jwt: support PrivateClaims in Config

This would help add extra claim for certain 2-leg JWT exchange.

For example, Google service account key can be used to generate an OIDC token, but Google TokenURL requires "target_audience" claims set.

See this example usage:
https://gist.github.com/wlhee/64bc518190053e2122ca1909c2977c67#file-exmaple-go-L29

Change-Id: Ic10b006e45a34210634c5a76261a7e3706066965
GitHub-Last-Rev: 7a6e247
GitHub-Pull-Request: #374
Reviewed-on: https://go-review.googlesource.com/c/oauth2/+/166220
Reviewed-by: Brad Fitzpatrick <[email protected]>
Run-TryBot: Brad Fitzpatrick <[email protected]>
TryBot-Result: Gobot Gobot <[email protected]>

Author: wlhee

jwt/jwt.go

@@ -66,6 +66,14 @@ type Config struct {
 	// request.  If empty, the value of TokenURL is used as the
 	// intended audience.
 	Audience string
+
+	// PrivateClaims optionally specifies custom private claims in the JWT.
+	// See http://tools.ietf.org/html/draft-jones-json-web-token-10#section-4.3
+	PrivateClaims map[string]interface{}
+
+	// UseIDToken optionally specifies whether ID token should be used instead
+	// of access token when the server returns both.
+	UseIDToken bool
 }
 
 // TokenSource returns a JWT TokenSource using the configuration
@@ -97,9 +105,10 @@ func (js jwtSource) Token() (*oauth2.Token, error) {
 	}
 	hc := oauth2.NewClient(js.ctx, nil)
 	claimSet := &jws.ClaimSet{
-		Iss:   js.conf.Email,
-		Scope: strings.Join(js.conf.Scopes, " "),
-		Aud:   js.conf.TokenURL,
+		Iss:           js.conf.Email,
+		Scope:         strings.Join(js.conf.Scopes, " "),
+		Aud:           js.conf.TokenURL,
+		PrivateClaims: js.conf.PrivateClaims,
 	}
 	if subject := js.conf.Subject; subject != "" {
 		claimSet.Sub = subject
@@ -166,5 +175,11 @@ func (js jwtSource) Token() (*oauth2.Token, error) {
 		}
 		token.Expiry = time.Unix(claimSet.Exp, 0)
 	}
+	if js.conf.UseIDToken {
+		if tokenRes.IDToken == "" {
+			return nil, fmt.Errorf("oauth2: response doesn't have JWT token")
+		}
+		token.AccessToken = tokenRes.IDToken
+	}
 	return token, nil
 }

jwt/jwt_test.go

@@ -11,6 +11,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"reflect"
 	"strings"
 	"testing"
 
@@ -221,6 +222,16 @@ func TestJWTFetch_AssertionPayload(t *testing.T) {
 			TokenURL:     ts.URL,
 			Audience:     "https://example.com",
 		},
+		{
+			Email:        "[email protected]",
+			PrivateKey:   dummyPrivateKey,
+			PrivateKeyID: "ABCDEFGHIJKLMNOPQRSTUVWXYZ",
+			TokenURL:     ts.URL,
+			PrivateClaims: map[string]interface{}{
+				"private0": "claim0",
+				"private1": "claim1",
+			},
+		},
 	} {
 		t.Run(conf.Email, func(t *testing.T) {
 			_, err := conf.TokenSource(context.Background()).Token()
@@ -261,6 +272,18 @@ func TestJWTFetch_AssertionPayload(t *testing.T) {
 			if got, want := claimSet.Prn, conf.Subject; got != want {
 				t.Errorf("payload prn = %q; want %q", got, want)
 			}
+			if len(conf.PrivateClaims) > 0 {
+				var got interface{}
+				if err := json.Unmarshal(gotjson, &got); err != nil {
+					t.Errorf("failed to parse payload; err = %q", err)
+				}
+				m := got.(map[string]interface{})
+				for v, k := range conf.PrivateClaims {
+					if !reflect.DeepEqual(m[v], k) {
+						t.Errorf("payload private claims key = %q: got %#v; want %#v", v, m[v], k)
+					}
+				}
+			}
 		})
 	}
 }