feat(transport): add support for API keys for gprc (#2326)

fixes: #485
googleapis · Jan 4, 2024 · 9dbfb73 · 9dbfb73
1 parent 858fb57
commit 9dbfb73
Show file tree

Hide file tree

Showing 2 changed files with 99 additions and 41 deletions.
diff --git a/transport/grpc/dial.go b/transport/grpc/dial.go
@@ -168,52 +168,56 @@ func dial(ctx context.Context, insecure bool, o *internal.DialSettings) (*grpc.C
 	// when dialing an insecure connection?
 	if !o.NoAuth && !insecure {
 		if o.APIKey != "" {
-			log.Print("API keys are not supported for gRPC APIs. Remove the WithAPIKey option from your client-creating call.")
-		}
-		creds, err := internal.Creds(ctx, o)
-		if err != nil {
-			return nil, err
-		}
-
-		grpcOpts = append(grpcOpts,
-			grpc.WithPerRPCCredentials(grpcTokenSource{
-				TokenSource:   oauth.TokenSource{creds.TokenSource},
-				quotaProject:  internal.GetQuotaProject(creds, o.QuotaProject),
+			grpcOpts = append(grpcOpts, grpc.WithPerRPCCredentials(grpcAPIKey{
+				apiKey:        o.APIKey,
 				requestReason: o.RequestReason,
-			}),
-		)
-
-		// Attempt Direct Path:
-		logRateLimiter.Do(func() {
-			logDirectPathMisconfig(endpoint, creds.TokenSource, o)
-		})
-		if isDirectPathEnabled(endpoint, o) && isTokenSourceDirectPathCompatible(creds.TokenSource, o) && metadata.OnGCE() {
-			// Overwrite all of the previously specific DialOptions, DirectPath uses its own set of credentials and certificates.
-			grpcOpts = []grpc.DialOption{
-				grpc.WithCredentialsBundle(grpcgoogle.NewDefaultCredentialsWithOptions(grpcgoogle.DefaultCredentialsOptions{oauth.TokenSource{creds.TokenSource}}))}
-			if timeoutDialerOption != nil {
-				grpcOpts = append(grpcOpts, timeoutDialerOption)
+			}))
+		} else {
+			creds, err := internal.Creds(ctx, o)
+			if err != nil {
+				return nil, err
 			}
-			// Check if google-c2p resolver is enabled for DirectPath
-			if isDirectPathXdsUsed(o) {
-				// google-c2p resolver target must not have a port number
-				if addr, _, err := net.SplitHostPort(endpoint); err == nil {
-					endpoint = "google-c2p:///" + addr
-				} else {
-					endpoint = "google-c2p:///" + endpoint
+			grpcOpts = append(grpcOpts, grpc.WithPerRPCCredentials(grpcTokenSource{
+				TokenSource:   oauth.TokenSource{TokenSource: creds.TokenSource},
+				quotaProject:  internal.GetQuotaProject(creds, o.QuotaProject),
+				requestReason: o.RequestReason,
+			}))
+			// Attempt Direct Path:
+			logRateLimiter.Do(func() {
+				logDirectPathMisconfig(endpoint, creds.TokenSource, o)
+			})
+			if isDirectPathEnabled(endpoint, o) && isTokenSourceDirectPathCompatible(creds.TokenSource, o) && metadata.OnGCE() {
+				// Overwrite all of the previously specific DialOptions, DirectPath uses its own set of credentials and certificates.
+				grpcOpts = []grpc.DialOption{
+					grpc.WithCredentialsBundle(grpcgoogle.NewDefaultCredentialsWithOptions(
+						grpcgoogle.DefaultCredentialsOptions{
+							PerRPCCreds: oauth.TokenSource{TokenSource: creds.TokenSource},
+						})),
 				}
-			} else {
-				if !strings.HasPrefix(endpoint, "dns:///") {
-					endpoint = "dns:///" + endpoint
+				if timeoutDialerOption != nil {
+					grpcOpts = append(grpcOpts, timeoutDialerOption)
 				}
-				grpcOpts = append(grpcOpts,
-					// For now all DirectPath go clients will be using the following lb config, but in future
-					// when different services need different configs, then we should change this to a
-					// per-service config.
-					grpc.WithDisableServiceConfig(),
-					grpc.WithDefaultServiceConfig(`{"loadBalancingConfig":[{"grpclb":{"childPolicy":[{"pick_first":{}}]}}]}`))
+				// Check if google-c2p resolver is enabled for DirectPath
+				if isDirectPathXdsUsed(o) {
+					// google-c2p resolver target must not have a port number
+					if addr, _, err := net.SplitHostPort(endpoint); err == nil {
+						endpoint = "google-c2p:///" + addr
+					} else {
+						endpoint = "google-c2p:///" + endpoint
+					}
+				} else {
+					if !strings.HasPrefix(endpoint, "dns:///") {
+						endpoint = "dns:///" + endpoint
+					}
+					grpcOpts = append(grpcOpts,
+						// For now all DirectPath go clients will be using the following lb config, but in future
+						// when different services need different configs, then we should change this to a
+						// per-service config.
+						grpc.WithDisableServiceConfig(),
+						grpc.WithDefaultServiceConfig(`{"loadBalancingConfig":[{"grpclb":{"childPolicy":[{"pick_first":{}}]}}]}`))
+				}
+				// TODO(cbro): add support for system parameters (quota project, request reason) via chained interceptor.
 			}
-			// TODO(cbro): add support for system parameters (quota project, request reason) via chained interceptor.
 		}
 	}
 
@@ -271,6 +275,31 @@ func (ts grpcTokenSource) GetRequestMetadata(ctx context.Context, uri ...string)
 	return metadata, nil
 }
 
+// grpcAPIKey supplies PerRPCCredentials from an API Key.
+type grpcAPIKey struct {
+	apiKey string
+
+	// Additional metadata attached as headers.
+	requestReason string
+}
+
+// GetRequestMetadata gets the request metadata as a map from a grpcAPIKey.
+func (ts grpcAPIKey) GetRequestMetadata(ctx context.Context, uri ...string) (
+	map[string]string, error) {
+	metadata := map[string]string{
+		"X-goog-api-key": ts.apiKey,
+	}
+	if ts.requestReason != "" {
+		metadata["X-goog-request-reason"] = ts.requestReason
+	}
+	return metadata, nil
+}
+
+// RequireTransportSecurity indicates whether the credentials requires transport security.
+func (ts grpcAPIKey) RequireTransportSecurity() bool {
+	return true
+}
+
 func isDirectPathEnabled(endpoint string, o *internal.DialSettings) bool {
 	if !o.EnableDirectPath {
 		return false

diff --git a/transport/grpc/dial_test.go b/transport/grpc/dial_test.go
@@ -12,6 +12,7 @@ import (
 	"testing"
 
 	"cloud.google.com/go/compute/metadata"
+	"github.com/google/go-cmp/cmp"
 	"golang.org/x/oauth2/google"
 	"google.golang.org/api/internal"
 	"google.golang.org/grpc"
@@ -143,3 +144,31 @@ func TestLogDirectPathMisconfigNotOnGCE(t *testing.T) {
 	}
 
 }
+
+func TestGRPCAPIKey_GetRequestMetadata(t *testing.T) {
+	for _, test := range []struct {
+		apiKey string
+		reason string
+	}{
+		{
+			apiKey: "MY_API_KEY",
+			reason: "MY_REQUEST_REASON",
+		},
+	} {
+		ts := grpcAPIKey{
+			apiKey:        test.apiKey,
+			requestReason: test.reason,
+		}
+		got, err := ts.GetRequestMetadata(context.Background())
+		if err != nil {
+			t.Fatal(err)
+		}
+		want := map[string]string{
+			"X-goog-api-key":        ts.apiKey,
+			"X-goog-request-reason": ts.requestReason,
+		}
+		if diff := cmp.Diff(want, got); diff != "" {
+			t.Errorf("mismatch (-want +got):\n%s", diff)
+		}
+	}
+}