From 339f3641d99846f2e452f723def18b577f431ca0 Mon Sep 17 00:00:00 2001
From: "Wenlei (Frank) He" <wlhe@google.com>
Date: Fri, 8 Mar 2019 09:20:43 -0800
Subject: [PATCH] Support using ID token as the ouath access token.

---
 jwt/jwt.go | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/jwt/jwt.go b/jwt/jwt.go
index 6d5fd61..ba71ab2 100644
--- a/jwt/jwt.go
+++ b/jwt/jwt.go
@@ -66,9 +66,12 @@ type Config struct {
 	// request.  If empty, the value of TokenURL is used as the
 	// intended audience.
 	Audience string
-	
+
 	// PrivateClaims optionally specifies private claims in the JWT.
 	PrivateClaims map[string]interface{}
+
+	// UseIDToken optionally uses ID token instead of access token.
+	UseIDToken bool
 }
 
 // TokenSource returns a JWT TokenSource using the configuration
@@ -100,10 +103,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,
-		PrivateClaims:  js.conf.PrivateClaims,
+		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
@@ -168,6 +171,9 @@ func (js jwtSource) Token() (*oauth2.Token, error) {
 		if err != nil {
 			return nil, fmt.Errorf("oauth2: error decoding JWT token: %v", err)
 		}
+		if js.conf.UseIDToken {
+			token.AccessToken = tokenRes.IDToken
+		}
 		token.Expiry = time.Unix(claimSet.Exp, 0)
 	}
 	return token, nil