Addressing PR comments

IABTechLab · Dec 10, 2024 · e741256 · e741256
1 parent 3a808a2
commit e741256
Show file tree

Hide file tree

Showing 9 changed files with 448 additions and 93 deletions.
diff --git a/src/main/java/com/uid2/shared/store/EncryptedRotatingSaltProvider.java b/src/main/java/com/uid2/shared/store/EncryptedRotatingSaltProvider.java
@@ -0,0 +1,33 @@
+package com.uid2.shared.store;
+
+import com.uid2.shared.cloud.DownloadCloudStorage;
+import com.uid2.shared.model.SaltEntry;
+import com.uid2.shared.store.reader.RotatingCloudEncryptionKeyProvider;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Collection;
+
+import static com.uid2.shared.util.CloudEncryptionHelpers.decryptInputStream;
+
+public class EncryptedRotatingSaltProvider extends RotatingSaltProvider {
+    private final RotatingCloudEncryptionKeyProvider cloudEncryptionKeyProvider;
+
+    public EncryptedRotatingSaltProvider(DownloadCloudStorage fileStreamProvider, String metadataPath, RotatingCloudEncryptionKeyProvider cloudEncryptionKeyProvider) {
+        super(fileStreamProvider, metadataPath);
+        this.cloudEncryptionKeyProvider = cloudEncryptionKeyProvider;
+    }
+
+    @Override
+    protected SaltEntry[] readInputStream(InputStream inputStream, SaltEntryBuilder entryBuilder, Integer size) throws IOException {
+        String decrypted = decryptInputStream(inputStream, cloudEncryptionKeyProvider);
+        SaltEntry[] entries = new SaltEntry[size];
+        int idx = 0;
+        for (String line : decrypted.split("\n")) {
+            final SaltEntry entry = entryBuilder.toEntry(line);
+            entries[idx] = entry;
+            idx++;
+        }
+        return entries;
+    }
+}
diff --git a/src/main/java/com/uid2/shared/store/EncryptedScopedStoreReader.java b/src/main/java/com/uid2/shared/store/EncryptedScopedStoreReader.java
@@ -1,22 +1,18 @@
 package com.uid2.shared.store;
 
 import com.uid2.shared.cloud.DownloadCloudStorage;
-import com.uid2.shared.model.CloudEncryptionKey;
 import com.uid2.shared.store.parser.Parser;
 import com.uid2.shared.store.parser.ParsingResult;
 import com.uid2.shared.store.scope.StoreScope;
 import com.uid2.shared.store.reader.RotatingCloudEncryptionKeyProvider;
-import io.vertx.core.json.JsonObject;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.*;
 
-import com.uid2.shared.encryption.AesGcm;
-
 import java.nio.charset.StandardCharsets;
-import java.util.Base64;
-import java.util.Map;
+
+import static com.uid2.shared.util.CloudEncryptionHelpers.decryptInputStream;
 
 public class EncryptedScopedStoreReader<T> extends ScopedStoreReader<T> {
     private static final Logger LOGGER = LoggerFactory.getLogger(EncryptedScopedStoreReader.class);
@@ -31,8 +27,7 @@ public EncryptedScopedStoreReader(DownloadCloudStorage fileStreamProvider, Store
     @Override
     protected long loadContent(String path) throws Exception {
         try (InputStream inputStream = this.contentStreamProvider.download(path)) {
-            String encryptedContent = inputStreamToString(inputStream);
-            String decryptedContent = getDecryptedContent(encryptedContent);
+            String decryptedContent = decryptInputStream(inputStream, cloudEncryptionKeyProvider);
             ParsingResult<T> parsed = this.parser.deserialize(new ByteArrayInputStream(decryptedContent.getBytes(StandardCharsets.UTF_8)));
             latestSnapshot.set(parsed.getData());
 
@@ -45,32 +40,4 @@ protected long loadContent(String path) throws Exception {
             throw e;
         }
     }
-
-    protected String getDecryptedContent(String encryptedContent) throws Exception {
-        JsonObject json = new JsonObject(encryptedContent);
-        int keyId = json.getInteger("key_id");
-        String encryptedPayload = json.getString("encrypted_payload");
-        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[] 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/main/java/com/uid2/shared/store/RotatingSaltProvider.java b/src/main/java/com/uid2/shared/store/RotatingSaltProvider.java
@@ -132,28 +132,24 @@ private SaltSnapshot loadSnapshot(JsonObject spec, String firstLevelSalt, SaltEn
         final Instant expires = Instant.ofEpochMilli(spec.getLong("expires", defaultExpires.toEpochMilli()));
 
         final String path = spec.getString("location");
-        int idx = 0;
-        final SaltEntry[] entries = new SaltEntry[spec.getInteger("size")];
-        Stream<String> stream = readInputStream(this.contentStreamProvider.download(path)).lines();
-        for (String l : stream.toList()) {
-            final SaltEntry entry = entryBuilder.toEntry(l);
-            entries[idx] = entry;
-            idx++;
-        }
+        Integer size = spec.getInteger("size");
+        SaltEntry[] entries = readInputStream(this.contentStreamProvider.download(path), entryBuilder, size);
 
-        LOGGER.info("Loaded " + idx + " salts");
+        LOGGER.info("Loaded " + size + " salts");
         return new SaltSnapshot(effective, expires, entries, firstLevelSalt);
     }
 
-    protected String readInputStream(InputStream inputStream) throws IOException {
+    protected SaltEntry[] readInputStream(InputStream inputStream, SaltEntryBuilder entryBuilder, Integer size) throws IOException {
         try (BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8))) {
-            StringBuilder stringBuilder = new StringBuilder();
             String line;
+            SaltEntry[] entries = new SaltEntry[size];
+            int idx = 0;
             while ((line = reader.readLine()) != null) {
-                stringBuilder.append(line);
-                stringBuilder.append(System.lineSeparator());
+                final SaltEntry entry = entryBuilder.toEntry(line);
+                entries[idx] = entry;
+                idx++;
             }
-            return stringBuilder.toString();
+            return entries;
         }
     }
 
@@ -225,7 +221,7 @@ public String encode(long id) {
         }
     }
 
-    static final class SaltEntryBuilder {
+    protected static final class SaltEntryBuilder {
         private final IdHashingScheme idHashingScheme;
 
         public SaltEntryBuilder(IdHashingScheme idHashingScheme) {

diff --git a/src/main/java/com/uid2/shared/store/reader/RotatingCloudEncryptionKeyProvider.java b/src/main/java/com/uid2/shared/store/reader/RotatingCloudEncryptionKeyProvider.java
@@ -62,7 +62,12 @@ public Map<Integer, CloudEncryptionKey> getAll() {
     }
 
     public CloudEncryptionKey getKey(int id) {
-        return reader.getSnapshot().get(id);
+        Map<Integer, CloudEncryptionKey> snapshot = reader.getSnapshot();
+        if(snapshot == null) {
+            return null;
+        }
+
+        return snapshot.get(id);
     }
 
     public void updateSiteToKeysMapping() {

diff --git a/.../store/RotatingEncryptedSaltProvider.java → ...2/shared/util/CloudEncryptionHelpers.java b/.../store/RotatingEncryptedSaltProvider.java → ...2/shared/util/CloudEncryptionHelpers.java
@@ -1,30 +1,19 @@
-package com.uid2.shared.store;
+package com.uid2.shared.util;
 
-import com.uid2.shared.cloud.DownloadCloudStorage;
+import java.io.InputStream;
 import com.uid2.shared.encryption.AesGcm;
 import com.uid2.shared.model.CloudEncryptionKey;
-import com.uid2.shared.store.reader.RotatingCloudEncryptionKeyProvider;
 
-import com.uid2.shared.store.reader.StoreReader;
+import com.uid2.shared.store.reader.RotatingCloudEncryptionKeyProvider;
 import io.vertx.core.json.JsonObject;
-import java.io.IOException;
-import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
 import java.util.Base64;
-import java.util.Collection;
 
-public class RotatingEncryptedSaltProvider extends RotatingSaltProvider implements StoreReader<Collection<RotatingSaltProvider.SaltSnapshot>> {
-    private final RotatingCloudEncryptionKeyProvider cloudEncryptionKeyProvider;
-
-    public RotatingEncryptedSaltProvider(DownloadCloudStorage fileStreamProvider, String metadataPath, RotatingCloudEncryptionKeyProvider cloudEncryptionKeyProvider) {
-        super(fileStreamProvider, metadataPath);
-        this.cloudEncryptionKeyProvider = cloudEncryptionKeyProvider;
-    }
-
-    @Override
-    protected String readInputStream(InputStream inputStream) throws IOException {
-        String encryptedContent = super.readInputStream(inputStream);
+import java.io.*;
 
+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");
@@ -41,8 +30,14 @@ protected String readInputStream(InputStream inputStream) throws IOException {
         return new String(decryptedBytes, StandardCharsets.UTF_8);
     }
 
-    @Override
-    public Collection<SaltSnapshot> getAll() {
-        return super.getSnapshots();
+    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();
+        }
     }
 }