authhandler.go raw

   1  // Copyright 2021 The Go Authors. All rights reserved.
   2  // Use of this source code is governed by a BSD-style
   3  // license that can be found in the LICENSE file.
   4  
   5  // Package authhandler implements a TokenSource to support
   6  // "three-legged OAuth 2.0" via a custom AuthorizationHandler.
   7  package authhandler
   8  
   9  import (
  10  	"context"
  11  	"errors"
  12  
  13  	"golang.org/x/oauth2"
  14  )
  15  
  16  const (
  17  	// Parameter keys for AuthCodeURL method to support PKCE.
  18  	codeChallengeKey       = "code_challenge"
  19  	codeChallengeMethodKey = "code_challenge_method"
  20  
  21  	// Parameter key for Exchange method to support PKCE.
  22  	codeVerifierKey = "code_verifier"
  23  )
  24  
  25  // PKCEParams holds parameters to support PKCE.
  26  type PKCEParams struct {
  27  	Challenge       string // The unpadded, base64-url-encoded string of the encrypted code verifier.
  28  	ChallengeMethod string // The encryption method (ex. S256).
  29  	Verifier        string // The original, non-encrypted secret.
  30  }
  31  
  32  // AuthorizationHandler is a 3-legged-OAuth helper that prompts
  33  // the user for OAuth consent at the specified auth code URL
  34  // and returns an auth code and state upon approval.
  35  type AuthorizationHandler func(authCodeURL string) (code string, state string, err error)
  36  
  37  // TokenSourceWithPKCE is an enhanced version of [oauth2.TokenSource] with PKCE support.
  38  //
  39  // The pkce parameter supports PKCE flow, which uses code challenge and code verifier
  40  // to prevent CSRF attacks. A unique code challenge and code verifier should be generated
  41  // by the caller at runtime. See https://www.oauth.com/oauth2-servers/pkce/ for more info.
  42  func TokenSourceWithPKCE(ctx context.Context, config *oauth2.Config, state string, authHandler AuthorizationHandler, pkce *PKCEParams) oauth2.TokenSource {
  43  	return oauth2.ReuseTokenSource(nil, authHandlerSource{config: config, ctx: ctx, authHandler: authHandler, state: state, pkce: pkce})
  44  }
  45  
  46  // TokenSource returns an [oauth2.TokenSource] that fetches access tokens
  47  // using 3-legged-OAuth flow.
  48  //
  49  // The provided [context.Context] is used for oauth2 Exchange operation.
  50  //
  51  // The provided [oauth2.Config] should be a full configuration containing AuthURL,
  52  // TokenURL, and Scope.
  53  //
  54  // An environment-specific AuthorizationHandler is used to obtain user consent.
  55  //
  56  // Per the OAuth protocol, a unique "state" string should be specified here.
  57  // This token source will verify that the "state" is identical in the request
  58  // and response before exchanging the auth code for OAuth token to prevent CSRF
  59  // attacks.
  60  func TokenSource(ctx context.Context, config *oauth2.Config, state string, authHandler AuthorizationHandler) oauth2.TokenSource {
  61  	return TokenSourceWithPKCE(ctx, config, state, authHandler, nil)
  62  }
  63  
  64  type authHandlerSource struct {
  65  	ctx         context.Context
  66  	config      *oauth2.Config
  67  	authHandler AuthorizationHandler
  68  	state       string
  69  	pkce        *PKCEParams
  70  }
  71  
  72  func (source authHandlerSource) Token() (*oauth2.Token, error) {
  73  	// Step 1: Obtain auth code.
  74  	var authCodeUrlOptions []oauth2.AuthCodeOption
  75  	if source.pkce != nil && source.pkce.Challenge != "" && source.pkce.ChallengeMethod != "" {
  76  		authCodeUrlOptions = []oauth2.AuthCodeOption{oauth2.SetAuthURLParam(codeChallengeKey, source.pkce.Challenge),
  77  			oauth2.SetAuthURLParam(codeChallengeMethodKey, source.pkce.ChallengeMethod)}
  78  	}
  79  	url := source.config.AuthCodeURL(source.state, authCodeUrlOptions...)
  80  	code, state, err := source.authHandler(url)
  81  	if err != nil {
  82  		return nil, err
  83  	}
  84  	if state != source.state {
  85  		return nil, errors.New("state mismatch in 3-legged-OAuth flow")
  86  	}
  87  
  88  	// Step 2: Exchange auth code for access token.
  89  	var exchangeOptions []oauth2.AuthCodeOption
  90  	if source.pkce != nil && source.pkce.Verifier != "" {
  91  		exchangeOptions = []oauth2.AuthCodeOption{oauth2.SetAuthURLParam(codeVerifierKey, source.pkce.Verifier)}
  92  	}
  93  	return source.config.Exchange(source.ctx, code, exchangeOptions...)
  94  }
  95