Support QuotaUser/QuotaProject/IAMAuthToken

- Attach fields into headers in both HTTP and gRPC
- Refactor the gRPC token source code for better readability
- Move HTTP transport code into internal/http.go

go/sgauth/default.go

@@ -15,7 +15,6 @@
 package sgauth
 
 import (
-	gRPCCredentials "google.golang.org/grpc/credentials"
 	"context"
 	"os"
 	"fmt"
@@ -34,7 +33,7 @@ import (
 // "Application Default Credentials".
 // It is a shortcut for FindDefaultCredentials(ctx, scope).TokenSource.
 func DefaultTokenSource(ctx context.Context, scope string) (internal.TokenSource, error) {
-	creds, err := applicationDefaultCredentials(ctx, &Settings{Scope:scope,})
+	creds, err := applicationDefaultCredentials(ctx, &Settings{Scope: scope,})
 	if err != nil {
 		return nil, err
 	}
@@ -59,26 +58,6 @@ func JWTTokenSource(ctx context.Context, settings *Settings) (internal.TokenSour
 	return ts, err
 }
 
-func NewGrpcApplicationDefault(ctx context.Context, settings *Settings) (gRPCCredentials.PerRPCCredentials, error) {
-	t, err := DefaultTokenSource(ctx, settings.Scope)
-	if err != nil {
-		return nil, err
-	}
-	return internal.GrpcTokenSource{t}, nil
-}
-
-func NewGrpcJWT(ctx context.Context, aud string) (gRPCCredentials.PerRPCCredentials, error) {
-	creds, err := applicationDefaultCredentials(ctx, &Settings{})
-	if creds != nil {
-		ts, err := credentials.JWTAccessTokenSourceFromJSON(creds.JSON, aud)
-		if (err != nil) {
-			return nil, err
-		}
-		return internal.GrpcTokenSource{ts}, nil
-	}
-	return nil, err
-}
-
 func findJSONCredentials(ctx context.Context, settings *Settings) (*credentials.Credentials, error) {
 	if settings.CredentialsJSON != "" {
 		return credentialsFromJSON(ctx, []byte(settings.CredentialsJSON),
@@ -139,7 +118,7 @@ func readCredentialsFile(ctx context.Context, filename string, settings *Setting
 }
 
 func credentialsFromJSON(ctx context.Context, jsonData []byte, scopes []string,
-	handler func(string)(string, error), state string) (*credentials.Credentials, error) {
+	handler func(string) (string, error), state string) (*credentials.Credentials, error) {
 	var f credentials.File
 	if err := json.Unmarshal(jsonData, &f); err != nil {
 		return nil, err

go/sgauth/grpc.go

@@ -26,22 +26,30 @@ import (
 
 func NewGrpcConn(ctx context.Context, settings *Settings, host string, port string) (*grpc.ClientConn, error) {
 	if settings == nil {
-		settings = &Settings {
+		settings = &Settings{
 			Scope: DefaultScope,
 		}
 	}
 
 	pool, _ := x509.SystemCertPool()
 	// error handling omitted
 	creds := credentials.NewClientTLSFromCert(pool, "")
-	var perRPC credentials.PerRPCCredentials
+	perRPC := internal.GrpcTokenSource{
+		QuotaUser:    settings.QuotaUser,
+		QuotaProject: settings.QuotaProject,
+		IAMAuthToken: settings.IAMAuthToken,
+	}
 
 	if settings.APIKey != "" {
-		perRPC = internal.GrpcApiKey{Value: settings.APIKey}
-	} else if settings.Scope != "" {
-		perRPC, _ = NewGrpcApplicationDefault(ctx, settings)
+		// API key
+		perRPC.ApiKey = settings.APIKey
 	} else {
-		perRPC, _ = NewGrpcJWT(ctx, settings.Audience)
+		// OAuth or JWT token
+		ts, err := newTokenSource(ctx, settings)
+		if err != nil {
+			return nil, err
+		}
+		perRPC.Source = *ts
 	}
 	return grpc.Dial(
 		fmt.Sprintf("%s:%s", host, port),

go/sgauth/http.go

@@ -20,46 +20,29 @@ import (
 	"golang.org/x/net/context"
 )
 
-// createClient creates an *http.Client from a TokenSource.
-// The returned client is not valid beyond the lifetime of the context.
-func createAuthTokenClient(src internal.TokenSource) *http.Client {
-	if src == nil {
-		return http.DefaultClient
-	}
-	return &http.Client{
-		Transport: &Transport{
-			Base:   http.DefaultClient.Transport,
-			Source: internal.ReuseTokenSource(nil, src),
-		},
-	}
-}
-
-// createClient creates an *http.Client from a TokenSource.
-// The returned client is not valid beyond the lifetime of the context.
-func createAPIKeyClient(key string) *http.Client {
-	if key == "" {
-		return http.DefaultClient
-	}
-	return &http.Client{
-		Transport: &Transport{
-			Base:   http.DefaultClient.Transport,
-			APIKey: key,
-		},
-	}
-}
-
 var DefaultScope = "https://www.googleapis.com/auth/cloud-platform"
 
-// DefaultClient returns an HTTP Client that uses the
-// DefaultTokenSource to obtain authentication credentials.
+// Returns the HTTP client using the given settings.
 func NewHTTPClient(ctx context.Context, settings *Settings) (*http.Client, error) {
+	transport := &internal.Transport{
+		Base:         http.DefaultClient.Transport,
+		QuotaUser:    settings.QuotaUser,
+		QuotaProject: settings.QuotaProject,
+		IAMAuthToken: settings.IAMAuthToken,
+	}
+
 	if settings.APIKey != "" {
-		return createAPIKeyClient(settings.APIKey), nil
+		// API key
+		transport.APIKey = settings.APIKey
 	} else {
+		// OAuth or JWT token
 		ts, err := newTokenSource(ctx, settings)
 		if err != nil {
 			return nil, err
 		}
-		return createAuthTokenClient(*ts), nil
+		transport.Source = ts
 	}
+	return &http.Client{
+		Transport: transport,
+	}, nil
 }

go/sgauth/internal/api_key.go → go/sgauth/internal/const.go

@@ -14,21 +14,11 @@
 // limitations under the License.
 package internal
 
-import "golang.org/x/net/context"
-
-// TokenSource supplies PerRPCCredentials from an oauth2.TokenSource.
-type GrpcApiKey struct {
-	Value string
-}
-
-// GetRequestMetadata gets the request metadata as a map from a TokenSource.
-func (key GrpcApiKey) GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error) {
-	return map[string]string{
-		"x-goog-api-key": key.Value,
-	}, nil
-}
-
-// RequireTransportSecurity indicates whether the credentials requires transport security.
-func (key GrpcApiKey) RequireTransportSecurity() bool {
-	return true
-}
+// Keys for HTTP google headers
+const (
+	headerAuth = "authorization"
+	headerApiKey = "X-Goog-Api-Key"
+	headerQuotaUser = "X-Goog-QuotaUser"
+	headerQuotaProject = "X-Goog-User-Project"
+	headerIAMAuthToken = "X-Goog-IAM-Authorization-Token"
+)

go/sgauth/internal/grpc.go

@@ -0,0 +1,61 @@
+//
+// Copyright 2018 Google Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package internal
+
+import "context"
+
+// TokenSource supplies PerRPCCredentials from an oauth2.TokenSource.
+type GrpcTokenSource struct {
+	Source TokenSource
+	ApiKey string
+
+	// Additional metadata attached as headers
+	QuotaUser string
+	QuotaProject string
+	IAMAuthToken string
+}
+
+// GetRequestMetadata gets the request metadata as a map from a TokenSource.
+func (ts GrpcTokenSource) GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error) {
+	metadata := map[string]string{}
+	if ts.ApiKey != "" {
+		metadata[headerApiKey] = ts.ApiKey
+	} else {
+		token, err := ts.Source.Token()
+		if err != nil {
+			return nil, err
+		}
+		metadata[headerAuth] = token.Type() + " " + token.AccessToken
+	}
+	attachAdditionalMetadata(metadata, ts)
+	return metadata, nil
+}
+
+// RequireTransportSecurity indicates whether the credentials requires transport security.
+func (ts GrpcTokenSource) RequireTransportSecurity() bool {
+	return true
+}
+
+func attachAdditionalMetadata(metadata map[string]string, ts GrpcTokenSource) {
+	if ts.QuotaUser != "" {
+		metadata[headerQuotaUser] = ts.QuotaUser
+	}
+	if ts.QuotaProject != "" {
+		metadata[headerQuotaProject] = ts.QuotaProject
+	}
+	if ts.IAMAuthToken != "" {
+		metadata[headerIAMAuthToken] = ts.IAMAuthToken
+	}
+}

go/sgauth/transport.go → go/sgauth/internal/http.go

@@ -12,14 +12,13 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
-package sgauth
+package internal
 
 import (
 	"errors"
 	"io"
 	"net/http"
 	"sync"
-	"github.com/google/oauth2l/go/sgauth/internal"
 )
 
 // Transport is an http.RoundTripper that makes OAuth 2.0 HTTP requests,
@@ -31,10 +30,14 @@ import (
 type Transport struct {
 	// Source supplies the token to add to outgoing requests'
 	// Authorization headers.
-	Source internal.TokenSource
-
+	Source TokenSource
 	APIKey string
 
+	// Additional metadata attached as headers
+	QuotaUser    string
+	QuotaProject string
+	IAMAuthToken string
+
 	// Base is the base RoundTripper used to make HTTP requests.
 	// If nil, http.DefaultTransport is used.
 	Base http.RoundTripper
@@ -54,10 +57,10 @@ func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
 			}
 		}()
 	}
-	
+
 	req2 := cloneRequest(req) // per RoundTripper contract
 	if t.APIKey != "" {
-		req2.Header.Set("X-Goog-Api-Key", t.APIKey)
+		req2.Header.Set(headerApiKey, t.APIKey)
 	} else {
 		if t.Source == nil {
 			return nil, errors.New("oauth2: Transport's Source is nil")
@@ -68,6 +71,9 @@ func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
 		}
 		token.SetAuthHeader(req2)
 	}
+
+	attachAdditionalHeaders(t, req2)
+
 	t.setModReq(req, req2)
 	res, err := t.base().RoundTrip(req2)
 
@@ -158,3 +164,15 @@ func (r *onEOFReader) runFunc() {
 		r.fn = nil
 	}
 }
+
+func attachAdditionalHeaders(t *Transport, req *http.Request) {
+	if (t.QuotaUser != "") {
+		req.Header.Set(headerQuotaUser, t.QuotaUser)
+	}
+	if (t.QuotaProject != "") {
+		req.Header.Set(headerQuotaProject, t.QuotaProject)
+	}
+	if (t.IAMAuthToken != "") {
+		req.Header.Set(headerIAMAuthToken, t.IAMAuthToken)
+	}
+}

go/sgauth/internal/token.go

@@ -195,25 +195,3 @@ func (tf *tokenRefresher) Token() (*Token, error) {
 	}
 	return tk, err
 }
-
-// TokenSource supplies PerRPCCredentials from an oauth2.TokenSource.
-type GrpcTokenSource struct {
-	TokenSource
-}
-
-// GetRequestMetadata gets the request metadata as a map from a TokenSource.
-func (ts GrpcTokenSource) GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error) {
-	token, err := ts.Token()
-	if err != nil {
-		return nil, err
-	}
-	return map[string]string{
-		"authorization": token.Type() + " " + token.AccessToken,
-	}, nil
-}
-
-// RequireTransportSecurity indicates whether the credentials requires transport security.
-func (ts GrpcTokenSource) RequireTransportSecurity() bool {
-	return true
-}
-

go/sgauth/settings.go

@@ -35,11 +35,11 @@ type Settings struct {
 	// We should have named them as quotaUser and quotaProject.
 	QuotaUser string
 	QuotaProject string
-
+	// IAM context attributes
+	IAMAuthToken string
 	// End-user OAuth Flow handler that redirects the user to the given URL
 	// and returns the token.
 	OAuthFlowHandler func(url string) (token string, err error)
-
 	// The state string used for 3LO session verification.
 	State string
 }