From a211f6aed92dabd6661da012b2f4bf97d481ad16 Mon Sep 17 00:00:00 2001
From: Moxie Marlinspike <moxie@thoughtcrime.org>
Date: Wed, 2 Dec 2015 15:06:09 -0800
Subject: [PATCH] Delete pending messages for an unlinked device

// FREEBIE
---
 .../textsecuregcm/WhisperServerService.java           |  2 +-
 .../textsecuregcm/controllers/DeviceController.java   |  5 +++++
 .../textsecuregcm/storage/Messages.java               |  3 +++
 .../textsecuregcm/storage/MessagesManager.java        |  4 ++++
 .../tests/controllers/DeviceControllerTest.java       | 11 +++++++++--
 5 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index 107b98bd3..9e7dead2e 100644
--- a/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -210,7 +210,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     environment.jersey().register(new AuthValueFactoryProvider.Binder());
 
     environment.jersey().register(new AccountController(pendingAccountsManager, accountsManager, rateLimiters, smsSender, messagesManager, new TimeProvider(), authorizationKey, config.getTestDevices()));
-    environment.jersey().register(new DeviceController(pendingDevicesManager, accountsManager, rateLimiters));
+    environment.jersey().register(new DeviceController(pendingDevicesManager, accountsManager, messagesManager, rateLimiters));
     environment.jersey().register(new DirectoryController(rateLimiters, directory));
     environment.jersey().register(new FederationControllerV1(accountsManager, attachmentController, messageController, keysControllerV1));
     environment.jersey().register(new FederationControllerV2(accountsManager, attachmentController, messageController, keysControllerV2));
diff --git a/src/main/java/org/whispersystems/textsecuregcm/controllers/DeviceController.java b/src/main/java/org/whispersystems/textsecuregcm/controllers/DeviceController.java
index f2e7d638d..473cebc07 100644
--- a/src/main/java/org/whispersystems/textsecuregcm/controllers/DeviceController.java
+++ b/src/main/java/org/whispersystems/textsecuregcm/controllers/DeviceController.java
@@ -32,6 +32,7 @@ import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.Device;
+import org.whispersystems.textsecuregcm.storage.MessagesManager;
 import org.whispersystems.textsecuregcm.storage.PendingDevicesManager;
 import org.whispersystems.textsecuregcm.util.Util;
 import org.whispersystems.textsecuregcm.util.VerificationCode;
@@ -65,14 +66,17 @@ public class DeviceController {
 
   private final PendingDevicesManager pendingDevices;
   private final AccountsManager       accounts;
+  private final MessagesManager       messages;
   private final RateLimiters          rateLimiters;
 
   public DeviceController(PendingDevicesManager pendingDevices,
                           AccountsManager accounts,
+                          MessagesManager messages,
                           RateLimiters rateLimiters)
   {
     this.pendingDevices  = pendingDevices;
     this.accounts        = accounts;
+    this.messages        = messages;
     this.rateLimiters    = rateLimiters;
   }
 
@@ -100,6 +104,7 @@ public class DeviceController {
 
     account.removeDevice(deviceId);
     accounts.update(account);
+    messages.clear(account.getNumber(), deviceId);
   }
 
   @Timed
diff --git a/src/main/java/org/whispersystems/textsecuregcm/storage/Messages.java b/src/main/java/org/whispersystems/textsecuregcm/storage/Messages.java
index 7f68bcca2..ae6e01fdf 100644
--- a/src/main/java/org/whispersystems/textsecuregcm/storage/Messages.java
+++ b/src/main/java/org/whispersystems/textsecuregcm/storage/Messages.java
@@ -60,6 +60,9 @@ public abstract class Messages {
   @SqlUpdate("DELETE FROM messages WHERE " + DESTINATION + " = :destination")
   abstract void clear(@Bind("destination") String destination);
 
+  @SqlUpdate("DELETE FROM messages WHERE " + DESTINATION + " = :destination AND " + DESTINATION_DEVICE + " = :destination_device")
+  abstract void clear(@Bind("destination") String destination, @Bind("destination_device") long destinationDevice);
+
   @SqlUpdate("DELETE FROM messages WHERE " + TIMESTAMP + " < :timestamp")
   public abstract void removeOld(@Bind("timestamp") long timestamp);
 
diff --git a/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java b/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java
index 8a9349f53..db38a03fc 100644
--- a/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java
+++ b/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java
@@ -29,6 +29,10 @@ public class MessagesManager {
     this.messages.clear(destination);
   }
 
+  public void clear(String destination, long deviceId) {
+    this.messages.clear(destination, deviceId);
+  }
+
   public Optional<OutgoingMessageEntity> delete(String destination, String source, long timestamp) {
     return Optional.fromNullable(this.messages.remove(destination, source, timestamp));
   }
diff --git a/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/DeviceControllerTest.java b/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/DeviceControllerTest.java
index 9aa0752e2..513afaa7f 100644
--- a/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/DeviceControllerTest.java
+++ b/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/DeviceControllerTest.java
@@ -30,6 +30,7 @@ import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.mappers.DeviceLimitExceededExceptionMapper;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
+import org.whispersystems.textsecuregcm.storage.MessagesManager;
 import org.whispersystems.textsecuregcm.storage.PendingDevicesManager;
 import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
 import org.whispersystems.textsecuregcm.util.VerificationCode;
@@ -48,8 +49,12 @@ import static org.mockito.Mockito.*;
 public class DeviceControllerTest {
   @Path("/v1/devices")
   static class DumbVerificationDeviceController extends DeviceController {
-    public DumbVerificationDeviceController(PendingDevicesManager pendingDevices, AccountsManager accounts, RateLimiters rateLimiters) {
-      super(pendingDevices, accounts, rateLimiters);
+    public DumbVerificationDeviceController(PendingDevicesManager pendingDevices,
+                                            AccountsManager accounts,
+                                            MessagesManager messages,
+                                            RateLimiters rateLimiters)
+    {
+      super(pendingDevices, accounts, messages, rateLimiters);
     }
 
     @Override
@@ -60,6 +65,7 @@ public class DeviceControllerTest {
 
   private PendingDevicesManager pendingDevicesManager = mock(PendingDevicesManager.class);
   private AccountsManager       accountsManager       = mock(AccountsManager.class       );
+  private MessagesManager       messagesManager       = mock(MessagesManager.class);
   private RateLimiters          rateLimiters          = mock(RateLimiters.class          );
   private RateLimiter           rateLimiter           = mock(RateLimiter.class           );
   private Account               account               = mock(Account.class               );
@@ -74,6 +80,7 @@ public class DeviceControllerTest {
                                                             .addProvider(new ConstraintViolationExceptionMapper())
                                                             .addResource(new DumbVerificationDeviceController(pendingDevicesManager,
                                                                                                               accountsManager,
+                                                                                                              messagesManager,
                                                                                                               rateLimiters))
                                                             .build();