From 4676163c3e22ad6696252611a9fef93180d7d1d0 Mon Sep 17 00:00:00 2001 From: Julie Qiu Date: Wed, 27 Dec 2023 15:49:22 -0500 Subject: [PATCH] feat(transport): add support for API keys for gprc --- transport/grpc/dial.go | 111 +++++++++++++++++++++++------------- transport/grpc/dial_test.go | 29 ++++++++++ 2 files changed, 99 insertions(+), 41 deletions(-) diff --git a/transport/grpc/dial.go b/transport/grpc/dial.go index 34f9ba8bdb7..2a7f0ad0d7d 100644 --- a/transport/grpc/dial.go +++ b/transport/grpc/dial.go @@ -146,52 +146,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. } } @@ -249,6 +253,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 index 812ee0863ff..cf87e144e19 100644 --- 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) + } + } +}