Making change to multiscope writer to not copy all salts for each site

IABTechLab · Dec 11, 2024 · f820ccc · f820ccc
1 parent 3883853
commit f820ccc
Show file tree

Hide file tree

Showing 6 changed files with 36 additions and 31 deletions.
diff --git a/src/main/java/com/uid2/admin/job/EncryptionJob/SaltEncryptionJob.java b/src/main/java/com/uid2/admin/job/EncryptionJob/SaltEncryptionJob.java
@@ -10,6 +10,7 @@
 import com.uid2.shared.store.RotatingSaltProvider;
 
 import java.util.Collection;
+import java.util.List;
 
 public class SaltEncryptionJob extends Job {
     private final Collection<OperatorKey> globalOperators;
@@ -32,9 +33,9 @@ public String getId() {
 
     @Override
     public void execute() throws Exception {
-        PrivateSiteDataMap<RotatingSaltProvider.SaltSnapshot> desiredPrivateState = PrivateSiteUtil.getPrivateSaltEntries(saltEntries, globalOperators);
-        multiScopeStoreWriter.uploadPrivateWithEncryption(desiredPrivateState, null);
-        PrivateSiteDataMap<RotatingSaltProvider.SaltSnapshot> desiredPublicState = PublicSiteUtil.getPublicSaltEntries(saltEntries, globalOperators);
-        multiScopeStoreWriter.uploadPublicWithEncryption(desiredPublicState, null);
+        List<Integer> desiredPrivateState = PrivateSiteUtil.getPrivateSaltSites(globalOperators);
+        multiScopeStoreWriter.uploadPrivateWithEncryption(desiredPrivateState, saltEntries, null);
+        List<Integer> desiredPublicState = PublicSiteUtil.getPublicSaltSites(globalOperators);
+        multiScopeStoreWriter.uploadPublicWithEncryption(desiredPublicState, saltEntries, null);
     }
 }
diff --git a/src/main/java/com/uid2/admin/store/MultiScopeStoreWriter.java b/src/main/java/com/uid2/admin/store/MultiScopeStoreWriter.java
@@ -69,6 +69,13 @@ public void uploadPrivateWithEncryption(Map<Integer, T> desiredState, JsonObject
         }
     }
 
+    public void uploadPrivateWithEncryption(List<Integer> siteIds, T desiredState, JsonObject extraMeta) throws Exception {
+        EncryptedStoreFactory<T> encryptedFactory = (EncryptedStoreFactory<T>) factory;
+        for (Integer siteId : siteIds) {
+            encryptedFactory.getEncryptedWriter(siteId,false).upload(desiredState, extraMeta);
+        }
+    }
+
     public void uploadPublicWithEncryption(Map<Integer, T> desiredPublicState, JsonObject extraMeta) throws Exception {
         EncryptedStoreFactory<T> encryptedFactory = (EncryptedStoreFactory<T>) factory;
         for (Map.Entry<Integer, T> entry : desiredPublicState.entrySet()) {
@@ -77,6 +84,13 @@ public void uploadPublicWithEncryption(Map<Integer, T> desiredPublicState, JsonO
         }
     }
 
+    public void uploadPublicWithEncryption(List<Integer> siteIds, T desiredState, JsonObject extraMeta) throws Exception {
+        EncryptedStoreFactory<T> encryptedFactory = (EncryptedStoreFactory<T>) factory;
+        for (Integer siteId : siteIds) {
+            encryptedFactory.getEncryptedWriter(siteId,true).upload(desiredState, extraMeta);
+        }
+    }
+
     public static <K, V> boolean areMapsEqual(Map<K, V> a, Map<K, V> b) {
         return a.size() == b.size() && a.entrySet().stream().allMatch(b.entrySet()::contains);
     }

diff --git a/src/main/java/com/uid2/admin/store/factory/SaltStoreFactory.java b/src/main/java/com/uid2/admin/store/factory/SaltStoreFactory.java
@@ -13,10 +13,14 @@
 import com.uid2.shared.store.reader.StoreReader;
 import com.uid2.shared.store.scope.EncryptedScope;
 import io.vertx.core.json.JsonObject;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.util.Collection;
 
 public class SaltStoreFactory implements EncryptedStoreFactory<Collection<RotatingSaltProvider.SaltSnapshot>> {
+    private static final Logger LOGGER = LoggerFactory.getLogger(SaltStoreFactory.class);
+
     JsonObject config;
     CloudPath rootMetadatapath;
     FileManager fileManager;
@@ -45,6 +49,7 @@ public StoreWriter<Collection<RotatingSaltProvider.SaltSnapshot>> getEncryptedWr
 
     @Override
     public StoreReader<Collection<RotatingSaltProvider.SaltSnapshot>> getEncryptedReader(Integer siteId, boolean isPublic) {
+        LOGGER.warn("getEncryptedReader called on SaltStoreFactory. This method is not implemented.");
         return null;
     }
 
@@ -55,11 +60,13 @@ public RotatingCloudEncryptionKeyProvider getCloudEncryptionProvider() {
 
     @Override
     public StoreReader<Collection<RotatingSaltProvider.SaltSnapshot>> getReader(Integer siteId) {
+        LOGGER.warn("getReader called on SaltStoreFactory. This method is not implemented.");
         return null;
     }
 
     @Override
     public StoreWriter<Collection<RotatingSaltProvider.SaltSnapshot>> getWriter(Integer siteId) {
+        LOGGER.warn("getWriter called on SaltStoreFactory. This method is not implemented.");
         return null;
     }
 }
diff --git a/src/main/java/com/uid2/admin/util/PrivateSiteUtil.java b/src/main/java/com/uid2/admin/util/PrivateSiteUtil.java
@@ -313,17 +313,8 @@ public static PrivateSiteDataMap<KeysetKey> getKeysetKeys(Collection<OperatorKey
         return result;
     }
 
-    public static PrivateSiteDataMap<RotatingSaltProvider.SaltSnapshot> getPrivateSaltEntries(
-            Collection<RotatingSaltProvider.SaltSnapshot> globalSaltEntries,
-            Collection<OperatorKey> operators) {
+    public static List<Integer> getPrivateSaltSites(Collection<OperatorKey> operators) {
         final PrivateSiteDataMap<RotatingSaltProvider.SaltSnapshot> result = getPrivateSites(operators);
-
-        globalSaltEntries.forEach(saltEntry -> {
-            result.forEach((publicSiteId, publicSiteData) -> {
-                publicSiteData.add(saltEntry);
-            });
-        });
-
-        return result;
+        return result.keySet().stream().toList();
     }
 }
diff --git a/src/main/java/com/uid2/admin/util/PublicSiteUtil.java b/src/main/java/com/uid2/admin/util/PublicSiteUtil.java
@@ -11,10 +11,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Map;
+import java.util.*;
 
 public class PublicSiteUtil {
     private static final Logger LOGGER = LoggerFactory.getLogger(PrivateSiteUtil.class);
@@ -128,18 +125,9 @@ public static PrivateSiteDataMap<KeysetKey> getPublicKeysetKeys(
         return result;
     }
 
-    public static PrivateSiteDataMap<RotatingSaltProvider.SaltSnapshot> getPublicSaltEntries(
-            Collection<RotatingSaltProvider.SaltSnapshot> globalSaltEntries,
-            Collection<OperatorKey> operators) {
+    public static List<Integer> getPublicSaltSites(Collection<OperatorKey> operators) {
         final PrivateSiteDataMap<RotatingSaltProvider.SaltSnapshot> result = getPublicSitesMap(operators);
-
-        globalSaltEntries.forEach(saltEntry -> {
-            result.forEach((publicSiteId, publicSiteData) -> {
-                publicSiteData.add(saltEntry);
-            });
-        });
-
-        return result;
+        return result.keySet().stream().toList();
     }
 
     public static PrivateSiteDataMap<ClientSideKeypair> getPublicClientKeypairs(

diff --git a/src/test/java/com/uid2/admin/store/writer/EncryptedSaltStoreWriterTest.java b/src/test/java/com/uid2/admin/store/writer/EncryptedSaltStoreWriterTest.java
@@ -59,6 +59,8 @@ public class EncryptedSaltStoreWriterTest {
 
     @Captor
     private ArgumentCaptor<String> pathCaptor;
+    @Captor
+    private ArgumentCaptor<String> cloudPathCaptor;
 
     @BeforeEach
     public void setUp() throws Exception {
@@ -118,7 +120,9 @@ public void testUploadNew() throws Exception {
 
         encryptedSaltStoreWriter.upload(snapshot);
 
-        verify(taggableCloudStorage).upload(pathCaptor.capture(), any(), any());
+        verify(taggableCloudStorage).upload(pathCaptor.capture(), cloudPathCaptor.capture(), any());
+        assertEquals(cloudPathCaptor.getValue(), "test/path");
+
         verifyFile(pathCaptor.getValue(), snapshot);
     }
 }