sch-UID2-5352 honouring refresh from in salt rotation

conf/default-config.json

@@ -1,3 +1,4 @@
 {
-  "enable_keysets": false
+  "enable_keysets": false,
+  "enable_salt_rotation_refresh_from": false
 }

conf/local-config.json

@@ -12,6 +12,7 @@
   "keys_acl_metadata_path": "keys_acl/metadata.json",
   "salts_metadata_path": "salts/metadata.json",
   "salt_snapshot_location_prefix": "salts/salts.txt.",
+  "enable_salt_rotation_refresh_from": false,
   "operators_metadata_path": "operators/metadata.json",
   "enclaves_metadata_path": "enclaves/metadata.json",
   "partners_metadata_path": "partners/metadata.json",

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

@@ -5,4 +5,5 @@ public class AdminConst {
     public static final String ROLE_OKTA_GROUP_MAP_MAINTAINER = "role_okta_group_map_maintainer";
     public static final String ROLE_OKTA_GROUP_MAP_PRIVILEGED = "role_okta_group_map_privileged";
     public static final String ROLE_OKTA_GROUP_MAP_SUPER_USER = "role_okta_group_map_super_user";
+    public static final String ENABLE_SALT_ROTATION_REFRESH_FROM = "enable_salt_rotation_refresh_from";
 }

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

@@ -232,7 +232,7 @@ public void run() {
             WriteLock writeLock = new WriteLock();
             KeyHasher keyHasher = new KeyHasher();
             IKeypairGenerator keypairGenerator = new SecureKeypairGenerator();
-            SaltRotation saltRotation = new SaltRotation(keyGenerator);
+            SaltRotation saltRotation = new SaltRotation(config, keyGenerator);
             EncryptionKeyService encryptionKeyService = new EncryptionKeyService(
                     config, auth, writeLock, encryptionKeyStoreWriter, keysetKeyStoreWriter, keyProvider, keysetKeysProvider, adminKeysetProvider, adminKeysetStoreWriter, keyGenerator, clock);
             KeysetManager keysetManager = new KeysetManager(

src/main/java/com/uid2/admin/secret/SaltRotation.java

@@ -1,10 +1,12 @@
 package com.uid2.admin.secret;
 
+import com.uid2.admin.AdminConst;
 import com.uid2.shared.model.SaltEntry;
 import com.uid2.shared.secret.IKeyGenerator;
 import com.uid2.shared.store.salt.RotatingSaltProvider;
 
 import com.uid2.shared.store.salt.RotatingSaltProvider.SaltSnapshot;
+import io.vertx.core.json.JsonObject;
 
 import java.time.Duration;
 import java.time.Instant;
@@ -17,12 +19,15 @@
 import static java.util.stream.Collectors.toList;
 
 public class SaltRotation {
+    private final static long THIRTY_DAYS_IN_MS = Duration.ofDays(30).toMillis();
+    private final static long DAY_IN_MS = Duration.ofDays(1).toMillis();
+
     private final IKeyGenerator keyGenerator;
-    private final long THIRTY_DAYS_IN_MS = Duration.ofDays(30).toMillis();
-    private final long DAY_IN_MS = Duration.ofDays(1).toMillis();
+    private final boolean isRefreshFromEnabled;
 
-    public SaltRotation(IKeyGenerator keyGenerator) {
+    public SaltRotation(JsonObject config, IKeyGenerator keyGenerator) {
         this.keyGenerator = keyGenerator;
+        this.isRefreshFromEnabled = config.getBoolean(AdminConst.ENABLE_SALT_ROTATION_REFRESH_FROM, false);
     }
 
     public Result rotateSalts(RotatingSaltProvider.SaltSnapshot lastSnapshot,
@@ -105,14 +110,15 @@ private List<Integer> pickSaltIndexesToRotate(
                 .sorted()
                 .toArray(Instant[]::new);
         final int maxSalts = (int) Math.ceil(lastSnapshot.getAllRotatingSalts().length * fraction);
+        final SaltEntry[] rotatableSalts = getRotatableSalts(lastSnapshot, nextEffective.toEpochMilli());
         final List<Integer> indexesToRotate = new ArrayList<>();
 
         Instant minLastUpdated = Instant.ofEpochMilli(0);
         for (Instant threshold : thresholds) {
             if (indexesToRotate.size() >= maxSalts) break;
             addIndexesToRotate(
                     indexesToRotate,
-                    lastSnapshot,
+                    rotatableSalts,
                     minLastUpdated.toEpochMilli(),
                     threshold.toEpochMilli(),
                     maxSalts - indexesToRotate.size()
@@ -122,12 +128,20 @@ private List<Integer> pickSaltIndexesToRotate(
         return indexesToRotate;
     }
 
+    private SaltEntry[] getRotatableSalts(SaltSnapshot lastSnapshot, long nextEffective) {
+        SaltEntry[] salts = lastSnapshot.getAllRotatingSalts();
+        if (isRefreshFromEnabled) {
+            return Arrays.stream(salts).filter(s -> s.refreshFrom() == nextEffective).toArray(SaltEntry[]::new);
+        }
+        return salts;
+    }
+
+
     private void addIndexesToRotate(List<Integer> entryIndexes,
-                                    SaltSnapshot lastSnapshot,
+                                    SaltEntry[] entries,
                                     long minLastUpdated,
                                     long maxLastUpdated,
                                     int maxIndexes) {
-        final SaltEntry[] entries = lastSnapshot.getAllRotatingSalts();
         final List<Integer> candidateIndexes = IntStream.range(0, entries.length)
                 .filter(i -> isBetween(entries[i].lastUpdated(), minLastUpdated, maxLastUpdated))
                 .boxed()

src/test/java/com/uid2/admin/secret/SaltRotationTest.java

@@ -1,8 +1,10 @@
 package com.uid2.admin.secret;
 
+import com.uid2.admin.AdminConst;
 import com.uid2.shared.model.SaltEntry;
 import com.uid2.shared.secret.IKeyGenerator;
 import com.uid2.shared.store.salt.RotatingSaltProvider;
+import io.vertx.core.json.JsonObject;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -38,7 +40,9 @@ private Instant daysLater(int days) {
     void setup() {
         MockitoAnnotations.openMocks(this);
 
-        saltRotation = new SaltRotation(keyGenerator);
+        JsonObject config = new JsonObject();
+
+        saltRotation = new SaltRotation(config, keyGenerator);
     }
 
     private static class SnapshotBuilder {
@@ -318,4 +322,44 @@ void rotateSaltsRemovePreviousSaltsOver90DaysOld() throws Exception {
         assertNull(salts[1].previousSalt());
     }
 
+
+    @Test
+    void rotateSaltsRotateWhenRefreshFromIsTargetDate() throws Exception {
+        JsonObject config = new JsonObject();
+        config.put(AdminConst.ENABLE_SALT_ROTATION_REFRESH_FROM, Boolean.TRUE);
+        saltRotation = new SaltRotation(config, keyGenerator);
+
+        final Duration[] minAges = {
+                Duration.ofDays(90),
+                Duration.ofDays(60),
+        };
+
+        var validForRotation1 = daysEarlier(120).toEpochMilli();
+        var validForRotation2 = daysEarlier(70).toEpochMilli();
+        var notValidForRotation = daysEarlier(30).toEpochMilli();
+        var refreshNow = targetDateAsInstant.toEpochMilli();
+        var refreshLater = daysLater(20).toEpochMilli();
+
+        var lastSnapshot = SnapshotBuilder.start()
+                .withEntries(
+                    new SaltEntry(1, "1", validForRotation1, "salt", refreshNow, null, null, null),
+                    new SaltEntry(2, "2", notValidForRotation, "salt", refreshNow, null, null, null),
+                    new SaltEntry(3, "3", validForRotation2, "salt", refreshLater, null, null, null)
+                )
+                .build(daysEarlier(1), daysLater(6));
+
+        var result = saltRotation.rotateSalts(lastSnapshot, minAges, 1, targetDate);
+        assertTrue(result.hasSnapshot());
+
+        var salts = result.getSnapshot().getAllRotatingSalts();
+
+        assertEquals(targetDateAsInstant.toEpochMilli(), salts[0].lastUpdated());
+        assertEquals(daysLater(30).toEpochMilli(), salts[0].refreshFrom());
+
+        assertEquals(notValidForRotation, salts[1].lastUpdated());
+        assertEquals(daysLater(30).toEpochMilli(), salts[1].refreshFrom());
+
+        assertEquals(validForRotation2, salts[2].lastUpdated());
+        assertEquals(refreshLater, salts[2].refreshFrom());
+    }
 }