From 950ef44c6e079baf075030377d90bf0c7e4b7b7a Mon Sep 17 00:00:00 2001
From: "Wenlei (Frank) He" <wlhe@google.com>
Date: Fri, 17 May 2019 17:56:57 +0000
Subject: [PATCH] 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: 7a6e247e68f742129ac9a5d5a5f1a8ad428ccb09
GitHub-Pull-Request: golang/oauth2#374
Reviewed-on: https://go-review.googlesource.com/c/oauth2/+/166220
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
---
 jwt/jwt.go      | 21 ++++++++++++++++++---
 jwt/jwt_test.go | 23 +++++++++++++++++++++++
 2 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/jwt/jwt.go b/jwt/jwt.go
index 99f3e0a3..b2bf1829 100644
--- a/jwt/jwt.go
+++ b/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
 }
diff --git a/jwt/jwt_test.go b/jwt/jwt_test.go
index 9dfa3b35..9772dc52 100644
--- a/jwt/jwt_test.go
+++ b/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:        "aaa2@xxx.com",
+			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)
+					}
+				}
+			}
 		})
 	}
 }