From 2052e62c01aed1633fd9b8dfe7b69bc685111b42 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Mon, 20 Mar 2023 10:44:33 -0400
Subject: [PATCH] Use a purpose-specific method when checking verification
 codes via the legacy registration API

---
 .../RegistrationServiceClient.java            | 27 +++++++++++++++++--
 .../src/main/proto/RegistrationService.proto  | 16 +++++++++++
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/registration/RegistrationServiceClient.java b/service/src/main/java/org/whispersystems/textsecuregcm/registration/RegistrationServiceClient.java
index f93da9398..779316bfc 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/registration/RegistrationServiceClient.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/registration/RegistrationServiceClient.java
@@ -181,8 +181,31 @@ public class RegistrationServiceClient implements Managed {
       final String verificationCode,
       final Duration timeout) {
 
-    return checkVerificationCodeSession(sessionId, verificationCode, timeout)
-        .thenApply(RegistrationServiceSession::verified);
+    return toCompletableFuture(stub.withDeadline(toDeadline(timeout))
+        .legacyCheckVerificationCode(CheckVerificationCodeRequest.newBuilder()
+            .setSessionId(ByteString.copyFrom(sessionId))
+            .setVerificationCode(verificationCode)
+            .build()))
+        .thenApply(response -> {
+          if (response.hasError()) {
+            switch (response.getError().getErrorType()) {
+              case CHECK_VERIFICATION_CODE_ERROR_TYPE_RATE_LIMITED ->
+                  throw new CompletionException(new RateLimitExceededException(response.getError().getMayRetry()
+                      ? Duration.ofSeconds(response.getError().getRetryAfterSeconds())
+                      : null, true));
+
+              case CHECK_VERIFICATION_CODE_ERROR_TYPE_NO_CODE_SENT,
+                  CHECK_VERIFICATION_CODE_ERROR_TYPE_ATTEMPT_EXPIRED,
+                  CHECK_VERIFICATION_CODE_ERROR_TYPE_SESSION_NOT_FOUND ->
+                  throw new CompletionException(new RegistrationServiceException(null));
+
+              default -> throw new CompletionException(
+                  new RuntimeException("Failed to check verification code: " + response.getError().getErrorType()));
+            }
+          } else {
+            return response.getVerified();
+          }
+        });
   }
 
   public CompletableFuture<RegistrationServiceSession> checkVerificationCodeSession(final byte[] sessionId,
diff --git a/service/src/main/proto/RegistrationService.proto b/service/src/main/proto/RegistrationService.proto
index 0e268d7e7..8f8c2ae7f 100644
--- a/service/src/main/proto/RegistrationService.proto
+++ b/service/src/main/proto/RegistrationService.proto
@@ -26,6 +26,8 @@ service RegistrationService {
    * session.
    */
   rpc check_verification_code (CheckVerificationCodeRequest) returns (CheckVerificationCodeResponse) {}
+
+  rpc legacy_check_verification_code (CheckVerificationCodeRequest) returns (LegacyCheckVerificationCodeResponse) {}
 }
 
 message CreateRegistrationSessionRequest {
@@ -343,6 +345,20 @@ message CheckVerificationCodeResponse {
   CheckVerificationCodeError error = 3;
 }
 
+message LegacyCheckVerificationCodeResponse {
+  /**
+   * Indicates whether the verification code given in the request that produced
+   * this response was correct.
+   */
+  bool verified = 1;
+
+  /**
+   * If a code could not be checked, explains the underlying error. Will be
+   * absent if no error occurred.
+   */
+  CheckVerificationCodeError error = 2;
+}
+
 message CheckVerificationCodeError {
   /**
    * The type of error that prevented a verification code from being checked.