Scope disconnection request listeners to a single connection

service/src/main/java/org/whispersystems/textsecuregcm/websocket/AuthenticatedConnectListener.java

@@ -13,7 +13,9 @@ import java.util.concurrent.ScheduledExecutorService;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.auth.AuthenticatedDevice;
+import org.whispersystems.textsecuregcm.auth.DisconnectionRequestManager;
 import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
+import org.whispersystems.textsecuregcm.identity.IdentityType;
 import org.whispersystems.textsecuregcm.limits.MessageDeliveryLoopMonitor;
 import org.whispersystems.textsecuregcm.metrics.MessageMetrics;
 import org.whispersystems.textsecuregcm.metrics.OpenWebSocketCounter;
@@ -48,6 +50,7 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
   private final PushNotificationManager pushNotificationManager;
   private final PushNotificationScheduler pushNotificationScheduler;
   private final WebSocketConnectionEventManager webSocketConnectionEventManager;
+  private final DisconnectionRequestManager disconnectionRequestManager;
   private final ScheduledExecutorService scheduledExecutorService;
   private final Scheduler messageDeliveryScheduler;
   private final ClientReleaseManager clientReleaseManager;
@@ -58,17 +61,18 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
   private final OpenWebSocketCounter openUnauthenticatedWebSocketCounter;
 
   public AuthenticatedConnectListener(
-      AccountsManager accountsManager,
-      ReceiptSender receiptSender,
-      MessagesManager messagesManager,
-      MessageMetrics messageMetrics,
-      PushNotificationManager pushNotificationManager,
-      PushNotificationScheduler pushNotificationScheduler,
-      WebSocketConnectionEventManager webSocketConnectionEventManager,
-      ScheduledExecutorService scheduledExecutorService,
-      Scheduler messageDeliveryScheduler,
-      ClientReleaseManager clientReleaseManager,
-      MessageDeliveryLoopMonitor messageDeliveryLoopMonitor,
+      final AccountsManager accountsManager,
+      final ReceiptSender receiptSender,
+      final MessagesManager messagesManager,
+      final MessageMetrics messageMetrics,
+      final PushNotificationManager pushNotificationManager,
+      final PushNotificationScheduler pushNotificationScheduler,
+      final WebSocketConnectionEventManager webSocketConnectionEventManager,
+      final DisconnectionRequestManager disconnectionRequestManager,
+      final ScheduledExecutorService scheduledExecutorService,
+      final Scheduler messageDeliveryScheduler,
+      final ClientReleaseManager clientReleaseManager,
+      final MessageDeliveryLoopMonitor messageDeliveryLoopMonitor,
       final ExperimentEnrollmentManager experimentEnrollmentManager) {
 
     this.accountsManager = accountsManager;
@@ -78,6 +82,7 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
     this.pushNotificationManager = pushNotificationManager;
     this.pushNotificationScheduler = pushNotificationScheduler;
     this.webSocketConnectionEventManager = webSocketConnectionEventManager;
+    this.disconnectionRequestManager = disconnectionRequestManager;
     this.scheduledExecutorService = scheduledExecutorService;
     this.messageDeliveryScheduler = messageDeliveryScheduler;
     this.clientReleaseManager = clientReleaseManager;
@@ -104,7 +109,7 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
       final AuthenticatedDevice auth = context.getAuthenticated(AuthenticatedDevice.class);
 
       final Optional<Account> maybeAuthenticatedAccount = accountsManager.getByAccountIdentifier(auth.accountIdentifier());
-      final Optional<Device> maybeAuthenticatedDevice = maybeAuthenticatedAccount.flatMap(account -> account.getDevice(auth.deviceId()));;
+      final Optional<Device> maybeAuthenticatedDevice = maybeAuthenticatedAccount.flatMap(account -> account.getDevice(auth.deviceId()));
 
       if (maybeAuthenticatedAccount.isEmpty() || maybeAuthenticatedDevice.isEmpty()) {
         log.warn("{}:{} not found when opening authenticated WebSocket", auth.accountIdentifier(), auth.deviceId());
@@ -127,7 +132,15 @@ public class AuthenticatedConnectListener implements WebSocketConnectListener {
           messageDeliveryLoopMonitor,
           experimentEnrollmentManager);
 
-      context.addWebsocketClosedListener((closingContext, statusCode, reason) -> {
+      disconnectionRequestManager.addListener(maybeAuthenticatedAccount.get().getIdentifier(IdentityType.ACI),
+          maybeAuthenticatedDevice.get().getId(),
+          connection);
+
+      context.addWebsocketClosedListener((_, _, _) -> {
+        disconnectionRequestManager.removeListener(maybeAuthenticatedAccount.get().getIdentifier(IdentityType.ACI),
+            maybeAuthenticatedDevice.get().getId(),
+            connection);
+
         // We begin the shutdown process by removing this client's "presence," which means it will again begin to
         // receive push notifications for inbound messages. We should do this first because, at this point, the
         // connection has already closed and attempts to actually deliver a message via the connection will not succeed.

service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java

@@ -37,6 +37,7 @@ import org.eclipse.jetty.util.StaticException;
 import org.reactivestreams.Publisher;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.whispersystems.textsecuregcm.auth.DisconnectionRequestListener;
 import org.whispersystems.textsecuregcm.controllers.MessageController;
 import org.whispersystems.textsecuregcm.entities.MessageProtos.Envelope;
 import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
@@ -65,7 +66,7 @@ import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Scheduler;
 
-public class WebSocketConnection implements WebSocketConnectionEventListener {
+public class WebSocketConnection implements WebSocketConnectionEventListener, DisconnectionRequestListener {
 
   private static final DistributionSummary messageTime = Metrics.summary(
       name(MessageController.class, "messageDeliveryDuration"));
@@ -506,24 +507,23 @@ public class WebSocketConnection implements WebSocketConnectionEventListener {
   }
 
   @Override
-  public void handleConnectionDisplaced(final boolean connectedElsewhere) {
+  public void handleConflictingMessageReader() {
+    closeConnection(4409, "Connected elsewhere");
+  }
+
+  @Override
+  public void handleDisconnectionRequest() {
+    closeConnection(4401, "Reauthentication required");
+  }
+
+  private void closeConnection(final int code, final String message) {
     final Tags tags = Tags.of(
         UserAgentTagUtil.getPlatformTag(client.getUserAgent()),
-        Tag.of("connectedElsewhere", String.valueOf(connectedElsewhere)));
+        // TODO We should probably just use the status code directly
+        Tag.of("connectedElsewhere", String.valueOf(code == 4409)));
 
     Metrics.counter(DISPLACEMENT_COUNTER_NAME, tags).increment();
 
-    final int code;
-    final String message;
-
-    if (connectedElsewhere) {
-      code = 4409;
-      message = "Connected elsewhere";
-    } else {
-      code = 4401;
-      message = "Reauthentication required";
-    }
-
     client.close(code, message);
   }