Add keepalive thread to provisioning socket

Closes #258
mollyim · May 10, 2024 · 53c2b06 · 53c2b06
1 parent 2be0f7e
commit 53c2b06
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 49 deletions.
diff --git a/app/src/main/java/org/thoughtcrime/securesms/registration/LinkDeviceRepository.kt b/app/src/main/java/org/thoughtcrime/securesms/registration/LinkDeviceRepository.kt
@@ -69,46 +69,44 @@ class LinkDeviceRepository(private val context: Application) {
     val deviceName: String? = progressImpl.deviceName
 
     return Single.fromCallable<ServiceResponse<LinkDeviceResponse>> {
-      try {
-        val ret = accountManager.getNewDeviceRegistration(tempIdentityKey)
-
-        val universalUnidentifiedAccess: Boolean = TextSecurePreferences.isUniversalUnidentifiedAccess(context)
-        val unidentifiedAccessKey: ByteArray = UnidentifiedAccess.deriveAccessKeyFrom(registrationData.profileKey)
-
-        val registrationLock: String? = ret.masterKey?.deriveRegistrationLock()
-        val encryptedDeviceName = deviceName?.let { DeviceNameCipher.encryptDeviceName(it.toByteArray(), ret.aciIdentity) }
-
-        val notDiscoverable =  SignalStore.phoneNumberPrivacy().phoneNumberDiscoverabilityMode == PhoneNumberDiscoverabilityMode.NOT_DISCOVERABLE
-
-        val accountAttributes = AccountAttributes(
-          signalingKey = null,
-          registrationId = registrationData.registrationId,
-          fetchesMessages = registrationData.isNotFcm,
-          registrationLock = registrationLock,
-          unidentifiedAccessKey = unidentifiedAccessKey,
-          unrestrictedUnidentifiedAccess = universalUnidentifiedAccess,
-          capabilities = AppCapabilities.getCapabilities(true),
-          discoverableByPhoneNumber = !notDiscoverable,
-          name = encryptedDeviceName?.let { Base64.encodeWithPadding(it) },
-          pniRegistrationId = registrationData.pniRegistrationId,
-          recoveryPassword = registrationData.recoveryPassword
-        )
-
-        val aciPreKeyCollection = RegistrationRepository.generateSignedAndLastResortPreKeys(ret.aciIdentity, SignalStore.account().aciPreKeys)
-        val pniPreKeyCollection = RegistrationRepository.generateSignedAndLastResortPreKeys(ret.pniIdentity, SignalStore.account().pniPreKeys)
-
-        val deviceId = accountManager.finishNewDeviceRegistration(
-          ret.provisioningCode,
-          accountAttributes,
-          aciPreKeyCollection, pniPreKeyCollection,
-          registrationData.fcmToken
-        )
-
-        ServiceResponse.forResult(LinkDeviceResponse(deviceName, ret, deviceId, aciPreKeyCollection, pniPreKeyCollection), 200, null)
-      } catch (e: IOException) {
-        ServiceResponse.forExecutionError(e)
-      }
-    }.subscribeOn(Schedulers.io()).map(::LinkDeviceResponseProcessor)
+      val ret = accountManager.getNewDeviceRegistration(tempIdentityKey)
+
+      val universalUnidentifiedAccess: Boolean = TextSecurePreferences.isUniversalUnidentifiedAccess(context)
+      val unidentifiedAccessKey: ByteArray = UnidentifiedAccess.deriveAccessKeyFrom(registrationData.profileKey)
+
+      val registrationLock: String? = ret.masterKey?.deriveRegistrationLock()
+      val encryptedDeviceName = deviceName?.let { DeviceNameCipher.encryptDeviceName(it.toByteArray(), ret.aciIdentity) }
+
+      val notDiscoverable =  SignalStore.phoneNumberPrivacy().phoneNumberDiscoverabilityMode == PhoneNumberDiscoverabilityMode.NOT_DISCOVERABLE
+
+      val accountAttributes = AccountAttributes(
+        signalingKey = null,
+        registrationId = registrationData.registrationId,
+        fetchesMessages = registrationData.isNotFcm,
+        registrationLock = registrationLock,
+        unidentifiedAccessKey = unidentifiedAccessKey,
+        unrestrictedUnidentifiedAccess = universalUnidentifiedAccess,
+        capabilities = AppCapabilities.getCapabilities(true),
+        discoverableByPhoneNumber = !notDiscoverable,
+        name = encryptedDeviceName?.let { Base64.encodeWithPadding(it) },
+        pniRegistrationId = registrationData.pniRegistrationId,
+        recoveryPassword = registrationData.recoveryPassword
+      )
+
+      val aciPreKeyCollection = RegistrationRepository.generateSignedAndLastResortPreKeys(ret.aciIdentity, SignalStore.account().aciPreKeys)
+      val pniPreKeyCollection = RegistrationRepository.generateSignedAndLastResortPreKeys(ret.pniIdentity, SignalStore.account().pniPreKeys)
+
+      val deviceId = accountManager.finishNewDeviceRegistration(
+        ret.provisioningCode,
+        accountAttributes,
+        aciPreKeyCollection, pniPreKeyCollection,
+        registrationData.fcmToken
+      )
+
+      ServiceResponse.forResult(LinkDeviceResponse(deviceName, ret, deviceId, aciPreKeyCollection, pniPreKeyCollection), 200, null)
+    }.subscribeOn(Schedulers.io())
+      .onErrorReturn { t -> ServiceResponse.forExecutionError(t) }
+      .map(::LinkDeviceResponseProcessor)
   }
 
   companion object {

diff --git a/...vice/src/main/java/org/whispersystems/signalservice/internal/push/ProvisioningSocket.java b/...vice/src/main/java/org/whispersystems/signalservice/internal/push/ProvisioningSocket.java
@@ -2,8 +2,11 @@
 
 import okio.ByteString;
 
+import org.signal.core.util.logging.Log;
 import org.signal.libsignal.protocol.IdentityKeyPair;
 import org.signal.libsignal.protocol.InvalidKeyException;
+import org.whispersystems.signalservice.api.util.SleepTimer;
+import org.whispersystems.signalservice.api.util.UptimeSleepTimer;
 import org.whispersystems.signalservice.api.websocket.HealthMonitor;
 import org.whispersystems.signalservice.internal.configuration.SignalServiceConfiguration;
 import org.whispersystems.signalservice.internal.crypto.PrimaryProvisioningCipher;
@@ -17,6 +20,8 @@
 
 public class ProvisioningSocket {
 
+  private static final String TAG = Log.tag(ProvisioningSocket.class);
+
   private final WebSocketConnection connection;
 
   private boolean connected = false;
@@ -54,20 +59,34 @@ public synchronized ProvisionMessage getProvisioningMessage(IdentityKeyPair temp
     if (!connected) {
       throw new IllegalStateException("No UUID requested yet!");
     }
-    ByteString bytes = readRequest();
-    connection.disconnect();
-    connected = false;
-    ProvisionMessage msg;
+    final Thread keepAliveThread = new Thread(() -> {
+      try {
+        while (connected) {
+          connection.sendKeepAlive();
+          //noinspection BusyWait
+          Thread.sleep(5_000);
+        }
+      } catch (InterruptedException | IOException ignored) {}
+    });
+    keepAliveThread.start();
     try {
-      msg = new PrimaryProvisioningCipher(null).decrypt(tempIdentity, bytes.toByteArray());
-    } catch (InvalidKeyException e) {
-      throw new AssertionError(e);
+      ByteString bytes = readRequest();
+      ProvisionMessage msg;
+      try {
+        msg = new PrimaryProvisioningCipher(null).decrypt(tempIdentity, bytes.toByteArray());
+      } catch (InvalidKeyException e) {
+        throw new AssertionError(e);
+      }
+      return msg;
+    } finally {
+      connection.disconnect();
+      connected = false;
+      keepAliveThread.interrupt();
     }
-    return msg;
   }
 
   private ByteString readRequest() throws TimeoutException, IOException {
-    WebSocketRequestMessage response = connection.readRequest(100000);
+    WebSocketRequestMessage response = connection.readRequest(500_000);
     return response.body;
   }
 }