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"
 
@@ -64,36 +65,30 @@ const (
 
 // credentialsFromJSON returns a google.Credentials based on the input.
 //
-// - 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 standard OAuth 2.0 flow will be executed if at least one of the
+//   following conditions is met:
+//   (1) the scope is non-empty and the scope for self-signed JWT flow is
+//       disabled.
+//   (2) Service Account Impersontation
 //
-// - Otherwise, executes a stanard OAuth 2.0 flow.
+// - Otherwise, executes a self-signed JWT flow (google.aip.dev/auth/4111)
 func credentialsFromJSON(ctx context.Context, data []byte, ds *DialSettings) (*google.Credentials, error) {
 	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 != "" {
+	if isOAuthFlow(data, ds) {
+		// Standard OAuth 2.0 Flow
 		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 {
+	isJWTFlow, err := isSelfSignedJWTFlow(cred.JSON)
+	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.GetAudience(), ds.GetScopes())
 		if err != nil {
 			return nil, err
 		}
@@ -102,16 +97,35 @@ func credentialsFromJSON(ctx context.Context, data []byte, ds *DialSettings) (*g
 	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")
-		}
-		audience = audiences[0]
+func isOAuthFlow(data []byte, ds *DialSettings) bool {
+	// Standard OAuth 2.0 Flow
+	return len(data) == 0 ||
+		(len(ds.GetScopes()) > 0 && !ds.EnableScopeForJWT) ||
+		ds.ImpersonationConfig != nil
+}
+
+func isSelfSignedJWTFlow(data []byte) (bool, error) {
+	// 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(data, &f); err != nil {
+		return false, err
+	}
+	return f.Type == serviceAccountKey, nil
+}
+
+func selfSignedJWTTokenSource(data []byte, audience string, scopes []string) (oauth2.TokenSource, error) {
+	if len(scopes) > 0 {
+		// Scopes are preferred in self-signed JWT
+		return google.JWTAccessTokenSourceWithScope(data, scopes...)
+	} else if audience != "" {
+		// Fallback to audience if scope is not provided
+		return google.JWTAccessTokenSourceFromJSON(data, audience)
+	} 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,32 @@ 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"},
+		EnableScopeForJWT: 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"},
+		EnableScopeForJWT: true,
+	}
+	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 +152,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 +172,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
+	EnableScopeForJWT   bool
 	TokenSource         oauth2.TokenSource
 	Credentials         *google.Credentials
 	CredentialsFile     string // if set, Token Source is ignored.
@@ -60,6 +61,14 @@ 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 len(ds.Audiences) > 0 {
+		return ds.Audiences[0]
+	}
+	return ds.DefaultAudience
+}
+
 // 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)
 }
+
+// EnableScopeForJWT returns a ClientOption that specifies if scope can be used
+// with self-signed JWT.
+func EnableScopeForJWT(scopeForJWT bool) ClientOption {
+	return enableScopeForJWT(audience)
+}
+
+type enableScopeForJWT bool
+
+func (w enableScopeForJWT) Apply(o *internal.DialSettings) {
+	o.EnableScopeForJWT = w
+}