From 976fbace2e0f507a8db50e4eb1852476fc7560ee Mon Sep 17 00:00:00 2001
From: Matthew Chen <matthew@signal.org>
Date: Thu, 9 Sep 2021 15:35:36 -0300
Subject: [PATCH] Whitelist sender key requests for websocket.

---
 .../MessageSendingPerformanceTest.swift        |  1 +
 .../src/Messages/MessageSender+SenderKey.swift |  9 ++++++++-
 .../src/Network/API/NetworkManager.swift       | 18 +++++++++++++-----
 .../src/Network/WebSockets/OWSWebSocket.swift  | 14 ++++++++++++++
 4 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/Signal/test/PerformanceTests/MessageSendingPerformanceTest.swift b/Signal/test/PerformanceTests/MessageSendingPerformanceTest.swift
index e4461c9731..8bb9997d95 100644
--- a/Signal/test/PerformanceTests/MessageSendingPerformanceTest.swift
+++ b/Signal/test/PerformanceTests/MessageSendingPerformanceTest.swift
@@ -211,6 +211,7 @@ class StubbableNetworkManager: NetworkManager {
     }
 
     public override func makePromise(request: TSRequest,
+                                     websocketSupportsRequest: Bool = false,
                                      remainingRetryCount: Int = 0) -> Promise<HTTPResponse> {
         Logger.info("Ignoring request: \(request)")
 
diff --git a/SignalServiceKit/src/Messages/MessageSender+SenderKey.swift b/SignalServiceKit/src/Messages/MessageSender+SenderKey.swift
index f0e384ba16..23045b9585 100644
--- a/SignalServiceKit/src/Messages/MessageSender+SenderKey.swift
+++ b/SignalServiceKit/src/Messages/MessageSender+SenderKey.swift
@@ -586,7 +586,14 @@ extension MessageSender {
             timestamp: timestamp,
             isOnline: isOnline)
 
-        return networkManager.makePromise(request: request)
+        // If we can use the websocket, try that and fail over to REST.
+        // If not, go straight to REST.
+        let remainingRetryCount: Int = (OWSWebSocket.canAppUseSocketsToMakeRequests
+                                            ? 1
+                                            : 0)
+        return networkManager.makePromise(request: request,
+                                          websocketSupportsRequest: true,
+                                          remainingRetryCount: remainingRetryCount)
     }
 }
 
diff --git a/SignalServiceKit/src/Network/API/NetworkManager.swift b/SignalServiceKit/src/Network/API/NetworkManager.swift
index 04ccce6d4a..9a5e619a75 100644
--- a/SignalServiceKit/src/Network/API/NetworkManager.swift
+++ b/SignalServiceKit/src/Network/API/NetworkManager.swift
@@ -18,11 +18,17 @@ public class NetworkManager: NSObject {
 
     // This method can be called from any thread.
     public func makePromise(request: TSRequest,
+                            websocketSupportsRequest: Bool = false,
                             remainingRetryCount: Int = 0) -> Promise<HTTPResponse> {
-        firstly {
-            FeatureFlags.deprecateREST
-                ? websocketRequestPromise(request: request)
-                : restRequestPromise(request: request)
+        firstly { () -> Promise<HTTPResponse> in
+            // Fail over to REST if websocket attempt fails.
+            if remainingRetryCount > 0,
+               OWSWebSocket.canAppUseSocketsToMakeRequests,
+               websocketSupportsRequest {
+                return websocketRequestPromise(request: request)
+            } else {
+                return restRequestPromise(request: request)
+            }
         }.recover(on: .global()) { error -> Promise<HTTPResponse> in
             if error.isRetryable,
                remainingRetryCount > 0 {
@@ -67,7 +73,9 @@ public class NetworkManager: NSObject {
 @objc
 public class OWSFakeNetworkManager: NetworkManager {
 
-    public override func makePromise(request: TSRequest, remainingRetryCount: Int = 0) -> Promise<HTTPResponse> {
+    public override func makePromise(request: TSRequest,
+                                     websocketSupportsRequest: Bool = false,
+                                     remainingRetryCount: Int = 0) -> Promise<HTTPResponse> {
         Logger.info("Ignoring request: \(request)")
         // Never resolve.
         let (promise, _) = Promise<HTTPResponse>.pending()
diff --git a/SignalServiceKit/src/Network/WebSockets/OWSWebSocket.swift b/SignalServiceKit/src/Network/WebSockets/OWSWebSocket.swift
index d07e174bab..c555fdaaca 100644
--- a/SignalServiceKit/src/Network/WebSockets/OWSWebSocket.swift
+++ b/SignalServiceKit/src/Network/WebSockets/OWSWebSocket.swift
@@ -638,6 +638,20 @@ public class OWSWebSocket: NSObject {
 
     // MARK: - Socket LifeCycle
 
+    public static var canAppUseSocketsToMakeRequests: Bool {
+        if signalService.isCensorshipCircumventionActive {
+            return false
+        } else if CurrentAppContext().isMainApp {
+            return true
+        } else if FeatureFlags.deprecateREST {
+            // When we deprecated REST, we _do_ want to open
+            // both websockets in the app extensions.
+            return true
+        } else {
+            return false
+        }
+    }
+
     // This method is thread-safe.
     public var shouldSocketBeOpen: Bool {