Had to adjust to work with 1mil salts

IABTechLab · Dec 11, 2024 · 1a23d0d · 1a23d0d
1 parent b2e9b68
commit 1a23d0d
Show file tree

Hide file tree

Showing 2 changed files with 77 additions and 16 deletions.
diff --git a/src/main/java/com/uid2/shared/util/CloudEncryptionHelpers.java b/src/main/java/com/uid2/shared/util/CloudEncryptionHelpers.java
@@ -1,6 +1,10 @@
 package com.uid2.shared.util;
 
 import java.io.InputStream;
+
+import com.fasterxml.jackson.core.JsonFactory;
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.core.JsonToken;
 import com.uid2.shared.encryption.AesGcm;
 import com.uid2.shared.model.CloudEncryptionKey;
 
@@ -13,31 +17,37 @@
 
 public class CloudEncryptionHelpers {
     public static String decryptInputStream(InputStream inputStream, RotatingCloudEncryptionKeyProvider cloudEncryptionKeyProvider) throws IOException {
-        String encryptedContent = inputStreamToString(inputStream);
-        JsonObject json = new JsonObject(encryptedContent);
-        int keyId = json.getInteger("key_id");
-        String encryptedPayload = json.getString("encrypted_payload");
+        JsonFactory factory = new JsonFactory();
+        JsonParser parser = factory.createParser(inputStream);
+        int keyId = -1;
+        byte[] encryptedPayload = null;
+        parser.nextToken();
+        while (parser.nextToken() != JsonToken.END_OBJECT) {
+                String fieldName = parser.getCurrentName();
+                if(fieldName.equals("key_id")) {
+                    parser.nextToken();
+                    keyId = parser.getIntValue();
+                }
+                if(fieldName.equals("encrypted_payload")) {
+                    parser.nextToken();
+                    encryptedPayload = parser.getBinaryValue();
+                }
+        }
+
+        if(keyId == -1 || encryptedPayload == null) {
+            throw new IllegalStateException("failed to parse json");
+        }
+
         CloudEncryptionKey decryptionKey = cloudEncryptionKeyProvider.getKey(keyId);
 
         if (decryptionKey == null) {
             throw new IllegalStateException("No matching S3 key found for decryption for key ID: " + keyId);
         }
 
         byte[] secret = Base64.getDecoder().decode(decryptionKey.getSecret());
-        byte[] encryptedBytes = Base64.getDecoder().decode(encryptedPayload);
+        byte[] encryptedBytes = encryptedPayload;
         byte[] decryptedBytes = AesGcm.decrypt(encryptedBytes, 0, secret);
 
         return new String(decryptedBytes, StandardCharsets.UTF_8);
     }
-
-    public static String inputStreamToString(InputStream inputStream) throws IOException {
-        try (BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8))) {
-            StringBuilder stringBuilder = new StringBuilder();
-            String line;
-            while ((line = reader.readLine()) != null) {
-                stringBuilder.append(line);
-            }
-            return stringBuilder.toString();
-        }
-    }
 }
diff --git a/src/test/java/com/uid2/shared/store/EncryptedRotatingSaltProviderTest.java b/src/test/java/com/uid2/shared/store/EncryptedRotatingSaltProviderTest.java
@@ -123,6 +123,57 @@ public void loadSaltSingleVersion() throws Exception {
         assertTrue(snapshot.getModifiedSince(Instant.now().minus(1, ChronoUnit.HOURS)).isEmpty());
     }
 
+    @Test
+    public void loadSaltSingleVersion1mil() throws Exception {
+        final String FIRST_LEVEL_SALT = "first_level_salt_value";
+        final String ID_PREFIX = "a";
+        final String ID_SECRET = "m3yMIcbg9vCaFLJsn4m4PfruZnvAZ72OxmFG5QsGMOw=";
+
+        final Instant generatedTime = Instant.now().minus(1, ChronoUnit.DAYS);
+        final Instant expireTime = Instant.now().plus(365, ChronoUnit.DAYS);
+
+        final JsonObject metadataJson = new JsonObject();
+        {
+            metadataJson.put("version", 2);
+            metadataJson.put("generated", generatedTime.getEpochSecond() * 1000L);
+            metadataJson.put("first_level", FIRST_LEVEL_SALT);
+            metadataJson.put("id_prefix", ID_PREFIX);
+            metadataJson.put("id_secret", ID_SECRET);
+            final JsonArray saltsRefList = new JsonArray();
+            {
+                final JsonObject saltsRef = new JsonObject();
+                saltsRef.put("effective", generatedTime.getEpochSecond() * 1000L);
+                saltsRef.put("expires", expireTime.getEpochSecond() * 1000L);
+                saltsRef.put("location", "salts.txt");
+                saltsRef.put("size", 1000000);
+                saltsRefList.add(saltsRef);
+            }
+            metadataJson.put("salts", saltsRefList);
+        }
+
+        final String effectiveTimeString = String.valueOf(generatedTime.getEpochSecond() * 1000L);
+        StringBuilder salts = new StringBuilder();
+        for (int i = 0; i < 1000000; i++) {
+            salts.append(i).append(",").append(effectiveTimeString).append(",").append("salt-string").append("\n");
+        }
+
+        when(cloudStorage.download("metadata"))
+                .thenReturn(new ByteArrayInputStream(metadataJson.toString().getBytes(StandardCharsets.US_ASCII)));
+        when(cloudStorage.download("salts.txt"))
+                .thenReturn(getEncryptedStream(salts.toString()));
+
+        EncryptedRotatingSaltProvider saltsProvider = new EncryptedRotatingSaltProvider(
+                cloudStorage, "metadata", keyProvider);
+
+        final JsonObject loadedMetadata = saltsProvider.getMetadata();
+        saltsProvider.loadContent(loadedMetadata);
+        assertEquals(2, saltsProvider.getVersion(loadedMetadata));
+
+        final ISaltProvider.ISaltSnapshot snapshot = saltsProvider.getSnapshot(Instant.now());
+        assertEquals(FIRST_LEVEL_SALT, snapshot.getFirstLevelSalt());
+        assertTrue(snapshot.getModifiedSince(Instant.now().minus(1, ChronoUnit.HOURS)).isEmpty());
+    }
+
     @Test
     public void loadSaltMultipleVersions() throws Exception {
         final String FIRST_LEVEL_SALT = "first_level_salt_value";