From dbfe4fd5aceb68426e551ae032cf9757dadeed82 Mon Sep 17 00:00:00 2001
From: Jeffrey Griffin <jeff@signal.org>
Date: Wed, 19 Dec 2018 12:08:24 -0800
Subject: [PATCH] replicate directory updates on GCM/APN token & signed prekey
 changes

---
 .../textsecuregcm/WhisperServerService.java   |  2 +-
 .../controllers/AccountController.java        | 21 ++++++++++++--
 .../controllers/KeysController.java           | 28 ++++++++++++++-----
 .../workers/DeleteUserCommand.java            |  3 ++
 .../controllers/AccountControllerTest.java    |  2 +-
 .../tests/controllers/KeyControllerTest.java  | 10 ++++---
 6 files changed, 50 insertions(+), 16 deletions(-)

diff --git a/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index aeb9f0487..f425fe869 100644
--- a/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -211,7 +211,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     environment.lifecycle().manage(directoryReconciler);
 
     AttachmentController attachmentController = new AttachmentController(rateLimiters, urlSigner);
-    KeysController       keysController       = new KeysController(rateLimiters, keys, accountsManager);
+    KeysController       keysController       = new KeysController(rateLimiters, keys, accountsManager, directoryQueue);
     MessageController    messageController    = new MessageController(rateLimiters, pushSender, receiptSender, accountsManager, messagesManager, apnFallbackManager);
     ProfileController    profileController    = new ProfileController(rateLimiters , accountsManager, config.getProfilesConfiguration());
 
diff --git a/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java b/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java
index c6bb2d3e5..29d14d3f2 100644
--- a/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java
+++ b/src/main/java/org/whispersystems/textsecuregcm/controllers/AccountController.java
@@ -270,7 +270,8 @@ public class AccountController {
   @Path("/gcm/")
   @Consumes(MediaType.APPLICATION_JSON)
   public void setGcmRegistrationId(@Auth Account account, @Valid GcmRegistrationId registrationId) {
-    Device device = account.getAuthenticatedDevice().get();
+    Device  device           = account.getAuthenticatedDevice().get();
+    boolean wasAccountActive = account.isActive();
 
     if (device.getGcmId() != null &&
         device.getGcmId().equals(registrationId.getGcmRegistrationId()))
@@ -286,6 +287,10 @@ public class AccountController {
     else                                     device.setFetchesMessages(false);
 
     accounts.update(account);
+
+    if (!wasAccountActive && account.isActive()) {
+      directoryQueue.addRegisteredUser(account.getNumber());
+    }
   }
 
   @Timed
@@ -308,12 +313,18 @@ public class AccountController {
   @Path("/apn/")
   @Consumes(MediaType.APPLICATION_JSON)
   public void setApnRegistrationId(@Auth Account account, @Valid ApnRegistrationId registrationId) {
-    Device device = account.getAuthenticatedDevice().get();
+    Device  device           = account.getAuthenticatedDevice().get();
+    boolean wasAccountActive = account.isActive();
+
     device.setApnId(registrationId.getApnRegistrationId());
     device.setVoipApnId(registrationId.getVoipRegistrationId());
     device.setGcmId(null);
     device.setFetchesMessages(true);
     accounts.update(account);
+
+    if (!wasAccountActive && account.isActive()) {
+      directoryQueue.addRegisteredUser(account.getNumber());
+    }
   }
 
   @Timed
@@ -413,7 +424,11 @@ public class AccountController {
       newUserMeter.mark();
     }
 
-    directoryQueue.addRegisteredUser(number);
+    if (account.isActive()) {
+      directoryQueue.addRegisteredUser(number);
+    } else {
+      directoryQueue.deleteRegisteredUser(number);
+    }
 
     messagesManager.clear(number);
     pendingAccounts.remove(number);
diff --git a/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java b/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
index a53b1507c..421cb7643 100644
--- a/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
+++ b/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
@@ -29,6 +29,7 @@ import org.whispersystems.textsecuregcm.entities.PreKeyResponseItem;
 import org.whispersystems.textsecuregcm.entities.PreKeyState;
 import org.whispersystems.textsecuregcm.entities.SignedPreKey;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
+import org.whispersystems.textsecuregcm.sqs.DirectoryQueue;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
@@ -61,11 +62,13 @@ public class KeysController {
   private final RateLimiters    rateLimiters;
   private final Keys            keys;
   private final AccountsManager accounts;
+  private final DirectoryQueue  directoryQueue;
 
-  public KeysController(RateLimiters rateLimiters, Keys keys, AccountsManager accounts) {
-    this.rateLimiters = rateLimiters;
-    this.keys         = keys;
-    this.accounts     = accounts;
+  public KeysController(RateLimiters rateLimiters, Keys keys, AccountsManager accounts, DirectoryQueue directoryQueue) {
+    this.rateLimiters   = rateLimiters;
+    this.keys           = keys;
+    this.accounts       = accounts;
+    this.directoryQueue = directoryQueue;
   }
 
   @GET
@@ -84,8 +87,9 @@ public class KeysController {
   @PUT
   @Consumes(MediaType.APPLICATION_JSON)
   public void setKeys(@Auth Account account, @Valid PreKeyState preKeys)  {
-    Device  device        = account.getAuthenticatedDevice().get();
-    boolean updateAccount = false;
+    Device  device           = account.getAuthenticatedDevice().get();
+    boolean wasAccountActive = account.isActive();
+    boolean updateAccount    = false;
 
     if (!preKeys.getSignedPreKey().equals(device.getSignedPreKey())) {
       device.setSignedPreKey(preKeys.getSignedPreKey());
@@ -99,6 +103,10 @@ public class KeysController {
 
     if (updateAccount) {
       accounts.update(account);
+
+      if (!wasAccountActive && account.isActive()) {
+        directoryQueue.addRegisteredUser(account.getNumber());
+      }
     }
 
     keys.store(account.getNumber(), device.getId(), preKeys.getPreKeys());
@@ -158,9 +166,15 @@ public class KeysController {
   @Path("/signed")
   @Consumes(MediaType.APPLICATION_JSON)
   public void setSignedKey(@Auth Account account, @Valid SignedPreKey signedPreKey) {
-    Device device = account.getAuthenticatedDevice().get();
+    Device  device           = account.getAuthenticatedDevice().get();
+    boolean wasAccountActive = account.isActive();
+
     device.setSignedPreKey(signedPreKey);
     accounts.update(account);
+
+    if (!wasAccountActive && account.isActive()) {
+      directoryQueue.addRegisteredUser(account.getNumber());
+    }
   }
 
   @Timed
diff --git a/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java b/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java
index a7f04a0e1..49f683d10 100644
--- a/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java
+++ b/src/main/java/org/whispersystems/textsecuregcm/workers/DeleteUserCommand.java
@@ -10,6 +10,7 @@ import org.whispersystems.textsecuregcm.WhisperServerConfiguration;
 import org.whispersystems.textsecuregcm.auth.AuthenticationCredentials;
 import org.whispersystems.textsecuregcm.providers.RedisClientFactory;
 import org.whispersystems.textsecuregcm.redis.ReplicatedJedisPool;
+import org.whispersystems.textsecuregcm.sqs.DirectoryQueue;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.Accounts;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
@@ -76,6 +77,7 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
       Accounts            accounts        = dbi.onDemand(Accounts.class);
       ReplicatedJedisPool cacheClient     = new RedisClientFactory(configuration.getCacheConfiguration().getUrl(), configuration.getCacheConfiguration().getReplicaUrls()).getRedisClientPool();
       ReplicatedJedisPool redisClient     = new RedisClientFactory(configuration.getDirectoryConfiguration().getRedisConfiguration().getUrl(), configuration.getDirectoryConfiguration().getRedisConfiguration().getReplicaUrls()).getRedisClientPool();
+      DirectoryQueue      directoryQueue  = new DirectoryQueue(configuration.getDirectoryConfiguration().getSqsConfiguration());
       DirectoryManager    directory       = new DirectoryManager(redisClient);
       AccountsManager     accountsManager = new AccountsManager(accounts, directory, cacheClient);
 
@@ -94,6 +96,7 @@ public class DeleteUserCommand extends EnvironmentCommand<WhisperServerConfigura
             device.get().setAuthenticationCredentials(new AuthenticationCredentials(Base64.encodeBytes(random)));
 
             accountsManager.update(account.get());
+            directoryQueue.deleteRegisteredUser(account.get().getNumber());
 
             logger.warn("Removed " + account.get().getNumber());
           } else {
diff --git a/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/AccountControllerTest.java b/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/AccountControllerTest.java
index d924e2903..2f110c79b 100644
--- a/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/AccountControllerTest.java
+++ b/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/AccountControllerTest.java
@@ -236,7 +236,7 @@ public class AccountControllerTest {
     assertThat(response.getStatus()).isEqualTo(204);
 
     verify(accountsManager, times(1)).create(isA(Account.class));
-    verify(directoryQueue, times(1)).addRegisteredUser(eq(SENDER));
+    verify(directoryQueue, times(1)).deleteRegisteredUser(eq(SENDER));
   }
 
   @Test
diff --git a/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeyControllerTest.java b/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeyControllerTest.java
index 90377a81c..7a980c008 100644
--- a/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeyControllerTest.java
+++ b/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeyControllerTest.java
@@ -14,6 +14,7 @@ import org.whispersystems.textsecuregcm.entities.PreKeyState;
 import org.whispersystems.textsecuregcm.entities.SignedPreKey;
 import org.whispersystems.textsecuregcm.limits.RateLimiter;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
+import org.whispersystems.textsecuregcm.sqs.DirectoryQueue;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
@@ -54,9 +55,10 @@ public class KeyControllerTest {
   private final SignedPreKey SAMPLE_SIGNED_KEY2 = new SignedPreKey(2222, "foobar", "sig22");
   private final SignedPreKey SAMPLE_SIGNED_KEY3 = new SignedPreKey(3333, "barfoo", "sig33");
 
-  private final Keys            keys          = mock(Keys.class           );
-  private final AccountsManager accounts      = mock(AccountsManager.class);
-  private final Account         existsAccount = mock(Account.class        );
+  private final Keys            keys           = mock(Keys.class           );
+  private final AccountsManager accounts       = mock(AccountsManager.class);
+  private final DirectoryQueue  directoryQueue = mock(DirectoryQueue.class);
+  private final Account         existsAccount  = mock(Account.class        );
 
   private RateLimiters          rateLimiters  = mock(RateLimiters.class);
   private RateLimiter           rateLimiter   = mock(RateLimiter.class );
@@ -66,7 +68,7 @@ public class KeyControllerTest {
                                                             .addProvider(AuthHelper.getAuthFilter())
                                                             .addProvider(new AuthValueFactoryProvider.Binder<>(Account.class))
                                                             .setTestContainerFactory(new GrizzlyWebTestContainerFactory())
-                                                            .addResource(new KeysController(rateLimiters, keys, accounts))
+                                                            .addResource(new KeysController(rateLimiters, keys, accounts, directoryQueue))
                                                             .build();
 
   @Before