From 209b4944feba1c62be2c9de4545e3b01a806b738 Mon Sep 17 00:00:00 2001
From: Mohan Li <67390330+mohanli-ml@users.noreply.github.com>
Date: Thu, 5 Aug 2021 12:03:16 -0700
Subject: [PATCH] feat: add allowNonDefaultServiceAccount option for DirectPath
 (#1433)

* feat: add allowNonDefaultServiceAccount option for DirectPath

* feat: add allowNonDefaultServiceAccount option for DirectPath
---
 .../InstantiatingGrpcChannelProvider.java     | 20 ++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java b/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java
index fe6ccb40e..225d76d6d 100644
--- a/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java
+++ b/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java
@@ -105,6 +105,7 @@ public final class InstantiatingGrpcChannelProvider implements TransportChannelP
   @Nullable private final Credentials credentials;
   @Nullable private final ChannelPrimer channelPrimer;
   @Nullable private final Boolean attemptDirectPath;
+  @Nullable private final Boolean allowNonDefaultServiceAccount;
   @VisibleForTesting final ImmutableMap<String, ?> directPathServiceConfig;
   @Nullable private final MtlsProvider mtlsProvider;
 
@@ -129,6 +130,7 @@ private InstantiatingGrpcChannelProvider(Builder builder) {
     this.credentials = builder.credentials;
     this.channelPrimer = builder.channelPrimer;
     this.attemptDirectPath = builder.attemptDirectPath;
+    this.allowNonDefaultServiceAccount = builder.allowNonDefaultServiceAccount;
     this.directPathServiceConfig =
         builder.directPathServiceConfig == null
             ? getDefaultDirectPathServiceConfig()
@@ -274,6 +276,13 @@ private boolean isDirectPathEnabled(String serviceAddress) {
     return false;
   }
 
+  private boolean isNonDefaultServiceAccountAllowed() {
+    if (allowNonDefaultServiceAccount != null && allowNonDefaultServiceAccount) {
+      return true;
+    }
+    return credentials instanceof ComputeEngineCredentials;
+  }
+
   // DirectPath should only be used on Compute Engine.
   // Notice Windows is supported for now.
   static boolean isOnComputeEngine() {
@@ -323,7 +332,7 @@ private ManagedChannel createSingleChannel() throws IOException, GeneralSecurity
 
     // TODO(weiranf): Add API in ComputeEngineCredentials to check default service account.
     if (isDirectPathEnabled(serviceAddress)
-        && credentials instanceof ComputeEngineCredentials
+        && isNonDefaultServiceAccountAllowed()
         && isOnComputeEngine()) {
       builder = ComputeEngineChannelBuilder.forAddress(serviceAddress, port);
       // Set default keepAliveTime and keepAliveTimeout when directpath environment is enabled.
@@ -435,6 +444,7 @@ public static final class Builder {
     @Nullable private Credentials credentials;
     @Nullable private ChannelPrimer channelPrimer;
     @Nullable private Boolean attemptDirectPath;
+    @Nullable private Boolean allowNonDefaultServiceAccount;
     @Nullable private ImmutableMap<String, ?> directPathServiceConfig;
 
     private Builder() {
@@ -459,6 +469,7 @@ private Builder(InstantiatingGrpcChannelProvider provider) {
       this.credentials = provider.credentials;
       this.channelPrimer = provider.channelPrimer;
       this.attemptDirectPath = provider.attemptDirectPath;
+      this.allowNonDefaultServiceAccount = provider.allowNonDefaultServiceAccount;
       this.directPathServiceConfig = provider.directPathServiceConfig;
       this.mtlsProvider = provider.mtlsProvider;
     }
@@ -654,6 +665,13 @@ public Builder setAttemptDirectPath(boolean attemptDirectPath) {
       return this;
     }
 
+    /** Whether allow non-default service account for DirectPath. */
+    @InternalApi("For internal use by google-cloud-java clients only")
+    public Builder setAllowNonDefaultServiceAccount(boolean allowNonDefaultServiceAccount) {
+      this.allowNonDefaultServiceAccount = allowNonDefaultServiceAccount;
+      return this;
+    }
+
     /**
      * Sets a service config for direct path. If direct path is not enabled, the provided service
      * config will be ignored.