From 38f9aef2af7ef7d987f4fe140d65193151c7a4b5 Mon Sep 17 00:00:00 2001
From: Scott Nonnenberg <scott@signal.org>
Date: Mon, 17 May 2021 11:29:37 -0700
Subject: [PATCH] On 401 response from Signal server, reconnect websocket

---
 ts/background.ts        | 32 ++++++++++++++++++++++++++++----
 ts/textsecure/WebAPI.ts | 31 ++++++++++++++++++++++++++-----
 2 files changed, 54 insertions(+), 9 deletions(-)

diff --git a/ts/background.ts b/ts/background.ts
index 7ca380f91..f2419fe33 100644
--- a/ts/background.ts
+++ b/ts/background.ts
@@ -21,6 +21,7 @@ import { removeStorageKeyJobQueue } from './jobs/removeStorageKeyJobQueue';
 import { ourProfileKeyService } from './services/ourProfileKey';
 import { shouldRespondWithProfileKey } from './util/shouldRespondWithProfileKey';
 import { setToExpire } from './services/MessageUpdater';
+import { LatestQueue } from './util/LatestQueue';
 
 const MAX_ATTACHMENT_DOWNLOAD_AGE = 3600 * 72 * 1000;
 
@@ -1442,6 +1443,29 @@ export async function startApp(): Promise<void> {
     }
   });
 
+  const reconnectToWebSocketQueue = new LatestQueue();
+
+  const enqueueReconnectToWebSocket = () => {
+    reconnectToWebSocketQueue.add(async () => {
+      if (!messageReceiver) {
+        window.log.info(
+          'reconnectToWebSocket: No messageReceiver. Early return.'
+        );
+        return;
+      }
+
+      window.log.info('reconnectToWebSocket starting...');
+      await disconnect();
+      connect();
+      window.log.info('reconnectToWebSocket complete.');
+    });
+  };
+
+  window.Whisper.events.on(
+    'mightBeUnlinked',
+    window._.debounce(enqueueReconnectToWebSocket, 1000, { maxWait: 5000 })
+  );
+
   function runStorageService() {
     window.Signal.Services.enableStorageService();
     window.textsecure.messaging.sendRequestKeySyncMessage();
@@ -1758,16 +1782,16 @@ export async function startApp(): Promise<void> {
     );
   }
 
-  function disconnect() {
+  async function disconnect() {
     window.log.info('disconnect');
 
     // Clear timer, since we're only called when the timer is expired
     disconnectTimer = null;
 
-    if (messageReceiver) {
-      messageReceiver.close();
-    }
     window.Signal.AttachmentDownloads.stop();
+    if (messageReceiver) {
+      await messageReceiver.close();
+    }
   }
 
   let connectCount = 0;
diff --git a/ts/textsecure/WebAPI.ts b/ts/textsecure/WebAPI.ts
index 9266c6637..462455ebd 100644
--- a/ts/textsecure/WebAPI.ts
+++ b/ts/textsecure/WebAPI.ts
@@ -330,6 +330,7 @@ type PromiseAjaxOptionsType = {
     | 'jsonwithdetails'
     | 'arraybuffer'
     | 'arraybufferwithdetails';
+  serverUrl?: string;
   stack?: string;
   timeout?: number;
   type: HTTPCodeType;
@@ -354,6 +355,11 @@ function isSuccess(status: number): boolean {
   return status >= 0 && status < 400;
 }
 
+function getHostname(url: string): string {
+  const urlObject = new URL(url);
+  return urlObject.hostname;
+}
+
 async function _promiseAjax(
   providedUrl: string | null,
   options: PromiseAjaxOptionsType
@@ -438,11 +444,25 @@ async function _promiseAjax(
 
     fetch(url, fetchOptions)
       .then(async response => {
-        // Build expired!
-        if (response.status === 499) {
-          window.log.error('Error: build expired');
-          await window.storage.put('remoteBuildExpiration', Date.now());
-          window.reduxActions.expiration.hydrateExpirationStatus(true);
+        if (options.serverUrl) {
+          if (
+            response.status === 499 &&
+            getHostname(options.serverUrl) === getHostname(url)
+          ) {
+            window.log.error('Got 499 from Signal Server. Build is expired.');
+            await window.storage.put('remoteBuildExpiration', Date.now());
+            window.reduxActions.expiration.hydrateExpirationStatus(true);
+          }
+          if (
+            !unauthenticated &&
+            response.status === 401 &&
+            getHostname(options.serverUrl) === getHostname(url)
+          ) {
+            window.log.error(
+              'Got 401 from Signal Server. We might be unlinked.'
+            );
+            window.Whisper.events.trigger('mightBeUnlinked');
+          }
         }
 
         let resultPromise;
@@ -1071,6 +1091,7 @@ export function initialize({
         type: param.httpType,
         user: param.username || username,
         redactUrl: param.redactUrl,
+        serverUrl: url,
         validateResponse: param.validateResponse,
         version,
         unauthenticated: param.unauthenticated,