Go: Oauth2 クライアント TokenSourceの実装

[golang.org/x/oauth2] (https://pkg.go.dev/golang.org/x/oauth2) の実装を追っていきます。

1. AuthCodeURL の発行
2. Exchange での authorization_code から、 oauth2.Token を取得

については触れるだけに留めて、取得した oauth2.Token のリフレッシュについて主に書いていきます。

API コールまで

初回のAPIコールまでを簡略化したコードで示します。

1. 認可ページの表示
2. コールバックで、Authorization Code を受け取る
3. トークンエンドポイントにリクエストを投げて、アクセストークン・リフレッシュトークンを取得した

package main

import (
	"context"
	"fmt"
	"golang.org/x/oauth2"
	"log"
)

func main() {
	ctx := context.Background()
	conf := &oauth2.Config{
		ClientID:     "YOUR_CLIENT_ID",
		ClientSecret: "YOUR_CLIENT_SECRET",
		Scopes:       []string{"SCOPE"},
		Endpoint: oauth2.Endpoint{
			AuthURL:  "https://provider.com/o/oauth2/auth",
			TokenURL: "https://provider.com/o/oauth2/token",
		},
	}
	// 1. 認可ページのURL
	url := conf.AuthCodeURL("your_state")
	// 2. Authorization Code を受け取る
	// ...
	// 3. トークンエンドポイントにリクエスト、アクセストークン・リフレッシュトークンを取得
	oauth2Token, err := conf.Exchange(ctx, code)
	if err != nil {
		log.Fatal(err)
	}

	client := conf.Client(ctx, oauth2Token)
	resp, err := client.Get("https://api.provider.com/v1/me")
	// ...
}

ここから本題に入っていきます。
この時取得した oauth2.Token を次回以降のリクエストにも使用する実装について考えていきます。

シナリオ

1. oauth2.Token を取得
2. oauth2.Token （アクセストークン, リフレッシュトークン） を DBなどのストレージに保存
3. 次回リクエスト時、ストレージから取得した oauth2.Token を使ってリクエストを行う

上のシナリオでアクセストークンの有効期限が切れていた場合、
grant_type=refresh_token として TokenURL にリクエストを行い、新たなアクセストークンを発行する必要があります。

シンプルな実装の例

// persistent storage から, oauth2.Token を取得
tokenSource := conf.TokenSource(ctx, token)
// アクセストークンの有効期限が切れていた場合、refresh_token から新たなアクセストークンを発行する
client := oauth2.NewClient(ctx, tokenSource)
resp, err := client.Get("https://api.provider.com/v1/me")
// ...
// 有効な oaauth2.Token を storage に保存するなど
token, err = tokenSrouce.Token()
// ...

oauth2.NewClient は、アクセストークンの有効期限が切れていた場合のリフレッシュを
Transport 層で行うよう実装された http.Client を返します。
この Client を使ってリクエストすれば、token のリフレッシュを気にしなくてよくなります。
RoundTripper がうまく使われていることが分かります。

ここで、Transport.RoundTrip() 内でどのようにして tokenの取得を行っているかを見ていきます。

type TokenSource

oauth2.Token とエラーを返すシンプルなメソッドを実装するインターフェイスです。
Transport 層でこのメソッドが呼ばれ、token の取得とセットが行われています。

type TokenSource interface {
	// Token returns a token or an error.
	// Token must be safe for concurrent use by multiple goroutines.
	// The returned Token must not be modified.
	Token() (*Token, error)
}

TokenSource の実装には以下のパターンがあります。

1. 上の実装例のように conf.TokenSource(ctx, token) を使う

2. oauth2.StaticTokenSource(token) を使う
   　※ これは常に同じ token を返すため、今回のシナリオにはそぐわない

3. インターフェイスを満たすようスクラッチで実装する
   　golang.org/x/oauth2/google の ComputeTokenSource など

TokenSource を実装してみる

例えば、リフレッシュしたタイミングでなにかの処理を実行するようにカスタムすることも可能です。
(Transport.RoundTrip() で呼ばれるため、Token()メソッドは concurrency-safe になるよう注意が必要です。)

type MyFunc func(*oauth2.Token) error

func Do(t *oauth2.Token) error {
	fmt.Println("refreshed!")
	return nil
}

type MyTokenSource struct {
	src oauth2.TokenSource
	f   MyFunc
}

func (s *MyTokenSource) Token() (*oauth2.Token, error) {
    t, err := s.src.Token()
    if err != nil {
        return nil, err
    }
    if err = s.f(t); err != nil {
        return t, err
    }
    return t, nil
}

src := conf.TokenSource(ctx, token)
mySrc := &MyTokenSource {
    src: src,
    f:   Do,
}
reuseSrc := oauth2.ReuseTokenSource(token, mySrc)
client := oauth2.NewClient(ctx, reuseSrc)

補足

oauth2.ReuseTokenSource で、mySrcをラップして使用しています。
下記のように NewClient は、 token を nil として初期化されます。
初回リクエスト時、 token が有効であるなら MyTokenSource.Token() が呼ばれるのを防ぐためにわざわざラップして NewClient に渡しています。

func NewClient(ctx context.Context, src TokenSource) *http.Client {
	if src == nil {
		return internal.ContextClient(ctx)
	}
	return &http.Client{
		Transport: &Transport{
			Base:   internal.ContextClient(ctx).Transport,
			Source: ReuseTokenSource(nil, src),
		},
	}
}