implemented v4 key rotation behind feature flag

Author: sophia-chen-ttd

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

@@ -8,4 +8,5 @@ private 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_V4_RAW_UID = "enable_v4_raw_uid";
 }

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

@@ -233,7 +233,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(keyGenerator, config);
             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/salt/KeyIdGenerator.java

@@ -0,0 +1,52 @@
+package com.uid2.admin.salt;
+
+import com.uid2.shared.model.SaltEntry;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Arrays;
+import java.util.Comparator;
+
+public class KeyIdGenerator {
+    private static final int MAX_KEY_ID = 16777215; // 3 bytes
+    private int lastActiveKeyId;
+
+    public KeyIdGenerator(SaltEntry[] salts) {
+        this.lastActiveKeyId = getLastActiveKeyId(salts);
+    }
+
+    private int getLastActiveKeyId(SaltEntry[] salts) {
+        SaltEntry[] sortedSaltsByLastUpdated = Arrays.stream(salts).sorted(Comparator.comparingLong(SaltEntry::lastUpdated)).toArray(SaltEntry[]::new);
+        var lastUpdated = sortedSaltsByLastUpdated[sortedSaltsByLastUpdated.length - 1].lastUpdated();
+
+        var sortedKeyIds = Arrays.stream(sortedSaltsByLastUpdated)
+                .filter(s -> s.lastUpdated() == lastUpdated)
+                .filter(s -> s.currentKey() != null)
+                .mapToInt(s -> s.currentKey().id())
+                .sorted()
+                .toArray();
+
+        if (sortedKeyIds.length == 0) return MAX_KEY_ID;
+
+        if (sortedKeyIds[sortedKeyIds.length - 1] == MAX_KEY_ID) {
+            if (sortedKeyIds[0] == 0) {
+                for (int i = 1; i < sortedKeyIds.length; i++) {
+                    if (sortedKeyIds[i] - sortedKeyIds[i - 1] > 1) {
+                        return sortedKeyIds[i - 1];
+                    }
+                }
+                return 0;
+            }
+            return MAX_KEY_ID;
+        }
+        return sortedKeyIds[sortedKeyIds.length - 1];
+    }
+
+    public int getNextKeyId() {
+        this.lastActiveKeyId += 1;
+        if (this.lastActiveKeyId > MAX_KEY_ID) {
+            this.lastActiveKeyId = 0;
+        }
+        return this.lastActiveKeyId;
+    }
+}

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

@@ -1,10 +1,12 @@
 package com.uid2.admin.salt;
 
+import com.uid2.admin.AdminConst;
 import com.uid2.shared.model.SaltEntry;
 import com.uid2.shared.secret.IKeyGenerator;
 
 import com.uid2.shared.store.salt.ISaltProvider.ISaltSnapshot;
 import com.uid2.shared.store.salt.RotatingSaltProvider.SaltSnapshot;
+import io.vertx.core.json.JsonObject;
 import lombok.Getter;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -17,12 +19,15 @@
 public class SaltRotation {
     private static final long THIRTY_DAYS_IN_MS = Duration.ofDays(30).toMillis();
     private static final double MAX_SALT_PERCENTAGE = 0.8;
+    private final boolean ENABLE_V4_RAW_UID;
 
     private final IKeyGenerator keyGenerator;
+
     private static final Logger LOGGER = LoggerFactory.getLogger(SaltRotation.class);
 
-    public SaltRotation(IKeyGenerator keyGenerator) {
+    public SaltRotation(IKeyGenerator keyGenerator, JsonObject config) {
         this.keyGenerator = keyGenerator;
+        this.ENABLE_V4_RAW_UID = config.getBoolean(AdminConst.ENABLE_V4_RAW_UID, false);
     }
 
     public Result rotateSalts(
@@ -99,21 +104,24 @@ private boolean isRefreshable(TargetDate targetDate, SaltEntry salt) {
     }
 
     private SaltEntry[] rotateSalts(SaltEntry[] oldSalts, List<SaltEntry> saltsToRotate, TargetDate targetDate) throws Exception {
+        var keyIdGenerator = new KeyIdGenerator(oldSalts);
         var saltIdsToRotate = saltsToRotate.stream().map(SaltEntry::id).collect(Collectors.toSet());
 
         var updatedSalts = new SaltEntry[oldSalts.length];
         for (int i = 0; i < oldSalts.length; i++) {
             var shouldRotate = saltIdsToRotate.contains(oldSalts[i].id());
-            updatedSalts[i] = updateSalt(oldSalts[i], targetDate, shouldRotate);
+            updatedSalts[i] = updateSalt(oldSalts[i], targetDate, shouldRotate, keyIdGenerator);
         }
         return updatedSalts;
     }
 
-    private SaltEntry updateSalt(SaltEntry oldSalt, TargetDate targetDate, boolean shouldRotate) throws Exception {
-        var currentSalt = shouldRotate ? this.keyGenerator.generateRandomKeyString(32) : oldSalt.currentSalt();
+    private SaltEntry updateSalt(SaltEntry oldSalt, TargetDate targetDate, boolean shouldRotate, KeyIdGenerator keyIdGenerator) throws Exception {
         var lastUpdated = shouldRotate ? targetDate.asEpochMs() : oldSalt.lastUpdated();
         var refreshFrom = calculateRefreshFrom(oldSalt, targetDate);
+        var currentSalt = calculateCurrentSalt(oldSalt, shouldRotate);
         var previousSalt = calculatePreviousSalt(oldSalt, shouldRotate, targetDate);
+        var currentKey = calculateCurrentKey(oldSalt, shouldRotate, keyIdGenerator);
+        var previousKey = calculatePreviousKey(oldSalt,shouldRotate, targetDate);
 
         return new SaltEntry(
                 oldSalt.id(),
@@ -122,8 +130,8 @@ private SaltEntry updateSalt(SaltEntry oldSalt, TargetDate targetDate, boolean s
                 currentSalt,
                 refreshFrom,
                 previousSalt,
-                null,
-                null
+                currentKey,
+                previousKey
         );
     }
 
@@ -132,6 +140,12 @@ private long calculateRefreshFrom(SaltEntry salt, TargetDate targetDate) {
         return Instant.ofEpochMilli(salt.lastUpdated()).truncatedTo(ChronoUnit.DAYS).toEpochMilli() + (multiplier * THIRTY_DAYS_IN_MS);
     }
 
+    private String calculateCurrentSalt(SaltEntry salt, boolean shouldRotate) throws Exception {
+        return shouldRotate ?
+                ENABLE_V4_RAW_UID ? null : this.keyGenerator.generateRandomKeyString(32)
+            : salt.currentSalt();
+    }
+
     private String calculatePreviousSalt(SaltEntry salt, boolean shouldRotate, TargetDate targetDate) {
         if (shouldRotate) {
             return salt.currentSalt();
@@ -142,6 +156,31 @@ private String calculatePreviousSalt(SaltEntry salt, boolean shouldRotate, Targe
         return null;
     }
 
+    private SaltEntry.KeyMaterial calculateCurrentKey(SaltEntry salt, boolean shouldRotate, KeyIdGenerator keyIdGenerator) throws Exception {
+        if (shouldRotate) {
+            if (ENABLE_V4_RAW_UID) {
+                return new SaltEntry.KeyMaterial(
+                        keyIdGenerator.getNextKeyId(),
+                        this.keyGenerator.generateRandomKeyString(32),
+                        this.keyGenerator.generateRandomKeyString(32)
+                );
+            } else {
+                return null;
+            }
+        }
+        return salt.currentKey();
+    }
+
+    private SaltEntry.KeyMaterial calculatePreviousKey(SaltEntry salt, boolean shouldRotate, TargetDate targetDate) {
+        if (shouldRotate) {
+            return salt.currentKey();
+        }
+        if (targetDate.saltAgeInDays(salt) < 90) {
+            return salt.previousKey();
+        }
+        return null;
+    }
+
     private List<SaltEntry> pickSaltsToRotate(
             Set<SaltEntry> refreshableSalts,
             TargetDate targetDate,

src/main/java/com/uid2/admin/store/writer/SaltSerializer.java

@@ -22,7 +22,7 @@ private static void addLine(SaltEntry entry, StringBuilder stringBuilder) {
                 .append(",")
                 .append(lastUpdated % 1000 == 0 ? lastUpdated + 1 : lastUpdated)
                 .append(",")
-                .append(entry.currentSalt());
+                .append(serializeNullable(entry.currentSalt()));
 
         stringBuilder.append(",");
         stringBuilder.append(serializeNullable(entry.refreshFrom()));

src/main/resources/localstack/s3/core/salts/metadata.json

@@ -11,10 +11,10 @@
       "location" : "salts/salts.txt.1670796729291",
       "size" : 2
     },{
-      "effective" : 1745907348982,
-      "expires" : 1766720293000,
-      "location" : "salts/salts.txt.1745907348982",
-      "size" : 2
+      "effective" : 1755648000000,
+      "expires" : 1756252800000,
+      "location" : "salts/salts.txt.1755648000000",
+      "size" : 1001
     }
   ]
 }