Allow for atomic account creation and activation

integration-tests/src/main/java/org/signal/integration/Operations.java

@@ -22,18 +22,24 @@ import java.util.ArrayList;
 import java.util.Base64;
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.Executors;
 import org.apache.commons.lang3.RandomUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.Validate;
 import org.apache.commons.lang3.tuple.Pair;
 import org.signal.integration.config.Config;
+import org.signal.libsignal.protocol.ecc.Curve;
+import org.signal.libsignal.protocol.ecc.ECKeyPair;
+import org.signal.libsignal.protocol.kem.KEMKeyPair;
+import org.signal.libsignal.protocol.kem.KEMKeyType;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
 import org.whispersystems.textsecuregcm.entities.AccountAttributes;
 import org.whispersystems.textsecuregcm.entities.AccountIdentityResponse;
 import org.whispersystems.textsecuregcm.entities.RegistrationRequest;
+import org.whispersystems.textsecuregcm.entities.SignedPreKey;
 import org.whispersystems.textsecuregcm.http.FaultTolerantHttpClient;
 import org.whispersystems.textsecuregcm.storage.Device;
 import org.whispersystems.textsecuregcm.util.HeaderUtils;
@@ -67,7 +73,8 @@ public final class Operations {
 
     // register account
     final RegistrationRequest registrationRequest = new RegistrationRequest(
-        null, registrationPassword, accountAttributes, true);
+        null, registrationPassword, accountAttributes, true,
+        Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty());
 
     final AccountIdentityResponse registrationResponse = apiPost("/v1/registration", registrationRequest)
         .authorized(number, accountPassword)
@@ -85,6 +92,42 @@ public final class Operations {
     return user;
   }
 
+  public static TestUser newRegisteredUserAtomic(final String number) {
+    final byte[] registrationPassword = RandomUtils.nextBytes(32);
+    final String accountPassword = Base64.getEncoder().encodeToString(RandomUtils.nextBytes(32));
+
+    final TestUser user = TestUser.create(number, accountPassword, registrationPassword);
+    final AccountAttributes accountAttributes = user.accountAttributes();
+
+    INTEGRATION_TOOLS.populateRecoveryPassword(number, registrationPassword).join();
+
+    final ECKeyPair aciIdentityKeyPair = Curve.generateKeyPair();
+    final ECKeyPair pniIdentityKeyPair = Curve.generateKeyPair();
+
+    // register account
+    final RegistrationRequest registrationRequest = new RegistrationRequest(null,
+        registrationPassword,
+        accountAttributes,
+        true,
+        Optional.of(Base64.getEncoder().encodeToString(aciIdentityKeyPair.getPublicKey().serialize())),
+        Optional.of(Base64.getEncoder().encodeToString(pniIdentityKeyPair.getPublicKey().serialize())),
+        Optional.of(generateSignedECPreKey(1, aciIdentityKeyPair)),
+        Optional.of(generateSignedECPreKey(2, pniIdentityKeyPair)),
+        Optional.of(generateSignedKEMPreKey(3, aciIdentityKeyPair)),
+        Optional.of(generateSignedKEMPreKey(4, pniIdentityKeyPair)),
+        Optional.empty(),
+        Optional.empty());
+
+    final AccountIdentityResponse registrationResponse = apiPost("/v1/registration", registrationRequest)
+        .authorized(number, accountPassword)
+        .executeExpectSuccess(AccountIdentityResponse.class);
+
+    user.setAciUuid(registrationResponse.uuid());
+    user.setPniUuid(registrationResponse.pni());
+
+    return user;
+  }
+
   public static void deleteUser(final TestUser user) {
     apiDelete("/v1/accounts/me").authorized(user).executeExpectSuccess();
   }
@@ -271,4 +314,16 @@ public final class Operations {
       throw new RuntimeException(e);
     }
   }
+
+  private static SignedPreKey generateSignedECPreKey(long id, final ECKeyPair identityKeyPair) {
+    final byte[] pubKey = Curve.generateKeyPair().getPublicKey().serialize();
+    final byte[] sig = identityKeyPair.getPrivateKey().calculateSignature(pubKey);
+    return new SignedPreKey(id, Base64.getEncoder().encodeToString(pubKey), Base64.getEncoder().encodeToString(sig));
+  }
+
+  private static SignedPreKey generateSignedKEMPreKey(long id, final ECKeyPair identityKeyPair) {
+    final byte[] pubKey = KEMKeyPair.generate(KEMKeyType.KYBER_1024).getPublicKey().serialize();
+    final byte[] sig = identityKeyPair.getPrivateKey().calculateSignature(pubKey);
+    return new SignedPreKey(id, Base64.getEncoder().encodeToString(pubKey), Base64.getEncoder().encodeToString(sig));
+  }
 }

integration-tests/src/test/java/org/signal/integration/IntegrationTest.java

@@ -14,6 +14,8 @@ import java.util.List;
 import org.apache.commons.lang3.tuple.Pair;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.whispersystems.textsecuregcm.entities.AccountIdentityResponse;
 import org.whispersystems.textsecuregcm.entities.CreateVerificationSessionRequest;
 import org.whispersystems.textsecuregcm.entities.IncomingMessage;
@@ -41,6 +43,19 @@ public class IntegrationTest {
     }
   }
 
+  @Test
+  public void testCreateAccountAtomic() throws Exception {
+    final TestUser user = Operations.newRegisteredUserAtomic("+19995550201");
+    try {
+      final Pair<Integer, AccountIdentityResponse> execute = Operations.apiGet("/v1/accounts/whoami")
+          .authorized(user)
+          .execute(AccountIdentityResponse.class);
+      assertEquals(200, execute.getLeft());
+    } finally {
+      Operations.deleteUser(user);
+    }
+  }
+
   @Test
   public void testRegistration() throws Exception {
     final UpdateVerificationSessionRequest originalRequest = new UpdateVerificationSessionRequest(
@@ -82,10 +97,19 @@ public class IntegrationTest {
     System.out.println("sms code supplied: " + codeVerified);
   }
 
-  @Test
+  @ParameterizedTest
-  public void testSendMessageUnsealed() throws Exception {
+  @ValueSource(booleans = {true, false})
-    final TestUser userA = Operations.newRegisteredUser("+19995550102");
+  public void testSendMessageUnsealed(final boolean atomicAccountCreation) throws Exception {
-    final TestUser userB = Operations.newRegisteredUser("+19995550103");
+    final TestUser userA;
+    final TestUser userB;
+
+    if (atomicAccountCreation) {
+      userA = Operations.newRegisteredUserAtomic("+19995550102");
+      userB = Operations.newRegisteredUserAtomic("+19995550103");
+    } else {
+      userA = Operations.newRegisteredUser("+19995550102");
+      userB = Operations.newRegisteredUser("+19995550103");
+    }
 
     try {
       final byte[] expectedContent = "Hello, World!".getBytes(StandardCharsets.UTF_8);

service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java

@@ -755,7 +755,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
             config.getCdnConfiguration().bucket(), zkProfileOperations, batchIdentityCheckExecutor),
         new ProvisioningController(rateLimiters, provisioningManager),
         new RegistrationController(accountsManager, phoneVerificationTokenManager, registrationLockVerificationManager,
-            rateLimiters),
+            keys, rateLimiters),
         new RemoteConfigController(remoteConfigsManager, adminEventLogger,
             config.getRemoteConfigConfiguration().authorizedTokens().value(),
             config.getRemoteConfigConfiguration().globalConfig()),

service/src/main/java/org/whispersystems/textsecuregcm/controllers/RegistrationController.java

@@ -21,6 +21,8 @@ import io.swagger.v3.oas.annotations.responses.ApiResponse;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Map;
 import java.util.Optional;
 import javax.validation.Valid;
 import javax.validation.constraints.NotNull;
@@ -46,6 +48,8 @@ import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
+import org.whispersystems.textsecuregcm.storage.Device;
+import org.whispersystems.textsecuregcm.storage.Keys;
 import org.whispersystems.textsecuregcm.util.HeaderUtils;
 import org.whispersystems.textsecuregcm.util.Util;
 
@@ -65,18 +69,23 @@ public class RegistrationController {
   private static final String COUNTRY_CODE_TAG_NAME = "countryCode";
   private static final String REGION_CODE_TAG_NAME = "regionCode";
   private static final String VERIFICATION_TYPE_TAG_NAME = "verification";
+  private static final String ACCOUNT_ACTIVATED_TAG_NAME = "accountActivated";
 
   private final AccountsManager accounts;
   private final PhoneVerificationTokenManager phoneVerificationTokenManager;
   private final RegistrationLockVerificationManager registrationLockVerificationManager;
+  private final Keys keys;
   private final RateLimiters rateLimiters;
 
   public RegistrationController(final AccountsManager accounts,
-      final PhoneVerificationTokenManager phoneVerificationTokenManager,
+                                final PhoneVerificationTokenManager phoneVerificationTokenManager,
-      final RegistrationLockVerificationManager registrationLockVerificationManager, final RateLimiters rateLimiters) {
+                                final RegistrationLockVerificationManager registrationLockVerificationManager,
+                                final Keys keys,
+                                final RateLimiters rateLimiters) {
     this.accounts = accounts;
     this.phoneVerificationTokenManager = phoneVerificationTokenManager;
     this.registrationLockVerificationManager = registrationLockVerificationManager;
+    this.keys = keys;
     this.rateLimiters = rateLimiters;
   }
 
@@ -134,13 +143,45 @@ public class RegistrationController {
       throw new WebApplicationException(Response.status(409, "device transfer available").build());
     }
 
-    final Account account = accounts.create(number, password, signalAgent, registrationRequest.accountAttributes(),
+    Account account = accounts.create(number, password, signalAgent, registrationRequest.accountAttributes(),
         existingAccount.map(Account::getBadges).orElseGet(ArrayList::new));
 
+    // If the request includes all the information we need to fully "activate" the account, we should do so
+    if (registrationRequest.supportsAtomicAccountCreation()) {
+      assert registrationRequest.aciIdentityKey().isPresent();
+      assert registrationRequest.pniIdentityKey().isPresent();
+      assert registrationRequest.aciSignedPreKey().isPresent();
+      assert registrationRequest.pniSignedPreKey().isPresent();
+      assert registrationRequest.aciPqLastResortPreKey().isPresent();
+      assert registrationRequest.pniPqLastResortPreKey().isPresent();
+
+      account = accounts.update(account, a -> {
+        a.setIdentityKey(registrationRequest.aciIdentityKey().get());
+        a.setPhoneNumberIdentityKey(registrationRequest.pniIdentityKey().get());
+      });
+
+      account = accounts.updateDevice(account, Device.MASTER_ID, device -> {
+        device.setSignedPreKey(registrationRequest.aciSignedPreKey().get());
+        device.setPhoneNumberIdentitySignedPreKey(registrationRequest.pniSignedPreKey().get());
+
+        registrationRequest.apnToken().ifPresent(apnRegistrationId -> {
+          device.setApnId(apnRegistrationId.apnRegistrationId());
+          device.setVoipApnId(apnRegistrationId.voipRegistrationId());
+        });
+
+        registrationRequest.gcmToken().ifPresent(gcmRegistrationId ->
+            device.setGcmId(gcmRegistrationId.gcmRegistrationId()));
+      });
+
+      keys.storePqLastResort(account.getUuid(), Map.of(Device.MASTER_ID, registrationRequest.aciPqLastResortPreKey().get()));
+      keys.storePqLastResort(account.getPhoneNumberIdentifier(), Map.of(Device.MASTER_ID, registrationRequest.pniPqLastResortPreKey().get()));
+    }
+
     Metrics.counter(ACCOUNT_CREATED_COUNTER_NAME, Tags.of(UserAgentTagUtil.getPlatformTag(userAgent),
             Tag.of(COUNTRY_CODE_TAG_NAME, Util.getCountryCode(number)),
             Tag.of(REGION_CODE_TAG_NAME, Util.getRegion(number)),
-            Tag.of(VERIFICATION_TYPE_TAG_NAME, verificationType.name())))
+            Tag.of(VERIFICATION_TYPE_TAG_NAME, verificationType.name()),
+            Tag.of(ACCOUNT_ACTIVATED_TAG_NAME, String.valueOf(account.isEnabled()))))
         .increment();
 
     return new AccountIdentityResponse(account.getUuid(),

service/src/main/java/org/whispersystems/textsecuregcm/entities/RegistrationRequest.java

@@ -6,14 +6,152 @@
 package org.whispersystems.textsecuregcm.entities;
 
 import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
-import javax.validation.Valid;
+import com.google.common.annotations.VisibleForTesting;
-import javax.validation.constraints.NotNull;
+import io.swagger.v3.oas.annotations.media.Schema;
 import org.whispersystems.textsecuregcm.util.ByteArrayAdapter;
 
-public record RegistrationRequest(String sessionId,
+import javax.validation.Valid;
-                                  @JsonDeserialize(using = ByteArrayAdapter.Deserializing.class) byte[] recoveryPassword,
+import javax.validation.constraints.AssertTrue;
-                                  @NotNull @Valid AccountAttributes accountAttributes,
+import javax.validation.constraints.NotNull;
-                                  boolean skipDeviceTransfer) implements PhoneVerificationRequest {
+import java.util.List;
+import java.util.Optional;
 
+public record RegistrationRequest(@Schema(requiredMode = Schema.RequiredMode.NOT_REQUIRED, description = """
+                                  The ID of an existing verification session as it appears in a verification session
+                                  metadata object. Must be provided if `recoveryPassword` is not provided; must not be
+                                  provided if `recoveryPassword` is provided.
+                                  """)
+                                  String sessionId,
 
+                                  @JsonDeserialize(using = ByteArrayAdapter.Deserializing.class)
+                                  @Schema(requiredMode = Schema.RequiredMode.NOT_REQUIRED, description = """
+                                  A base64-encoded registration recovery password. Must be provided if `sessionId` is
+                                  not provided; must not be provided if `sessionId` is provided
+                                  """)
+                                  byte[] recoveryPassword,
+
+                                  @NotNull
+                                  @Valid
+                                  @Schema(requiredMode = Schema.RequiredMode.REQUIRED)
+                                  AccountAttributes accountAttributes,
+
+                                  @Schema(requiredMode = Schema.RequiredMode.REQUIRED, description = """
+                                  If true, indicates that the end user has elected not to transfer data from another
+                                  device even though a device transfer is technically possible given the capabilities of
+                                  the calling device and the device associated with the existing account (if any). If
+                                  false and if a device transfer is technically possible, the registration request will
+                                  fail with an HTTP/409 response indicating that the client should prompt the user to
+                                  transfer data from an existing device.
+                                  """)
+                                  boolean skipDeviceTransfer,
+
+                                  @Schema(requiredMode = Schema.RequiredMode.NOT_REQUIRED, description = """
+                                  The ACI-associated identity key for the account, encoded as a base64 string. If
+                                  provided, an account will be created "atomically," and all other properties needed for
+                                  atomic account creation must also be present.
+                                  """)
+                                  Optional<String> aciIdentityKey,
+
+                                  @Schema(requiredMode = Schema.RequiredMode.NOT_REQUIRED, description = """
+                                  The PNI-associated identity key for the account, encoded as a base64 string. If
+                                  provided, an account will be created "atomically," and all other properties needed for
+                                  atomic account creation must also be present.
+                                  """)
+                                  Optional<String> pniIdentityKey,
+
+                                  @Schema(requiredMode = Schema.RequiredMode.NOT_REQUIRED, description = """
+                                  A signed EC pre-key to be associated with this account's ACI. If provided, an account
+                                  will be created "atomically," and all other properties needed for atomic account
+                                  creation must also be present.
+                                  """)
+                                  Optional<@Valid SignedPreKey> aciSignedPreKey,
+
+                                  @Schema(requiredMode = Schema.RequiredMode.NOT_REQUIRED, description = """
+                                  A signed EC pre-key to be associated with this account's PNI. If provided, an account
+                                  will be created "atomically," and all other properties needed for atomic account
+                                  creation must also be present.
+                                  """)
+                                  Optional<@Valid SignedPreKey> pniSignedPreKey,
+
+                                  @Schema(requiredMode = Schema.RequiredMode.NOT_REQUIRED, description = """
+                                  A signed Kyber-1024 "last resort" pre-key to be associated with this account's ACI. If
+                                  provided, an account will be created "atomically," and all other properties needed for
+                                  atomic account creation must also be present.
+                                  """)
+                                  Optional<@Valid SignedPreKey> aciPqLastResortPreKey,
+
+                                  @Schema(requiredMode = Schema.RequiredMode.NOT_REQUIRED, description = """
+                                  A signed Kyber-1024 "last resort" pre-key to be associated with this account's PNI. If
+                                  provided, an account will be created "atomically," and all other properties needed for
+                                  atomic account creation must also be present.
+                                  """)
+                                  Optional<@Valid SignedPreKey> pniPqLastResortPreKey,
+
+                                  @Schema(requiredMode = Schema.RequiredMode.NOT_REQUIRED, description = """
+                                  An APNs token set for the account's primary device. If provided, the account's primary
+                                  device will be notified of new messages via push notifications to the given token. If
+                                  creating an account "atomically," callers must provide exactly one of an APNs token
+                                  set, an FCM token, or an `AccountAttributes` entity with `fetchesMessages` set to
+                                  `true`.
+                                  """)
+                                  Optional<@Valid ApnRegistrationId> apnToken,
+
+                                  @Schema(requiredMode = Schema.RequiredMode.NOT_REQUIRED, description = """
+                                  An FCM/GCM token for the account's primary device. If provided, the account's primary
+                                  device will be notified of new messages via push notifications to the given token. If
+                                  creating an account "atomically," callers must provide exactly one of an APNs token
+                                  set, an FCM token, or an `AccountAttributes` entity with `fetchesMessages` set to
+                                  `true`.
+                                  """)
+                                  Optional<@Valid GcmRegistrationId> gcmToken) implements PhoneVerificationRequest {
+
+  @AssertTrue
+  public boolean isEverySignedKeyValid() {
+    return validatePreKeySignature(aciIdentityKey(), aciSignedPreKey())
+        && validatePreKeySignature(pniIdentityKey(), pniSignedPreKey())
+        && validatePreKeySignature(aciIdentityKey(), aciPqLastResortPreKey())
+        && validatePreKeySignature(pniIdentityKey(), pniPqLastResortPreKey());
+  }
+
+  @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
+  private static boolean validatePreKeySignature(final Optional<String> maybeIdentityKey,
+      final Optional<SignedPreKey> maybeSignedPreKey) {
+
+    return maybeSignedPreKey.map(signedPreKey -> maybeIdentityKey
+            .map(identityKey -> PreKeySignatureValidator.validatePreKeySignatures(identityKey, List.of(signedPreKey)))
+            .orElse(false))
+        .orElse(true);
+  }
+
+  @AssertTrue
+  public boolean isCompleteRequest() {
+    final boolean hasNoAtomicAccountCreationParameters =
+        aciIdentityKey().isEmpty()
+            && pniIdentityKey().isEmpty()
+            && aciSignedPreKey().isEmpty()
+            && pniSignedPreKey().isEmpty()
+            && aciPqLastResortPreKey().isEmpty()
+            && pniPqLastResortPreKey().isEmpty();
+
+    return supportsAtomicAccountCreation() || hasNoAtomicAccountCreationParameters;
+  }
+
+  public boolean supportsAtomicAccountCreation() {
+    return hasExactlyOneMessageDeliveryChannel()
+        && aciIdentityKey().isPresent()
+        && pniIdentityKey().isPresent()
+        && aciSignedPreKey().isPresent()
+        && pniSignedPreKey().isPresent()
+        && aciPqLastResortPreKey().isPresent()
+        && pniPqLastResortPreKey().isPresent();
+  }
+
+  @VisibleForTesting
+  boolean hasExactlyOneMessageDeliveryChannel() {
+    if (accountAttributes.getFetchesMessages()) {
+      return apnToken.isEmpty() && gcmToken.isEmpty();
+    } else {
+      return apnToken.isPresent() ^ gcmToken.isPresent();
+    }
+  }
 }

service/src/test/java/org/whispersystems/textsecuregcm/controllers/AccountControllerTest.java

@@ -1667,7 +1667,7 @@ class AccountControllerTest {
     var deviceMessages = List.of(
         new IncomingMessage(1, 2, 2, "content2"),
         new IncomingMessage(1, 3, 3, "content3"));
-    var deviceKeys = List.of(1L, 2L, 3L).stream().collect(Collectors.toMap(Function.identity(), n -> KeysHelper.signedPreKey(n + 100, pniIdentityKeyPair)));
+    var deviceKeys = List.of(1L, 2L, 3L).stream().collect(Collectors.toMap(Function.identity(), n -> KeysHelper.signedECPreKey(n + 100, pniIdentityKeyPair)));
 
     final Map<Long, Integer> registrationIds = Map.of(1L, 17, 2L, 47, 3L, 89);
 
@@ -1721,7 +1721,7 @@ class AccountControllerTest {
     var deviceMessages = List.of(
         new IncomingMessage(1, 2, 2, "content2"),
         new IncomingMessage(1, 3, 3, "content3"));
-    var deviceKeys = List.of(1L, 2L, 3L).stream().collect(Collectors.toMap(Function.identity(), n -> KeysHelper.signedPreKey(n + 100, pniIdentityKeyPair)));
+    var deviceKeys = List.of(1L, 2L, 3L).stream().collect(Collectors.toMap(Function.identity(), n -> KeysHelper.signedECPreKey(n + 100, pniIdentityKeyPair)));
 
     final Map<Long, Integer> registrationIds = Map.of(1L, 17, 2L, 47, 3L, 89);
 

service/src/test/java/org/whispersystems/textsecuregcm/controllers/RegistrationControllerTest.java

@@ -9,8 +9,11 @@ import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import com.google.i18n.phonenumbers.PhoneNumberUtil;
@@ -19,8 +22,11 @@ import io.dropwizard.testing.junit5.ResourceExtension;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.Base64;
+import java.util.Map;
 import java.util.Optional;
+import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
+import java.util.function.Consumer;
 import java.util.stream.Stream;
 import javax.annotation.Nullable;
 import javax.ws.rs.WebApplicationException;
@@ -39,12 +45,17 @@ import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.CsvSource;
 import org.junit.jupiter.params.provider.EnumSource;
 import org.junit.jupiter.params.provider.MethodSource;
+import org.signal.libsignal.protocol.ecc.Curve;
+import org.signal.libsignal.protocol.ecc.ECKeyPair;
 import org.whispersystems.textsecuregcm.auth.PhoneVerificationTokenManager;
 import org.whispersystems.textsecuregcm.auth.RegistrationLockError;
 import org.whispersystems.textsecuregcm.auth.RegistrationLockVerificationManager;
 import org.whispersystems.textsecuregcm.entities.AccountAttributes;
+import org.whispersystems.textsecuregcm.entities.ApnRegistrationId;
+import org.whispersystems.textsecuregcm.entities.GcmRegistrationId;
 import org.whispersystems.textsecuregcm.entities.RegistrationRequest;
 import org.whispersystems.textsecuregcm.entities.RegistrationServiceSession;
+import org.whispersystems.textsecuregcm.entities.SignedPreKey;
 import org.whispersystems.textsecuregcm.limits.RateLimiter;
 import org.whispersystems.textsecuregcm.limits.RateLimiters;
 import org.whispersystems.textsecuregcm.mappers.ImpossiblePhoneNumberExceptionMapper;
@@ -53,8 +64,12 @@ import org.whispersystems.textsecuregcm.mappers.RateLimitExceededExceptionMapper
 import org.whispersystems.textsecuregcm.registration.RegistrationServiceClient;
 import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
+import org.whispersystems.textsecuregcm.storage.Device;
+import org.whispersystems.textsecuregcm.storage.Keys;
 import org.whispersystems.textsecuregcm.storage.RegistrationRecoveryPasswordsManager;
 import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
+import org.whispersystems.textsecuregcm.tests.util.KeysHelper;
+import org.whispersystems.textsecuregcm.util.MockUtils;
 import org.whispersystems.textsecuregcm.util.SystemMapper;
 
 @ExtendWith(DropwizardExtensionsSupport.class)
@@ -73,6 +88,7 @@ class RegistrationControllerTest {
       RegistrationLockVerificationManager.class);
   private final RegistrationRecoveryPasswordsManager registrationRecoveryPasswordsManager = mock(
       RegistrationRecoveryPasswordsManager.class);
+  private final Keys keys = mock(Keys.class);
   private final RateLimiters rateLimiters = mock(RateLimiters.class);
 
   private final RateLimiter registrationLimiter = mock(RateLimiter.class);
@@ -87,7 +103,7 @@ class RegistrationControllerTest {
       .addResource(
           new RegistrationController(accountsManager,
               new PhoneVerificationTokenManager(registrationServiceClient, registrationRecoveryPasswordsManager),
-              registrationLockVerificationManager, rateLimiters))
+              registrationLockVerificationManager, keys, rateLimiters))
       .build();
 
   @BeforeEach
@@ -97,10 +113,10 @@ class RegistrationControllerTest {
 
   @Test
   public void testRegistrationRequest() throws Exception {
-    assertFalse(new RegistrationRequest("", new byte[0], new AccountAttributes(), true).isValid());
+    assertFalse(new RegistrationRequest("", new byte[0], new AccountAttributes(), true, Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty()).isValid());
-    assertFalse(new RegistrationRequest("some", new byte[32], new AccountAttributes(), true).isValid());
+    assertFalse(new RegistrationRequest("some", new byte[32], new AccountAttributes(), true, Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty()).isValid());
-    assertTrue(new RegistrationRequest("", new byte[32], new AccountAttributes(), true).isValid());
+    assertTrue(new RegistrationRequest("", new byte[32], new AccountAttributes(), true, Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty()).isValid());
-    assertTrue(new RegistrationRequest("some", new byte[0], new AccountAttributes(), true).isValid());
+    assertTrue(new RegistrationRequest("some", new byte[0], new AccountAttributes(), true, Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty()).isValid());
   }
 
   @Test
@@ -332,6 +348,402 @@ class RegistrationControllerTest {
     }
   }
 
+  @ParameterizedTest
+  @MethodSource
+  void atomicAccountCreationConflictingChannel(final RegistrationRequest conflictingChannelRequest) {
+    when(registrationServiceClient.getSession(any(), any()))
+        .thenReturn(
+            CompletableFuture.completedFuture(
+                Optional.of(new RegistrationServiceSession(new byte[16], NUMBER, true, null, null, null,
+                    SESSION_EXPIRATION_SECONDS))));
+
+    try (final Response response = resources.getJerseyTest()
+        .target("/v1/registration")
+        .request()
+        .header(HttpHeaders.AUTHORIZATION, AuthHelper.getProvisioningAuthHeader(NUMBER, PASSWORD))
+        .post(Entity.json(conflictingChannelRequest))) {
+
+      assertEquals(422, response.getStatus());
+    }
+  }
+
+  static Stream<Arguments> atomicAccountCreationConflictingChannel() {
+    final Optional<String> aciIdentityKey;
+    final Optional<String> pniIdentityKey;
+    final Optional<SignedPreKey> aciSignedPreKey;
+    final Optional<SignedPreKey> pniSignedPreKey;
+    final Optional<SignedPreKey> aciPqLastResortPreKey;
+    final Optional<SignedPreKey> pniPqLastResortPreKey;
+    {
+      final ECKeyPair aciIdentityKeyPair = Curve.generateKeyPair();
+      final ECKeyPair pniIdentityKeyPair = Curve.generateKeyPair();
+
+      aciIdentityKey = Optional.of(KeysHelper.serializeIdentityKey(aciIdentityKeyPair));
+      pniIdentityKey = Optional.of(KeysHelper.serializeIdentityKey(pniIdentityKeyPair));
+      aciSignedPreKey = Optional.of(KeysHelper.signedECPreKey(1, aciIdentityKeyPair));
+      pniSignedPreKey = Optional.of(KeysHelper.signedECPreKey(2, pniIdentityKeyPair));
+      aciPqLastResortPreKey = Optional.of(KeysHelper.signedKEMPreKey(3, aciIdentityKeyPair));
+      pniPqLastResortPreKey = Optional.of(KeysHelper.signedKEMPreKey(4, pniIdentityKeyPair));
+    }
+
+    final AccountAttributes fetchesMessagesAccountAttributes =
+        new AccountAttributes(true, 1, "test", null, true, new Device.DeviceCapabilities());
+
+    final AccountAttributes pushAccountAttributes =
+        new AccountAttributes(false, 1, "test", null, true, new Device.DeviceCapabilities());
+
+    return Stream.of(
+        // "Fetches messages" is true, but an APNs token is provided
+        Arguments.of(new RegistrationRequest("session-id",
+            new byte[0],
+            fetchesMessagesAccountAttributes,
+            true,
+            aciIdentityKey,
+            pniIdentityKey,
+            aciSignedPreKey,
+            pniSignedPreKey,
+            aciPqLastResortPreKey,
+            pniPqLastResortPreKey,
+            Optional.of(new ApnRegistrationId("apns-token", null)),
+            Optional.empty())),
+
+        // "Fetches messages" is true, but an FCM (GCM) token is provided
+        Arguments.of(new RegistrationRequest("session-id",
+            new byte[0],
+            fetchesMessagesAccountAttributes,
+            true,
+            aciIdentityKey,
+            pniIdentityKey,
+            aciSignedPreKey,
+            pniSignedPreKey,
+            aciPqLastResortPreKey,
+            pniPqLastResortPreKey,
+            Optional.empty(),
+            Optional.of(new GcmRegistrationId("gcm-token")))),
+
+        // "Fetches messages" is false, but multiple types of push tokens are provided
+        Arguments.of(new RegistrationRequest("session-id",
+            new byte[0],
+            pushAccountAttributes,
+            true,
+            aciIdentityKey,
+            pniIdentityKey,
+            aciSignedPreKey,
+            pniSignedPreKey,
+            aciPqLastResortPreKey,
+            pniPqLastResortPreKey,
+            Optional.of(new ApnRegistrationId("apns-token", null)),
+            Optional.of(new GcmRegistrationId("gcm-token"))))
+    );
+  }
+
+  @ParameterizedTest
+  @MethodSource
+  void atomicAccountCreationPartialSignedPreKeys(final RegistrationRequest partialSignedPreKeyRequest) {
+    when(registrationServiceClient.getSession(any(), any()))
+        .thenReturn(
+            CompletableFuture.completedFuture(
+                Optional.of(new RegistrationServiceSession(new byte[16], NUMBER, true, null, null, null,
+                    SESSION_EXPIRATION_SECONDS))));
+
+    final Invocation.Builder request = resources.getJerseyTest()
+        .target("/v1/registration")
+        .request()
+        .header(HttpHeaders.AUTHORIZATION, AuthHelper.getProvisioningAuthHeader(NUMBER, PASSWORD));
+
+    try (final Response response = request.post(Entity.json(partialSignedPreKeyRequest))) {
+      assertEquals(422, response.getStatus());
+    }
+  }
+
+  static Stream<Arguments> atomicAccountCreationPartialSignedPreKeys() {
+    final Optional<String> aciIdentityKey;
+    final Optional<String> pniIdentityKey;
+    final Optional<SignedPreKey> aciSignedPreKey;
+    final Optional<SignedPreKey> pniSignedPreKey;
+    final Optional<SignedPreKey> aciPqLastResortPreKey;
+    final Optional<SignedPreKey> pniPqLastResortPreKey;
+    {
+      final ECKeyPair aciIdentityKeyPair = Curve.generateKeyPair();
+      final ECKeyPair pniIdentityKeyPair = Curve.generateKeyPair();
+
+      aciIdentityKey = Optional.of(KeysHelper.serializeIdentityKey(aciIdentityKeyPair));
+      pniIdentityKey = Optional.of(KeysHelper.serializeIdentityKey(pniIdentityKeyPair));
+      aciSignedPreKey = Optional.of(KeysHelper.signedECPreKey(1, aciIdentityKeyPair));
+      pniSignedPreKey = Optional.of(KeysHelper.signedECPreKey(2, pniIdentityKeyPair));
+      aciPqLastResortPreKey = Optional.of(KeysHelper.signedKEMPreKey(3, aciIdentityKeyPair));
+      pniPqLastResortPreKey = Optional.of(KeysHelper.signedKEMPreKey(4, pniIdentityKeyPair));
+    }
+
+    final AccountAttributes accountAttributes =
+        new AccountAttributes(true, 1, "test", null, true, new Device.DeviceCapabilities());
+
+    return Stream.of(
+        // Signed PNI EC pre-key is missing
+        Arguments.of(new RegistrationRequest("session-id",
+            new byte[0],
+            accountAttributes,
+            true,
+            aciIdentityKey,
+            pniIdentityKey,
+            aciSignedPreKey,
+            Optional.empty(),
+            aciPqLastResortPreKey,
+            pniPqLastResortPreKey,
+            Optional.empty(),
+            Optional.empty())),
+
+        // Signed ACI EC pre-key is missing
+        Arguments.of(new RegistrationRequest("session-id",
+            new byte[0],
+            accountAttributes,
+            true,
+            aciIdentityKey,
+            pniIdentityKey,
+            Optional.empty(),
+            pniSignedPreKey,
+            aciPqLastResortPreKey,
+            pniPqLastResortPreKey,
+            Optional.empty(),
+            Optional.empty())),
+
+        // Signed PNI KEM pre-key is missing
+        Arguments.of(new RegistrationRequest("session-id",
+            new byte[0],
+            accountAttributes,
+            true,
+            aciIdentityKey,
+            pniIdentityKey,
+            aciSignedPreKey,
+            pniSignedPreKey,
+            aciPqLastResortPreKey,
+            Optional.empty(),
+            Optional.empty(),
+            Optional.empty())),
+
+        // Signed ACI KEM pre-key is missing
+        Arguments.of(new RegistrationRequest("session-id",
+            new byte[0],
+            accountAttributes,
+            true,
+            aciIdentityKey,
+            pniIdentityKey,
+            aciSignedPreKey,
+            pniSignedPreKey,
+            Optional.empty(),
+            pniPqLastResortPreKey,
+            Optional.empty(),
+            Optional.empty())),
+
+        // All signed pre-keys are present, but ACI identity key is missing
+        Arguments.of(new RegistrationRequest("session-id",
+            new byte[0],
+            accountAttributes,
+            true,
+            Optional.empty(),
+            pniIdentityKey,
+            aciSignedPreKey,
+            pniSignedPreKey,
+            aciPqLastResortPreKey,
+            pniPqLastResortPreKey,
+            Optional.empty(),
+            Optional.empty())),
+
+        // All signed pre-keys are present, but PNI identity key is missing
+        Arguments.of(new RegistrationRequest("session-id",
+            new byte[0],
+            accountAttributes,
+            true,
+            aciIdentityKey,
+            Optional.empty(),
+            aciSignedPreKey,
+            pniSignedPreKey,
+            aciPqLastResortPreKey,
+            pniPqLastResortPreKey,
+            Optional.empty(),
+            Optional.empty()))
+    );
+  }
+
+
+  @ParameterizedTest
+  @MethodSource
+  @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
+  void atomicAccountCreationSuccess(final RegistrationRequest registrationRequest,
+      final String expectedAciIdentityKey,
+      final String expectedPniIdentityKey,
+      final SignedPreKey expectedAciSignedPreKey,
+      final SignedPreKey expectedPniSignedPreKey,
+      final SignedPreKey expectedAciPqLastResortPreKey,
+      final SignedPreKey expectedPniPqLastResortPreKey,
+      final Optional<String> expectedApnsToken,
+      final Optional<String> expectedApnsVoipToken,
+      final Optional<String> expectedGcmToken) throws InterruptedException {
+
+    when(registrationServiceClient.getSession(any(), any()))
+        .thenReturn(
+            CompletableFuture.completedFuture(
+                Optional.of(new RegistrationServiceSession(new byte[16], NUMBER, true, null, null, null,
+                    SESSION_EXPIRATION_SECONDS))));
+
+    final UUID accountIdentifier = UUID.randomUUID();
+    final UUID phoneNumberIdentifier = UUID.randomUUID();
+
+    final Account account = MockUtils.buildMock(Account.class, a -> {
+      when(a.getUuid()).thenReturn(accountIdentifier);
+      when(a.getPhoneNumberIdentifier()).thenReturn(phoneNumberIdentifier);
+    });
+
+    final Device device = mock(Device.class);
+
+    when(accountsManager.create(any(), any(), any(), any(), any()))
+        .thenReturn(account);
+
+    when(accountsManager.update(eq(account), any())).thenAnswer(invocation -> {
+      final Consumer<Account> accountUpdater = invocation.getArgument(1);
+      accountUpdater.accept(account);
+
+      return invocation.getArgument(0);
+    });
+
+    when(accountsManager.updateDevice(eq(account), eq(Device.MASTER_ID), any())).thenAnswer(invocation -> {
+      final Consumer<Device> deviceUpdater = invocation.getArgument(2);
+      deviceUpdater.accept(device);
+
+      return invocation.getArgument(0);
+    });
+
+    final Invocation.Builder request = resources.getJerseyTest()
+        .target("/v1/registration")
+        .request()
+        .header(HttpHeaders.AUTHORIZATION, AuthHelper.getProvisioningAuthHeader(NUMBER, PASSWORD));
+
+    try (Response response = request.post(Entity.json(registrationRequest))) {
+      assertEquals(200, response.getStatus());
+    }
+
+    verify(accountsManager).create(any(), any(), any(), any(), any());
+    verify(accountsManager).updateDevice(eq(account), eq(Device.MASTER_ID), any());
+
+    verify(account).setIdentityKey(expectedAciIdentityKey);
+    verify(account).setPhoneNumberIdentityKey(expectedPniIdentityKey);
+
+    verify(device).setSignedPreKey(expectedAciSignedPreKey);
+    verify(device).setPhoneNumberIdentitySignedPreKey(expectedPniSignedPreKey);
+
+    verify(keys).storePqLastResort(accountIdentifier, Map.of(Device.MASTER_ID, expectedAciPqLastResortPreKey));
+    verify(keys).storePqLastResort(phoneNumberIdentifier, Map.of(Device.MASTER_ID, expectedPniPqLastResortPreKey));
+
+    expectedApnsToken.ifPresentOrElse(expectedToken -> verify(device).setApnId(expectedToken),
+        () -> verify(device, never()).setApnId(any()));
+
+    expectedApnsVoipToken.ifPresentOrElse(expectedToken -> verify(device).setVoipApnId(expectedToken),
+        () -> verify(device, never()).setVoipApnId(any()));
+
+    expectedGcmToken.ifPresentOrElse(expectedToken -> verify(device).setGcmId(expectedToken),
+        () -> verify(device, never()).setGcmId(any()));
+  }
+
+  private static Stream<Arguments> atomicAccountCreationSuccess() {
+    final Optional<String> aciIdentityKey;
+    final Optional<String> pniIdentityKey;
+    final Optional<SignedPreKey> aciSignedPreKey;
+    final Optional<SignedPreKey> pniSignedPreKey;
+    final Optional<SignedPreKey> aciPqLastResortPreKey;
+    final Optional<SignedPreKey> pniPqLastResortPreKey;
+    {
+      final ECKeyPair aciIdentityKeyPair = Curve.generateKeyPair();
+      final ECKeyPair pniIdentityKeyPair = Curve.generateKeyPair();
+
+      aciIdentityKey = Optional.of(KeysHelper.serializeIdentityKey(aciIdentityKeyPair));
+      pniIdentityKey = Optional.of(KeysHelper.serializeIdentityKey(pniIdentityKeyPair));
+      aciSignedPreKey = Optional.of(KeysHelper.signedECPreKey(1, aciIdentityKeyPair));
+      pniSignedPreKey = Optional.of(KeysHelper.signedECPreKey(2, pniIdentityKeyPair));
+      aciPqLastResortPreKey = Optional.of(KeysHelper.signedKEMPreKey(3, aciIdentityKeyPair));
+      pniPqLastResortPreKey = Optional.of(KeysHelper.signedKEMPreKey(4, pniIdentityKeyPair));
+    }
+
+    final AccountAttributes fetchesMessagesAccountAttributes =
+        new AccountAttributes(true, 1, "test", null, true, new Device.DeviceCapabilities());
+
+    final AccountAttributes pushAccountAttributes =
+        new AccountAttributes(false, 1, "test", null, true, new Device.DeviceCapabilities());
+
+    final String apnsToken = "apns-token";
+    final String apnsVoipToken = "apns-voip-token";
+    final String gcmToken = "gcm-token";
+
+    return Stream.of(
+        // Fetches messages; no push tokens
+        Arguments.of(new RegistrationRequest("session-id",
+                new byte[0],
+                fetchesMessagesAccountAttributes,
+                true,
+                aciIdentityKey,
+                pniIdentityKey,
+                aciSignedPreKey,
+                pniSignedPreKey,
+                aciPqLastResortPreKey,
+                pniPqLastResortPreKey,
+                Optional.empty(),
+                Optional.empty()),
+            aciIdentityKey.get(),
+            pniIdentityKey.get(),
+            aciSignedPreKey.get(),
+            pniSignedPreKey.get(),
+            aciPqLastResortPreKey.get(),
+            pniPqLastResortPreKey.get(),
+            Optional.empty(),
+            Optional.empty(),
+            Optional.empty()),
+
+        // Has APNs tokens
+        Arguments.of(new RegistrationRequest("session-id",
+                new byte[0],
+                pushAccountAttributes,
+                true,
+                aciIdentityKey,
+                pniIdentityKey,
+                aciSignedPreKey,
+                pniSignedPreKey,
+                aciPqLastResortPreKey,
+                pniPqLastResortPreKey,
+                Optional.of(new ApnRegistrationId(apnsToken, apnsVoipToken)),
+                Optional.empty()),
+            aciIdentityKey.get(),
+            pniIdentityKey.get(),
+            aciSignedPreKey.get(),
+            pniSignedPreKey.get(),
+            aciPqLastResortPreKey.get(),
+            pniPqLastResortPreKey.get(),
+            Optional.of(apnsToken),
+            Optional.of(apnsVoipToken),
+            Optional.empty()),
+
+        // Fetches messages; no push tokens
+        Arguments.of(new RegistrationRequest("session-id",
+                new byte[0],
+                pushAccountAttributes,
+                true,
+                aciIdentityKey,
+                pniIdentityKey,
+                aciSignedPreKey,
+                pniSignedPreKey,
+                aciPqLastResortPreKey,
+                pniPqLastResortPreKey,
+                Optional.empty(),
+                Optional.of(new GcmRegistrationId(gcmToken))),
+            aciIdentityKey.get(),
+            pniIdentityKey.get(),
+            aciSignedPreKey.get(),
+            pniSignedPreKey.get(),
+            aciPqLastResortPreKey.get(),
+            pniPqLastResortPreKey.get(),
+            Optional.empty(),
+            Optional.empty(),
+            Optional.of(gcmToken)));
+  }
+
   /**
    * Valid request JSON with the give session ID and skipDeviceTransfer
    */

service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerChangeNumberIntegrationTest.java

@@ -148,7 +148,7 @@ class AccountsManagerChangeNumberIntegrationTest {
     final String secondNumber = "+18005552222";
     final int rotatedPniRegistrationId = 17;
     final ECKeyPair pniIdentityKeyPair = Curve.generateKeyPair();
-    final SignedPreKey rotatedSignedPreKey = KeysHelper.signedPreKey(1L, pniIdentityKeyPair);
+    final SignedPreKey rotatedSignedPreKey = KeysHelper.signedECPreKey(1L, pniIdentityKeyPair);
 
     final AccountAttributes accountAttributes = new AccountAttributes(true, rotatedPniRegistrationId + 1, "test", null, true, new Device.DeviceCapabilities());
     final Account account = accountsManager.create(originalNumber, "password", null, accountAttributes, new ArrayList<>());

service/src/test/java/org/whispersystems/textsecuregcm/tests/controllers/KeysControllerTest.java

@@ -25,7 +25,6 @@ import io.dropwizard.auth.PolymorphicAuthValueFactoryProvider;
 import io.dropwizard.testing.junit5.DropwizardExtensionsSupport;
 import io.dropwizard.testing.junit5.ResourceExtension;
 import java.time.Duration;
-import java.util.Base64;
 import java.util.Collections;
 import java.util.LinkedList;
 import java.util.List;
@@ -100,14 +99,14 @@ class KeysControllerTest {
 
   private final SignedPreKey SAMPLE_PQ_KEY_PNI = new SignedPreKey(8888, "test7", "sig");
 
-  private final SignedPreKey SAMPLE_SIGNED_KEY = KeysHelper.signedPreKey(1111, IDENTITY_KEY_PAIR);
+  private final SignedPreKey SAMPLE_SIGNED_KEY = KeysHelper.signedECPreKey(1111, IDENTITY_KEY_PAIR);
-  private final SignedPreKey SAMPLE_SIGNED_KEY2 = KeysHelper.signedPreKey(2222, IDENTITY_KEY_PAIR);
+  private final SignedPreKey SAMPLE_SIGNED_KEY2 = KeysHelper.signedECPreKey(2222, IDENTITY_KEY_PAIR);
-  private final SignedPreKey SAMPLE_SIGNED_KEY3 = KeysHelper.signedPreKey(3333, IDENTITY_KEY_PAIR);
+  private final SignedPreKey SAMPLE_SIGNED_KEY3 = KeysHelper.signedECPreKey(3333, IDENTITY_KEY_PAIR);
-  private final SignedPreKey SAMPLE_SIGNED_PNI_KEY = KeysHelper.signedPreKey(4444, PNI_IDENTITY_KEY_PAIR);
+  private final SignedPreKey SAMPLE_SIGNED_PNI_KEY = KeysHelper.signedECPreKey(4444, PNI_IDENTITY_KEY_PAIR);
-  private final SignedPreKey SAMPLE_SIGNED_PNI_KEY2 = KeysHelper.signedPreKey(5555, PNI_IDENTITY_KEY_PAIR);
+  private final SignedPreKey SAMPLE_SIGNED_PNI_KEY2 = KeysHelper.signedECPreKey(5555, PNI_IDENTITY_KEY_PAIR);
-  private final SignedPreKey SAMPLE_SIGNED_PNI_KEY3 = KeysHelper.signedPreKey(6666, PNI_IDENTITY_KEY_PAIR);
+  private final SignedPreKey SAMPLE_SIGNED_PNI_KEY3 = KeysHelper.signedECPreKey(6666, PNI_IDENTITY_KEY_PAIR);
-  private final SignedPreKey VALID_DEVICE_SIGNED_KEY = KeysHelper.signedPreKey(89898, IDENTITY_KEY_PAIR);
+  private final SignedPreKey VALID_DEVICE_SIGNED_KEY = KeysHelper.signedECPreKey(89898, IDENTITY_KEY_PAIR);
-  private final SignedPreKey VALID_DEVICE_PNI_SIGNED_KEY = KeysHelper.signedPreKey(7777, PNI_IDENTITY_KEY_PAIR);
+  private final SignedPreKey VALID_DEVICE_PNI_SIGNED_KEY = KeysHelper.signedECPreKey(7777, PNI_IDENTITY_KEY_PAIR);
 
   private final static Keys KEYS = mock(Keys.class               );
   private final static AccountsManager             accounts                    = mock(AccountsManager.class            );
@@ -252,7 +251,7 @@ class KeysControllerTest {
 
   @Test
   void putSignedPreKeyV2() {
-    SignedPreKey   test     = KeysHelper.signedPreKey(9998, IDENTITY_KEY_PAIR);
+    SignedPreKey   test     = KeysHelper.signedECPreKey(9998, IDENTITY_KEY_PAIR);
     Response response = resources.getJerseyTest()
                                  .target("/v2/keys/signed")
                                  .request()
@@ -268,7 +267,7 @@ class KeysControllerTest {
 
   @Test
   void putPhoneNumberIdentitySignedPreKeyV2() {
-    final SignedPreKey replacementKey = KeysHelper.signedPreKey(9998, PNI_IDENTITY_KEY_PAIR);
+    final SignedPreKey replacementKey = KeysHelper.signedECPreKey(9998, PNI_IDENTITY_KEY_PAIR);
 
     Response response = resources.getJerseyTest()
         .target("/v2/keys/signed")
@@ -286,7 +285,7 @@ class KeysControllerTest {
 
   @Test
   void disabledPutSignedPreKeyV2() {
-    SignedPreKey   test     = KeysHelper.signedPreKey(9999, IDENTITY_KEY_PAIR);
+    SignedPreKey   test     = KeysHelper.signedECPreKey(9999, IDENTITY_KEY_PAIR);
     Response response = resources.getJerseyTest()
                                  .target("/v2/keys/signed")
                                  .request()
@@ -659,7 +658,7 @@ class KeysControllerTest {
   void putKeysTestV2() {
     final PreKey preKey = new PreKey(31337, "foobar");
     final ECKeyPair identityKeyPair = Curve.generateKeyPair();
-    final SignedPreKey signedPreKey = KeysHelper.signedPreKey(31338, identityKeyPair);
+    final SignedPreKey signedPreKey = KeysHelper.signedECPreKey(31338, identityKeyPair);
     final String identityKey = KeysHelper.serializeIdentityKey(identityKeyPair);
 
     PreKeyState preKeyState = new PreKeyState(identityKey, signedPreKey, List.of(preKey));
@@ -687,9 +686,9 @@ class KeysControllerTest {
   void putKeysPqTestV2() {
     final PreKey preKey = new PreKey(31337, "foobar");
     final ECKeyPair identityKeyPair = Curve.generateKeyPair();
-    final SignedPreKey signedPreKey = KeysHelper.signedPreKey(31338, identityKeyPair);
+    final SignedPreKey signedPreKey = KeysHelper.signedECPreKey(31338, identityKeyPair);
-    final SignedPreKey pqPreKey = KeysHelper.signedPreKey(31339, identityKeyPair);
+    final SignedPreKey pqPreKey = KeysHelper.signedECPreKey(31339, identityKeyPair);
-    final SignedPreKey pqLastResortPreKey = KeysHelper.signedPreKey(31340, identityKeyPair);
+    final SignedPreKey pqLastResortPreKey = KeysHelper.signedECPreKey(31340, identityKeyPair);
     final String identityKey = KeysHelper.serializeIdentityKey(identityKeyPair);
 
     PreKeyState preKeyState = new PreKeyState(identityKey, signedPreKey, List.of(preKey), List.of(pqPreKey), pqLastResortPreKey);
@@ -719,7 +718,7 @@ class KeysControllerTest {
   void putKeysByPhoneNumberIdentifierTestV2() {
     final PreKey preKey = new PreKey(31337, "foobar");
     final ECKeyPair identityKeyPair = Curve.generateKeyPair();
-    final SignedPreKey signedPreKey = KeysHelper.signedPreKey(31338, identityKeyPair);
+    final SignedPreKey signedPreKey = KeysHelper.signedECPreKey(31338, identityKeyPair);
     final String identityKey = KeysHelper.serializeIdentityKey(identityKeyPair);
 
     PreKeyState preKeyState = new PreKeyState(identityKey, signedPreKey, List.of(preKey));
@@ -748,9 +747,9 @@ class KeysControllerTest {
   void putKeysByPhoneNumberIdentifierPqTestV2() {
     final PreKey preKey = new PreKey(31337, "foobar");
     final ECKeyPair identityKeyPair = Curve.generateKeyPair();
-    final SignedPreKey signedPreKey = KeysHelper.signedPreKey(31338, identityKeyPair);
+    final SignedPreKey signedPreKey = KeysHelper.signedECPreKey(31338, identityKeyPair);
-    final SignedPreKey pqPreKey = KeysHelper.signedPreKey(31339, identityKeyPair);
+    final SignedPreKey pqPreKey = KeysHelper.signedECPreKey(31339, identityKeyPair);
-    final SignedPreKey pqLastResortPreKey = KeysHelper.signedPreKey(31340, identityKeyPair);
+    final SignedPreKey pqLastResortPreKey = KeysHelper.signedECPreKey(31340, identityKeyPair);
     final String identityKey = KeysHelper.serializeIdentityKey(identityKeyPair);
 
     PreKeyState preKeyState = new PreKeyState(identityKey, signedPreKey, List.of(preKey), List.of(pqPreKey), pqLastResortPreKey);
@@ -796,7 +795,7 @@ class KeysControllerTest {
   void disabledPutKeysTestV2() {
     final PreKey       preKey       = new PreKey(31337, "foobar");
     final ECKeyPair identityKeyPair = Curve.generateKeyPair();
-    final SignedPreKey signedPreKey = KeysHelper.signedPreKey(31338, identityKeyPair);
+    final SignedPreKey signedPreKey = KeysHelper.signedECPreKey(31338, identityKeyPair);
     final String       identityKey  = KeysHelper.serializeIdentityKey(identityKeyPair);
 
     List<PreKey> preKeys = new LinkedList<PreKey>() {{
@@ -830,7 +829,7 @@ class KeysControllerTest {
   @Test
   void putIdentityKeyNonPrimary() {
     final PreKey       preKey       = new PreKey(31337, "foobar");
-    final SignedPreKey signedPreKey = KeysHelper.signedPreKey(31338, IDENTITY_KEY_PAIR);
+    final SignedPreKey signedPreKey = KeysHelper.signedECPreKey(31338, IDENTITY_KEY_PAIR);
 
     List<PreKey> preKeys = List.of(preKey);
 

service/src/test/java/org/whispersystems/textsecuregcm/tests/util/KeysHelper.java

@@ -8,6 +8,8 @@ package org.whispersystems.textsecuregcm.tests.util;
 import java.util.Base64;
 import org.signal.libsignal.protocol.ecc.Curve;
 import org.signal.libsignal.protocol.ecc.ECKeyPair;
+import org.signal.libsignal.protocol.kem.KEMKeyPair;
+import org.signal.libsignal.protocol.kem.KEMKeyType;
 import org.whispersystems.textsecuregcm.entities.SignedPreKey;
 
 public final class KeysHelper {
@@ -15,9 +17,15 @@ public final class KeysHelper {
     return Base64.getEncoder().encodeToString(keyPair.getPublicKey().serialize());
   }
   
-  public static SignedPreKey signedPreKey(long id, final ECKeyPair signingKey) {
+  public static SignedPreKey signedECPreKey(long id, final ECKeyPair identityKeyPair) {
     final byte[] pubKey = Curve.generateKeyPair().getPublicKey().serialize();
-    final byte[] sig = signingKey.getPrivateKey().calculateSignature(pubKey);
+    final byte[] sig = identityKeyPair.getPrivateKey().calculateSignature(pubKey);
+    return new SignedPreKey(id, Base64.getEncoder().encodeToString(pubKey), Base64.getEncoder().encodeToString(sig));
+  }
+
+  public static SignedPreKey signedKEMPreKey(long id, final ECKeyPair identityKeyPair) {
+    final byte[] pubKey = KEMKeyPair.generate(KEMKeyType.KYBER_1024).getPublicKey().serialize();
+    final byte[] sig = identityKeyPair.getPrivateKey().calculateSignature(pubKey);
     return new SignedPreKey(id, Base64.getEncoder().encodeToString(pubKey), Base64.getEncoder().encodeToString(sig));
   }
 }