From eaa868cf0692e7ddc4ea8d79fe65b0b4bc9c64f1 Mon Sep 17 00:00:00 2001
From: Jon Chambers <jon@signal.org>
Date: Wed, 4 Oct 2023 15:16:35 -0400
Subject: [PATCH] Add a remote address interceptor to base gRPC tests

---
 .../textsecuregcm/grpc/GrpcTestUtils.java           | 11 ++++++-----
 .../grpc/MockRemoteAddressInterceptor.java          |  2 +-
 .../textsecuregcm/grpc/SimpleBaseGrpcTest.java      | 13 +++++++++----
 3 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/GrpcTestUtils.java b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/GrpcTestUtils.java
index 2432b88bb..a3f2d0789 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/GrpcTestUtils.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/GrpcTestUtils.java
@@ -26,23 +26,24 @@ public final class GrpcTestUtils {
     // noop
   }
 
-  public static MockAuthenticationInterceptor setupAuthenticatedExtension(
+  public static void setupAuthenticatedExtension(
       final GrpcServerExtension extension,
+      final MockAuthenticationInterceptor mockAuthenticationInterceptor,
+      final MockRemoteAddressInterceptor mockRemoteAddressInterceptor,
       final UUID authenticatedAci,
       final long authenticatedDeviceId,
       final BindableService service) {
-    final MockAuthenticationInterceptor mockAuthenticationInterceptor = new MockAuthenticationInterceptor();
     mockAuthenticationInterceptor.setAuthenticatedDevice(authenticatedAci, authenticatedDeviceId);
     extension.getServiceRegistry()
-        .addService(ServerInterceptors.intercept(service, mockAuthenticationInterceptor, new ErrorMappingInterceptor()));
-    return mockAuthenticationInterceptor;
+        .addService(ServerInterceptors.intercept(service, mockRemoteAddressInterceptor, mockAuthenticationInterceptor, new ErrorMappingInterceptor()));
   }
 
   public static void setupUnauthenticatedExtension(
       final GrpcServerExtension extension,
+      final MockRemoteAddressInterceptor mockRemoteAddressInterceptor,
       final BindableService service) {
     extension.getServiceRegistry()
-        .addService(ServerInterceptors.intercept(service, new ErrorMappingInterceptor()));
+        .addService(ServerInterceptors.intercept(service, mockRemoteAddressInterceptor, new ErrorMappingInterceptor()));
   }
 
   public static void assertStatusException(final Status expected, final Executable serviceCall) {
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/MockRemoteAddressInterceptor.java b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/MockRemoteAddressInterceptor.java
index 11e316b42..a30aa79db 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/MockRemoteAddressInterceptor.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/MockRemoteAddressInterceptor.java
@@ -28,7 +28,7 @@ public class MockRemoteAddressInterceptor implements ServerInterceptor {
       final Metadata headers,
       final ServerCallHandler<ReqT, RespT> next) {
 
-    return remoteAddress != null
+    return remoteAddress == null
         ? next.startCall(serverCall, headers)
         : Contexts.interceptCall(
             Context.current().withValue(RemoteAddressUtil.REMOTE_ADDRESS_CONTEXT_KEY, remoteAddress),
diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/SimpleBaseGrpcTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/SimpleBaseGrpcTest.java
index 8d3907022..f9064fae6 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/grpc/SimpleBaseGrpcTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/grpc/SimpleBaseGrpcTest.java
@@ -59,7 +59,8 @@ public abstract class SimpleBaseGrpcTest<SERVICE extends BindableService, STUB e
 
   private AutoCloseable mocksCloseable;
 
-  private MockAuthenticationInterceptor mockAuthenticationInterceptor;
+  private final MockAuthenticationInterceptor mockAuthenticationInterceptor = new MockAuthenticationInterceptor();
+  private final MockRemoteAddressInterceptor mockRemoteAddressInterceptor = new MockRemoteAddressInterceptor();
 
   private SERVICE service;
 
@@ -112,9 +113,9 @@ public abstract class SimpleBaseGrpcTest<SERVICE extends BindableService, STUB e
   protected void baseSetup() {
     mocksCloseable = MockitoAnnotations.openMocks(this);
     service = requireNonNull(createServiceBeforeEachTest(), "created service must not be `null`");
-    mockAuthenticationInterceptor = GrpcTestUtils.setupAuthenticatedExtension(
-        GRPC_SERVER_EXTENSION_AUTHENTICATED, AUTHENTICATED_ACI, AUTHENTICATED_DEVICE_ID, service);
-    GrpcTestUtils.setupUnauthenticatedExtension(GRPC_SERVER_EXTENSION_UNAUTHENTICATED, service);
+    GrpcTestUtils.setupAuthenticatedExtension(
+        GRPC_SERVER_EXTENSION_AUTHENTICATED, mockAuthenticationInterceptor, mockRemoteAddressInterceptor, AUTHENTICATED_ACI, AUTHENTICATED_DEVICE_ID, service);
+    GrpcTestUtils.setupUnauthenticatedExtension(GRPC_SERVER_EXTENSION_UNAUTHENTICATED, mockRemoteAddressInterceptor, service);
     try {
       authenticatedServiceStub = createStub(GRPC_SERVER_EXTENSION_AUTHENTICATED.getChannel());
       unauthenticatedServiceStub = createStub(GRPC_SERVER_EXTENSION_UNAUTHENTICATED.getChannel());
@@ -143,4 +144,8 @@ public abstract class SimpleBaseGrpcTest<SERVICE extends BindableService, STUB e
   protected STUB unauthenticatedServiceStub() {
     return unauthenticatedServiceStub;
   }
+
+  protected MockRemoteAddressInterceptor getMockRemoteAddressInterceptor() {
+    return mockRemoteAddressInterceptor;
+  }
 }