Fix bug with stale linked devices when changing number.

app/src/main/java/org/thoughtcrime/securesms/components/settings/app/changenumber/ChangeNumberRepository.kt

@@ -6,12 +6,12 @@ import io.reactivex.rxjava3.core.Single
 import io.reactivex.rxjava3.schedulers.Schedulers
 import org.signal.core.util.logging.Log
 import org.signal.libsignal.protocol.IdentityKeyPair
+import org.signal.libsignal.protocol.SignalProtocolAddress
 import org.signal.libsignal.protocol.state.SignalProtocolStore
 import org.signal.libsignal.protocol.util.KeyHelper
 import org.signal.libsignal.protocol.util.Medium
 import org.thoughtcrime.securesms.crypto.IdentityKeyUtil
 import org.thoughtcrime.securesms.crypto.PreKeyUtil
-import org.thoughtcrime.securesms.crypto.storage.PreKeyMetadataStore
 import org.thoughtcrime.securesms.database.IdentityDatabase
 import org.thoughtcrime.securesms.database.SignalDatabase
 import org.thoughtcrime.securesms.database.model.databaseprotos.PendingChangeNumberMetadata
@@ -30,7 +30,6 @@ import org.whispersystems.signalservice.api.KeyBackupSystemNoDataException
 import org.whispersystems.signalservice.api.SignalServiceAccountManager
 import org.whispersystems.signalservice.api.SignalServiceMessageSender
 import org.whispersystems.signalservice.api.account.ChangePhoneNumberRequest
-import org.whispersystems.signalservice.api.messages.multidevice.DeviceInfo
 import org.whispersystems.signalservice.api.push.PNI
 import org.whispersystems.signalservice.api.push.ServiceId
 import org.whispersystems.signalservice.api.push.ServiceIdType
@@ -41,13 +40,17 @@ import org.whispersystems.signalservice.internal.push.OutgoingPushMessage
 import org.whispersystems.signalservice.internal.push.SignalServiceProtos.SyncMessage
 import org.whispersystems.signalservice.internal.push.VerifyAccountResponse
 import org.whispersystems.signalservice.internal.push.WhoAmIResponse
+import org.whispersystems.signalservice.internal.push.exceptions.MismatchedDevicesException
 import java.io.IOException
 import java.security.MessageDigest
 import java.security.SecureRandom
 
 private val TAG: String = Log.tag(ChangeNumberRepository::class.java)
 
-class ChangeNumberRepository(private val accountManager: SignalServiceAccountManager = ApplicationDependencies.getSignalServiceAccountManager()) {
+class ChangeNumberRepository(
+  private val accountManager: SignalServiceAccountManager = ApplicationDependencies.getSignalServiceAccountManager(),
+  private val messageSender: SignalServiceMessageSender = ApplicationDependencies.getSignalServiceMessageSender()
+) {
 
   fun ensureDecryptionsDrained(): Completable {
     return Completable.create { emitter ->
@@ -61,9 +64,26 @@ class ChangeNumberRepository(private val accountManager: SignalServiceAccountMan
 
   fun changeNumber(code: String, newE164: String): Single<ServiceResponse<VerifyAccountResponse>> {
     return Single.fromCallable {
-      val (request: ChangePhoneNumberRequest, metadata: PendingChangeNumberMetadata) = createChangeNumberRequest(code, newE164, null)
-      SignalStore.misc().setPendingChangeNumberMetadata(metadata)
-      accountManager.changeNumber(request)
+      var completed = false
+      var attempts = 0
+      lateinit var changeNumberResponse: ServiceResponse<VerifyAccountResponse>
+
+      while (!completed && attempts < 5) {
+        val (request: ChangePhoneNumberRequest, metadata: PendingChangeNumberMetadata) = createChangeNumberRequest(code, newE164, null)
+        SignalStore.misc().setPendingChangeNumberMetadata(metadata)
+
+        changeNumberResponse = accountManager.changeNumber(request)
+
+        val possibleError: Throwable? = changeNumberResponse.applicationError.orElse(null)
+        if (possibleError is MismatchedDevicesException) {
+          messageSender.handleChangeNumberMismatchDevices(possibleError.mismatchedDevices)
+          attempts++
+        } else {
+          completed = true
+        }
+      }
+
+      changeNumberResponse
     }.subscribeOn(Schedulers.io())
       .onErrorReturn { t -> ServiceResponse.forExecutionError(t) }
   }
@@ -75,22 +95,40 @@ class ChangeNumberRepository(private val accountManager: SignalServiceAccountMan
     tokenData: TokenData
   ): Single<ServiceResponse<VerifyAccountRepository.VerifyAccountWithRegistrationLockResponse>> {
     return Single.fromCallable {
-      try {
-        val kbsData: KbsPinData = KbsRepository.restoreMasterKey(pin, tokenData.enclave, tokenData.basicAuth, tokenData.tokenResponse)!!
-        val registrationLock: String = kbsData.masterKey.deriveRegistrationLock()
-        val (request: ChangePhoneNumberRequest, metadata: PendingChangeNumberMetadata) = createChangeNumberRequest(code, newE164, registrationLock)
+      val kbsData: KbsPinData
+      val registrationLock: String
 
+      try {
+        kbsData = KbsRepository.restoreMasterKey(pin, tokenData.enclave, tokenData.basicAuth, tokenData.tokenResponse)!!
+        registrationLock = kbsData.masterKey.deriveRegistrationLock()
+      } catch (e: KeyBackupSystemWrongPinException) {
+        return@fromCallable ServiceResponse.forExecutionError(e)
+      } catch (e: KeyBackupSystemNoDataException) {
+        return@fromCallable ServiceResponse.forExecutionError(e)
+      } catch (e: IOException) {
+        return@fromCallable ServiceResponse.forExecutionError(e)
+      }
+
+      var completed = false
+      var attempts = 0
+      lateinit var changeNumberResponse: ServiceResponse<VerifyAccountResponse>
+
+      while (!completed && attempts < 5) {
+        val (request: ChangePhoneNumberRequest, metadata: PendingChangeNumberMetadata) = createChangeNumberRequest(code, newE164, registrationLock)
         SignalStore.misc().setPendingChangeNumberMetadata(metadata)
 
-        val response: ServiceResponse<VerifyAccountResponse> = accountManager.changeNumber(request)
-        VerifyAccountRepository.VerifyAccountWithRegistrationLockResponse.from(response, kbsData)
-      } catch (e: KeyBackupSystemWrongPinException) {
-        ServiceResponse.forExecutionError(e)
-      } catch (e: KeyBackupSystemNoDataException) {
-        ServiceResponse.forExecutionError(e)
-      } catch (e: IOException) {
-        ServiceResponse.forExecutionError(e)
+        changeNumberResponse = accountManager.changeNumber(request)
+
+        val possibleError: Throwable? = changeNumberResponse.applicationError.orElse(null)
+        if (possibleError is MismatchedDevicesException) {
+          messageSender.handleChangeNumberMismatchDevices(possibleError.mismatchedDevices)
+          attempts++
+        } else {
+          completed = true
+        }
       }
+
+      VerifyAccountRepository.VerifyAccountWithRegistrationLockResponse.from(changeNumberResponse, kbsData)
     }.subscribeOn(Schedulers.io())
       .onErrorReturn { t -> ServiceResponse.forExecutionError(t) }
   }
@@ -199,49 +237,57 @@ class ChangeNumberRepository(private val accountManager: SignalServiceAccountMan
     newE164: String,
     registrationLock: String?
   ): ChangeNumberRequestData {
-    val messageSender: SignalServiceMessageSender = ApplicationDependencies.getSignalServiceMessageSender()
-    val pniProtocolStore: SignalProtocolStore = ApplicationDependencies.getProtocolStore().pni()
-    val pniMetadataStore: PreKeyMetadataStore = SignalStore.account().pniPreKeys
-
-    val devices: List<DeviceInfo> = accountManager.getDevices()
+    val selfIdentifier: String = SignalStore.account().requireAci().toString()
+    val aciProtocolStore: SignalProtocolStore = ApplicationDependencies.getProtocolStore().aci()
 
     val pniIdentity: IdentityKeyPair = IdentityKeyUtil.generateIdentityKeyPair()
     val deviceMessages = mutableListOf<OutgoingPushMessage>()
-    val devicePniSignedPreKeys = mutableMapOf<String, SignedPreKeyEntity>()
-    val pniRegistrationIds = mutableMapOf<String, Int>()
-    val primaryDeviceId = SignalServiceAddress.DEFAULT_DEVICE_ID.toString()
+    val devicePniSignedPreKeys = mutableMapOf<Int, SignedPreKeyEntity>()
+    val pniRegistrationIds = mutableMapOf<Int, Int>()
+    val primaryDeviceId: Int = SignalServiceAddress.DEFAULT_DEVICE_ID
 
-    for (device in devices) {
-      val deviceId = device.id.toString()
+    val devices: List<Int> = listOf(primaryDeviceId) + aciProtocolStore.getSubDeviceSessions(selfIdentifier)
 
-      // Signed Prekeys
-      val signedPreKeyRecord = if (deviceId == primaryDeviceId) {
-        PreKeyUtil.generateAndStoreSignedPreKey(pniProtocolStore, pniMetadataStore, pniIdentity.privateKey)
-      } else {
-        PreKeyUtil.generateSignedPreKey(SecureRandom().nextInt(Medium.MAX_VALUE), pniIdentity.privateKey)
+    devices
+      .filter { it == primaryDeviceId || aciProtocolStore.containsSession(SignalProtocolAddress(selfIdentifier, it)) }
+      .forEach { deviceId ->
+        // Signed Prekeys
+        val signedPreKeyRecord = if (deviceId == primaryDeviceId) {
+          PreKeyUtil.generateAndStoreSignedPreKey(ApplicationDependencies.getProtocolStore().pni(), SignalStore.account().pniPreKeys, pniIdentity.privateKey)
+        } else {
+          PreKeyUtil.generateSignedPreKey(SecureRandom().nextInt(Medium.MAX_VALUE), pniIdentity.privateKey)
+        }
+        devicePniSignedPreKeys[deviceId] = SignedPreKeyEntity(signedPreKeyRecord.id, signedPreKeyRecord.keyPair.publicKey, signedPreKeyRecord.signature)
+
+        // Registration Ids
+        var pniRegistrationId = -1
+        while (pniRegistrationId < 0 || pniRegistrationIds.values.contains(pniRegistrationId)) {
+          pniRegistrationId = KeyHelper.generateRegistrationId(false)
+        }
+        pniRegistrationIds[deviceId] = pniRegistrationId
+
+        // Device Messages
+        if (deviceId != primaryDeviceId) {
+          val pniChangeNumber = SyncMessage.PniChangeNumber.newBuilder()
+            .setIdentityKeyPair(pniIdentity.serialize().toProtoByteString())
+            .setSignedPreKey(signedPreKeyRecord.serialize().toProtoByteString())
+            .setRegistrationId(pniRegistrationId)
+            .build()
+
+          deviceMessages += messageSender.getEncryptedSyncPniChangeNumberMessage(deviceId, pniChangeNumber)
+        }
       }
-      devicePniSignedPreKeys[deviceId] = SignedPreKeyEntity(signedPreKeyRecord.id, signedPreKeyRecord.keyPair.publicKey, signedPreKeyRecord.signature)
 
-      // Registration Ids
-      var pniRegistrationId = -1
-      while (pniRegistrationId < 0 || pniRegistrationIds.values.contains(pniRegistrationId)) {
-        pniRegistrationId = KeyHelper.generateRegistrationId(false)
-      }
-      pniRegistrationIds[deviceId] = pniRegistrationId
+    val request = ChangePhoneNumberRequest(
+      newE164,
+      code,
+      registrationLock,
+      pniIdentity.publicKey,
+      deviceMessages,
+      devicePniSignedPreKeys.mapKeys { it.key.toString() },
+      pniRegistrationIds.mapKeys { it.key.toString() }
+    )
 
-      // Device Messages
-      if (deviceId != primaryDeviceId) {
-        val pniChangeNumber = SyncMessage.PniChangeNumber.newBuilder()
-          .setIdentityKeyPair(pniIdentity.serialize().toProtoByteString())
-          .setSignedPreKey(signedPreKeyRecord.serialize().toProtoByteString())
-          .setRegistrationId(pniRegistrationId)
-          .build()
-
-        deviceMessages += messageSender.getEncryptedSyncPniChangeNumberMessage(device.id, pniChangeNumber)
-      }
-    }
-
-    val request = ChangePhoneNumberRequest(newE164, code, registrationLock, pniIdentity.publicKey, deviceMessages, devicePniSignedPreKeys, pniRegistrationIds)
     val metadata = PendingChangeNumberMetadata.newBuilder()
       .setPreviousPni(SignalStore.account().pni!!.toByteString())
       .setPniIdentityKeyPair(pniIdentity.serialize().toProtoByteString())

app/src/main/java/org/thoughtcrime/securesms/push/SignalServiceNetworkAccess.kt

@@ -211,7 +211,7 @@ open class SignalServiceNetworkAccess(context: Context) {
     COUNTRY_CODE_UZBEKISTAN,
   )
 
-  val uncensoredConfiguration: SignalServiceConfiguration = SignalServiceConfiguration(
+  open val uncensoredConfiguration: SignalServiceConfiguration = SignalServiceConfiguration(
     arrayOf(SignalServiceUrl(BuildConfig.SIGNAL_URL, serviceTrustStore)),
     mapOf(
       0 to arrayOf(SignalCdnUrl(BuildConfig.SIGNAL_CDN_URL, serviceTrustStore)),