Create s3 keys for newly added operators

pom.xml

@@ -6,7 +6,7 @@
 
     <groupId>com.uid2</groupId>
     <artifactId>uid2-admin</artifactId>
-    <version>5.10.0</version>
+    <version>5.10.1-alpha-81-SNAPSHOT</version>
 
     <properties>
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
@@ -16,7 +16,7 @@
         <!-- check micrometer.version vertx-micrometer-metrics consumes before bumping up -->
         <micrometer.version>1.1.0</micrometer.version>
         <junit-jupiter.version>5.7.0</junit-jupiter.version>
-        <uid2-shared.version>7.15.0</uid2-shared.version>
+        <uid2-shared.version>7.16.0</uid2-shared.version>
         <okta-jwt.version>0.5.8</okta-jwt.version>
         <image.version>${project.version}</image.version>
     </properties>

src/main/java/com/uid2/admin/Main.java

@@ -255,7 +255,7 @@ public void run() {
                     clientSideKeypairService,
                     new ServiceService(auth, writeLock, serviceStoreWriter, serviceProvider, siteProvider, serviceLinkProvider),
                     new ServiceLinkService(auth, writeLock, serviceLinkStoreWriter, serviceLinkProvider, serviceProvider, siteProvider),
-                    new OperatorKeyService(config, auth, writeLock, operatorKeyStoreWriter, operatorKeyProvider, siteProvider, keyGenerator, keyHasher),
+                    new OperatorKeyService(config, auth, writeLock, operatorKeyStoreWriter, operatorKeyProvider, siteProvider, keyGenerator, keyHasher, s3KeyManager),
                     new SaltService(auth, writeLock, saltStoreWriter, saltProvider, saltRotation),
                     new SiteService(auth, writeLock, siteStoreWriter, siteProvider, clientKeyProvider),
                     new PartnerConfigService(auth, writeLock, partnerStoreWriter, partnerConfigProvider),

src/main/java/com/uid2/admin/managers/S3KeyManager.java

@@ -36,7 +36,7 @@ S3Key generateS3Key(int siteId, long activates, long created) throws Exception {
     }
 
     String generateSecret() throws Exception {
-        // Assuming want to generate a 32-byte key
+        //Generate a 32-byte key for AesGcm
         return keyGenerator.generateRandomKeyString(32);
     }
 
@@ -46,15 +46,6 @@ void addS3Key(S3Key s3Key) throws Exception {
         s3KeyStoreWriter.upload(s3Keys, null);
     }
 
-    // Method to create and add an S3 key that activates immediately for a specific site
-    public S3Key createAndAddImmediate3Key(int siteId) throws Exception {
-        int newKeyId = getNextKeyId();
-        long created = Instant.now().getEpochSecond();
-        S3Key newKey = new S3Key(newKeyId, siteId, created, created, generateSecret());
-        addS3Key(newKey);
-        return newKey;
-    }
-
     int getNextKeyId() {
         Map<Integer, S3Key> s3Keys = s3KeyProvider.getAll();
         if (s3Keys == null || s3Keys.isEmpty()) {
@@ -63,8 +54,17 @@ int getNextKeyId() {
         return s3Keys.keySet().stream().max(Integer::compareTo).orElse(0) + 1;
     }
 
-    public S3Key getS3Key(int id) {
-        return s3KeyProvider.getAll().get(id);
+    // Method to create and upload an S3 key that activates immediately for a specific site, for emergency rotation
+    public S3Key createAndAddImmediate3Key(int siteId) throws Exception {
+        int newKeyId = getNextKeyId();
+        long created = Instant.now().getEpochSecond();
+        S3Key newKey = new S3Key(newKeyId, siteId, created, created, generateSecret());
+        addS3Key(newKey);
+        return newKey;
+    }
+
+    public S3Key getS3KeyByKeyIdentifier(int keyIdentifier) {
+        return s3KeyProvider.getAll().get(keyIdentifier);
     }
 
     public Optional<S3Key> getS3KeyBySiteId(int siteId) {
@@ -98,6 +98,7 @@ int countKeysForSite(int siteId) {
 
     public void generateKeysForOperators(Collection<OperatorKey> operatorKeys, long keyActivateInterval, int keyCountPerSite) throws Exception {
         this.s3KeyProvider.loadContent();
+
         if (operatorKeys == null || operatorKeys.isEmpty()) {
             throw new IllegalArgumentException("Operator keys collection must not be null or empty");
         }
@@ -108,12 +109,14 @@ public void generateKeysForOperators(Collection<OperatorKey> operatorKeys, long
             throw new IllegalArgumentException("Key count per site must be greater than zero");
         }
 
+        // Extract all the unique site IDs from input operator keys collection
         Set<Integer> uniqueSiteIds = new HashSet<>();
         for (OperatorKey operatorKey : operatorKeys) {
             uniqueSiteIds.add(operatorKey.getSiteId());
         }
 
         for (Integer siteId : uniqueSiteIds) {
+            // Check if the site ID already exists in the S3 key provider and has fewer than the required number of keys
             int currentKeyCount = countKeysForSite(siteId);
             if (currentKeyCount < keyCountPerSite) {
                 int keysToGenerate = keyCountPerSite - currentKeyCount;
@@ -123,6 +126,9 @@ public void generateKeysForOperators(Collection<OperatorKey> operatorKeys, long
                     S3Key s3Key = generateS3Key(siteId, activated, created);
                     addS3Key(s3Key);
                 }
+                LOGGER.info("Generated " + keysToGenerate + " keys for site ID " + siteId);
+            } else {
+                LOGGER.info("Site ID " + siteId + " already has the required number of keys. Skipping key generation.");
             }
         }
     }

src/main/java/com/uid2/admin/vertx/service/OperatorKeyService.java

@@ -3,6 +3,7 @@
 import com.fasterxml.jackson.databind.ObjectWriter;
 import com.uid2.admin.auth.AdminAuthMiddleware;
 import com.uid2.admin.auth.RevealedKey;
+import com.uid2.admin.managers.S3KeyManager;
 import com.uid2.shared.model.Site;
 import com.uid2.shared.secret.IKeyGenerator;
 import com.uid2.admin.store.writer.OperatorKeyStoreWriter;
@@ -45,6 +46,9 @@ public class OperatorKeyService implements IService {
     private final IKeyGenerator keyGenerator;
     private final KeyHasher keyHasher;
     private final String operatorKeyPrefix;
+    private final S3KeyManager s3KeyManager;
+    private final long s3KeyActivatesInSeconds;
+    private final int s3KeyCountPerSite;
 
     public OperatorKeyService(JsonObject config,
                               AdminAuthMiddleware auth,
@@ -53,16 +57,20 @@ public OperatorKeyService(JsonObject config,
                               RotatingOperatorKeyProvider operatorKeyProvider,
                               RotatingSiteStore siteProvider,
                               IKeyGenerator keyGenerator,
-                              KeyHasher keyHasher) {
+                              KeyHasher keyHasher,
+                              S3KeyManager s3KeyManager) {
         this.auth = auth;
         this.writeLock = writeLock;
         this.operatorKeyStoreWriter = operatorKeyStoreWriter;
         this.operatorKeyProvider = operatorKeyProvider;
         this.siteProvider = siteProvider;
         this.keyGenerator = keyGenerator;
         this.keyHasher = keyHasher;
+        this.s3KeyManager = s3KeyManager;
 
         this.operatorKeyPrefix = config.getString("operator_key_prefix");
+        this.s3KeyActivatesInSeconds = config.getLong("s3_key_activates_in_seconds",0L);
+        this.s3KeyCountPerSite = config.getInteger("s3_key_count_per_site",0);
     }
 
     @Override
@@ -266,6 +274,8 @@ private void handleOperatorAdd(RoutingContext rc) {
             // upload to storage
             operatorKeyStoreWriter.upload(operators);
 
+            s3KeyManager.generateKeysForOperators(Collections.singletonList(newOperator), s3KeyActivatesInSeconds, s3KeyCountPerSite);
+
             // respond with new key
             rc.response().end(JSON_WRITER.writeValueAsString(new RevealedKey<>(newOperator, key)));
         } catch (Exception e) {
@@ -372,12 +382,14 @@ private void handleOperatorUpdate(RoutingContext rc) {
                 return;
             }
 
+            boolean siteIdChanged = false;
             if (!rc.queryParam("site_id").isEmpty()) {
                 final Site site = RequestUtil.getSiteFromParam(rc, "site_id", this.siteProvider);
                 if (site == null) {
                     ResponseUtil.error(rc, 404, "site id not found");
                     return;
                 }
+                siteIdChanged = true;
                 existingOperator.setSiteId(site.getId());
             }
 
@@ -400,6 +412,10 @@ private void handleOperatorUpdate(RoutingContext rc) {
             // upload to storage
             operatorKeyStoreWriter.upload(operators);
 
+            if (siteIdChanged) {
+                s3KeyManager.generateKeysForOperators(Collections.singletonList(existingOperator), s3KeyActivatesInSeconds, s3KeyCountPerSite);
+            }
+
             // return the updated client
             rc.response().end(JSON_WRITER.writeValueAsString(existingOperator));
         } catch (Exception e) {

src/test/java/com/uid2/admin/managers/S3KeyManagerTest.java

@@ -11,6 +11,7 @@
 import org.mockito.ArgumentCaptor;
 
 import java.util.*;
+import java.util.stream.Collectors;
 
 import static org.junit.jupiter.api.Assertions.*;
 import static org.mockito.ArgumentMatchers.any;
@@ -104,7 +105,7 @@ void testGetS3Key() {
         existingKeys.put(1, s3Key);
         when(s3KeyProvider.getAll()).thenReturn(existingKeys);
 
-        S3Key result = s3KeyManager.getS3Key(1);
+        S3Key result = s3KeyManager.getS3KeyByKeyIdentifier(1);
 
         assertEquals(s3Key, result);
     }
@@ -292,4 +293,148 @@ void testCountKeysForSite() {
         assertEquals(1, countForSite2);
         assertEquals(0, countForSite3);
     }
+
+    @Test
+    void testGenerateKeysForOperators() throws Exception {
+        Collection<OperatorKey> operatorKeys = Arrays.asList(
+                createOperatorKey("hash1", 100),
+                createOperatorKey("hash2", 100),
+                createOperatorKey("hash3", 200)
+        );
+        long keyActivateInterval = 3600; // 1 hour
+        int keyCountPerSite = 3;
+
+        Map<Integer, S3Key> existingKeys = new HashMap<>();
+        existingKeys.put(1, new S3Key(1, 100, 1000L, 900L, "existingKey1"));
+        when(s3KeyProvider.getAll()).thenReturn(existingKeys);
+
+        doReturn("generatedSecret").when(s3KeyManager).generateSecret();
+
+        s3KeyManager.generateKeysForOperators(operatorKeys, keyActivateInterval, keyCountPerSite);
+
+        verify(s3KeyProvider, times(1)).loadContent();
+
+        ArgumentCaptor<S3Key> keyCaptor = ArgumentCaptor.forClass(S3Key.class);
+        verify(s3KeyManager, times(5)).addS3Key(keyCaptor.capture());
+
+        List<S3Key> capturedKeys = keyCaptor.getAllValues();
+        assertEquals(5, capturedKeys.size());
+
+        List<S3Key> site100Keys = capturedKeys.stream()
+                .filter(key -> key.getSiteId() == 100)
+                .sorted(Comparator.comparingLong(S3Key::getActivates))
+                .collect(Collectors.toList());
+        assertEquals(2, site100Keys.size());
+        assertTrue(site100Keys.get(1).getActivates() - site100Keys.get(0).getActivates() >= keyActivateInterval);
+
+        List<S3Key> site200Keys = capturedKeys.stream()
+                .filter(key -> key.getSiteId() == 200)
+                .sorted(Comparator.comparingLong(S3Key::getActivates))
+                .collect(Collectors.toList());
+        assertEquals(3, site200Keys.size());
+        assertTrue(site200Keys.get(1).getActivates() - site200Keys.get(0).getActivates() >= keyActivateInterval);
+    }
+
+    @Test
+    void testGenerateKeysForOperators_NoNewKeysNeeded() throws Exception {
+        Collection<OperatorKey> operatorKeys = Collections.singletonList(
+                createOperatorKey("hash1", 100)
+        );
+        long keyActivateInterval = 3600;
+        int keyCountPerSite = 3;
+
+        Map<Integer, S3Key> existingKeys = new HashMap<>();
+        existingKeys.put(1, new S3Key(1, 100, 1000L, 900L, "existingKey1"));
+        existingKeys.put(2, new S3Key(2, 100, 2000L, 1900L, "existingKey2"));
+        existingKeys.put(3, new S3Key(3, 100, 3000L, 2900L, "existingKey3"));
+        when(s3KeyProvider.getAll()).thenReturn(existingKeys);
+
+        s3KeyManager.generateKeysForOperators(operatorKeys, keyActivateInterval, keyCountPerSite);
+
+        verify(s3KeyManager, never()).addS3Key(any());
+    }
+
+    @Test
+    void testGenerateKeysForOperators_EmptyOperatorKeys() {
+        Collection<OperatorKey> operatorKeys = Collections.emptyList();
+        long keyActivateInterval = 3600;
+        int keyCountPerSite = 3;
+
+        assertThrows(IllegalArgumentException.class, () ->
+                s3KeyManager.generateKeysForOperators(operatorKeys, keyActivateInterval, keyCountPerSite)
+        );
+    }
+
+    @Test
+    void testGenerateKeysForOperators_InvalidKeyActivateInterval() {
+        Collection<OperatorKey> operatorKeys = Collections.singletonList(
+                createOperatorKey("hash1", 100)
+        );
+        long keyActivateInterval = 0;
+        int keyCountPerSite = 3;
+
+        assertThrows(IllegalArgumentException.class, () ->
+                s3KeyManager.generateKeysForOperators(operatorKeys, keyActivateInterval, keyCountPerSite)
+        );
+    }
+
+    @Test
+    void testGenerateKeysForOperators_InvalidKeyCountPerSite() {
+        Collection<OperatorKey> operatorKeys = Collections.singletonList(
+                createOperatorKey("hash1", 100)
+        );
+        long keyActivateInterval = 3600;
+        int keyCountPerSite = 0;
+
+        assertThrows(IllegalArgumentException.class, () ->
+                s3KeyManager.generateKeysForOperators(operatorKeys, keyActivateInterval, keyCountPerSite)
+        );
+    }
+
+    @Test
+    void testGenerateKeysForOperators_MultipleSitesWithVaryingExistingKeys() throws Exception {
+        Collection<OperatorKey> operatorKeys = Arrays.asList(
+                createOperatorKey("hash1", 100),
+                createOperatorKey("hash2", 200),
+                createOperatorKey("hash3", 300)
+        );
+        long keyActivateInterval = 3600;
+        int keyCountPerSite = 3;
+
+        Map<Integer, S3Key> existingKeys = new HashMap<>();
+        existingKeys.put(1, new S3Key(1, 100, 1000L, 900L, "existingKey1"));
+        existingKeys.put(2, new S3Key(2, 200, 2000L, 1900L, "existingKey2"));
+        existingKeys.put(3, new S3Key(3, 200, 3000L, 2900L, "existingKey3"));
+        when(s3KeyProvider.getAll()).thenReturn(existingKeys);
+
+        doReturn("generatedSecret").when(s3KeyManager).generateSecret();
+
+        s3KeyManager.generateKeysForOperators(operatorKeys, keyActivateInterval, keyCountPerSite);
+
+        ArgumentCaptor<S3Key> keyCaptor = ArgumentCaptor.forClass(S3Key.class);
+        verify(s3KeyManager, times(6)).addS3Key(keyCaptor.capture());
+
+        List<S3Key> capturedKeys = keyCaptor.getAllValues();
+        assertEquals(6, capturedKeys.size());
+
+        assertEquals(2, capturedKeys.stream().filter(key -> key.getSiteId() == 100).count());
+        assertEquals(1, capturedKeys.stream().filter(key -> key.getSiteId() == 200).count());
+        assertEquals(3, capturedKeys.stream().filter(key -> key.getSiteId() == 300).count());
+    }
+
+    private OperatorKey createOperatorKey(String keyHash, int siteId) {
+        return new OperatorKey(
+                keyHash,
+                "salt",
+                "name",
+                "contact",
+                "protocol",
+                System.currentTimeMillis(),
+                false,
+                siteId,
+                Collections.emptySet(),
+                null,
+                "keyId"
+        );
+    }
 }