From bf05e47e260602a3e4a89a341fec14695055acd1 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Wed, 10 Jan 2024 16:59:01 -0500
Subject: [PATCH] Get accounts asynchronously when checking SVR credentials via
 gRPC

---
 ...erviceCredentialsAnonymousGrpcService.java | 19 +++++++++++--------
 ...ceCredentialsAnonymousGrpcServiceTest.java |  4 +++-
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/ExternalServiceCredentialsAnonymousGrpcService.java b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/ExternalServiceCredentialsAnonymousGrpcService.java
index f76bf3d02..8536b7de8 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/ExternalServiceCredentialsAnonymousGrpcService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/ExternalServiceCredentialsAnonymousGrpcService.java
@@ -24,6 +24,7 @@ import org.whispersystems.textsecuregcm.storage.Account;
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.util.function.Tuples;
 
 public class ExternalServiceCredentialsAnonymousGrpcService extends
     ReactorExternalServiceCredentialsAnonymousGrpc.ExternalServiceCredentialsAnonymousImplBase {
@@ -60,15 +61,17 @@ public class ExternalServiceCredentialsAnonymousGrpcService extends
         svrCredentialsGenerator,
         MAX_SVR_PASSWORD_AGE_SECONDS);
 
-    // the username associated with the provided number
-    final Optional<String> matchingUsername = accountsManager
-        .getByE164(request.getNumber())
-        .map(Account::getUuid)
-        .map(svrCredentialsGenerator::generateForUuid)
-        .map(ExternalServiceCredentials::username);
+    return Mono.fromFuture(() -> accountsManager.getByE164Async(request.getNumber()))
+        // the username associated with the provided number
+        .map(maybeAccount -> maybeAccount.map(Account::getUuid)
+            .map(svrCredentialsGenerator::generateForUuid)
+            .map(ExternalServiceCredentials::username))
+        .flatMapMany(maybeMatchingUsername -> Flux.fromIterable(credentials)
+            .map(credential -> Tuples.of(maybeMatchingUsername, credential)))
+        .reduce(CheckSvrCredentialsResponse.newBuilder(), ((builder, usernameAndCredentialInfo) -> {
+          final Optional<String> matchingUsername = usernameAndCredentialInfo.getT1();
+          final ExternalServiceCredentialsSelector.CredentialInfo credentialInfo = usernameAndCredentialInfo.getT2();
 
-    return Flux.fromIterable(credentials)
-        .reduce(CheckSvrCredentialsResponse.newBuilder(), ((builder, credentialInfo) -> {
           final AuthCheckResult authCheckResult;
           if (!credentialInfo.valid()) {
             authCheckResult = AuthCheckResult.AUTH_CHECK_RESULT_INVALID;
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/ExternalServiceCredentialsAnonymousGrpcServiceTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/ExternalServiceCredentialsAnonymousGrpcServiceTest.java
index 328527cf5..f1994c8fe 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/ExternalServiceCredentialsAnonymousGrpcServiceTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/ExternalServiceCredentialsAnonymousGrpcServiceTest.java
@@ -12,6 +12,7 @@ import java.time.Duration;
 import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mock;
@@ -58,7 +59,8 @@ class ExternalServiceCredentialsAnonymousGrpcServiceTest extends
 
   @BeforeEach
   public void setup() {
-    Mockito.when(accountsManager.getByE164(USER_E164)).thenReturn(Optional.of(account(USER_UUID)));
+    Mockito.when(accountsManager.getByE164Async(USER_E164))
+        .thenReturn(CompletableFuture.completedFuture(Optional.of(account(USER_UUID))));
   }
 
   @Test