Ease use of JWT by passing URI to auth library

The URI no longer needs to be provided to the Credential explicitly,
which prevents needing to know a magic string and allows using the same
Credential with multiple services.

2 files changed, 72 insertions, 10 deletions

diff --git a/auth/src/main/java/io/grpc/auth/ClientAuthInterceptor.java b/auth/src/main/java/io/grpc/auth/ClientAuthInterceptor.java
index 4ed477d84..d5f566182 100644
--- a/auth/src/main/java/io/grpc/auth/ClientAuthInterceptor.java
+++ b/auth/src/main/java/io/grpc/auth/ClientAuthInterceptor.java
@@ -42,8 +42,11 @@ import io.grpc.ClientInterceptors.CheckedForwardingClientCall;
 import io.grpc.Metadata;
 import io.grpc.MethodDescriptor;
 import io.grpc.Status;
+import io.grpc.StatusException;
 
 import java.io.IOException;
+import java.net.URI;
+import java.net.URISyntaxException;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Executor;
@@ -70,22 +73,27 @@ public class ClientAuthInterceptor implements ClientInterceptor {
   }
 
   @Override
-  public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(MethodDescriptor<ReqT, RespT> method,
-      CallOptions callOptions, Channel next) {
+  public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
+      final MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, final Channel next) {
     // TODO(ejona86): If the call fails for Auth reasons, this does not properly propagate info that
     // would be in WWW-Authenticate, because it does not yet have access to the header.
     return new CheckedForwardingClientCall<ReqT, RespT>(next.newCall(method, callOptions)) {
       @Override
       protected void checkedStart(Listener<RespT> responseListener, Metadata headers)
-          throws Exception {
+          throws StatusException {
         Metadata cachedSaved;
+        URI uri = serviceUri(next, method);
         synchronized (ClientAuthInterceptor.this) {
           // TODO(louiscryan): This is icky but the current auth library stores the same
           // metadata map until the next refresh cycle. This will be fixed once
           // https://github.com/google/google-auth-library-java/issues/3
           // is resolved.
-          if (lastMetadata == null || lastMetadata != getRequestMetadata()) {
-            lastMetadata = getRequestMetadata();
+          // getRequestMetadata() may return a different map based on the provided URI, i.e., for
+          // JWT. However, today it does not cache JWT and so we won't bother tring to cache its
+          // return value based on the URI.
+          Map<String, List<String>> latestMetadata = getRequestMetadata(uri);
+          if (lastMetadata == null || lastMetadata != latestMetadata) {
+            lastMetadata = latestMetadata;
             cached = toHeaders(lastMetadata);
           }
           cachedSaved = cached;
@@ -96,11 +104,37 @@ public class ClientAuthInterceptor implements ClientInterceptor {
     };
   }
 
-  private Map<String, List<String>> getRequestMetadata() {
+  /**
+   * Generate a JWT-specific service URI. The URI is simply an identifier with enough information
+   * for a service to know that the JWT was intended for it. The URI will commonly be verified with
+   * a simple string equality check.
+   */
+  private URI serviceUri(Channel channel, MethodDescriptor<?, ?> method) throws StatusException {
+    String authority = channel.authority();
+    if (authority == null) {
+      throw Status.UNAUTHENTICATED.withDescription("Channel has no authority").asException();
+    }
+    // Always use HTTPS, by definition.
+    final String scheme = "https";
+    // The default port must not be present. Alternative ports should be present.
+    final String suffixToStrip = ":443";
+    if (authority.endsWith(suffixToStrip)) {
+      authority = authority.substring(0, authority.length() - suffixToStrip.length());
+    }
+    String path = "/" + MethodDescriptor.extractFullServiceName(method.getFullMethodName());
+    try {
+      return new URI(scheme, authority, path, null, null);
+    } catch (URISyntaxException e) {
+      throw Status.UNAUTHENTICATED.withDescription("Unable to construct service URI for auth")
+          .withCause(e).asException();
+    }
+  }
+
+  private Map<String, List<String>> getRequestMetadata(URI uri) throws StatusException {
     try {
-      return credentials.getRequestMetadata();
+      return credentials.getRequestMetadata(uri);
     } catch (IOException e) {
-      throw Status.UNAUTHENTICATED.withCause(e).asRuntimeException();
+      throw Status.UNAUTHENTICATED.withCause(e).asException();
     }
   }
 
diff --git a/auth/src/test/java/io/grpc/auth/ClientAuthInterceptorTests.java b/auth/src/test/java/io/grpc/auth/ClientAuthInterceptorTests.java
index f0010423e..6736bd096 100644
--- a/auth/src/test/java/io/grpc/auth/ClientAuthInterceptorTests.java
+++ b/auth/src/test/java/io/grpc/auth/ClientAuthInterceptorTests.java
@@ -34,6 +34,7 @@ package io.grpc.auth;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.isA;
 import static org.mockito.Matchers.same;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -51,6 +52,7 @@ import io.grpc.Channel;
 import io.grpc.ClientCall;
 import io.grpc.Metadata;
 import io.grpc.MethodDescriptor;
+import io.grpc.MethodDescriptor.Marshaller;
 import io.grpc.Status;
 
 import org.junit.Assert;
@@ -64,6 +66,7 @@ import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
 import java.io.IOException;
+import java.net.URI;
 import java.util.Date;
 import java.util.concurrent.Executors;
 
@@ -82,6 +85,11 @@ public class ClientAuthInterceptorTests {
   Credentials credentials;
 
   @Mock
+  Marshaller<String> stringMarshaller;
+
+  @Mock
+  Marshaller<Integer> intMarshaller;
+
   MethodDescriptor<String, Integer> descriptor;
 
   @Mock
@@ -99,7 +107,10 @@ public class ClientAuthInterceptorTests {
   @Before
   public void startUp() throws IOException {
     MockitoAnnotations.initMocks(this);
+    descriptor = MethodDescriptor.create(
+        MethodDescriptor.MethodType.UNKNOWN, "a.service/method", stringMarshaller, intMarshaller);
     when(channel.newCall(same(descriptor), any(CallOptions.class))).thenReturn(call);
+    doReturn("localhost:443").when(channel).authority();
     interceptor = new ClientAuthInterceptor(credentials,
         Executors.newSingleThreadExecutor());
   }
@@ -111,7 +122,7 @@ public class ClientAuthInterceptorTests {
     values.put("Authorization", "token2");
     values.put("Extra-Authorization", "token3");
     values.put("Extra-Authorization", "token4");
-    when(credentials.getRequestMetadata()).thenReturn(Multimaps.asMap(values));
+    when(credentials.getRequestMetadata(any(URI.class))).thenReturn(Multimaps.asMap(values));
     ClientCall<String, Integer> interceptedCall =
         interceptor.interceptCall(descriptor, CallOptions.DEFAULT, channel);
     Metadata headers = new Metadata();
@@ -128,7 +139,7 @@ public class ClientAuthInterceptorTests {
 
   @Test
   public void testCredentialsThrows() throws IOException {
-    when(credentials.getRequestMetadata()).thenThrow(new IOException("Broken"));
+    when(credentials.getRequestMetadata(any(URI.class))).thenThrow(new IOException("Broken"));
     ClientCall<String, Integer> interceptedCall =
         interceptor.interceptCall(descriptor, CallOptions.DEFAULT, channel);
     Metadata headers = new Metadata();
@@ -160,4 +171,21 @@ public class ClientAuthInterceptorTests {
     Assert.assertArrayEquals(new String[]{"Bearer allyourbase"},
         Iterables.toArray(authorization, String.class));
   }
+
+  @Test
+  public void verifyServiceUri() throws IOException {
+    ClientCall<String, Integer> interceptedCall;
+
+    doReturn("example.com:443").when(channel).authority();
+    interceptedCall = interceptor.interceptCall(descriptor, CallOptions.DEFAULT, channel);
+    interceptedCall.start(listener, new Metadata());
+    verify(credentials).getRequestMetadata(URI.create("https://example.com/a.service"));
+    interceptedCall.cancel();
+
+    doReturn("example.com:123").when(channel).authority();
+    interceptedCall = interceptor.interceptCall(descriptor, CallOptions.DEFAULT, channel);
+    interceptedCall.start(listener, new Metadata());
+    verify(credentials).getRequestMetadata(URI.create("https://example.com:123/a.service"));
+    interceptedCall.cancel();
+  }
 }