Close unidentified websocket when the app is in background

app/src/main/java/org/thoughtcrime/securesms/messages/IncomingMessageObserver.kt

@@ -172,6 +172,13 @@ class IncomingMessageObserver(private val context: Application) {
     }
   }
 
+  private fun shouldKeepAliveUnidentified(): Boolean {
+    val timeIdle = lock.withLock {
+      if (appVisible) 0 else System.currentTimeMillis() - lastInteractionTime
+    }
+    return timeIdle <= websocketReadTimeout
+  }
+
   private fun isConnectionNecessary(): Boolean {
     val timeIdle: Long
     val keepAliveEntries: Set<Pair<String, Long>>
@@ -379,7 +386,7 @@ class IncomingMessageObserver(private val context: Application) {
           decryptionDrained = false
         }
 
-        signalWebSocket.connect()
+        signalWebSocket.connect(shouldKeepAliveUnidentified())
         try {
           while (isConnectionNecessary()) {
             try {
@@ -430,11 +437,12 @@ class IncomingMessageObserver(private val context: Application) {
               }
             } catch (e: WebSocketUnavailableException) {
               Log.i(TAG, "Pipe unexpectedly unavailable, connecting")
-              signalWebSocket.connect()
+              signalWebSocket.connect(shouldKeepAliveUnidentified())
             } catch (e: TimeoutException) {
               Log.w(TAG, "Application level read timeout...")
               attempts = 0
             }
+            signalWebSocket.setKeepAliveUnidentified(shouldKeepAliveUnidentified())
           }
 
           if (!appVisible) {

app/src/main/java/org/thoughtcrime/securesms/net/SignalWebSocketHealthMonitor.java

@@ -110,12 +110,14 @@ private void onStateChange(WebSocketConnectionState connectionState, HealthState
   }
 
   @Override
-  public void onKeepAliveResponse(long sentTimestamp, boolean isIdentifiedWebSocket) {
+  public void onKeepAliveResponse(long sentTimestamp, boolean isIdentifiedWebSocket, boolean keepMonitoring) {
     final long keepAliveTime = System.currentTimeMillis();
     executor.execute(() -> {
       if (isIdentifiedWebSocket) {
+        identified.needsKeepAlive = keepMonitoring;
         identified.lastKeepAliveReceived = keepAliveTime;
       } else {
+        unidentified.needsKeepAlive = keepMonitoring;
         unidentified.lastKeepAliveReceived = keepAliveTime;
       }
     });
@@ -143,6 +145,13 @@ private static class HealthState {
 
     private volatile boolean needsKeepAlive;
     private volatile long    lastKeepAliveReceived;
+
+    /**
+     * Checks if a keep-alive response has not been received yet based on the provided send time.
+     */
+    private boolean isKeepAlivePending(long sendTime) {
+      return needsKeepAlive && lastKeepAliveReceived < sendTime;
+    }
   }
 
   /**
@@ -172,9 +181,8 @@ public void run() {
           sleepUntil(responseRequiredTime);
 
           if (shouldKeepRunning && isKeepAliveNecessary()) {
-            if (identified.lastKeepAliveReceived < keepAliveSendTime || unidentified.lastKeepAliveReceived < keepAliveSendTime) {
+            if (identified.isKeepAlivePending(keepAliveSendTime)) {
               Log.w(TAG, "Missed keep alives, identified last: " + identified.lastKeepAliveReceived +
-                         " unidentified last: " + unidentified.lastKeepAliveReceived +
                          " needed by: " + responseRequiredTime);
               signalWebSocket.forceNewWebSockets();
             }

libsignal-service/src/main/java/org/whispersystems/signalservice/api/SignalWebSocket.java

@@ -47,6 +47,7 @@ public final class SignalWebSocket {
   private       CompositeDisposable                       unidentifiedWebSocketStateDisposable;
 
   private boolean canConnect;
+  private boolean keepAliveUnidentified;
 
   public SignalWebSocket(WebSocketFactory webSocketFactory) {
     this.webSocketFactory                     = webSocketFactory;
@@ -75,16 +76,26 @@ public Observable<WebSocketConnectionState> getUnidentifiedWebSocketState() {
   /**
    * Indicate that WebSocketConnections can now be made and attempt to connect both of them.
    */
-  public synchronized void connect() {
+  public synchronized void connect(boolean includeUnidentified) {
     canConnect = true;
+    keepAliveUnidentified = includeUnidentified;
     try {
       getWebSocket();
-      getUnidentifiedWebSocket();
+      if (includeUnidentified) {
+        getUnidentifiedWebSocket();
+      }
     } catch (WebSocketUnavailableException e) {
       throw new AssertionError(e);
     }
   }
 
+  public synchronized void setKeepAliveUnidentified(boolean keepAlive) {
+    if (unidentifiedWebSocket != null) {
+      unidentifiedWebSocket.setKeepAlive(keepAlive);
+    }
+    keepAliveUnidentified = keepAlive;
+  }
+
   /**
    * Indicate that WebSocketConnections can no longer be made and disconnect both of them.
    */
@@ -183,7 +194,11 @@ public synchronized void sendKeepAlive() throws IOException {
     if (canConnect) {
       try {
         getWebSocket().sendKeepAlive();
-        getUnidentifiedWebSocket().sendKeepAlive();
+        // Avoid creating the unidentified socket only for sending keep-alives,
+        // unless it's already open.
+        if (keepAliveUnidentified && unidentifiedWebSocket != null && !unidentifiedWebSocket.isDead()) {
+          getUnidentifiedWebSocket().sendKeepAlive();
+        }
       } catch (WebSocketUnavailableException e) {
         throw new AssertionError(e);
       }

libsignal-service/src/main/java/org/whispersystems/signalservice/api/websocket/HealthMonitor.java

@@ -4,7 +4,7 @@
  * Callbacks to provide WebSocket health information to a monitor.
  */
 public interface HealthMonitor {
-  void onKeepAliveResponse(long sentTimestamp, boolean isIdentifiedWebSocket);
+  void onKeepAliveResponse(long sentTimestamp, boolean isIdentifiedWebSocket, boolean keepMonitoring);
 
   void onMessageError(int status, boolean isIdentifiedWebSocket);
 }

libsignal-service/src/main/java/org/whispersystems/signalservice/internal/websocket/WebSocketConnection.java

@@ -76,6 +76,8 @@ public class WebSocketConnection extends WebSocketListener {
 
   private WebSocket client;
 
+  private boolean keepAlive;
+
   public WebSocketConnection(String name,
                              SignalServiceConfiguration serviceConfiguration,
                              Optional<CredentialsProvider> credentialsProvider,
@@ -101,6 +103,7 @@ public WebSocketConnection(String name,
     this.dns                 = serviceConfiguration.getDns();
     this.signalProxy         = serviceConfiguration.getSignalProxy();
     this.healthMonitor       = healthMonitor;
+    this.keepAlive           = true;
     this.webSocketState      = BehaviorSubject.createDefault(WebSocketConnectionState.DISCONNECTED);
     this.allowStories        = allowStories;
     this.serviceUrls         = serviceConfiguration.getSignalServiceUrls();
@@ -112,6 +115,10 @@ public String getName() {
     return name;
   }
 
+  public void setKeepAlive(boolean keepAlive) {
+    this.keepAlive = keepAlive;
+  }
+
   private Pair<SignalServiceUrl, String> getConnectionInfo() {
     SignalServiceUrl serviceUrl = serviceUrls[random.nextInt(serviceUrls.length)];
     String           uri        = serviceUrl.getUrl().replace("https://", "wss://").replace("http://", "ws://");
@@ -266,7 +273,7 @@ public synchronized void sendResponse(WebSocketResponseMessage response) throws
   }
 
   public synchronized void sendKeepAlive() throws IOException {
-    if (client != null) {
+    if (client != null && keepAlive) {
       log( "Sending keep alive...");
       long id = System.currentTimeMillis();
       byte[] message = new WebSocketMessage.Builder()
@@ -311,7 +318,7 @@ public synchronized void onMessage(WebSocket webSocket, ByteString payload) {
             healthMonitor.onMessageError(message.response.status, credentialsProvider.isPresent());
           }
         } else if (keepAlives.remove(message.response.id)) {
-          healthMonitor.onKeepAliveResponse(message.response.id, credentialsProvider.isPresent());
+          healthMonitor.onKeepAliveResponse(message.response.id, credentialsProvider.isPresent(), keepAlive);
         }
       }
 
@@ -333,7 +340,9 @@ public synchronized void onClosed(WebSocket webSocket, int code, String reason)
 
   @Override
   public synchronized void onFailure(WebSocket webSocket, Throwable t, Response response) {
-    warn("onFailure()", t);
+    if (keepAlive || response != null) {
+      warn("onFailure()", t);
+    }
 
     if (response != null && (response.code() == 401 || response.code() == 403)) {
       webSocketState.onNext(WebSocketConnectionState.AUTHENTICATION_FAILED);