feat: use CallCredentials to add request header for better error messages

PiperOrigin-RevId: 468318869

Author: On Demand Rides and Deliveries

src/main/java/com/google/fleetengine/auth/client/FleetEngineAuthCallCredentials.java

@@ -0,0 +1,57 @@
+// Copyright 2021 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.fleetengine.auth.client;
+
+import com.google.fleetengine.auth.token.FleetEngineToken;
+import io.grpc.CallCredentials;
+import io.grpc.Metadata;
+import io.grpc.Status;
+import java.util.concurrent.Executor;
+
+/** Adds an athorization header containing a Fleet Engine JWT to the request metadata. */
+public class FleetEngineAuthCallCredentials extends CallCredentials {
+
+  /** Authorization header name. */
+  private static final Metadata.Key<String> AUTHORIZATION_HEADER =
+      Metadata.Key.of("authorization", Metadata.ASCII_STRING_MARSHALLER);
+
+  private final FleetEngineTokenProvider tokenProvider;
+
+  public static FleetEngineAuthCallCredentials create(FleetEngineTokenProvider tokenProvider) {
+    return new FleetEngineAuthCallCredentials(tokenProvider);
+  }
+
+  private FleetEngineAuthCallCredentials(FleetEngineTokenProvider tokenProvider) {
+    this.tokenProvider = tokenProvider;
+  }
+
+  @Override
+  public void applyRequestMetadata(
+      RequestInfo requestInfo, Executor appExecutor, CallCredentials.MetadataApplier applier) {
+    Metadata headers = new Metadata();
+    try {
+      FleetEngineToken token = tokenProvider.getSignedToken();
+      headers.put(AUTHORIZATION_HEADER, String.format("Bearer %s", token.jwt()));
+    } catch (Exception e) {
+      applier.fail(Status.UNAUTHENTICATED.withDescription("Unable to create token").withCause(e));
+      return;
+    }
+
+    applier.apply(headers);
+  }
+
+  @Override
+  public void thisUsesUnstableApi() {}
+}

src/main/java/com/google/fleetengine/auth/client/FleetEngineAuthClientInterceptor.java

@@ -14,19 +14,15 @@
 
 package com.google.fleetengine.auth.client;
 
-import com.google.common.annotations.VisibleForTesting;
-import com.google.fleetengine.auth.token.FleetEngineToken;
-import com.google.fleetengine.auth.token.factory.signer.SigningTokenException;
 import io.grpc.CallOptions;
 import io.grpc.Channel;
 import io.grpc.ClientCall;
 import io.grpc.ClientInterceptor;
-import io.grpc.ForwardingClientCall;
-import io.grpc.Metadata;
 import io.grpc.MethodDescriptor;
 
 /**
- * Intercepts an outgoing gRPC request and attaches a valid Fleet Engine JWT to the header.
+ * Intercepts an outgoing gRPC request and attaches a valid Fleet Engine JWT to the header by using
+ * {@link FleetEngineAuthCallCredentials}.
  *
  * <p>Works with generated gRPC stubby classes:
  *
@@ -37,11 +33,7 @@
  */
 public class FleetEngineAuthClientInterceptor implements ClientInterceptor {
 
-  /** Authorization header name. */
-  private static final Metadata.Key<String> AUTHORIZATION_HEADER =
-      Metadata.Key.of("authorization", Metadata.ASCII_STRING_MARSHALLER);
-
-  private final FleetEngineTokenProvider tokenProvider;
+  private final FleetEngineAuthCallCredentials callCredentials;
 
   /**
    * Creates a gRPC client interceptor that attaches tokens from {@code tokenProvider} to outgoing
@@ -50,39 +42,22 @@ public class FleetEngineAuthClientInterceptor implements ClientInterceptor {
    * @param tokenProvider provides valid Fleet Engine JWTs for an outgoing gRPC request.
    */
   public static FleetEngineAuthClientInterceptor create(FleetEngineTokenProvider tokenProvider) {
-    return new FleetEngineAuthClientInterceptor(tokenProvider);
+    return new FleetEngineAuthClientInterceptor(
+        FleetEngineAuthCallCredentials.create(tokenProvider));
   }
 
   /** Constructor. */
-  private FleetEngineAuthClientInterceptor(FleetEngineTokenProvider tokenProvider) {
-    this.tokenProvider = tokenProvider;
+  private FleetEngineAuthClientInterceptor(FleetEngineAuthCallCredentials callCredentials) {
+    this.callCredentials = callCredentials;
   }
 
   /**
-   * Overrides {@link ClientInterceptor#interceptCall(MethodDescriptor, CallOptions, Channel)} and
-   * attaches a Fleet Engine JWT to the header of the outgoing gRPC request.
+   * Returns a new call with the {@link FleetEngineAuthCallCredentials} added to the {@code
+   * callOptions}.
    */
   @Override
   public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
       MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {
-    ClientCall<ReqT, RespT> call = next.newCall(method, callOptions);
-    return new ForwardingClientCall.SimpleForwardingClientCall<ReqT, RespT>(call) {
-      @Override
-      public void start(Listener<RespT> responseListener, Metadata headers) {
-        addAuthorizationHeader(headers);
-        super.start(responseListener, headers);
-      }
-    };
-  }
-
-  /** Adds the signed base64 encode JWT to the header. */
-  @VisibleForTesting
-  void addAuthorizationHeader(Metadata headers) {
-    try {
-      FleetEngineToken token = tokenProvider.getSignedToken();
-      headers.put(AUTHORIZATION_HEADER, String.format("Bearer %s", token.jwt()));
-    } catch (SigningTokenException e) {
-      throw new WritingAuthorizationHeaderException("Exception while getting token.", e);
-    }
+    return next.newCall(method, callOptions.withCallCredentials(this.callCredentials));
   }
 }

src/main/java/com/google/fleetengine/auth/client/WritingAuthorizationHeaderException.java

@@ -15,6 +15,7 @@
 package com.google.fleetengine.auth.client;
 
 /** Signals that an exception occurred while writing the authorization header. */
+@Deprecated
 public class WritingAuthorizationHeaderException extends RuntimeException {
 
   /**

…leetEngineAuthClientInterceptorTest.java …/FleetEngineAuthCallCredentialsTest.javasrc/test/java/com/google/fleetengine/auth/client/FleetEngineAuthClientInterceptorTest.java renamed to src/test/java/com/google/fleetengine/auth/client/FleetEngineAuthCallCredentialsTest.java src/test/java/com/google/fleetengine/auth/client/FleetEngineAuthClientInterceptorTest.java renamed to src/test/java/com/google/fleetengine/auth/client/FleetEngineAuthCallCredentialsTest.java

@@ -15,39 +15,41 @@
 package com.google.fleetengine.auth.client;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import com.google.fleetengine.auth.EmptyFleetEngineTokenClaims;
 import com.google.fleetengine.auth.token.FleetEngineToken;
 import com.google.fleetengine.auth.token.FleetEngineTokenType;
 import com.google.fleetengine.auth.token.factory.signer.SigningTokenException;
+import io.grpc.CallCredentials;
 import io.grpc.Metadata;
 import java.time.Instant;
 import java.util.Date;
-import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 @RunWith(JUnit4.class)
-public class FleetEngineAuthClientInterceptorTest {
+public class FleetEngineAuthCallCredentialsTest {
 
   private static final Metadata.Key<String> AUTHORIZATION_HEADER =
       Metadata.Key.of("authorization", Metadata.ASCII_STRING_MARSHALLER);
   private static final String FAKE_JWT = "fake.jwt.token";
   private static final String FAKE_AUTHORIZATION_HEADER = String.format("Bearer %s", FAKE_JWT);
-  private Metadata headers = new Metadata();
   @Mock private FleetEngineTokenProvider tokenProvider;
+  @Mock private CallCredentials.MetadataApplier applier;
   private FleetEngineToken fleetEngineToken;
 
   @Before
   public void setUp() {
     MockitoAnnotations.openMocks(this);
-    headers = new Metadata();
 
     fleetEngineToken =
         FleetEngineToken.builder()
@@ -59,27 +61,28 @@ public void setUp() {
   }
 
   @Test
-  public void addAuthorizationHeader_addsHeaderCorrectly() throws SigningTokenException {
+  public void applyRequestMetadata_addsHeaderCorrectly() throws SigningTokenException {
     fleetEngineToken = fleetEngineToken.toBuilder().setJwt(FAKE_JWT).build();
     when(tokenProvider.getSignedToken()).thenReturn(fleetEngineToken);
-    FleetEngineAuthClientInterceptor clientInterceptor =
-        FleetEngineAuthClientInterceptor.create(this.tokenProvider);
+    FleetEngineAuthCallCredentials callCredentials =
+        FleetEngineAuthCallCredentials.create(this.tokenProvider);
 
-    clientInterceptor.addAuthorizationHeader(headers);
+    callCredentials.applyRequestMetadata(null, null, applier);
 
-    assertThat(headers.get(AUTHORIZATION_HEADER)).isEqualTo(FAKE_AUTHORIZATION_HEADER);
+    ArgumentCaptor<Metadata> metadataCaptor = ArgumentCaptor.forClass(Metadata.class);
+    verify(applier).apply(metadataCaptor.capture());
+    assertThat(metadataCaptor.getValue().get(AUTHORIZATION_HEADER))
+        .isEqualTo(FAKE_AUTHORIZATION_HEADER);
   }
 
   @Test
-  public void addAuthorizationHeader_onException_throwsException() throws SigningTokenException {
+  public void applyRequestMetadata_onException_callsDail() throws SigningTokenException {
     fleetEngineToken = fleetEngineToken.toBuilder().setJwt(FAKE_JWT).build();
     when(tokenProvider.getSignedToken()).thenThrow(mock(SigningTokenException.class));
+    FleetEngineAuthCallCredentials callCredentials =
+        FleetEngineAuthCallCredentials.create(this.tokenProvider);
 
-    FleetEngineAuthClientInterceptor clientInterceptor =
-        FleetEngineAuthClientInterceptor.create(this.tokenProvider);
-
-    Assert.assertThrows(
-        WritingAuthorizationHeaderException.class,
-        () -> clientInterceptor.addAuthorizationHeader(headers));
+    callCredentials.applyRequestMetadata(null, null, applier);
+    verify(applier).fail(any());
   }
 }