From 5e34823a4989c670855a57fd49520b27bb76821d Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Fri, 4 Sep 2020 12:10:24 -0400
Subject: [PATCH] Optionally send online-only messages via keyspace
 notifications.

---
 .../textsecuregcm/WhisperServerService.java   |  2 +-
 .../textsecuregcm/push/WebsocketSender.java   | 61 +++++++++++++------
 .../storage/MessagesManager.java              |  8 +++
 .../websocket/WebSocketConnection.java        |  6 ++
 4 files changed, 56 insertions(+), 21 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index 995b43ef8..8753e1fbc 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -300,7 +300,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     PubSubManager              pubSubManager              = new PubSubManager(pubsubClient, dispatchManager);
     APNSender                  apnSender                  = new APNSender(accountsManager, config.getApnConfiguration());
     GCMSender                  gcmSender                  = new GCMSender(accountsManager, config.getGcmConfiguration().getApiKey());
-    WebsocketSender            websocketSender            = new WebsocketSender(messagesManager, pubSubManager, clientPresenceManager);
+    WebsocketSender            websocketSender            = new WebsocketSender(messagesManager, pubSubManager, clientPresenceManager, featureFlagsManager);
     RateLimiters               rateLimiters               = new RateLimiters(config.getLimitsConfiguration(), cacheCluster);
 
     AccountAuthenticator                  accountAuthenticator                  = new AccountAuthenticator(accountsManager);
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/push/WebsocketSender.java b/service/src/main/java/org/whispersystems/textsecuregcm/push/WebsocketSender.java
index 9124253fe..66d9a3d05 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/push/WebsocketSender.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/push/WebsocketSender.java
@@ -20,10 +20,13 @@ import com.codahale.metrics.Meter;
 import com.codahale.metrics.MetricRegistry;
 import com.codahale.metrics.SharedMetricRegistries;
 import com.google.protobuf.ByteString;
+import io.micrometer.core.instrument.Counter;
+import io.micrometer.core.instrument.Metrics;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.Device;
+import org.whispersystems.textsecuregcm.storage.FeatureFlagsManager;
 import org.whispersystems.textsecuregcm.storage.MessagesManager;
 import org.whispersystems.textsecuregcm.storage.PubSubManager;
 import org.whispersystems.textsecuregcm.util.Constants;
@@ -60,38 +63,56 @@ public class WebsocketSender {
   private final Meter provisioningOnlineMeter  = metricRegistry.meter(name(getClass(), "provisioning_online" ));
   private final Meter provisioningOfflineMeter = metricRegistry.meter(name(getClass(), "provisioning_offline"));
 
+  private final Counter ephemeralOnlineCounter  = Metrics.counter(name(getClass(), "ephemeral"), "online", "true");
+  private final Counter ephemeralOfflineCounter = Metrics.counter(name(getClass(), "ephemeral"), "offline", "true");
+
   private final MessagesManager       messagesManager;
   private final PubSubManager         pubSubManager;
   private final ClientPresenceManager clientPresenceManager;
+  private final FeatureFlagsManager   featureFlagsManager;
 
-  public WebsocketSender(MessagesManager messagesManager, PubSubManager pubSubManager, ClientPresenceManager clientPresenceManager) {
+  private static final String KEYSPACE_DELIVERY_FEATURE_FLAG = "keyspace-delivery-for-online-messages";
+
+  public WebsocketSender(MessagesManager messagesManager, PubSubManager pubSubManager, ClientPresenceManager clientPresenceManager, final FeatureFlagsManager featureFlagsManager) {
     this.messagesManager       = messagesManager;
     this.pubSubManager         = pubSubManager;
     this.clientPresenceManager = clientPresenceManager;
+    this.featureFlagsManager   = featureFlagsManager;
   }
 
   public boolean sendMessage(Account account, Device device, Envelope message, Type channel, boolean online) {
-    WebsocketAddress address       = new WebsocketAddress(account.getNumber(), device.getId());
-    PubSubMessage    pubSubMessage = PubSubMessage.newBuilder()
-                                                  .setType(PubSubMessage.Type.DELIVER)
-                                                  .setContent(message.toByteString())
-                                                  .build();
-
-    pubSubManager.publish(address, pubSubMessage);
-
-    if (clientPresenceManager.isPresent(account.getUuid(), device.getId())) {
-      if      (channel == Type.APN) apnOnlineMeter.mark();
-      else if (channel == Type.GCM) gcmOnlineMeter.mark();
-      else                          websocketOnlineMeter.mark();
-
-      return true;
+    if (online && featureFlagsManager.isFeatureFlagActive(KEYSPACE_DELIVERY_FEATURE_FLAG)) {
+      if (clientPresenceManager.isPresent(account.getUuid(), device.getId())) {
+        ephemeralOnlineCounter.increment();
+        messagesManager.insertEphemeral(account.getUuid(), device.getId(), message);
+        return true;
+      } else {
+        ephemeralOfflineCounter.increment();
+        return false;
+      }
     } else {
-      if      (channel == Type.APN) apnOfflineMeter.mark();
-      else if (channel == Type.GCM) gcmOfflineMeter.mark();
-      else                          websocketOfflineMeter.mark();
+      WebsocketAddress address = new WebsocketAddress(account.getNumber(), device.getId());
+      PubSubMessage pubSubMessage = PubSubMessage.newBuilder()
+              .setType(PubSubMessage.Type.DELIVER)
+              .setContent(message.toByteString())
+              .build();
 
-      if (!online) queueMessage(account, device, message);
-      return false;
+      pubSubManager.publish(address, pubSubMessage);
+
+      if (clientPresenceManager.isPresent(account.getUuid(), device.getId())) {
+        if (channel == Type.APN) apnOnlineMeter.mark();
+        else if (channel == Type.GCM) gcmOnlineMeter.mark();
+        else websocketOnlineMeter.mark();
+
+        return true;
+      } else {
+        if (channel == Type.APN) apnOfflineMeter.mark();
+        else if (channel == Type.GCM) gcmOfflineMeter.mark();
+        else websocketOfflineMeter.mark();
+
+        if (!online) queueMessage(account, device, message);
+        return false;
+      }
     }
   }
 
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java
index c67771898..1e8065046 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/MessagesManager.java
@@ -41,6 +41,14 @@ public class MessagesManager {
     messagesCache.insert(UUID.randomUUID(), destinationUuid, destinationDevice, message);
   }
 
+  public void insertEphemeral(final UUID destinationUuid, final long destinationDevice, final Envelope message) {
+    messagesCache.insertEphemeral(destinationUuid, destinationDevice, message);
+  }
+
+  public Optional<Envelope> takeEphemeralMessage(final UUID destinationUuid, final long destinationDevice) {
+    return messagesCache.takeEphemeralMessage(destinationUuid, destinationDevice);
+  }
+
   public OutgoingMessageEntityList getMessagesForDevice(String destination, UUID destinationUuid, long destinationDevice, final String userAgent) {
     RedisOperation.unchecked(() -> pushLatencyManager.recordQueueRead(destinationUuid, destinationDevice, userAgent));
 
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java b/service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java
index 0367a0c55..571b6ab7e 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/websocket/WebSocketConnection.java
@@ -224,6 +224,12 @@ public class WebSocketConnection implements DispatchChannel, MessageAvailability
   @Override
   public void handleNewEphemeralMessageAvailable() {
     ephemeralMessageAvailableMeter.mark();
+
+    final Optional<Envelope> maybeMessage = messagesManager.takeEphemeralMessage(account.getUuid(), device.getId());
+
+    if (maybeMessage.isPresent()) {
+      sendMessage(maybeMessage.get(), Optional.empty(), false);
+    }
   }
 
   @Override