authhandler: Add support for PKCE

- Added new TokenSourceWithPKCE function to authhandler package.
- Updated Token method to support PKCE flow, sending code challenge and challenge method on the auth-code request, and sending code verifier on the exchange request.
- Updated google/default.go to support PKCE param.

Change-Id: Iab895bc01407c4742706061753f5329a772068ec
GitHub-Last-Rev: c1fddd2
GitHub-Pull-Request: #568
Reviewed-on: https://go-review.googlesource.com/c/oauth2/+/410515
Run-TryBot: Cody Oss <[email protected]>
TryBot-Result: Gopher Robot <[email protected]>
Reviewed-by: Shin Fan <[email protected]>
Reviewed-by: Cody Oss <[email protected]>

authhandler/authhandler.go

@@ -13,11 +13,36 @@ import (
 	"golang.org/x/oauth2"
 )
 
+const (
+	// Parameter keys for AuthCodeURL method to support PKCE.
+	codeChallengeKey       = "code_challenge"
+	codeChallengeMethodKey = "code_challenge_method"
+
+	// Parameter key for Exchange method to support PKCE.
+	codeVerifierKey = "code_verifier"
+)
+
+// PKCEParams holds parameters to support PKCE.
+type PKCEParams struct {
+	Challenge       string // The unpadded, base64-url-encoded string of the encrypted code verifier.
+	ChallengeMethod string // The encryption method (ex. S256).
+	Verifier        string // The original, non-encrypted secret.
+}
+
 // AuthorizationHandler is a 3-legged-OAuth helper that prompts
 // the user for OAuth consent at the specified auth code URL
 // and returns an auth code and state upon approval.
 type AuthorizationHandler func(authCodeURL string) (code string, state string, err error)
 
+// TokenSourceWithPKCE is an enhanced version of TokenSource with PKCE support.
+//
+// The pkce parameter supports PKCE flow, which uses code challenge and code verifier
+// to prevent CSRF attacks. A unique code challenge and code verifier should be generated
+// by the caller at runtime. See https://www.oauth.com/oauth2-servers/pkce/ for more info.
+func TokenSourceWithPKCE(ctx context.Context, config *oauth2.Config, state string, authHandler AuthorizationHandler, pkce *PKCEParams) oauth2.TokenSource {
+	return oauth2.ReuseTokenSource(nil, authHandlerSource{config: config, ctx: ctx, authHandler: authHandler, state: state, pkce: pkce})
+}
+
 // TokenSource returns an oauth2.TokenSource that fetches access tokens
 // using 3-legged-OAuth flow.
 //
@@ -33,24 +58,37 @@ type AuthorizationHandler func(authCodeURL string) (code string, state string, e
 // and response before exchanging the auth code for OAuth token to prevent CSRF
 // attacks.
 func TokenSource(ctx context.Context, config *oauth2.Config, state string, authHandler AuthorizationHandler) oauth2.TokenSource {
-	return oauth2.ReuseTokenSource(nil, authHandlerSource{config: config, ctx: ctx, authHandler: authHandler, state: state})
+	return TokenSourceWithPKCE(ctx, config, state, authHandler, nil)
 }
 
 type authHandlerSource struct {
 	ctx         context.Context
 	config      *oauth2.Config
 	authHandler AuthorizationHandler
 	state       string
+	pkce        *PKCEParams
 }
 
 func (source authHandlerSource) Token() (*oauth2.Token, error) {
-	url := source.config.AuthCodeURL(source.state)
+	// Step 1: Obtain auth code.
+	var authCodeUrlOptions []oauth2.AuthCodeOption
+	if source.pkce != nil && source.pkce.Challenge != "" && source.pkce.ChallengeMethod != "" {
+		authCodeUrlOptions = []oauth2.AuthCodeOption{oauth2.SetAuthURLParam(codeChallengeKey, source.pkce.Challenge),
+			oauth2.SetAuthURLParam(codeChallengeMethodKey, source.pkce.ChallengeMethod)}
+	}
+	url := source.config.AuthCodeURL(source.state, authCodeUrlOptions...)
 	code, state, err := source.authHandler(url)
 	if err != nil {
 		return nil, err
 	}
 	if state != source.state {
 		return nil, errors.New("state mismatch in 3-legged-OAuth flow")
 	}
-	return source.config.Exchange(source.ctx, code)
+
+	// Step 2: Exchange auth code for access token.
+	var exchangeOptions []oauth2.AuthCodeOption
+	if source.pkce != nil && source.pkce.Verifier != "" {
+		exchangeOptions = []oauth2.AuthCodeOption{oauth2.SetAuthURLParam(codeVerifierKey, source.pkce.Verifier)}
+	}
+	return source.config.Exchange(source.ctx, code, exchangeOptions...)
 }

authhandler/authhandler_test.go

@@ -97,3 +97,61 @@ func TestTokenExchange_StateMismatch(t *testing.T) {
 		t.Errorf("err = %q; want %q", err, want_err)
 	}
 }
+
+func TestTokenExchangeWithPKCE_Success(t *testing.T) {
+	authhandler := func(authCodeURL string) (string, string, error) {
+		if authCodeURL == "testAuthCodeURL?client_id=testClientID&code_challenge=codeChallenge&code_challenge_method=plain&response_type=code&scope=pubsub&state=testState" {
+			return "testCode", "testState", nil
+		}
+		return "", "", fmt.Errorf("invalid authCodeURL: %q", authCodeURL)
+	}
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		r.ParseForm()
+		if r.Form.Get("code") == "testCode" && r.Form.Get("code_verifier") == "codeChallenge" {
+			w.Header().Set("Content-Type", "application/json")
+			w.Write([]byte(`{
+				"access_token": "90d64460d14870c08c81352a05dedd3465940a7c",
+				"scope": "pubsub",
+				"token_type": "bearer",
+				"expires_in": 3600
+			}`))
+		}
+	}))
+	defer ts.Close()
+
+	conf := &oauth2.Config{
+		ClientID: "testClientID",
+		Scopes:   []string{"pubsub"},
+		Endpoint: oauth2.Endpoint{
+			AuthURL:  "testAuthCodeURL",
+			TokenURL: ts.URL,
+		},
+	}
+	pkce := PKCEParams{
+		Challenge:       "codeChallenge",
+		ChallengeMethod: "plain",
+		Verifier:        "codeChallenge",
+	}
+
+	tok, err := TokenSourceWithPKCE(context.Background(), conf, "testState", authhandler, &pkce).Token()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !tok.Valid() {
+		t.Errorf("got invalid token: %v", tok)
+	}
+	if got, want := tok.AccessToken, "90d64460d14870c08c81352a05dedd3465940a7c"; got != want {
+		t.Errorf("access token = %q; want %q", got, want)
+	}
+	if got, want := tok.TokenType, "bearer"; got != want {
+		t.Errorf("token type = %q; want %q", got, want)
+	}
+	if got := tok.Expiry.IsZero(); got {
+		t.Errorf("token expiry is zero = %v, want false", got)
+	}
+	scope := tok.Extra("scope")
+	if got, want := scope, "pubsub"; got != want {
+		t.Errorf("scope = %q; want %q", got, want)
+	}
+}

google/default.go

@@ -54,11 +54,14 @@ type CredentialsParams struct {
 	// Optional.
 	Subject string
 
-	// AuthHandler is the AuthorizationHandler used for 3-legged OAuth flow. Optional.
+	// AuthHandler is the AuthorizationHandler used for 3-legged OAuth flow. Required for 3LO flow.
 	AuthHandler authhandler.AuthorizationHandler
 
-	// State is a unique string used with AuthHandler. Optional.
+	// State is a unique string used with AuthHandler. Required for 3LO flow.
 	State string
+
+	// PKCE is used to support PKCE flow. Optional for 3LO flow.
+	PKCE *authhandler.PKCEParams
 }
 
 func (params CredentialsParams) deepCopy() CredentialsParams {
@@ -176,7 +179,7 @@ func CredentialsFromJSONWithParams(ctx context.Context, jsonData []byte, params
 	if config != nil {
 		return &Credentials{
 			ProjectID:   "",
-			TokenSource: authhandler.TokenSource(ctx, config, params.State, params.AuthHandler),
+			TokenSource: authhandler.TokenSourceWithPKCE(ctx, config, params.State, params.AuthHandler, params.PKCE),
 			JSON:        jsonData,
 		}, nil
 	}