From a58f3f0fe369bb9575b47fd0be926a8b2d74038b Mon Sep 17 00:00:00 2001
From: Moxie Marlinspike <moxie@thoughtcrime.org>
Date: Tue, 24 Mar 2015 10:48:14 -0700
Subject: [PATCH] Check subscription status on websocket keepalive.

// FREEBIE
---
 .../dispatch/DispatchManager.java             |  4 +++
 .../textsecuregcm/WhisperServerService.java   |  4 +--
 .../controllers/KeepAliveController.java      | 27 +++++++++++++++++--
 .../textsecuregcm/storage/PubSubManager.java  |  4 +++
 4 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/src/main/java/org/whispersystems/dispatch/DispatchManager.java b/src/main/java/org/whispersystems/dispatch/DispatchManager.java
index 2bea48cd0..e52f104da 100644
--- a/src/main/java/org/whispersystems/dispatch/DispatchManager.java
+++ b/src/main/java/org/whispersystems/dispatch/DispatchManager.java
@@ -75,6 +75,10 @@ public class DispatchManager extends Thread {
     }
   }
 
+  public boolean hasSubscription(String name) {
+    return subscriptions.containsKey(name);
+  }
+
   @Override
   public void run() {
     while (running) {
diff --git a/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index cb712e1bc..414758eae 100644
--- a/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -206,11 +206,11 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
       WebSocketEnvironment webSocketEnvironment = new WebSocketEnvironment(environment, config, 90000);
       webSocketEnvironment.setAuthenticator(new WebSocketAccountAuthenticator(deviceAuthenticator));
       webSocketEnvironment.setConnectListener(new AuthenticatedConnectListener(accountsManager, pushSender, messagesManager, pubSubManager));
-      webSocketEnvironment.jersey().register(new KeepAliveController());
+      webSocketEnvironment.jersey().register(new KeepAliveController(pubSubManager));
 
       WebSocketEnvironment provisioningEnvironment = new WebSocketEnvironment(environment, config);
       provisioningEnvironment.setConnectListener(new ProvisioningConnectListener(pubSubManager));
-      provisioningEnvironment.jersey().register(new KeepAliveController());
+      provisioningEnvironment.jersey().register(new KeepAliveController(pubSubManager));
       
       WebSocketResourceProviderFactory webSocketServlet    = new WebSocketResourceProviderFactory(webSocketEnvironment   );
       WebSocketResourceProviderFactory provisioningServlet = new WebSocketResourceProviderFactory(provisioningEnvironment);
diff --git a/src/main/java/org/whispersystems/textsecuregcm/controllers/KeepAliveController.java b/src/main/java/org/whispersystems/textsecuregcm/controllers/KeepAliveController.java
index f4fa3a6c0..903d56b30 100644
--- a/src/main/java/org/whispersystems/textsecuregcm/controllers/KeepAliveController.java
+++ b/src/main/java/org/whispersystems/textsecuregcm/controllers/KeepAliveController.java
@@ -1,19 +1,42 @@
 package org.whispersystems.textsecuregcm.controllers;
 
 import com.codahale.metrics.annotation.Timed;
+import org.whispersystems.textsecuregcm.storage.Account;
+import org.whispersystems.textsecuregcm.storage.PubSubManager;
+import org.whispersystems.textsecuregcm.websocket.WebsocketAddress;
+import org.whispersystems.websocket.session.WebSocketSession;
+import org.whispersystems.websocket.session.WebSocketSessionContext;
 
 import javax.ws.rs.GET;
-import javax.ws.rs.PUT;
 import javax.ws.rs.Path;
 import javax.ws.rs.core.Response;
 
+import io.dropwizard.auth.Auth;
+
 
 @Path("/v1/keepalive")
 public class KeepAliveController {
 
+  private final PubSubManager pubSubManager;
+
+  public KeepAliveController(PubSubManager pubSubManager) {
+    this.pubSubManager = pubSubManager;
+  }
+
   @Timed
   @GET
-  public Response getKeepAlive() {
+  public Response getKeepAlive(@Auth(required = false) Account account,
+                               @WebSocketSession       WebSocketSessionContext context)
+  {
+    if (account != null) {
+      WebsocketAddress address = new WebsocketAddress(account.getNumber(),
+                                                      account.getAuthenticatedDevice().get().getId());
+
+      if (!pubSubManager.hasLocalSubscription(address)) {
+        context.getClient().close(1000, "OK");
+      }
+    }
+
     return Response.ok().build();
   }
 
diff --git a/src/main/java/org/whispersystems/textsecuregcm/storage/PubSubManager.java b/src/main/java/org/whispersystems/textsecuregcm/storage/PubSubManager.java
index af5f0d975..7eef751e4 100644
--- a/src/main/java/org/whispersystems/textsecuregcm/storage/PubSubManager.java
+++ b/src/main/java/org/whispersystems/textsecuregcm/storage/PubSubManager.java
@@ -59,6 +59,10 @@ public class PubSubManager implements Managed {
     dispatchManager.unsubscribe(serializedAddress, dispatchChannel);
   }
 
+  public boolean hasLocalSubscription(WebsocketAddress address) {
+    return dispatchManager.hasSubscription(address.serialize());
+  }
+
   public boolean publish(WebsocketAddress address, PubSubMessage message) {
     return publish(address.serialize().getBytes(), message);
   }