Fix experiments in BackupAuthManager

signalapp · Mar 20, 2024 · e7d1ead · e7d1ead
1 parent a9b5359
commit e7d1ead
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 13 deletions.
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -700,7 +700,7 @@ public void run(WhisperServerConfiguration config, Environment environment) thro
     ServerZkReceiptOperations zkReceiptOperations = new ServerZkReceiptOperations(zkSecretParams);
 
     Cdn3BackupCredentialGenerator cdn3BackupCredentialGenerator = new Cdn3BackupCredentialGenerator(config.getTus());
-    BackupAuthManager backupAuthManager = new BackupAuthManager(dynamicConfigurationManager, rateLimiters, accountsManager, backupsGenericZkSecretParams, clock);
+    BackupAuthManager backupAuthManager = new BackupAuthManager(experimentEnrollmentManager, rateLimiters, accountsManager, backupsGenericZkSecretParams, clock);
     BackupsDb backupsDb = new BackupsDb(
         dynamoDbAsyncClient,
         config.getDynamoDbTables().getBackups().getTableName(),

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/backup/BackupAuthManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/backup/BackupAuthManager.java
@@ -21,6 +21,7 @@
 import org.signal.libsignal.zkgroup.backups.BackupAuthCredentialResponse;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
 import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
+import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
@@ -43,19 +44,19 @@ public class BackupAuthManager {
   final static String BACKUP_EXPERIMENT_NAME = "backup";
   final static String BACKUP_MEDIA_EXPERIMENT_NAME = "backupMedia";
 
-  private final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager;
+  private final ExperimentEnrollmentManager experimentEnrollmentManager;
   private final GenericServerSecretParams serverSecretParams;
   private final Clock clock;
   private final RateLimiters rateLimiters;
   private final AccountsManager accountsManager;
 
   public BackupAuthManager(
-      final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager,
+      final ExperimentEnrollmentManager experimentEnrollmentManager,
       final RateLimiters rateLimiters,
       final AccountsManager accountsManager,
       final GenericServerSecretParams serverSecretParams,
       final Clock clock) {
-    this.dynamicConfigurationManager = dynamicConfigurationManager;
+    this.experimentEnrollmentManager = experimentEnrollmentManager;
     this.rateLimiters = rateLimiters;
     this.accountsManager = accountsManager;
     this.serverSecretParams = serverSecretParams;
@@ -157,9 +158,6 @@ private Optional<Long> receiptLevel(final Account account) {
   }
 
   private boolean inExperiment(final String experimentName, final Account account) {
-    return dynamicConfigurationManager.getConfiguration()
-        .getExperimentEnrollmentConfiguration(experimentName)
-        .map(config -> config.getEnrolledUuids().contains(account.getUuid()))
-        .orElse(false);
+    return this.experimentEnrollmentManager.isEnrolled(account.getUuid(), experimentName);
   }
 }
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/tests/util/ExperimentHelper.java b/service/src/test/java/org/whispersystems/textsecuregcm/tests/util/ExperimentHelper.java
@@ -7,6 +7,7 @@
 
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
 import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicExperimentEnrollmentConfiguration;
+import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
 import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
 import java.util.Optional;
 import java.util.Set;
@@ -17,7 +18,7 @@
 
 public class ExperimentHelper {
 
-  public static DynamicConfigurationManager<DynamicConfiguration> withEnrollment(
+  private static DynamicConfigurationManager<DynamicConfiguration> withEnrollment(
       final String experimentName,
       final Set<UUID> enrolledUuids,
       final int enrollmentPercentage) {
@@ -31,11 +32,11 @@ public static DynamicConfigurationManager<DynamicConfiguration> withEnrollment(
     return dcm;
   }
 
-  public static DynamicConfigurationManager<DynamicConfiguration> withEnrollment(final String experimentName, final Set<UUID> enrolledUuids) {
-    return withEnrollment(experimentName, enrolledUuids, 0);
+  public static ExperimentEnrollmentManager withEnrollment(final String experimentName, final Set<UUID> enrolledUuids) {
+    return new ExperimentEnrollmentManager(withEnrollment(experimentName, enrolledUuids, 0));
   }
 
-  public static DynamicConfigurationManager<DynamicConfiguration> withEnrollment(final String experimentName, final UUID enrolledUuid) {
-    return withEnrollment(experimentName, Set.of(enrolledUuid), 0);
+  public static ExperimentEnrollmentManager withEnrollment(final String experimentName, final UUID enrolledUuid) {
+    return new ExperimentEnrollmentManager(withEnrollment(experimentName, Set.of(enrolledUuid), 0));
   }
 }