auth: Promote OAuth2 service accounts to JWT

JWT needs less configuration and zero round-trips to initialize.

Fixes #785

2 files changed, 165 insertions, 1 deletions

diff --git a/auth/src/main/java/io/grpc/auth/GoogleAuthLibraryCallCredentials.java b/auth/src/main/java/io/grpc/auth/GoogleAuthLibraryCallCredentials.java
index 9e40da482..4564b693e 100644
--- a/auth/src/main/java/io/grpc/auth/GoogleAuthLibraryCallCredentials.java
+++ b/auth/src/main/java/io/grpc/auth/GoogleAuthLibraryCallCredentials.java
@@ -45,17 +45,27 @@ import io.grpc.MethodDescriptor;
 import io.grpc.Status;
 import io.grpc.StatusException;
 
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Method;
 import java.net.URI;
 import java.net.URISyntaxException;
+import java.security.PrivateKey;
+import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Executor;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 import javax.annotation.Nullable;
 
 /**
  * Wraps {@link Credentials} as a {@link CallCredentials}.
  */
 final class GoogleAuthLibraryCallCredentials implements CallCredentials {
+  private static final Logger log
+      = Logger.getLogger(GoogleAuthLibraryCallCredentials.class.getName());
+  private static final JwtHelper jwtHelper
+      = createJwtHelperOrNull(GoogleAuthLibraryCallCredentials.class.getClassLoader());
 
   @VisibleForTesting
   final Credentials creds;
@@ -64,7 +74,16 @@ final class GoogleAuthLibraryCallCredentials implements CallCredentials {
   private Map<String, List<String>> lastMetadata;
 
   public GoogleAuthLibraryCallCredentials(Credentials creds) {
-    this.creds = checkNotNull(creds, "creds");
+    this(creds, jwtHelper);
+  }
+
+  @VisibleForTesting
+  GoogleAuthLibraryCallCredentials(Credentials creds, JwtHelper jwtHelper) {
+    checkNotNull(creds, "creds");
+    if (jwtHelper != null) {
+      creds = jwtHelper.tryServiceAccountToJwt(creds);
+    }
+    this.creds = creds;
   }
 
   @Override
@@ -168,4 +187,74 @@ final class GoogleAuthLibraryCallCredentials implements CallCredentials {
     }
     return headers;
   }
+
+  @VisibleForTesting
+  @Nullable
+  static JwtHelper createJwtHelperOrNull(ClassLoader loader) {
+    Class<?> rawServiceAccountClass;
+    try {
+      // Specify loader so it can be overriden in tests
+      rawServiceAccountClass
+          = Class.forName("com.google.auth.oauth2.ServiceAccountCredentials", false, loader);
+    } catch (ClassNotFoundException ex) {
+      return null;
+    }
+    try {
+      return new JwtHelper(rawServiceAccountClass, loader);
+    } catch (ReflectiveOperationException ex) {
+      // Failure is a bug in this class, but we still choose to gracefully recover
+      log.log(Level.WARNING, "Failed to create JWT helper. This is unexpected", ex);
+      return null;
+    }
+  }
+
+  @VisibleForTesting
+  static class JwtHelper {
+    private final Class<? extends Credentials> serviceAccountClass;
+    private final Constructor<? extends Credentials> jwtConstructor;
+    private final Method getScopes;
+    private final Method getClientId;
+    private final Method getClientEmail;
+    private final Method getPrivateKey;
+    private final Method getPrivateKeyId;
+
+    public JwtHelper(Class<?> rawServiceAccountClass, ClassLoader loader)
+        throws ReflectiveOperationException {
+      serviceAccountClass = rawServiceAccountClass.asSubclass(Credentials.class);
+      getScopes = serviceAccountClass.getMethod("getScopes");
+      getClientId = serviceAccountClass.getMethod("getClientId");
+      getClientEmail = serviceAccountClass.getMethod("getClientEmail");
+      getPrivateKey = serviceAccountClass.getMethod("getPrivateKey");
+      getPrivateKeyId = serviceAccountClass.getMethod("getPrivateKeyId");
+      Class<? extends Credentials> jwtClass = Class.forName(
+          "com.google.auth.oauth2.ServiceAccountJwtAccessCredentials", false, loader)
+          .asSubclass(Credentials.class);
+      jwtConstructor
+          = jwtClass.getConstructor(String.class, String.class, PrivateKey.class, String.class);
+    }
+
+    public Credentials tryServiceAccountToJwt(Credentials creds) {
+      if (!serviceAccountClass.isInstance(creds)) {
+        return creds;
+      }
+      try {
+        creds = serviceAccountClass.cast(creds);
+        Collection<?> scopes = (Collection<?>) getScopes.invoke(creds);
+        if (scopes.size() != 0) {
+          // Leave as-is, since the scopes may limit access within the service.
+          return creds;
+        }
+        return jwtConstructor.newInstance(
+            getClientId.invoke(creds),
+            getClientEmail.invoke(creds),
+            getPrivateKey.invoke(creds),
+            getPrivateKeyId.invoke(creds));
+      } catch (ReflectiveOperationException ex) {
+        // Failure is a bug in this class, but we still choose to gracefully recover
+        log.log(Level.WARNING,
+            "Failed converting service account credential to JWT. This is unexpected", ex);
+        return creds;
+      }
+    }
+  }
 }
diff --git a/auth/src/test/java/io/grpc/auth/GoogleAuthLibraryCallCredentialsTests.java b/auth/src/test/java/io/grpc/auth/GoogleAuthLibraryCallCredentialsTests.java
index c993625e4..50b3c21e9 100644
--- a/auth/src/test/java/io/grpc/auth/GoogleAuthLibraryCallCredentialsTests.java
+++ b/auth/src/test/java/io/grpc/auth/GoogleAuthLibraryCallCredentialsTests.java
@@ -34,6 +34,8 @@ package io.grpc.auth;
 import static com.google.common.base.Charsets.US_ASCII;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.doAnswer;
@@ -44,6 +46,7 @@ import static org.mockito.Mockito.when;
 import com.google.auth.Credentials;
 import com.google.auth.oauth2.AccessToken;
 import com.google.auth.oauth2.OAuth2Credentials;
+import com.google.auth.oauth2.ServiceAccountCredentials;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.LinkedListMultimap;
 import com.google.common.collect.ListMultimap;
@@ -70,7 +73,10 @@ import org.mockito.stubbing.Answer;
 
 import java.io.IOException;
 import java.net.URI;
+import java.security.KeyPair;
+import java.security.KeyPairGenerator;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Date;
 import java.util.List;
 import java.util.concurrent.Executor;
@@ -269,6 +275,75 @@ public class GoogleAuthLibraryCallCredentialsTests {
     verify(credentials).getRequestMetadata(eq(new URI("https://example.com:123/a.service")));
   }
 
+  @Test
+  public void serviceAccountToJwt() throws Exception {
+    KeyPair pair = KeyPairGenerator.getInstance("RSA").generateKeyPair();
+    ServiceAccountCredentials credentials = new ServiceAccountCredentials(
+        null, "email@example.com", pair.getPrivate(), null, null) {
+      @Override
+      public AccessToken refreshAccessToken() {
+        throw new AssertionError();
+      }
+    };
+
+    GoogleAuthLibraryCallCredentials callCredentials =
+        new GoogleAuthLibraryCallCredentials(credentials);
+    callCredentials.applyRequestMetadata(method, attrs, executor, applier);
+    assertEquals(1, runPendingRunnables());
+
+    verify(applier).apply(headersCaptor.capture());
+    Metadata headers = headersCaptor.getValue();
+    String[] authorization = Iterables.toArray(headers.getAll(AUTHORIZATION), String.class);
+    assertEquals(1, authorization.length);
+    assertTrue(authorization[0], authorization[0].startsWith("Bearer "));
+    // JWT is reasonably long. Normal tokens aren't.
+    assertTrue(authorization[0], authorization[0].length() > 300);
+  }
+
+  @Test
+  public void serviceAccountWithScopeNotToJwt() throws Exception {
+    final AccessToken token = new AccessToken("allyourbase", new Date(Long.MAX_VALUE));
+    KeyPair pair = KeyPairGenerator.getInstance("RSA").generateKeyPair();
+    ServiceAccountCredentials credentials = new ServiceAccountCredentials(
+        null, "email@example.com", pair.getPrivate(), null, Arrays.asList("somescope")) {
+      @Override
+      public AccessToken refreshAccessToken() {
+        return token;
+      }
+    };
+
+    GoogleAuthLibraryCallCredentials callCredentials =
+        new GoogleAuthLibraryCallCredentials(credentials);
+    callCredentials.applyRequestMetadata(method, attrs, executor, applier);
+    assertEquals(1, runPendingRunnables());
+
+    verify(applier).apply(headersCaptor.capture());
+    Metadata headers = headersCaptor.getValue();
+    Iterable<String> authorization = headers.getAll(AUTHORIZATION);
+    assertArrayEquals(new String[]{"Bearer allyourbase"},
+        Iterables.toArray(authorization, String.class));
+  }
+
+  @Test
+  public void oauthClassesNotInClassPath() throws Exception {
+    ListMultimap<String, String> values = LinkedListMultimap.create();
+    values.put("Authorization", "token1");
+    when(credentials.getRequestMetadata(eq(expectedUri))).thenReturn(Multimaps.asMap(values));
+
+    assertNull(GoogleAuthLibraryCallCredentials.createJwtHelperOrNull(null));
+    GoogleAuthLibraryCallCredentials callCredentials =
+        new GoogleAuthLibraryCallCredentials(credentials, null);
+    callCredentials.applyRequestMetadata(method, attrs, executor, applier);
+    assertEquals(1, runPendingRunnables());
+
+    verify(credentials).getRequestMetadata(eq(expectedUri));
+    verify(applier).apply(headersCaptor.capture());
+    Metadata headers = headersCaptor.getValue();
+    Iterable<String> authorization = headers.getAll(AUTHORIZATION);
+    assertArrayEquals(new String[]{"token1"},
+        Iterables.toArray(authorization, String.class));
+  }
+
   private int runPendingRunnables() {
     ArrayList<Runnable> savedPendingRunnables = pendingRunnables;
     pendingRunnables = new ArrayList<Runnable>();