switched a sequence

IABTechLab · Jun 26, 2024 · 16cfa2a · 16cfa2a
1 parent 1676d64
commit 16cfa2a
Show file tree

Hide file tree

Showing 4 changed files with 109 additions and 9 deletions.
diff --git a/src/main/java/com/uid2/admin/Main.java b/src/main/java/com/uid2/admin/Main.java
@@ -201,17 +201,21 @@ public void run() {
             RotatingS3KeyProvider s3KeyProvider = new RotatingS3KeyProvider(cloudStorage, s3KeyGlobalScope);
             S3KeyStoreWriter s3KeyStoreWriter = new S3KeyStoreWriter(s3KeyProvider, fileManager, jsonWriter, versionGenerator, clock, s3KeyGlobalScope);
             S3KeyManager s3KeyManager = new S3KeyManager(s3KeyProvider, s3KeyStoreWriter);
-            s3KeyManager.generateKeysForOperators(operatorKeyProvider.getAll(), config.getLong("s3_key_activates_in_seconds"), config.getInteger("s3_key_count_per_site"));
             try {
-                s3KeyProvider.loadContent(s3KeyProvider.getMetadata());
+                s3KeyProvider.loadContent();
             } catch (CloudStorageException e) {
                 if (e.getMessage().contains("The specified key does not exist")) {
                     s3KeyStoreWriter.upload(new HashMap<>(), null);
                     s3KeyProvider.loadContent(s3KeyProvider.getMetadata());
-                } else {
+                } else if (e.getMessage().contains("AmazonS3Exception: test-core-bucket")){
+                    s3KeyStoreWriter.upload(new HashMap<>(), null);
+                    s3KeyProvider.loadContent(s3KeyProvider.getMetadata());
+                }else{
                     throw e;
                 }
             }
+            s3KeyManager.generateKeysForOperators(operatorKeyProvider.getAll(), config.getLong("s3_key_activates_in_seconds"), config.getInteger("s3_key_count_per_site"));
+
 
             String enclaveMetadataPath = config.getString(EnclaveIdentifierProvider.ENCLAVES_METADATA_PATH);
             EnclaveIdentifierProvider enclaveIdProvider = new EnclaveIdentifierProvider(cloudStorage, enclaveMetadataPath);

diff --git a/src/main/java/com/uid2/admin/managers/S3KeyManager.java b/src/main/java/com/uid2/admin/managers/S3KeyManager.java
@@ -101,19 +101,29 @@ public void generateKeysForOperators(Collection<OperatorKey> operatorKeys, long
         if (keyCountPerSite <= 0) {
             throw new IllegalArgumentException("Key count per site must be greater than zero");
         }
+
         Set<Integer> uniqueSiteIds = new HashSet<>();
         for (OperatorKey operatorKey : operatorKeys) {
             uniqueSiteIds.add(operatorKey.getSiteId());
         }
 
+        System.out.println("Unique Site IDs: " + uniqueSiteIds); // Log the unique site IDs
+
         for (Integer siteId : uniqueSiteIds) {
+            System.out.println("Checking if site has keys: " + siteId); // Log the current site ID being checked
             if (!doesSiteHaveKeys(siteId)) {
+                System.out.println("Site does not have keys: " + siteId); // Log if the site does not have keys
                 for (int i = 0; i < keyCountPerSite; i++) {
                     long created = Instant.now().getEpochSecond();
                     long activated = created + (i * keyActivateInterval);
-                    addS3Key(generateS3Key(siteId, activated, created));
+                    S3Key s3Key = generateS3Key(siteId, activated, created);
+                    addS3Key(s3Key);
+                    System.out.println("Generated and added S3Key: " + s3Key); // Log the generated key
                 }
+            } else {
+                System.out.println("Site already has keys: " + siteId); // Log if the site already has keys
             }
         }
     }
+
 }
diff --git a/src/main/java/com/uid2/admin/store/writer/S3KeyStoreWriter.java b/src/main/java/com/uid2/admin/store/writer/S3KeyStoreWriter.java
@@ -33,7 +33,7 @@ public S3KeyStoreWriter(StoreReader<Map<Integer, S3Key>> provider, FileManager f
     public void upload(Map<Integer, S3Key> data, JsonObject extraMeta) throws Exception {
         JsonArray jsonS3Keys = new JsonArray();
         for (Map.Entry<Integer, S3Key> s3KeyEntry : data.entrySet()) {
-            jsonS3Keys.add(JsonObject.mapFrom(s3KeyEntry.getValue()));
+            jsonS3Keys.add(s3KeyEntry.getValue());
         }
         writer.upload(jsonS3Keys.encodePrettily(), extraMeta);
     }

diff --git a/src/test/java/com/uid2/admin/managers/S3KeyManagerTest.java b/src/test/java/com/uid2/admin/managers/S3KeyManagerTest.java
@@ -2,17 +2,15 @@
 
 import ch.qos.logback.classic.Logger;
 import com.uid2.admin.store.writer.S3KeyStoreWriter;
+import com.uid2.shared.auth.OperatorKey;
 import com.uid2.shared.model.S3Key;
 import com.uid2.shared.secret.IKeyGenerator;
 import com.uid2.shared.store.reader.RotatingS3KeyProvider;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.ArgumentCaptor;
 
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
+import java.util.*;
 
 import static org.junit.jupiter.api.Assertions.*;
 import static org.mockito.ArgumentMatchers.any;
@@ -189,4 +187,92 @@ void testCreateAndAddImmediateS3Key() throws Exception {
 
         verify(s3KeyStoreWriter, times(1)).upload(any(Map.class), eq(null));
     }
+
+    @Test
+    public void testDoesSiteHaveKeys_SiteHasKeys() {
+        int siteId = 1;
+        S3Key s3Key = new S3Key(siteId, siteId, 0L, 0L, "key");
+        Map<Integer, S3Key> allKeys = new HashMap<>();
+        allKeys.put(1, s3Key);
+
+        when(s3KeyProvider.getAll()).thenReturn(allKeys);
+
+        boolean result = s3KeyManager.doesSiteHaveKeys(siteId);
+        assertTrue(result);
+    }
+    @Test
+    public void testDoesSiteHaveKeys_SiteDoesNotHaveKeys() {
+        int siteId = 1;
+        Map<Integer, S3Key> allKeys = new HashMap<>();
+
+        when(s3KeyProvider.getAll()).thenReturn(allKeys);
+
+        boolean result = s3KeyManager.doesSiteHaveKeys(siteId);
+        assertFalse(result);
+    }
+
+    @Test
+    public void testDoesSiteHaveKeys_AllKeysNull() {
+        int siteId = 1;
+
+        when(s3KeyProvider.getAll()).thenReturn(null);
+
+        boolean result = s3KeyManager.doesSiteHaveKeys(siteId);
+        assertFalse(result);
+    }
+
+    @Test
+    public void testDoesSiteHaveKeys_MultipleKeysDifferentSiteIds() {
+        S3Key s3Key1 = new S3Key(1, 1, 0L, 0L, "key1");
+        S3Key s3Key2 = new S3Key(2, 2, 0L, 0L, "key2");
+        Map<Integer, S3Key> allKeys = new HashMap<>();
+        allKeys.put(1, s3Key1);
+        allKeys.put(2, s3Key2);
+
+        when(s3KeyProvider.getAll()).thenReturn(allKeys);
+
+        assertTrue(s3KeyManager.doesSiteHaveKeys(1));
+        assertTrue(s3KeyManager.doesSiteHaveKeys(2));
+        assertFalse(s3KeyManager.doesSiteHaveKeys(3)); // Site ID 3 does not exist
+    }
+
+    @Test
+    public void testDoesSiteHaveKeys_SameSiteIdMultipleKeys() {
+        int siteId = 1;
+        S3Key s3Key1 = new S3Key(siteId, siteId, 0L, 0L, "key1");
+        S3Key s3Key2 = new S3Key(siteId, siteId, 0L, 0L, "key2");
+        Map<Integer, S3Key> allKeys = new HashMap<>();
+        allKeys.put(1, s3Key1);
+        allKeys.put(2, s3Key2);
+
+        when(s3KeyProvider.getAll()).thenReturn(allKeys);
+
+        boolean result = s3KeyManager.doesSiteHaveKeys(siteId);
+        assertTrue(result);
+    }
+
+    @Test
+    public void testDoesSiteHaveKeys_LargeNumberOfKeys() {
+        Map<Integer, S3Key> allKeys = new HashMap<>();
+        for (int i = 1; i <= 1000; i++) {
+            S3Key s3Key = new S3Key(i, i, 0L, 0L, "key" + i);
+            allKeys.put(i, s3Key);
+        }
+
+        when(s3KeyProvider.getAll()).thenReturn(allKeys);
+
+        for (int i = 1; i <= 1000; i++) {
+            assertTrue(s3KeyManager.doesSiteHaveKeys(i));
+        }
+        assertFalse(s3KeyManager.doesSiteHaveKeys(1001)); // Site ID 1001 does not exist
+    }
+
+    @Test
+    public void testDoesSiteHaveKeys_EmptyKeys() {
+        when(s3KeyProvider.getAll()).thenReturn(new HashMap<>());
+
+        assertFalse(s3KeyManager.doesSiteHaveKeys(1));
+    }
+
+
 }