Fix atomic registrations when not using session ID.

app/src/main/java/org/thoughtcrime/securesms/registration/RegistrationData.kt

@@ -1,7 +1,7 @@
 package org.thoughtcrime.securesms.registration
 
 import org.signal.libsignal.zkgroup.profiles.ProfileKey
-import org.whispersystems.signalservice.api.account.PreKeyCollections
+import org.whispersystems.signalservice.api.account.PreKeyCollection
 
 data class RegistrationData(
   val code: String,
@@ -9,7 +9,8 @@ data class RegistrationData(
   val password: String,
   val registrationId: Int,
   val profileKey: ProfileKey,
-  val preKeyCollections: PreKeyCollections,
+  val aciPreKeyCollection: PreKeyCollection,
+  val pniPreKeyCollection: PreKeyCollection,
   val fcmToken: String?,
   val pniRegistrationId: Int,
   val recoveryPassword: String?

app/src/main/java/org/thoughtcrime/securesms/registration/RegistrationRepository.java

@@ -15,7 +15,6 @@ import org.signal.libsignal.protocol.state.PreKeyRecord;
 import org.signal.libsignal.protocol.state.SignedPreKeyRecord;
 import org.signal.libsignal.protocol.util.KeyHelper;
 import org.signal.libsignal.zkgroup.profiles.ProfileKey;
-import org.thoughtcrime.securesms.crypto.IdentityKeyUtil;
 import org.thoughtcrime.securesms.crypto.PreKeyUtil;
 import org.thoughtcrime.securesms.crypto.ProfileKeyUtil;
 import org.thoughtcrime.securesms.crypto.SenderKeyUtil;
@@ -40,8 +39,6 @@ import org.thoughtcrime.securesms.util.TextSecurePreferences;
 import org.whispersystems.signalservice.api.KbsPinData;
 import org.whispersystems.signalservice.api.SignalServiceAccountManager;
 import org.whispersystems.signalservice.api.account.PreKeyCollection;
-import org.whispersystems.signalservice.api.account.PreKeyCollections;
-import org.whispersystems.signalservice.api.account.PreKeyUpload;
 import org.whispersystems.signalservice.api.push.ACI;
 import org.whispersystems.signalservice.api.push.PNI;
 import org.whispersystems.signalservice.api.push.ServiceIdType;
@@ -148,11 +145,11 @@ public final class RegistrationRepository {
     SenderKeyUtil.clearAllState();
 
     SignalServiceAccountDataStoreImpl aciProtocolStore    = ApplicationDependencies.getProtocolStore().aci();
-    PreKeyCollection                  aciPreKeyCollection = registrationData.getPreKeyCollections().getAciPreKeyCollection();
+    PreKeyCollection                  aciPreKeyCollection = registrationData.getAciPreKeyCollection();
     PreKeyMetadataStore               aciMetadataStore    = SignalStore.account().aciPreKeys();
 
     SignalServiceAccountDataStoreImpl pniProtocolStore    = ApplicationDependencies.getProtocolStore().pni();
-    PreKeyCollection                  pniPreKeyCollection = registrationData.getPreKeyCollections().getPniPreKeyCollection();
+    PreKeyCollection                  pniPreKeyCollection = registrationData.getPniPreKeyCollection();
     PreKeyMetadataStore               pniMetadataStore    = SignalStore.account().pniPreKeys();
 
     storePreKeys(aciProtocolStore, aciMetadataStore, aciPreKeyCollection);
@@ -189,23 +186,24 @@ public final class RegistrationRepository {
     ApplicationDependencies.getIncomingMessageObserver();
   }
 
-  public static @Nullable PreKeyCollections generatePreKeys() {
-    final IdentityKeyPair     keyPair          = IdentityKeyUtil.generateIdentityKeyPair();
-    final PreKeyMetadataStore aciMetadataStore = SignalStore.account().aciPreKeys();
-    final PreKeyMetadataStore pniMetadataStore = SignalStore.account().pniPreKeys();
-
-    try {
-      return new PreKeyCollections(keyPair,
-                                   generatePreKeysForType(ServiceIdType.ACI, keyPair, aciMetadataStore),
-                                   generatePreKeysForType(ServiceIdType.PNI, keyPair, pniMetadataStore)
-      );
-    } catch (IOException e) {
-      Log.e(TAG, "Failed to generate prekeys!", e);
-      return null;
+  public static PreKeyCollection generatePreKeysForType(ServiceIdType serviceIdType) {
+    IdentityKeyPair keyPair;
+    PreKeyMetadataStore metadataStore;
+    if (serviceIdType == ServiceIdType.ACI) {
+      if (!SignalStore.account().hasAciIdentityKey()) {
+        SignalStore.account().generateAciIdentityKeyIfNecessary();
+      }
+      keyPair = SignalStore.account().getAciIdentityKey();
+      metadataStore = SignalStore.account().aciPreKeys();
+    } else if (serviceIdType == ServiceIdType.PNI) {
+      if (!SignalStore.account().hasPniIdentityKey()) {
+        SignalStore.account().generatePniIdentityKeyIfNecessary();
+      }
+      keyPair = SignalStore.account().getPniIdentityKey();
+      metadataStore = SignalStore.account().pniPreKeys();
+    } else {
+      throw new IllegalArgumentException("serviceIdType is not one of {ACI, PNI}");
     }
-  }
-
-  private static PreKeyCollection generatePreKeysForType(ServiceIdType serviceIdType, IdentityKeyPair keyPair, PreKeyMetadataStore metadataStore) throws IOException {
     int                nextSignedPreKeyId = metadataStore.getNextSignedPreKeyId();
     SignedPreKeyRecord signedPreKey       = PreKeyUtil.generateSignedPreKey(nextSignedPreKeyId, keyPair.getPrivateKey());
     metadataStore.setActiveSignedPreKeyId(signedPreKey.getId());
@@ -222,6 +220,7 @@ public final class RegistrationRepository {
     List<KyberPreKeyRecord> oneTimeKyberPreKeys        = PreKeyUtil.generateOneTimeKyberPreKeyRecords(oneTimeKyberPreKeyIdOffset, keyPair.getPrivateKey());
 
     return new PreKeyCollection(
+        keyPair,
         nextSignedPreKeyId,
         ecOneTimePreKeyIdOffset,
         nextKyberPreKeyId,

app/src/main/java/org/thoughtcrime/securesms/registration/VerifyAccountRepository.kt

@@ -186,7 +186,7 @@ class VerifyAccountRepository(private val context: Application) {
     )
 
     return Single.fromCallable {
-      val response = accountManager.registerAccount(sessionId, registrationData.recoveryPassword, accountAttributes, registrationData.preKeyCollections, registrationData.fcmToken, true)
+      val response = accountManager.registerAccount(sessionId, registrationData.recoveryPassword, accountAttributes, registrationData.aciPreKeyCollection, registrationData.pniPreKeyCollection, registrationData.fcmToken, true)
       VerifyResponse.from(response, kbsData, pin)
     }.subscribeOn(Schedulers.io())
   }

app/src/main/java/org/thoughtcrime/securesms/registration/viewmodel/RegistrationViewModel.java

@@ -32,6 +32,7 @@ import org.thoughtcrime.securesms.util.Util;
 import org.whispersystems.signalservice.api.KbsPinData;
 import org.whispersystems.signalservice.api.KeyBackupSystemNoDataException;
 import org.whispersystems.signalservice.api.kbs.PinHashUtil;
+import org.whispersystems.signalservice.api.push.ServiceIdType;
 import org.whispersystems.signalservice.api.push.exceptions.IncorrectCodeException;
 import org.whispersystems.signalservice.internal.ServiceResponse;
 import org.whispersystems.signalservice.internal.contacts.entities.TokenResponse;
@@ -235,7 +236,8 @@ public final class RegistrationViewModel extends BaseRegistrationViewModel {
                                 getRegistrationSecret(),
                                 registrationRepository.getRegistrationId(),
                                 registrationRepository.getProfileKey(getNumber().getE164Number()),
-                                Objects.requireNonNull(RegistrationRepository.generatePreKeys()),
+                                RegistrationRepository.generatePreKeysForType(ServiceIdType.ACI),
+                                RegistrationRepository.generatePreKeysForType(ServiceIdType.PNI),
                                 getFcmToken(),
                                 registrationRepository.getPniRegistrationId(),
                                 getSessionId() != null ? null : getRecoveryPassword());

libsignal/service/src/main/java/org/whispersystems/signalservice/api/SignalServiceAccountManager.java

@@ -19,7 +19,7 @@ import org.signal.libsignal.zkgroup.profiles.ProfileKey;
 import org.whispersystems.signalservice.api.account.AccountAttributes;
 import org.whispersystems.signalservice.api.account.ChangePhoneNumberRequest;
 import org.whispersystems.signalservice.api.account.PniKeyDistributionRequest;
-import org.whispersystems.signalservice.api.account.PreKeyCollections;
+import org.whispersystems.signalservice.api.account.PreKeyCollection;
 import org.whispersystems.signalservice.api.account.PreKeyUpload;
 import org.whispersystems.signalservice.api.crypto.ProfileCipher;
 import org.whispersystems.signalservice.api.crypto.ProfileCipherOutputStream;
@@ -317,9 +317,9 @@ public class SignalServiceAccountManager {
     }
   }
 
-  public @Nonnull ServiceResponse<VerifyAccountResponse> registerAccount(@Nullable String sessionId, @Nullable String recoveryPassword, AccountAttributes attributes, PreKeyCollections preKeys, String fcmToken, boolean skipDeviceTransfer) {
+  public @Nonnull ServiceResponse<VerifyAccountResponse> registerAccount(@Nullable String sessionId, @Nullable String recoveryPassword, AccountAttributes attributes, PreKeyCollection aciPreKeys, PreKeyCollection pniPreKeys, String fcmToken, boolean skipDeviceTransfer) {
     try {
-      VerifyAccountResponse response = pushServiceSocket.submitRegistrationRequest(sessionId, recoveryPassword, attributes, preKeys, fcmToken, skipDeviceTransfer);
+      VerifyAccountResponse response = pushServiceSocket.submitRegistrationRequest(sessionId, recoveryPassword, attributes, aciPreKeys, pniPreKeys, fcmToken, skipDeviceTransfer);
       return ServiceResponse.forResult(response, 200, null);
     } catch (IOException e) {
       return ServiceResponse.forUnknownError(e);

libsignal/service/src/main/java/org/whispersystems/signalservice/api/account/PreKeyCollection.kt

@@ -6,6 +6,7 @@
 package org.whispersystems.signalservice.api.account
 
 import org.signal.libsignal.protocol.IdentityKey
+import org.signal.libsignal.protocol.IdentityKeyPair
 import org.signal.libsignal.protocol.state.KyberPreKeyRecord
 import org.signal.libsignal.protocol.state.PreKeyRecord
 import org.signal.libsignal.protocol.state.SignedPreKeyRecord
@@ -17,6 +18,7 @@ import org.whispersystems.signalservice.api.push.ServiceIdType
  * the service approves the keys we have a local copy to persist.
  */
 data class PreKeyCollection(
+  val identityKeyPair: IdentityKeyPair,
   val nextSignedPreKeyId: Int,
   val ecOneTimePreKeyIdOffset: Int,
   val lastResortKyberPreKeyId: Int,

libsignal/service/src/main/java/org/whispersystems/signalservice/api/account/PreKeyCollections.kt

@@ -1,17 +0,0 @@
-/*
- * Copyright 2023 Signal Messenger, LLC
- * SPDX-License-Identifier: AGPL-3.0-only
- */
-
-package org.whispersystems.signalservice.api.account
-
-import org.signal.libsignal.protocol.IdentityKeyPair
-
-/**
- * A holder class to hold onto two [PreKeyCollection] that are very similar but share an [IdentityKeyPair]
- */
-data class PreKeyCollections(
-  val identityKeyPair: IdentityKeyPair,
-  val aciPreKeyCollection: PreKeyCollection,
-  val pniPreKeyCollection: PreKeyCollection
-)

libsignal/service/src/main/java/org/whispersystems/signalservice/internal/push/PushServiceSocket.java

@@ -40,10 +40,9 @@ import org.signal.storageservice.protos.groups.GroupExternalCredential;
 import org.signal.storageservice.protos.groups.GroupJoinInfo;
 import org.whispersystems.signalservice.api.account.AccountAttributes;
 import org.whispersystems.signalservice.api.account.ChangePhoneNumberRequest;
-import org.whispersystems.signalservice.api.account.PreKeyCollections;
 import org.whispersystems.signalservice.api.account.PniKeyDistributionRequest;
-import org.whispersystems.signalservice.api.account.PreKeyUpload;
 import org.whispersystems.signalservice.api.account.PreKeyCollection;
+import org.whispersystems.signalservice.api.account.PreKeyUpload;
 import org.whispersystems.signalservice.api.crypto.UnidentifiedAccess;
 import org.whispersystems.signalservice.api.groupsv2.CredentialResponse;
 import org.whispersystems.signalservice.api.groupsv2.GroupsV2AuthorizationString;
@@ -394,7 +393,7 @@ public class PushServiceSocket {
     }
   }
 
-  public VerifyAccountResponse submitRegistrationRequest(@Nullable String sessionId, @Nullable String recoveryPassword, AccountAttributes attributes, PreKeyCollections preKeys, @Nullable String fcmToken, boolean skipDeviceTransfer) throws IOException {
+  public VerifyAccountResponse submitRegistrationRequest(@Nullable String sessionId, @Nullable String recoveryPassword, AccountAttributes attributes, PreKeyCollection aciPreKeys, PreKeyCollection pniPreKeys, @Nullable String fcmToken, boolean skipDeviceTransfer) throws IOException {
     String path = REGISTRATION_PATH;
     if (sessionId == null && recoveryPassword == null) {
       throw new IllegalArgumentException("Neither Session ID nor Recovery Password provided.");
@@ -411,36 +410,29 @@ public class PushServiceSocket {
       gcmRegistrationId = new GcmRegistrationId(fcmToken, true);
     }
 
-    RegistrationSessionRequestBody body;
-    if (sessionId != null) {
-      final PreKeyCollection aciPreKeys = preKeys.getAciPreKeyCollection();
-      final PreKeyCollection pniPreKeys = preKeys.getPniPreKeyCollection();
-      final SignedPreKeyEntity aciSignedPreKey = new SignedPreKeyEntity(Objects.requireNonNull(aciPreKeys.getSignedPreKey()).getId(),
-                                                                        aciPreKeys.getSignedPreKey().getKeyPair().getPublicKey(),
-                                                                        aciPreKeys.getSignedPreKey().getSignature());
-      final SignedPreKeyEntity pniSignedPreKey = new SignedPreKeyEntity(Objects.requireNonNull(pniPreKeys.getSignedPreKey()).getId(),
-                                                                        pniPreKeys.getSignedPreKey().getKeyPair().getPublicKey(),
-                                                                        pniPreKeys.getSignedPreKey().getSignature());
-      final KyberPreKeyEntity aciLastResortKyberPreKey = new KyberPreKeyEntity(Objects.requireNonNull(aciPreKeys.getLastResortKyberPreKey()).getId(),
-                                                                               aciPreKeys.getLastResortKyberPreKey().getKeyPair().getPublicKey(),
-                                                                               aciPreKeys.getLastResortKyberPreKey().getSignature());
-      final KyberPreKeyEntity pniLastResortKyberPreKey = new KyberPreKeyEntity(Objects.requireNonNull(pniPreKeys.getLastResortKyberPreKey()).getId(),
-                                                                               pniPreKeys.getLastResortKyberPreKey().getKeyPair().getPublicKey(),
-                                                                               pniPreKeys.getLastResortKyberPreKey().getSignature());
-      body = new RegistrationSessionRequestBody(sessionId,
-                                                null,
-                                                attributes,
-                                                Base64.encodeBytesWithoutPadding(aciPreKeys.getIdentityKey().serialize()),
-                                                Base64.encodeBytesWithoutPadding(pniPreKeys.getIdentityKey().serialize()),
-                                                aciSignedPreKey,
-                                                pniSignedPreKey,
-                                                aciLastResortKyberPreKey,
-                                                pniLastResortKyberPreKey,
-                                                gcmRegistrationId,
-                                                skipDeviceTransfer);
-    } else {
-      body = new RegistrationSessionRequestBody(null, recoveryPassword, attributes, null, null, null, null, null, null, null, skipDeviceTransfer);
-    }
+    final SignedPreKeyEntity aciSignedPreKey = new SignedPreKeyEntity(Objects.requireNonNull(aciPreKeys.getSignedPreKey()).getId(),
+                                                                      aciPreKeys.getSignedPreKey().getKeyPair().getPublicKey(),
+                                                                      aciPreKeys.getSignedPreKey().getSignature());
+    final SignedPreKeyEntity pniSignedPreKey = new SignedPreKeyEntity(Objects.requireNonNull(pniPreKeys.getSignedPreKey()).getId(),
+                                                                      pniPreKeys.getSignedPreKey().getKeyPair().getPublicKey(),
+                                                                      pniPreKeys.getSignedPreKey().getSignature());
+    final KyberPreKeyEntity aciLastResortKyberPreKey = new KyberPreKeyEntity(Objects.requireNonNull(aciPreKeys.getLastResortKyberPreKey()).getId(),
+                                                                             aciPreKeys.getLastResortKyberPreKey().getKeyPair().getPublicKey(),
+                                                                             aciPreKeys.getLastResortKyberPreKey().getSignature());
+    final KyberPreKeyEntity pniLastResortKyberPreKey = new KyberPreKeyEntity(Objects.requireNonNull(pniPreKeys.getLastResortKyberPreKey()).getId(),
+                                                                             pniPreKeys.getLastResortKyberPreKey().getKeyPair().getPublicKey(),
+                                                                             pniPreKeys.getLastResortKyberPreKey().getSignature());
+    RegistrationSessionRequestBody body = new RegistrationSessionRequestBody(sessionId,
+                                                                             recoveryPassword,
+                                                                             attributes,
+                                                                             Base64.encodeBytesWithoutPadding(aciPreKeys.getIdentityKey().serialize()),
+                                                                             Base64.encodeBytesWithoutPadding(pniPreKeys.getIdentityKey().serialize()),
+                                                                             aciSignedPreKey,
+                                                                             pniSignedPreKey,
+                                                                             aciLastResortKyberPreKey,
+                                                                             pniLastResortKyberPreKey,
+                                                                             gcmRegistrationId,
+                                                                             skipDeviceTransfer);
 
     String response = makeServiceRequest(path, "POST", JsonUtil.toJson(body), NO_HEADERS, new RegistrationSessionResponseHandler(), Optional.empty());
     return JsonUtil.fromJson(response, VerifyAccountResponse.class);