Explicitly create registration sessions

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

@@ -29,6 +29,7 @@ import java.util.ArrayList;
 import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
+import java.util.concurrent.CompletionException;
 import javax.annotation.Nullable;
 import javax.servlet.http.HttpServletRequest;
 import javax.validation.Valid;
@@ -232,7 +233,7 @@ public class AccountController {
                              @PathParam("token") String pushToken,
                              @PathParam("number") String number,
                              @QueryParam("voip") @DefaultValue("true") boolean useVoip)
-      throws ImpossiblePhoneNumberException, NonNormalizedPhoneNumberException {
+      throws ImpossiblePhoneNumberException, NonNormalizedPhoneNumberException, RateLimitExceededException {
 
     final PushNotification.TokenType tokenType = switch(pushType) {
       case "apn" -> useVoip ? PushNotification.TokenType.APN_VOIP : PushNotification.TokenType.APN;
@@ -242,9 +243,18 @@ public class AccountController {
 
     Util.requireNormalizedNumber(number);
 
-    String pushChallenge = generatePushChallenge();
-    StoredVerificationCode storedVerificationCode =
-        new StoredVerificationCode(null, clock.millis(), pushChallenge, null);
+    final Phonenumber.PhoneNumber phoneNumber;
+    try {
+      phoneNumber = PhoneNumberUtil.getInstance().parse(number, null);
+    } catch (final NumberParseException e) {
+      // This should never happen since we just verified that the number is already normalized
+      throw new BadRequestException("Bad phone number");
+    }
+
+    final String pushChallenge = generatePushChallenge();
+    final byte[] sessionId = createRegistrationSession(phoneNumber);
+    final StoredVerificationCode storedVerificationCode =
+        new StoredVerificationCode(null, clock.millis(), pushChallenge, sessionId);
 
     pendingAccounts.store(number, storedVerificationCode);
     pushNotificationManager.sendRegistrationChallengeNotification(pushToken, tokenType, storedVerificationCode.pushCode());
@@ -346,8 +356,16 @@ public class AccountController {
       }
     }).orElse(ClientType.UNKNOWN);
 
-    final byte[] sessionId = registrationServiceClient.sendRegistrationCode(phoneNumber,
-        messageTransport, clientType, acceptLanguage.orElse(null), REGISTRATION_RPC_TIMEOUT).join();
+    // During the transition to explicit session creation, some previously-stored records may not have a session ID;
+    // after the transition, we can assume that any existing record has an associated session ID.
+    final byte[] sessionId =  maybeStoredVerificationCode.isPresent() && maybeStoredVerificationCode.get().sessionId() != null ?
+        maybeStoredVerificationCode.get().sessionId() : createRegistrationSession(phoneNumber);
+
+    registrationServiceClient.sendRegistrationCode(sessionId,
+        messageTransport,
+        clientType,
+        acceptLanguage.orElse(null),
+        REGISTRATION_RPC_TIMEOUT).join();
 
     final StoredVerificationCode storedVerificationCode = new StoredVerificationCode(null,
         clock.millis(),
@@ -940,4 +958,23 @@ public class AccountController {
 
     return Hex.toStringCondensed(challenge);
   }
+
+  private byte[] createRegistrationSession(final Phonenumber.PhoneNumber phoneNumber) throws RateLimitExceededException {
+
+    try {
+      return registrationServiceClient.createRegistrationSession(phoneNumber, REGISTRATION_RPC_TIMEOUT).join();
+    } catch (final CompletionException e) {
+      Throwable cause = e;
+
+      while (cause instanceof CompletionException) {
+        cause = cause.getCause();
+      }
+
+      if (cause instanceof RateLimitExceededException rateLimitExceededException) {
+        throw rateLimitExceededException;
+      }
+
+      throw e;
+    }
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/registration/RegistrationServiceClient.java

@@ -16,15 +16,19 @@ import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
+import java.time.Instant;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
 import org.apache.commons.lang3.StringUtils;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.signal.registration.rpc.CheckVerificationCodeRequest;
 import org.signal.registration.rpc.CheckVerificationCodeResponse;
+import org.signal.registration.rpc.CreateRegistrationSessionRequest;
 import org.signal.registration.rpc.RegistrationServiceGrpc;
 import org.signal.registration.rpc.SendVerificationCodeRequest;
+import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
 
 public class RegistrationServiceClient implements Managed {
 
@@ -52,17 +56,36 @@ public class RegistrationServiceClient implements Managed {
     this.callbackExecutor = callbackExecutor;
   }
 
-  public CompletableFuture<byte[]> sendRegistrationCode(final Phonenumber.PhoneNumber phoneNumber,
+  public CompletableFuture<byte[]> createRegistrationSession(final Phonenumber.PhoneNumber phoneNumber, final Duration timeout) {
+    final long e164 = Long.parseLong(
+        PhoneNumberUtil.getInstance().format(phoneNumber, PhoneNumberUtil.PhoneNumberFormat.E164).substring(1));
+
+    return toCompletableFuture(stub.withDeadline(toDeadline(timeout))
+        .createSession(CreateRegistrationSessionRequest.newBuilder()
+            .setE164(e164)
+            .build()))
+        .thenApply(response -> switch (response.getResponseCase()) {
+          case SESSION_METADATA -> response.getSessionMetadata().getSessionId().toByteArray();
+
+          case ERROR -> {
+            switch (response.getError().getErrorType()) {
+              case ERROR_TYPE_RATE_LIMITED -> throw new CompletionException(new RateLimitExceededException(Duration.ofSeconds(response.getError().getRetryAfterSeconds())));
+              default -> throw new RuntimeException("Unrecognized error type from registration service: " + response.getError().getErrorType());
+            }
+          }
+
+          case RESPONSE_NOT_SET -> throw new RuntimeException("No response from registration service");
+        });
+  }
+
+  public CompletableFuture<byte[]> sendRegistrationCode(final byte[] sessionId,
       final MessageTransport messageTransport,
       final ClientType clientType,
       @Nullable final String acceptLanguage,
       final Duration timeout) {
 
-    final long e164 = Long.parseLong(
-        PhoneNumberUtil.getInstance().format(phoneNumber, PhoneNumberUtil.PhoneNumberFormat.E164).substring(1));
-
     final SendVerificationCodeRequest.Builder requestBuilder = SendVerificationCodeRequest.newBuilder()
-        .setE164(e164)
+        .setSessionId(ByteString.copyFrom(sessionId))
         .setTransport(getRpcMessageTransport(messageTransport))
         .setClientType(getRpcClientType(clientType));
 

service/src/main/proto/RegistrationService.proto

@@ -6,8 +6,13 @@ package org.signal.registration.rpc;
 
 service RegistrationService {
   /**
-   * Sends a verification code to a destination phone number and returns the
-   * ID of the newly-created registration session.
+   * Create a new registration session for a given destination phone number.
+   */
+  rpc create_session (CreateRegistrationSessionRequest) returns (CreateRegistrationSessionResponse) {}
+
+  /**
+   * Sends a verification code to a destination phone number within the context
+   * of a previously-created registration session.
    */
   rpc send_verification_code (SendVerificationCodeRequest) returns (SendVerificationCodeResponse) {}
 
@@ -18,9 +23,71 @@ service RegistrationService {
   rpc check_verification_code (CheckVerificationCodeRequest) returns (CheckVerificationCodeResponse) {}
 }
 
+message CreateRegistrationSessionRequest {
+  /**
+   * The phone number for which to create a new registration session.
+   */
+  uint64 e164 = 1;
+}
+
+message CreateRegistrationSessionResponse {
+  oneof response {
+    /**
+     * Metadata for the newly-created session.
+     */
+    RegistrationSessionMetadata session_metadata = 1;
+
+    /**
+     * A response explaining why a session could not be created as requested.
+     */
+    CreateRegistrationSessionError error = 2;
+  }
+}
+
+message RegistrationSessionMetadata {
+  /**
+   * An opaque sequence of bytes that uniquely identifies the registration
+   * session associated with this registration attempt.
+   */
+  bytes session_id = 1;
+}
+
+message CreateRegistrationSessionError {
+  /**
+   * The type of error that prevented a session from being created.
+   */
+  CreateRegistrationSessionErrorType error_type = 1;
+
+  /**
+   * Indicates that this error is fatal and should not be retried without
+   * modification. Non-fatal errors may be retried without modification after
+   * the duration indicated by `retry_after_seconds`.
+   */
+  bool fatal = 2;
+
+  /**
+   * If this error is not fatal (see `fatal`), indicates the duration in seconds
+   * from the present after which the request may be retried without
+   * modification. This value has no meaning otherwise.
+   */
+  uint64 retry_after_seconds = 3;
+}
+
+enum CreateRegistrationSessionErrorType {
+  ERROR_TYPE_UNSPECIFIED = 0;
+
+  /**
+   * Indicates that a session could not be created because too many requests to
+   * create a session for the given phone number have been received in some
+   * window of time. Callers should wait and try again later.
+   */
+  ERROR_TYPE_RATE_LIMITED = 1;
+}
+
 message SendVerificationCodeRequest {
   /**
-   * The phone number to which to send a verification code.
+   * The phone number to which to send a verification code. Ignored (and may be
+   * null if `session_id` is set.
    */
   uint64 e164 = 1;
 
@@ -31,8 +98,8 @@ message SendVerificationCodeRequest {
   MessageTransport transport = 2;
 
   /**
-   * The value of the `Accept-Language` header provided by remote clients (if
-   * any).
+   * A prioritized list of languages accepted by the destination; should be
+   * provided in the same format as the value of an HTTP Accept-Language header.
    */
   string accept_language = 3;
 
@@ -40,6 +107,11 @@ message SendVerificationCodeRequest {
    * The type of client requesting a verification code.
    */
   ClientType client_type = 4;
+
+  /**
+   * The ID of a session within which to send (or re-send) a verification code.
+   */
+  bytes session_id = 5;
 }
 
 enum MessageTransport {