Implement a libsignal-net shadowing web socket.

libsignal-service/src/main/java/org/whispersystems/signalservice/internal/util/CompletableFutureExtensions.kt

@@ -0,0 +1,25 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.signalservice.internal.util
+
+import org.signal.libsignal.internal.CompletableFuture
+
+/**
+ * A Kotlin friendly adapter for [org.signal.libsignal.internal.CompletableFuture.whenComplete]
+ * taking two callbacks ([onSuccess] and [onFailure]) instead of a [java.util.function.BiConsumer].
+ */
+fun <T> CompletableFuture<T>.whenComplete(
+  onSuccess: ((T?) -> Unit),
+  onFailure: ((Throwable) -> Unit)
+): CompletableFuture<T> {
+  return this.whenComplete { value, throwable ->
+    if (throwable != null) {
+      onFailure(throwable)
+    } else {
+      onSuccess(value)
+    }
+  }
+}

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

@@ -11,10 +11,10 @@ import io.reactivex.rxjava3.schedulers.Schedulers
 import io.reactivex.rxjava3.subjects.BehaviorSubject
 import io.reactivex.rxjava3.subjects.SingleSubject
 import org.signal.core.util.logging.Log
-import org.signal.libsignal.internal.CompletableFuture
 import org.signal.libsignal.net.ChatService
 import org.whispersystems.signalservice.api.websocket.HealthMonitor
 import org.whispersystems.signalservice.api.websocket.WebSocketConnectionState
+import org.whispersystems.signalservice.internal.util.whenComplete
 import java.time.Instant
 import java.util.Optional
 import kotlin.time.Duration.Companion.seconds
@@ -52,6 +52,31 @@ class LibSignalChatConnection(
       ByteArray(0),
       SEND_TIMEOUT.toInt()
     )
+
+    private fun WebSocketRequestMessage.toLibSignalRequest(timeout: Long = SEND_TIMEOUT): LibSignalRequest {
+      return LibSignalRequest(
+        this.verb?.uppercase() ?: "GET",
+        this.path ?: "",
+        this.headers.associate {
+          val parts = it.split(':', limit = 2)
+          if (parts.size != 2) {
+            throw IllegalArgumentException("Headers must contain at least one colon")
+          }
+          parts[0] to parts[1]
+        },
+        this.body?.toByteArray() ?: byteArrayOf(),
+        timeout.toInt()
+      )
+    }
+
+    private fun LibSignalResponse.toWebsocketResponse(isUnidentified: Boolean): WebsocketResponse {
+      return WebsocketResponse(
+        this.status,
+        this.body.decodeToString(),
+        this.headers,
+        isUnidentified
+      )
+    }
   }
 
   override val name = "[$name:${System.identityHashCode(this)}]"
@@ -176,42 +201,4 @@ class LibSignalChatConnection(
   override fun sendResponse(response: WebSocketResponseMessage?) {
     throw NotImplementedError()
   }
-
-  private fun WebSocketRequestMessage.toLibSignalRequest(timeout: Long = SEND_TIMEOUT): LibSignalRequest {
-    return LibSignalRequest(
-      this.verb?.uppercase() ?: "GET",
-      this.path ?: "",
-      this.headers.associate {
-        val parts = it.split(':', limit = 2)
-        if (parts.size != 2) {
-          throw IllegalArgumentException("Headers must contain at least one colon")
-        }
-        parts[0] to parts[1]
-      },
-      this.body?.toByteArray() ?: byteArrayOf(),
-      timeout.toInt()
-    )
-  }
-
-  private fun LibSignalResponse.toWebsocketResponse(isUnidentified: Boolean): WebsocketResponse {
-    return WebsocketResponse(
-      this.status,
-      this.body.decodeToString(),
-      this.headers,
-      isUnidentified
-    )
-  }
-
-  private fun <T> CompletableFuture<T>.whenComplete(
-    onSuccess: ((T?) -> Unit),
-    onFailure: ((Throwable) -> Unit)
-  ): CompletableFuture<T> {
-    return this.whenComplete { value, throwable ->
-      if (throwable != null) {
-        onFailure(throwable)
-      } else {
-        onSuccess(value)
-      }
-    }
-  }
 }

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

@@ -0,0 +1,227 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.signalservice.internal.websocket
+
+import io.reactivex.rxjava3.core.Observable
+import io.reactivex.rxjava3.core.Single
+import okhttp3.Response
+import okhttp3.WebSocket
+import org.signal.core.util.logging.Log
+import org.signal.libsignal.net.ChatService
+import org.whispersystems.signalservice.api.util.CredentialsProvider
+import org.whispersystems.signalservice.api.websocket.HealthMonitor
+import org.whispersystems.signalservice.api.websocket.WebSocketConnectionState
+import org.whispersystems.signalservice.internal.configuration.SignalServiceConfiguration
+import org.whispersystems.signalservice.internal.util.whenComplete
+import java.util.Optional
+import java.util.concurrent.ExecutorService
+import java.util.concurrent.Executors
+import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.atomic.AtomicInteger
+import java.util.concurrent.atomic.AtomicLong
+import kotlin.random.Random
+import kotlin.time.Duration
+import kotlin.time.Duration.Companion.days
+import kotlin.time.Duration.Companion.minutes
+import kotlin.time.Duration.Companion.seconds
+
+/**
+ * A wrapper on top of [OkHttpWebSocketConnection] that sends a keep alive via
+ * libsignal-net for a configurable percentage of the _successful_ web socket requests.
+ *
+ * Stats are collected for the shadowed requests and persisted across app restarts
+ * using [org.thoughtcrime.securesms.keyvalue.InternalValues].
+ *
+ * When a hardcoded error threshold is reached, the user is notified to submit debug logs.
+ *
+ * @see [org.thoughtcrime.securesms.util.FeatureFlags.libSignalWebSocketShadowingPercentage]
+ */
+class ShadowingWebSocketConnection(
+  name: String,
+  serviceConfiguration: SignalServiceConfiguration,
+  credentialsProvider: Optional<CredentialsProvider>,
+  signalAgent: String,
+  healthMonitor: HealthMonitor,
+  allowStories: Boolean,
+  private val chatService: ChatService,
+  private val shadowPercentage: Int,
+  private val bridge: WebSocketShadowingBridge
+) : OkHttpWebSocketConnection(
+  name,
+  serviceConfiguration,
+  credentialsProvider,
+  signalAgent,
+  healthMonitor,
+  allowStories
+) {
+  private var stats: Stats = try {
+    bridge.readStatsSnapshot()?.let {
+      Stats.fromSnapshot(it)
+    } ?: Stats()
+  } catch (ex: Exception) {
+    Log.w(TAG, "Failed to restore Stats from a snapshot.", ex)
+    Stats()
+  }
+  private val canShadow: AtomicBoolean = AtomicBoolean(false)
+  private val executor: ExecutorService = Executors.newSingleThreadExecutor()
+
+  override fun connect(): Observable<WebSocketConnectionState> {
+    executor.submit {
+      chatService.connectUnauthenticated().whenComplete(
+        onSuccess = {
+          canShadow.set(true)
+          Log.i(TAG, "Shadow socket connected.")
+        },
+        onFailure = {
+          canShadow.set(false)
+          Log.i(TAG, "Shadow socket failed to connect.")
+        }
+      )
+    }
+    return super.connect()
+  }
+
+  override fun onClosing(webSocket: WebSocket, code: Int, reason: String) {
+    saveSnapshot()
+    super.onClosing(webSocket, code, reason)
+  }
+
+  override fun onFailure(webSocket: WebSocket, t: Throwable, response: Response?) {
+    saveSnapshot()
+    super.onFailure(webSocket, t, response)
+  }
+
+  override fun disconnect() {
+    executor.submit {
+      chatService.disconnect().thenApply {
+        canShadow.set(false)
+        Log.i(TAG, "Shadow socket disconnected.")
+      }
+    }
+    super.disconnect()
+  }
+
+  override fun sendRequest(request: WebSocketRequestMessage): Single<WebsocketResponse> {
+    return super.sendRequest(request).doOnSuccess(::sendShadow)
+  }
+
+  private fun sendShadow(actualResponse: WebsocketResponse) {
+    executor.submit {
+      if (canShadow.get() && Random.nextInt(100) < this.shadowPercentage) {
+        libsignalKeepAlive(actualResponse)
+        val snapshotAge = System.currentTimeMillis() - stats.lastSnapshot.get()
+        if (snapshotAge > SNAPSHOT_INTERVAL.inWholeMilliseconds) {
+          saveSnapshot()
+        }
+      }
+    }
+  }
+
+  private fun shouldSubmitLogs(): Boolean {
+    val requestsCompared = stats.requestsCompared.get()
+    // Should not happen in practice, but helps avoid a division by zero later if it does.
+    if (requestsCompared == 0) {
+      return false
+    }
+    val timeSinceLastNotificationMs = System.currentTimeMillis() - stats.lastNotified.get()
+    val percentFailed = stats.failures.get() * 100 / requestsCompared
+    return timeSinceLastNotificationMs > FULL_DAY.inWholeMilliseconds &&
+      percentFailed > FAILURE_PERCENTAGE
+  }
+
+  private fun libsignalKeepAlive(actualResponse: WebsocketResponse) {
+    val request = ChatService.Request(
+      "GET",
+      "/v1/keepalive",
+      emptyMap(),
+      ByteArray(0),
+      KEEP_ALIVE_TIMEOUT.inWholeMilliseconds.toInt()
+    )
+    chatService.unauthenticatedSendAndDebug(request)
+      .whenComplete(
+        onSuccess = {
+          stats.requestsCompared.incrementAndGet()
+          val goodStatus = (it?.response?.status ?: -1) in 200..299
+          if (!goodStatus) {
+            stats.badStatuses.incrementAndGet()
+          }
+          stats.reconnects.addAndGet(it?.debugInfo?.reconnectCount ?: 0)
+          Log.i(TAG, "$it")
+        },
+        onFailure = {
+          stats.requestsCompared.incrementAndGet()
+          stats.failures.incrementAndGet()
+          Log.w(TAG, "Shadow request failed: reason=[$it]")
+          Log.i(TAG, "Actual response status=${actualResponse.status}")
+          if (shouldSubmitLogs()) {
+            Log.i(TAG, "Notification to submit logs triggered.")
+            bridge.triggerFailureNotification("Experimental websocket transport failures!")
+            stats.reset()
+          }
+        }
+      )
+  }
+
+  private fun saveSnapshot() {
+    executor.submit {
+      Log.d(TAG, "Persisting shadowing stats snapshot.")
+      bridge.writeStatsSnapshot(stats.snapshot())
+    }
+  }
+
+  companion object {
+    private val TAG: String = Log.tag(ShadowingWebSocketConnection::class.java)
+    private val FULL_DAY: Duration = 1.days
+
+    // If more than this percentage of shadow requests fail, the
+    // notification to submit logs will be triggered.
+    private const val FAILURE_PERCENTAGE: Int = 10
+    private val KEEP_ALIVE_TIMEOUT: Duration = 3.seconds
+    private val SNAPSHOT_INTERVAL: Duration = 10.minutes
+  }
+
+  class Stats(
+    requestsCompared: Int = 0,
+    failures: Int = 0,
+    badStatuses: Int = 0,
+    reconnects: Int = 0,
+    lastNotified: Long = 0
+  ) {
+    val requestsCompared: AtomicInteger = AtomicInteger(requestsCompared)
+    val failures: AtomicInteger = AtomicInteger(failures)
+    val badStatuses: AtomicInteger = AtomicInteger(badStatuses)
+    val reconnects: AtomicInteger = AtomicInteger(reconnects)
+    val lastNotified: AtomicLong = AtomicLong(lastNotified)
+    val lastSnapshot: AtomicLong = AtomicLong(0)
+
+    fun reset() {
+      requestsCompared.set(0)
+      failures.set(0)
+      badStatuses.set(0)
+      reconnects.set(0)
+      // Do not reset lastNotified nor lastSnapshot
+    }
+
+    companion object {
+      fun fromSnapshot(bytes: ByteArray): Stats {
+        val snapshot = Snapshot.ADAPTER.decode(bytes)
+        return Stats(snapshot.requestsCompared, snapshot.failures, snapshot.badStatuses, snapshot.reconnects, snapshot.lastNotified)
+      }
+    }
+
+    fun snapshot(): ByteArray {
+      lastSnapshot.set(System.currentTimeMillis())
+      return Snapshot.Builder()
+        .requestsCompared(requestsCompared.get())
+        .failures(failures.get())
+        .badStatuses(badStatuses.get())
+        .reconnects(reconnects.get())
+        .lastNotified(lastNotified.get())
+        .build()
+        .encode()
+    }
+  }
+}

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

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.whispersystems.signalservice.internal.websocket
+
+/**
+ * An interface to support app<->signal-service interop for the purposes of web socket shadowing.
+ */
+interface WebSocketShadowingBridge {
+  /**
+   * Persist shadowing stats snapshot.
+   */
+  fun writeStatsSnapshot(bytes: ByteArray)
+
+  /**
+   * Restore shadowing stats snapshot.
+   */
+  fun readStatsSnapshot(): ByteArray?
+
+  /**
+   * Display a notification the user to submit debug logs, with a custom message.
+   */
+  fun triggerFailureNotification(message: String)
+}