From ff1ef90a6dd9e2978e45bcad34ee7716b6366fc5 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Fri, 18 Aug 2023 13:22:57 -0400
Subject: [PATCH] Defer actions taken after rate limit checks

---
 .../whispersystems/textsecuregcm/grpc/CallingGrpcService.java  | 2 +-
 .../org/whispersystems/textsecuregcm/grpc/KeysGrpcService.java | 2 +-
 .../textsecuregcm/grpc/CallingGrpcServiceTest.java             | 3 +++
 .../whispersystems/textsecuregcm/grpc/KeysGrpcServiceTest.java | 3 +++
 4 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/CallingGrpcService.java b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/CallingGrpcService.java
index f981f094d..4f60f9647 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/CallingGrpcService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/CallingGrpcService.java
@@ -34,7 +34,7 @@ public class CallingGrpcService extends ReactorCallingGrpc.CallingImplBase {
     final AuthenticatedDevice authenticatedDevice = AuthenticationUtil.requireAuthenticatedDevice();
 
     return rateLimiters.getTurnLimiter().validateReactive(authenticatedDevice.accountIdentifier())
-        .then(Mono.fromSupplier(() -> turnTokenGenerator.generate(authenticatedDevice.accountIdentifier())))
+        .then(Mono.defer(() -> Mono.fromSupplier(() -> turnTokenGenerator.generate(authenticatedDevice.accountIdentifier()))))
         .map(turnToken -> GetTurnCredentialsResponse.newBuilder()
             .setUsername(turnToken.username())
             .setPassword(turnToken.password())
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcService.java b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcService.java
index 89d885a76..4f8277557 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcService.java
@@ -135,7 +135,7 @@ public class KeysGrpcService extends ReactorKeysGrpc.KeysImplBase {
         request.getDeviceId();
 
     return rateLimiters.getPreKeysLimiter().validateReactive(rateLimitKey)
-        .then(Mono.fromFuture(accountsManager.getByServiceIdentifierAsync(targetIdentifier))
+        .then(Mono.fromFuture(() -> accountsManager.getByServiceIdentifierAsync(targetIdentifier))
             .flatMap(Mono::justOrEmpty))
         .switchIfEmpty(Mono.error(Status.NOT_FOUND.asException()))
         .flatMap(targetAccount ->
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/CallingGrpcServiceTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/CallingGrpcServiceTest.java
index 122448fd6..0753dfadf 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/CallingGrpcServiceTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/CallingGrpcServiceTest.java
@@ -12,6 +12,7 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
 import static org.mockito.Mockito.when;
 
 import io.grpc.ServerInterceptors;
@@ -102,5 +103,7 @@ class CallingGrpcServiceTest {
     assertEquals(Status.Code.RESOURCE_EXHAUSTED, exception.getStatus().getCode());
     assertNotNull(exception.getTrailers());
     assertEquals(retryAfter, exception.getTrailers().get(RateLimitUtil.RETRY_AFTER_DURATION_KEY));
+
+    verifyNoInteractions(turnTokenGenerator);
   }
 }
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcServiceTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcServiceTest.java
index 52860fc19..1ae1a443f 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcServiceTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/KeysGrpcServiceTest.java
@@ -14,6 +14,7 @@ import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
 import static org.mockito.Mockito.when;
 
 import com.google.protobuf.ByteString;
@@ -669,5 +670,7 @@ class KeysGrpcServiceTest {
     assertEquals(Status.Code.RESOURCE_EXHAUSTED, exception.getStatus().getCode());
     assertNotNull(exception.getTrailers());
     assertEquals(retryAfterDuration, exception.getTrailers().get(RateLimitUtil.RETRY_AFTER_DURATION_KEY));
+
+    verifyNoInteractions(accountsManager);
   }
 }