Adding the start of what is needed for salt encryption

Author: cody-constine-ttd

src/main/java/com/uid2/shared/store/EncryptedScopedStoreReader.java

@@ -50,14 +50,7 @@ protected String getDecryptedContent(String encryptedContent) throws Exception {
         JsonObject json = new JsonObject(encryptedContent);
         int keyId = json.getInteger("key_id");
         String encryptedPayload = json.getString("encrypted_payload");
-        Map<Integer, CloudEncryptionKey> cloudEncryptionKeys = cloudEncryptionKeyProvider.getAll();
-        CloudEncryptionKey decryptionKey = null;
-        for (CloudEncryptionKey key : cloudEncryptionKeys.values()) {
-            if (key.getId() == keyId) {
-                decryptionKey = key;
-                break;
-            }
-        }
+        CloudEncryptionKey decryptionKey = cloudEncryptionKeyProvider.getKey(keyId);
 
         if (decryptionKey == null) {
             throw new IllegalStateException("No matching S3 key found for decryption for key ID: " + keyId);

src/main/java/com/uid2/shared/store/RotatingEncryptedSaltProvider.java

@@ -0,0 +1,41 @@
+package com.uid2.shared.store;
+
+import com.uid2.shared.cloud.DownloadCloudStorage;
+import com.uid2.shared.encryption.AesGcm;
+import com.uid2.shared.model.CloudEncryptionKey;
+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;
+
+public class RotatingEncryptedSaltProvider extends RotatingSaltProvider{
+    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);
+
+        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);
+    }
+}

src/main/java/com/uid2/shared/store/RotatingSaltProvider.java

@@ -13,6 +13,7 @@
 import org.hashids.Hashids;
 
 import java.io.BufferedReader;
+import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.nio.charset.StandardCharsets;
@@ -21,6 +22,7 @@
 import java.util.*;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 /*
   1. metadata.json format
@@ -132,20 +134,29 @@ private SaltSnapshot loadSnapshot(JsonObject spec, String firstLevelSalt, SaltEn
         final String path = spec.getString("location");
         int idx = 0;
         final SaltEntry[] entries = new SaltEntry[spec.getInteger("size")];
-
-        try (InputStream inputStream = this.contentStreamProvider.download(path);
-             InputStreamReader inputStreamReader = new InputStreamReader(inputStream, StandardCharsets.UTF_8);
-             BufferedReader reader = new BufferedReader(inputStreamReader)) {
-            for (String l; (l = reader.readLine()) != null; ++idx) {
-                final SaltEntry entry = entryBuilder.toEntry(l);
-                entries[idx] = entry;
-            }
+        Stream<String> stream = readInputStream(this.contentStreamProvider.download(path)).lines();
+        for (String l : stream.toList()) {
+            final SaltEntry entry = entryBuilder.toEntry(l);
+            entries[idx] = entry;
+            idx++;
         }
 
         LOGGER.info("Loaded " + idx + " salts");
         return new SaltSnapshot(effective, expires, entries, firstLevelSalt);
     }
 
+    protected String readInputStream(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);
+                stringBuilder.append(System.lineSeparator());
+            }
+            return stringBuilder.toString();
+        }
+    }
+
     public static class SaltSnapshot implements ISaltSnapshot {
         private final Instant effective;
         private final Instant expires;

src/main/java/com/uid2/shared/store/reader/RotatingCloudEncryptionKeyProvider.java

@@ -61,6 +61,10 @@ public Map<Integer, CloudEncryptionKey> getAll() {
         return keys != null ? keys : new HashMap<>();
     }
 
+    public CloudEncryptionKey getKey(int id) {
+        return reader.getSnapshot().get(id);
+    }
+
     public void updateSiteToKeysMapping() {
         Map<Integer, CloudEncryptionKey> allKeys = getAll();
         siteToKeysMap.clear();

src/test/java/com/uid2/shared/store/EncryptedScopedStoreReaderTest.java

@@ -20,6 +20,8 @@
 import static com.uid2.shared.TestUtilites.toInputStream;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -52,6 +54,7 @@ void setUp() {
         Map<Integer, CloudEncryptionKey> mockKeyMap = new HashMap<>();
         mockKeyMap.put(encryptionKey.getId(), encryptionKey);
         when(keyProvider.getAll()).thenReturn(mockKeyMap);
+        when(keyProvider.getKey(1)).thenReturn(mockKeyMap.get(1));
     }
 
     @Test
@@ -127,6 +130,7 @@ void testDecryptionOfEncryptedContent() throws Exception {
     void testHandlingInvalidEncryptionKey() throws Exception {
         // Set key provider to return an empty map
         when(keyProvider.getAll()).thenReturn(new HashMap<>());
+        when(keyProvider.getKey(anyInt())).thenReturn(null);
 
         String secretKey = encryptionKey.getSecret();
         byte[] secretKeyBytes = Base64.getDecoder().decode(secretKey);
@@ -159,6 +163,7 @@ void testLoadWithMultipleEncryptionKeys() throws Exception {
         mockKeyMap.put(encryptionKey.getId(), encryptionKey);
         mockKeyMap.put(newKey.getId(), newKey);
         when(keyProvider.getAll()).thenReturn(mockKeyMap);
+        when(keyProvider.getKey(2)).thenReturn(mockKeyMap.get(2));
 
         byte[] encryptedPayload = AesGcm.encrypt("value1,value2".getBytes(StandardCharsets.UTF_8), newKeyBytes);
         String encryptedPayloadBase64 = Base64.getEncoder().encodeToString(encryptedPayload);