Merge pull request #567 from IABTechLab/sch-UID2-5851-migration-to-key-rotation

sch-UID2-5851 migration from salt to key rotation

Author: sophia-chen-ttd

pom.xml

@@ -16,7 +16,7 @@
         <!-- check micrometer.version vertx-micrometer-metrics consumes before bumping up -->
         <micrometer.version>1.12.2</micrometer.version>
         <junit-jupiter.version>5.11.2</junit-jupiter.version>
-        <uid2-shared.version>10.9.4</uid2-shared.version>
+        <uid2-shared.version>11.0.0</uid2-shared.version>
         <okta-jwt.version>0.5.10</okta-jwt.version>
         <image.version>${project.version}</image.version>
     </properties>

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,56 @@
+package com.uid2.admin.salt;
+
+import com.uid2.shared.model.SaltEntry;
+
+import java.util.Arrays;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ *  Assumptions:
+ *  - The latest assigned key ids are from the latest updated buckets
+ *  - Key ids from these buckets will always be monotonically increasing (apart from wraparound) as they have not rotated again after last assignment
+ *
+ * Intended outcomes of KeyIdGenerator:
+ * - Key ids are always monotonically increasing, starting from 0
+ * - When the last allocated key id reaches 16777215, the next key id will wrap around to 0
+ * - Continuing to increment from the highest key id will result in monotonic incrementation of key ids for all newly rotated buckets
+ **/
+public class KeyIdGenerator {
+    private static final int MAX_KEY_ID = 16777215; // 3 bytes
+    private final AtomicInteger nextActiveKeyId;
+
+    public KeyIdGenerator(SaltEntry[] buckets) {
+        this.nextActiveKeyId = new AtomicInteger(getNextActiveKeyId(buckets));
+    }
+
+    private static int getNextActiveKeyId(SaltEntry[] buckets) {
+        long lastUpdatedTimestampWithKey = Arrays.stream(buckets).filter(s -> s.currentKeySalt() != null).mapToLong(SaltEntry::lastUpdated).max().orElse(0);
+        if (lastUpdatedTimestampWithKey == 0) return 0;
+
+        int[] lastActiveKeyIdsSorted = Arrays.stream(buckets)
+                .filter(s -> s.lastUpdated() == lastUpdatedTimestampWithKey && s.currentKeySalt() != null)
+                .mapToInt(s -> s.currentKeySalt().id())
+                .sorted()
+                .toArray();
+
+        int highestId = lastActiveKeyIdsSorted[lastActiveKeyIdsSorted.length - 1];
+
+        int nextKeyId = highestId + 1;
+        if (nextKeyId <= MAX_KEY_ID) return nextKeyId;
+
+        // Wrapped case - find the last consecutive ID from 0
+        for (int i = 0; i < lastActiveKeyIdsSorted.length - 1; i++) {
+            if (lastActiveKeyIdsSorted[i + 1] - lastActiveKeyIdsSorted[i] > 1) {
+                return lastActiveKeyIdsSorted[i] + 1;
+            }
+        }
+
+        return 0;
+    }
+
+    public int getNextKeyId() {
+        return nextActiveKeyId.getAndUpdate(id ->
+                id + 1 > MAX_KEY_ID ? 0 : id + 1
+        );
+    }
+}

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(
@@ -57,6 +62,7 @@ public Result rotateSalts(
         logSaltAges("refreshable-salts", targetDate, refreshableSalts);
         logSaltAges("rotated-salts", targetDate, saltsToRotate);
         logSaltAges("total-salts", targetDate, Arrays.asList(postRotationSalts));
+        logBucketFormatCount(targetDate, postRotationSalts);
 
         var nextSnapshot = new SaltSnapshot(
                 nextEffective,
@@ -99,45 +105,85 @@ 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();
-        var lastUpdated = shouldRotate ? targetDate.asEpochMs() : oldSalt.lastUpdated();
-        var refreshFrom = calculateRefreshFrom(oldSalt, targetDate);
-        var previousSalt = calculatePreviousSalt(oldSalt, shouldRotate, targetDate);
+    private SaltEntry updateSalt(SaltEntry oldBucket, TargetDate targetDate, boolean shouldRotate, KeyIdGenerator keyIdGenerator) throws Exception {
+        var lastUpdated = shouldRotate ? targetDate.asEpochMs() : oldBucket.lastUpdated();
+        var refreshFrom = calculateRefreshFrom(oldBucket, targetDate);
+        var currentSalt = calculateCurrentSalt(oldBucket, shouldRotate);
+        var previousSalt = calculatePreviousSalt(oldBucket, shouldRotate, targetDate);
+        var currentKeySalt = calculateCurrentKeySalt(oldBucket, shouldRotate, keyIdGenerator);
+        var previousKeySalt = calculatePreviousKeySalt(oldBucket,shouldRotate, targetDate);
 
         return new SaltEntry(
-                oldSalt.id(),
-                oldSalt.hashedId(),
+                oldBucket.id(),
+                oldBucket.hashedId(),
                 lastUpdated,
                 currentSalt,
                 refreshFrom,
                 previousSalt,
-                null,
-                null
+                currentKeySalt,
+                previousKeySalt
         );
     }
 
-    private long calculateRefreshFrom(SaltEntry salt, TargetDate targetDate) {
-        long multiplier = targetDate.saltAgeInDays(salt) / 30 + 1;
-        return Instant.ofEpochMilli(salt.lastUpdated()).truncatedTo(ChronoUnit.DAYS).toEpochMilli() + (multiplier * THIRTY_DAYS_IN_MS);
+    private long calculateRefreshFrom(SaltEntry bucket, TargetDate targetDate) {
+        long multiplier = targetDate.saltAgeInDays(bucket) / 30 + 1;
+        return Instant.ofEpochMilli(bucket.lastUpdated()).truncatedTo(ChronoUnit.DAYS).toEpochMilli() + (multiplier * THIRTY_DAYS_IN_MS);
+    }
+
+    private String calculateCurrentSalt(SaltEntry bucket, boolean shouldRotate) throws Exception {
+        if (shouldRotate) {
+            if (ENABLE_V4_RAW_UID) {
+                return null;
+            }
+            else {
+                return this.keyGenerator.generateRandomKeyString(32);
+            }
+        }
+        return bucket.currentSalt();
     }
 
-    private String calculatePreviousSalt(SaltEntry salt, boolean shouldRotate, TargetDate targetDate) {
+    private String calculatePreviousSalt(SaltEntry bucket, boolean shouldRotate, TargetDate targetDate) {
         if (shouldRotate) {
-            return salt.currentSalt();
+            return bucket.currentSalt();
         }
-        if (targetDate.saltAgeInDays(salt) < 90) {
-            return salt.previousSalt();
+        if (targetDate.saltAgeInDays(bucket) < 90) {
+            return bucket.previousSalt();
+        }
+        return null;
+    }
+
+    private SaltEntry.KeyMaterial calculateCurrentKeySalt(SaltEntry bucket, 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 bucket.currentKeySalt();
+    }
+
+    private SaltEntry.KeyMaterial calculatePreviousKeySalt(SaltEntry bucket, boolean shouldRotate, TargetDate targetDate) {
+        if (shouldRotate) {
+            return bucket.currentKeySalt();
+        }
+        if (targetDate.saltAgeInDays(bucket) < 90) {
+            return bucket.previousKeySalt();
         }
         return null;
     }
@@ -207,6 +253,24 @@ private void logSaltAges(String saltCountType, TargetDate targetDate, Collection
         }
     }
 
+
+    /** Logging to monitor migration of buckets from salts (old format - v2/v3) to encryption keys (new format - v4) **/
+    private void logBucketFormatCount(TargetDate targetDate, SaltEntry[] postRotationBuckets) {
+        int totalKeys = 0, totalSalts = 0, totalPreviousKeys = 0, totalPreviousSalts = 0;
+
+        for (SaltEntry bucket : postRotationBuckets) {
+            if (bucket.currentKeySalt() != null) totalKeys++;
+            if (bucket.currentSalt() != null) totalSalts++;
+            if (bucket.previousKeySalt() != null) totalPreviousKeys++;
+            if (bucket.previousSalt() != null) totalPreviousSalts++;
+        }
+
+        LOGGER.info("UID bucket format: target_date={} bucket_format={} bucket_count={}", targetDate, "total-current-key-buckets", totalKeys);
+        LOGGER.info("UID bucket format: target_date={} bucket_format={} bucket_count={}", targetDate, "total-current-salt-buckets", totalSalts);
+        LOGGER.info("UID bucket format: target_date={} bucket_format={} bucket_count={}", targetDate, "total-previous-key-buckets", totalPreviousKeys);
+        LOGGER.info("UID bucket format: target_date={} bucket_format={} bucket_count={}", targetDate, "total-previous-salt-buckets", totalPreviousSalts);
+    }
+
     @Getter
     public static final class Result {
         private final SaltSnapshot snapshot; // can be null if new snapshot is not needed

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

@@ -22,16 +22,16 @@ 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()));
 
         stringBuilder.append(",");
         stringBuilder.append(serializeNullable(entry.previousSalt()));
 
-        appendKey(stringBuilder, entry.currentKey());
-        appendKey(stringBuilder, entry.previousKey());
+        appendKey(stringBuilder, entry.currentKeySalt());
+        appendKey(stringBuilder, entry.previousKeySalt());
 
         stringBuilder.append("\n");
     }

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

@@ -4,7 +4,6 @@
 import com.uid2.admin.store.version.VersionGenerator;
 import com.uid2.shared.cloud.CloudStorageException;
 import com.uid2.shared.cloud.TaggableCloudStorage;
-import com.uid2.shared.model.SaltEntry;
 import com.uid2.shared.store.CloudPath;
 import com.uid2.shared.store.salt.RotatingSaltProvider;
 import io.vertx.core.json.JsonArray;

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
     }
   ]
 }