feat(internal): [AIP-4111] support scopes for self-signed JWT auth flow

go.mod

@@ -10,7 +10,7 @@ require (
 	golang.org/x/lint v0.0.0-20210508222113-6edffad5e616
 	// TODO(codyoss): unfreeze after min version of 1.14
 	golang.org/x/net v0.0.0-20210503060351-7fd8e65b6420
-	golang.org/x/oauth2 v0.0.0-20210514164344-f6687ab2804c
+	golang.org/x/oauth2 v0.0.0-20210615190721-d04028783cf1
 	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
 	golang.org/x/sys v0.0.0-20210611083646-a4fc73990273
 	golang.org/x/tools v0.1.3

go.sum

@@ -247,8 +247,9 @@ golang.org/x/oauth2 v0.0.0-20201208152858-08078c50e5b5/go.mod h1:KelEdhl1UZF7XfJ
 golang.org/x/oauth2 v0.0.0-20210218202405-ba52d332ba99/go.mod h1:KelEdhl1UZF7XfJ4dDtk6s++YSgaE7mD/BuKKDLBl4A=
 golang.org/x/oauth2 v0.0.0-20210220000619-9bb904979d93/go.mod h1:KelEdhl1UZF7XfJ4dDtk6s++YSgaE7mD/BuKKDLBl4A=
 golang.org/x/oauth2 v0.0.0-20210313182246-cd4f82c27b84/go.mod h1:KelEdhl1UZF7XfJ4dDtk6s++YSgaE7mD/BuKKDLBl4A=
-golang.org/x/oauth2 v0.0.0-20210514164344-f6687ab2804c h1:pkQiBZBvdos9qq4wBAHqlzuZHEXo07pqV06ef90u1WI=
 golang.org/x/oauth2 v0.0.0-20210514164344-f6687ab2804c/go.mod h1:KelEdhl1UZF7XfJ4dDtk6s++YSgaE7mD/BuKKDLBl4A=
+golang.org/x/oauth2 v0.0.0-20210615190721-d04028783cf1 h1:x622Z2o4hgCr/4CiKWc51jHVKaWdtVpBNmEI8wI9Qns=
+golang.org/x/oauth2 v0.0.0-20210615190721-d04028783cf1/go.mod h1:KelEdhl1UZF7XfJ4dDtk6s++YSgaE7mD/BuKKDLBl4A=
 golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=

internal/creds.go

@@ -7,6 +7,7 @@ package internal
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io/ioutil"
 
@@ -62,56 +63,68 @@ const (
 	serviceAccountKey = "service_account"
 )
 
-// credentialsFromJSON returns a google.Credentials based on the input.
+// credentialsFromJSON returns a google.Credentials from the JSON data
 //
-// - A self-signed JWT auth flow will be executed if: the data file is a service
-//   account, no user are scopes provided, an audience is provided, a user
-//   specified endpoint is not provided, and credentials will not be
-//   impersonated.
+// - A self-signed JWT flow will be executed if the following conditions are
+// met:
+//   (1) At least one of the following is true:
+//       (a) Scope for self-signed JWT flow is enabled
+//       (b) Audiences are explicitly provided by users
+//       (b) No scope is provided
+//   (2) No service account impersontation
 //
-// - Otherwise, executes a stanard OAuth 2.0 flow.
+// - Otherwise, executes standard OAuth 2.0 flow
+// More details: google.aip.dev/auth/4111
 func credentialsFromJSON(ctx context.Context, data []byte, ds *DialSettings) (*google.Credentials, error) {
+	// By default, a standard OAuth 2.0 token source is created
 	cred, err := google.CredentialsFromJSON(ctx, data, ds.GetScopes()...)
 	if err != nil {
 		return nil, err
 	}
-	// Standard OAuth 2.0 Flow
-	if len(data) == 0 ||
-		len(ds.Scopes) > 0 ||
-		(ds.DefaultAudience == "" && len(ds.Audiences) == 0) ||
-		ds.ImpersonationConfig != nil ||
-		ds.Endpoint != "" {
-		return cred, nil
-	}
 
-	// Check if JSON is a service account and if so create a self-signed JWT.
-	var f struct {
-		Type string `json:"type"`
-		// The rest JSON fields are omitted because they are not used.
-	}
-	if err := json.Unmarshal(cred.JSON, &f); err != nil {
+	// Override the token source to use self-signed JWT if conditions are met
+	isJWTFlow, err := isSelfSignedJWTFlow(data, ds)
+	if err != nil {
 		return nil, err
 	}
-	if f.Type == serviceAccountKey {
-		ts, err := selfSignedJWTTokenSource(data, ds.DefaultAudience, ds.Audiences)
+	if isJWTFlow {
+		ts, err := selfSignedJWTTokenSource(data, ds)
 		if err != nil {
 			return nil, err
 		}
 		cred.TokenSource = ts
 	}
+
 	return cred, err
 }
 
-func selfSignedJWTTokenSource(data []byte, defaultAudience string, audiences []string) (oauth2.TokenSource, error) {
-	audience := defaultAudience
-	if len(audiences) > 0 {
-		// TODO(shinfan): Update golang oauth to support multiple audiences.
-		if len(audiences) > 1 {
-			return nil, fmt.Errorf("multiple audiences support is not implemented")
+func isSelfSignedJWTFlow(data []byte, ds *DialSettings) (bool, error) {
+	if (ds.EnableJwtWithScope || ds.HasCustomAudience() || len(ds.GetScopes()) == 0) &&
+		ds.ImpersonationConfig == nil {
+		// Check if JSON is a service account and if so create a self-signed JWT.
+		var f struct {
+			Type string `json:"type"`
+			// The rest JSON fields are omitted because they are not used.
 		}
-		audience = audiences[0]
+		if err := json.Unmarshal(data, &f); err != nil {
+			return false, err
+		}
+		return f.Type == serviceAccountKey, nil
+	}
+	return false, nil
+}
+
+func selfSignedJWTTokenSource(data []byte, ds *DialSettings) (oauth2.TokenSource, error) {
+	if len(ds.GetScopes()) > 0 && !ds.HasCustomAudience() {
+		// Scopes are preferred in self-signed JWT unless the scope is not available
+		// or a custom audience is used.
+		return google.JWTAccessTokenSourceWithScope(data, ds.GetScopes()...)
+	} else if ds.GetAudience() != "" {
+		// Fallback to audience if scope is not provided
+		return google.JWTAccessTokenSourceFromJSON(data, ds.GetAudience())
+	} else {
+		return nil, errors.New("neither scopes or audience are provided for the self-signed JWT")
 	}
-	return google.JWTAccessTokenSourceFromJSON(data, audience)
 }
 
 // QuotaProjectFromCreds returns the quota project from the JSON blob in the provided credentials.

internal/creds_test.go

@@ -38,6 +38,7 @@ func TestTokenSource(t *testing.T) {
 	ds = &DialSettings{
 		TokenSource:     ts,
 		CredentialsFile: "testdata/service-account.json",
+		DefaultScopes:   []string{"foo"},
 	}
 	got, err = Creds(ctx, ds)
 	if err != nil {
@@ -54,14 +55,20 @@ func TestDefaultServiceAccount(t *testing.T) {
 
 	// Load a valid JSON file. No way to really test the contents; we just
 	// verify that there is no error.
-	ds := &DialSettings{CredentialsFile: "testdata/service-account.json"}
+	ds := &DialSettings{
+		CredentialsFile: "testdata/service-account.json",
+		DefaultScopes:   []string{"foo"},
+	}
 	if _, err := Creds(ctx, ds); err != nil {
 		t.Errorf("got %v, wanted no error", err)
 	}
 
 	// Load valid JSON. No way to really test the contents; we just
 	// verify that there is no error.
-	ds = &DialSettings{CredentialsJSON: []byte(validServiceAccountJSON)}
+	ds = &DialSettings{
+		CredentialsJSON: []byte(validServiceAccountJSON),
+		DefaultScopes:   []string{"foo"},
+	}
 	if _, err := Creds(ctx, ds); err != nil {
 		t.Errorf("got %v, wanted no error", err)
 	}
@@ -85,6 +92,82 @@ func TestJWTWithAudience(t *testing.T) {
 	}
 }
 
+func TestJWTWithScope(t *testing.T) {
+	ctx := context.Background()
+
+	// Load a valid JSON file. No way to really test the contents; we just
+	// verify that there is no error.
+	ds := &DialSettings{
+		CredentialsFile:    "testdata/service-account.json",
+		Scopes:             []string{"foo"},
+		EnableJwtWithScope: true,
+	}
+	if _, err := Creds(ctx, ds); err != nil {
+		t.Errorf("got %v, wanted no error", err)
+	}
+
+	// Load valid JSON. No way to really test the contents; we just
+	// verify that there is no error.
+	ds = &DialSettings{
+		CredentialsJSON:    []byte(validServiceAccountJSON),
+		Scopes:             []string{"foo"},
+		EnableJwtWithScope: true,
+	}
+	if _, err := Creds(ctx, ds); err != nil {
+		t.Errorf("got %v, wanted no error", err)
+	}
+}
+
+func TestJWTWithDefaultScopes(t *testing.T) {
+	ctx := context.Background()
+
+	// Load a valid JSON file. No way to really test the contents; we just
+	// verify that there is no error.
+	ds := &DialSettings{
+		CredentialsFile:    "testdata/service-account.json",
+		DefaultScopes:      []string{"foo"},
+		EnableJwtWithScope: true,
+	}
+	if _, err := Creds(ctx, ds); err != nil {
+		t.Errorf("got %v, wanted no error", err)
+	}
+
+	// Load valid JSON. No way to really test the contents; we just
+	// verify that there is no error.
+	ds = &DialSettings{
+		CredentialsJSON:    []byte(validServiceAccountJSON),
+		DefaultScopes:      []string{"foo"},
+		EnableJwtWithScope: true,
+	}
+	if _, err := Creds(ctx, ds); err != nil {
+		t.Errorf("got %v, wanted no error", err)
+	}
+}
+
+func TestJWTWithDefaultAudience(t *testing.T) {
+	ctx := context.Background()
+
+	// Load a valid JSON file. No way to really test the contents; we just
+	// verify that there is no error.
+	ds := &DialSettings{
+		CredentialsFile: "testdata/service-account.json",
+		DefaultAudience: "foo",
+	}
+	if _, err := Creds(ctx, ds); err != nil {
+		t.Errorf("got %v, wanted no error", err)
+	}
+
+	// Load valid JSON. No way to really test the contents; we just
+	// verify that there is no error.
+	ds = &DialSettings{
+		CredentialsJSON: []byte(validServiceAccountJSON),
+		DefaultAudience: "foo",
+	}
+	if _, err := Creds(ctx, ds); err != nil {
+		t.Errorf("got %v, wanted no error", err)
+	}
+}
+
 func TestOAuth(t *testing.T) {
 	ctx := context.Background()
 
@@ -119,7 +202,13 @@ const validServiceAccountJSON = `{
 func TestQuotaProjectFromCreds(t *testing.T) {
 	ctx := context.Background()
 
-	cred, err := credentialsFromJSON(ctx, []byte(validServiceAccountJSON), &DialSettings{Endpoint: "foo.googleapis.com"})
+	cred, err := credentialsFromJSON(
+		ctx,
+		[]byte(validServiceAccountJSON),
+		&DialSettings{
+			Endpoint:      "foo.googleapis.com",
+			DefaultScopes: []string{"foo"},
+		})
 	if err != nil {
 		t.Fatalf("got %v, wanted no error", err)
 	}
@@ -133,7 +222,13 @@ func TestQuotaProjectFromCreds(t *testing.T) {
 	"quota_project_id": "foobar"
 }`)
 
-	cred, err = credentialsFromJSON(ctx, []byte(quotaProjectJSON), &DialSettings{Endpoint: "foo.googleapis.com"})
+	cred, err = credentialsFromJSON(
+		ctx,
+		[]byte(quotaProjectJSON),
+		&DialSettings{
+			Endpoint:      "foo.googleapis.com",
+			DefaultScopes: []string{"foo"},
+		})
 	if err != nil {
 		t.Fatalf("got %v, wanted no error", err)
 	}

internal/settings.go

@@ -24,6 +24,7 @@ type DialSettings struct {
 	DefaultMTLSEndpoint string
 	Scopes              []string
 	DefaultScopes       []string
+	EnableJwtWithScope  bool
 	TokenSource         oauth2.TokenSource
 	Credentials         *google.Credentials
 	CredentialsFile     string // if set, Token Source is ignored.
@@ -60,6 +61,19 @@ func (ds *DialSettings) GetScopes() []string {
 	return ds.DefaultScopes
 }
 
+// GetAudience returns the user-provided audience, if set, or else falls back to the default audience.
+func (ds *DialSettings) GetAudience() string {
+	if ds.HasCustomAudience() {
+		return ds.Audiences[0]
+	}
+	return ds.DefaultAudience
+}
+
+// HasCustomAudience returns true if a custom audience is provided by users.
+func (ds *DialSettings) HasCustomAudience() bool {
+	return len(ds.Audiences) > 0
+}
+
 // Validate reports an error if ds is invalid.
 func (ds *DialSettings) Validate() error {
 	if ds.SkipValidation {

option/internaloption/internaloption.go

@@ -94,3 +94,15 @@ func (w withDefaultScopes) Apply(o *internal.DialSettings) {
 	o.DefaultScopes = make([]string, len(w))
 	copy(o.DefaultScopes, w)
 }
+
+// EnableJwtWithScope returns a ClientOption that specifies if scope can be used
+// with self-signed JWT.
+func EnableJwtWithScope(enableScope bool) option.ClientOption {
+	return enableJwtWithScope(enableScope)
+}
+
+type enableJwtWithScope bool
+
+func (w enableJwtWithScope) Apply(o *internal.DialSettings) {
+	o.EnableJwtWithScope = bool(w)
+}