Make AppEngineAuthCredentials Restorable

Author: Ajay Kannan

gcloud-java-core/src/main/java/com/google/gcloud/AuthCredentials.java

@@ -21,6 +21,7 @@
 import com.google.auth.oauth2.AccessToken;
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.auth.oauth2.ServiceAccountCredentials;
+import com.google.common.collect.ImmutableSet;
 
 import java.io.IOException;
 import java.io.InputStream;
@@ -37,32 +38,42 @@ public abstract class AuthCredentials implements Restorable<AuthCredentials> {
 
   private static class AppEngineAuthCredentials extends AuthCredentials {
 
-    private static final AuthCredentials INSTANCE = new AppEngineAuthCredentials();
-    private static final AppEngineAuthCredentialsState STATE = new AppEngineAuthCredentialsState();
+    private final Collection<String> scopes;
 
     private static class AppEngineCredentials extends GoogleCredentials {
 
       private final Object appIdentityService;
+      private final Method getAccessToken;
+      private final Method getAccessTokenResult;
       private final Collection<String> scopes;
       private final boolean scopesRequired;
 
-      AppEngineCredentials() {
+      AppEngineCredentials(Collection<String> scopes) {
         try {
           Class<?> factoryClass =
               Class.forName("com.google.appengine.api.appidentity.AppIdentityServiceFactory");
           Method method = factoryClass.getMethod("getAppIdentityService");
           this.appIdentityService = method.invoke(null);
-          this.scopes = null;
-          this.scopesRequired = true;
+          Class<?> serviceClass =
+              Class.forName("com.google.appengine.api.appidentity.AppIdentityService");
+          Class<?> tokenResultClass = Class.forName(
+              "com.google.appengine.api.appidentity.AppIdentityService$GetAccessTokenResult");
+          this.getAccessTokenResult = serviceClass.getMethod("getAccessToken", Iterable.class);
+          this.getAccessToken = tokenResultClass.getMethod("getAccessToken");
+          this.scopes = scopes;
+          this.scopesRequired = (scopes == null || scopes.isEmpty());
         } catch (Exception e) {
           throw new RuntimeException("Could not create AppEngineCredentials using reflection.");
         }
       }
 
-      AppEngineCredentials(Collection<String> scopes, Object appIdentityService) {
-          this.appIdentityService = appIdentityService;
-          this.scopes = scopes;
-          this.scopesRequired = (scopes == null || scopes.isEmpty());
+      AppEngineCredentials(Collection<String> scopes, Object appIdentityService,
+          Method getAccessToken, Method getAccessTokenResult) {
+        this.appIdentityService = appIdentityService;
+        this.getAccessToken = getAccessToken;
+        this.getAccessTokenResult = getAccessTokenResult;
+        this.scopes = scopes;
+        this.scopesRequired = (scopes == null || scopes.isEmpty());
       }
 
       /**
@@ -74,13 +85,7 @@ public AccessToken refreshAccessToken() throws IOException {
           throw new IOException("AppEngineCredentials requires createScoped call before use.");
         }
         try {
-          Class<?> serviceClass =
-              Class.forName("com.google.appengine.api.appidentity.AppIdentityService");
-          Class<?> tokenResultClass = Class.forName(
-              "com.google.appengine.api.appidentity.AppIdentityService$GetAccessTokenResult");
-          Method getAccessTokenResult = serviceClass.getMethod("getAccessToken", Iterable.class);
           Object accessTokenResult = getAccessTokenResult.invoke(appIdentityService, scopes);
-          Method getAccessToken = tokenResultClass.getMethod("getAccessToken");
           String accessToken = (String) getAccessToken.invoke(accessTokenResult);
           return new AccessToken(accessToken, null);
         } catch (Exception e) {
@@ -95,7 +100,8 @@ public boolean createScopedRequired() {
 
       @Override
       public GoogleCredentials createScoped(Collection<String> scopes) {
-        return new AppEngineCredentials(scopes, appIdentityService);
+        return new AppEngineCredentials(
+            scopes, appIdentityService, getAccessToken, getAccessTokenResult);
       }
     }
 
@@ -104,9 +110,15 @@ private static class AppEngineAuthCredentialsState
 
       private static final long serialVersionUID = 3558563960848658928L;
 
+      private final Collection<String> scopes;
+
+      private AppEngineAuthCredentialsState(Collection<String> scopes) {
+        this.scopes = scopes;
+      }
+
       @Override
       public AuthCredentials restore() {
-        return INSTANCE;
+        return new AppEngineAuthCredentials(scopes);
       }
 
       @Override
@@ -120,14 +132,18 @@ public boolean equals(Object obj) {
       }
     }
 
+    AppEngineAuthCredentials(Collection<String> scopes) {
+      this.scopes = scopes != null ? ImmutableSet.copyOf(scopes) : null;
+    }
+
     @Override
-    protected GoogleCredentials credentials() {
-      return new AppEngineCredentials();
+    public GoogleCredentials credentials() {
+      return new AppEngineCredentials(null);
     }
 
     @Override
     public RestorableState<AuthCredentials> capture() {
-      return STATE;
+      return new AppEngineAuthCredentialsState(scopes);
     }
   }
 
@@ -176,7 +192,7 @@ public boolean equals(Object obj) {
     }
 
     @Override
-    protected GoogleCredentials credentials() {
+    public GoogleCredentials credentials() {
       return new ServiceAccountCredentials(null, account, privateKey, null, null);
     }
 
@@ -232,29 +248,20 @@ public boolean equals(Object obj) {
     }
 
     @Override
-    protected GoogleCredentials credentials() {
+    public GoogleCredentials credentials() {
       return googleCredentials;
     }
 
-    public ServiceAccountAuthCredentials toServiceAccountCredentials() {
-      if (googleCredentials instanceof ServiceAccountCredentials) {
-        ServiceAccountCredentials credentials = (ServiceAccountCredentials) googleCredentials;
-        return new ServiceAccountAuthCredentials(credentials.getClientEmail(),
-            credentials.getPrivateKey());
-      }
-      return null;
-    }
-
     @Override
     public RestorableState<AuthCredentials> capture() {
       return STATE;
     }
   }
 
-  protected abstract GoogleCredentials credentials();
+  public abstract GoogleCredentials credentials();
 
   public static AuthCredentials createForAppEngine() {
-    return AppEngineAuthCredentials.INSTANCE;
+    return new AppEngineAuthCredentials(null);
   }
 
   /**
@@ -297,11 +304,12 @@ public static ServiceAccountAuthCredentials createFor(String account, PrivateKey
    * Account Authentication</a>.
    * </p>
    *
-   * @param jsonCredentialStream stream for Service Account Credentials in JSON format
+   * @param jsonCredentialStream stream for Service Account Credentials or User Credentials in JSON
+   *     format
    * @return the credentials instance.
    * @throws IOException if the credentials cannot be created from the stream.
    */
-  public static ServiceAccountAuthCredentials createForJson(InputStream jsonCredentialStream)
+  public static AuthCredentials createForJson(InputStream jsonCredentialStream)
       throws IOException {
     GoogleCredentials tempCredentials = GoogleCredentials.fromStream(jsonCredentialStream);
     if (tempCredentials instanceof ServiceAccountCredentials) {
@@ -310,9 +318,9 @@ public static ServiceAccountAuthCredentials createForJson(InputStream jsonCreden
       return new ServiceAccountAuthCredentials(
           tempServiceAccountCredentials.getClientEmail(),
           tempServiceAccountCredentials.getPrivateKey());
-    } else {
-      throw new IOException(
-          "The given JSON Credentials Stream is not a service account credential.");
     }
+    throw new IOException(
+        "The given JSON credentials stream could not be parsed as service account credentials or"
+        + " user credentials.");
   }
 }

gcloud-java-storage/src/main/java/com/google/gcloud/storage/StorageImpl.java

@@ -31,26 +31,26 @@
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.google.api.services.storage.model.StorageObject;
+import com.google.auth.oauth2.GoogleCredentials;
+import com.google.auth.oauth2.ServiceAccountCredentials;
 import com.google.common.base.Function;
 import com.google.common.base.Functions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
-import com.google.common.collect.Sets;
 import com.google.common.hash.Hashing;
 import com.google.common.io.BaseEncoding;
 import com.google.common.primitives.Ints;
 import com.google.gcloud.AuthCredentials;
-import com.google.gcloud.AuthCredentials.ApplicationDefaultAuthCredentials;
 import com.google.gcloud.AuthCredentials.ServiceAccountAuthCredentials;
-import com.google.gcloud.PageImpl;
 import com.google.gcloud.BaseService;
 import com.google.gcloud.ExceptionHandler;
 import com.google.gcloud.ExceptionHandler.Interceptor;
-import com.google.gcloud.RetryHelper.RetryHelperException;
 import com.google.gcloud.Page;
+import com.google.gcloud.PageImpl;
+import com.google.gcloud.RetryHelper.RetryHelperException;
 import com.google.gcloud.spi.StorageRpc;
 import com.google.gcloud.spi.StorageRpc.RewriteResponse;
 import com.google.gcloud.spi.StorageRpc.Tuple;
@@ -71,7 +71,6 @@
 import java.util.EnumMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.Callable;
 import java.util.concurrent.TimeUnit;
 
@@ -566,15 +565,13 @@ public URL signUrl(BlobInfo blobInfo, long duration, TimeUnit unit, SignUrlOptio
     ServiceAccountAuthCredentials cred =
         (ServiceAccountAuthCredentials) optionMap.get(SignUrlOption.Option.SERVICE_ACCOUNT_CRED);
     if (cred == null) {
-      AuthCredentials serviceCred = this.options().authCredentials();
-      if (serviceCred instanceof ServiceAccountAuthCredentials) {
-        cred = (ServiceAccountAuthCredentials) serviceCred;
-      } else {
-        if (serviceCred instanceof ApplicationDefaultAuthCredentials) {
-          cred = ((ApplicationDefaultAuthCredentials) serviceCred).toServiceAccountCredentials();
-        }
-      }
-      checkArgument(cred != null, "Signing key was not provided and could not be derived");
+      AuthCredentials authCredentials = this.options().authCredentials();
+      GoogleCredentials serviceCred =
+          authCredentials != null ? authCredentials.credentials() : null;
+      checkArgument(
+          serviceCred instanceof ServiceAccountCredentials,
+          "Signing key was not provided and could not be derived");
+      cred = (ServiceAccountAuthCredentials) authCredentials;
     }
     // construct signature - see https://cloud.google.com/storage/docs/access-control#Signed-URLs
     StringBuilder stBuilder = new StringBuilder();