add context support

Author: mkaczanowski

generate.go

@@ -1,6 +1,7 @@
 package oauth2
 
 import (
+	"context"
 	"net/http"
 	"time"
 )
@@ -17,11 +18,11 @@ type (
 
 	// AuthorizeGenerate generate the authorization code interface
 	AuthorizeGenerate interface {
-		Token(data *GenerateBasic) (code string, err error)
+		Token(ctx context.Context, data *GenerateBasic) (code string, err error)
 	}
 
 	// AccessGenerate generate the access and refresh tokens interface
 	AccessGenerate interface {
-		Token(data *GenerateBasic, isGenRefresh bool) (access, refresh string, err error)
+		Token(ctx context.Context, data *GenerateBasic, isGenRefresh bool) (access, refresh string, err error)
 	}
 )

generates/access.go

@@ -2,6 +2,7 @@ package generates
 
 import (
 	"bytes"
+	"context"
 	"encoding/base64"
 	"strconv"
 	"strings"
@@ -20,7 +21,7 @@ type AccessGenerate struct {
 }
 
 // Token based on the UUID generated token
-func (ag *AccessGenerate) Token(data *oauth2.GenerateBasic, isGenRefresh bool) (string, string, error) {
+func (ag *AccessGenerate) Token(ctx context.Context, data *oauth2.GenerateBasic, isGenRefresh bool) (string, string, error) {
 	buf := bytes.NewBufferString(data.Client.GetID())
 	buf.WriteString(data.UserID)
 	buf.WriteString(strconv.FormatInt(data.CreateAt.UnixNano(), 10))

generates/access_test.go

@@ -1,6 +1,7 @@
 package generates_test
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -22,7 +23,7 @@ func TestAccess(t *testing.T) {
 			CreateAt: time.Now(),
 		}
 		gen := generates.NewAccessGenerate()
-		access, refresh, err := gen.Token(data, true)
+		access, refresh, err := gen.Token(context.Background(), data, true)
 		So(err, ShouldBeNil)
 		So(access, ShouldNotBeEmpty)
 		So(refresh, ShouldNotBeEmpty)

generates/authorize.go

@@ -2,6 +2,7 @@ package generates
 
 import (
 	"bytes"
+	"context"
 	"encoding/base64"
 	"strings"
 
@@ -18,7 +19,7 @@ func NewAuthorizeGenerate() *AuthorizeGenerate {
 type AuthorizeGenerate struct{}
 
 // Token based on the UUID generated token
-func (ag *AuthorizeGenerate) Token(data *oauth2.GenerateBasic) (string, error) {
+func (ag *AuthorizeGenerate) Token(ctx context.Context, data *oauth2.GenerateBasic) (string, error) {
 	buf := bytes.NewBufferString(data.Client.GetID())
 	buf.WriteString(data.UserID)
 	token := uuid.NewMD5(uuid.Must(uuid.NewRandom()), buf.Bytes())

generates/authorize_test.go

@@ -1,6 +1,7 @@
 package generates_test
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -22,7 +23,7 @@ func TestAuthorize(t *testing.T) {
 			CreateAt: time.Now(),
 		}
 		gen := generates.NewAuthorizeGenerate()
-		code, err := gen.Token(data)
+		code, err := gen.Token(context.Background(), data)
 		So(err, ShouldBeNil)
 		So(code, ShouldNotBeEmpty)
 		Println("\nAuthorize Code:" + code)

generates/jwt_access.go

@@ -1,6 +1,7 @@
 package generates
 
 import (
+	"context"
 	"encoding/base64"
 	"strings"
 	"time"
@@ -41,7 +42,7 @@ type JWTAccessGenerate struct {
 }
 
 // Token based on the UUID generated token
-func (a *JWTAccessGenerate) Token(data *oauth2.GenerateBasic, isGenRefresh bool) (string, string, error) {
+func (a *JWTAccessGenerate) Token(ctx context.Context, data *oauth2.GenerateBasic, isGenRefresh bool) (string, string, error) {
 	claims := &JWTAccessClaims{
 		StandardClaims: jwt.StandardClaims{
 			Audience:  data.Client.GetID(),

generates/jwt_access_test.go

@@ -1,6 +1,7 @@
 package generates_test
 
 import (
+	"context"
 	"fmt"
 	"testing"
 	"time"
@@ -28,7 +29,7 @@ func TestJWTAccess(t *testing.T) {
 		}
 
 		gen := generates.NewJWTAccessGenerate([]byte("00000000"), jwt.SigningMethodHS512)
-		access, refresh, err := gen.Token(data, true)
+		access, refresh, err := gen.Token(context.Background(), data, true)
 		So(err, ShouldBeNil)
 		So(access, ShouldNotBeEmpty)
 		So(refresh, ShouldNotBeEmpty)

manage.go

@@ -1,6 +1,7 @@
 package oauth2
 
 import (
+	"context"
 	"net/http"
 	"time"
 )
@@ -21,26 +22,26 @@ type TokenGenerateRequest struct {
 // Manager authorization management interface
 type Manager interface {
 	// get the client information
-	GetClient(clientID string) (cli ClientInfo, err error)
+	GetClient(ctx context.Context, clientID string) (cli ClientInfo, err error)
 
 	// generate the authorization token(code)
-	GenerateAuthToken(rt ResponseType, tgr *TokenGenerateRequest) (authToken TokenInfo, err error)
+	GenerateAuthToken(ctx context.Context, rt ResponseType, tgr *TokenGenerateRequest) (authToken TokenInfo, err error)
 
 	// generate the access token
-	GenerateAccessToken(rt GrantType, tgr *TokenGenerateRequest) (accessToken TokenInfo, err error)
+	GenerateAccessToken(ctx context.Context, rt GrantType, tgr *TokenGenerateRequest) (accessToken TokenInfo, err error)
 
 	// refreshing an access token
-	RefreshAccessToken(tgr *TokenGenerateRequest) (accessToken TokenInfo, err error)
+	RefreshAccessToken(ctx context.Context, tgr *TokenGenerateRequest) (accessToken TokenInfo, err error)
 
 	// use the access token to delete the token information
-	RemoveAccessToken(access string) (err error)
+	RemoveAccessToken(ctx context.Context, access string) (err error)
 
 	// use the refresh token to delete the token information
-	RemoveRefreshToken(refresh string) (err error)
+	RemoveRefreshToken(ctx context.Context, refresh string) (err error)
 
 	// according to the access token for corresponding token information
-	LoadAccessToken(access string) (ti TokenInfo, err error)
+	LoadAccessToken(ctx context.Context, access string) (ti TokenInfo, err error)
 
 	// according to the refresh token for corresponding token information
-	LoadRefreshToken(refresh string) (ti TokenInfo, err error)
+	LoadRefreshToken(ctx context.Context, refresh string) (ti TokenInfo, err error)
 }

manage/manage_test.go

@@ -1,6 +1,7 @@
 package manage_test
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -15,6 +16,7 @@ import (
 func TestManager(t *testing.T) {
 	Convey("Manager test", t, func() {
 		manager := manage.NewDefaultManager()
+		ctx := context.Background()
 
 		manager.MustTokenStorage(store.NewMemoryTokenStore())
 
@@ -34,7 +36,7 @@ func TestManager(t *testing.T) {
 		}
 
 		Convey("GetClient test", func() {
-			cli, err := manager.GetClient("1")
+			cli, err := manager.GetClient(ctx, "1")
 			So(err, ShouldBeNil)
 			So(cli.GetSecret(), ShouldEqual, "11")
 		})
@@ -55,7 +57,8 @@ func TestManager(t *testing.T) {
 }
 
 func testManager(tgr *oauth2.TokenGenerateRequest, manager oauth2.Manager) {
-	cti, err := manager.GenerateAuthToken(oauth2.Code, tgr)
+	ctx := context.Background()
+	cti, err := manager.GenerateAuthToken(ctx, oauth2.Code, tgr)
 	So(err, ShouldBeNil)
 
 	code := cti.GetCode()
@@ -67,58 +70,59 @@ func testManager(tgr *oauth2.TokenGenerateRequest, manager oauth2.Manager) {
 		RedirectURI:  tgr.RedirectURI,
 		Code:         code,
 	}
-	ati, err := manager.GenerateAccessToken(oauth2.AuthorizationCode, atParams)
+	ati, err := manager.GenerateAccessToken(ctx, oauth2.AuthorizationCode, atParams)
 	So(err, ShouldBeNil)
 
 	accessToken, refreshToken := ati.GetAccess(), ati.GetRefresh()
 	So(accessToken, ShouldNotBeEmpty)
 	So(refreshToken, ShouldNotBeEmpty)
 
-	ainfo, err := manager.LoadAccessToken(accessToken)
+	ainfo, err := manager.LoadAccessToken(ctx, accessToken)
 	So(err, ShouldBeNil)
 	So(ainfo.GetClientID(), ShouldEqual, atParams.ClientID)
 
-	arinfo, err := manager.LoadRefreshToken(accessToken)
+	arinfo, err := manager.LoadRefreshToken(ctx, accessToken)
 	So(err, ShouldNotBeNil)
 	So(arinfo, ShouldBeNil)
 
-	rainfo, err := manager.LoadAccessToken(refreshToken)
+	rainfo, err := manager.LoadAccessToken(ctx, refreshToken)
 	So(err, ShouldNotBeNil)
 	So(rainfo, ShouldBeNil)
 
-	rinfo, err := manager.LoadRefreshToken(refreshToken)
+	rinfo, err := manager.LoadRefreshToken(ctx, refreshToken)
 	So(err, ShouldBeNil)
 	So(rinfo.GetClientID(), ShouldEqual, atParams.ClientID)
 
 	atParams.Refresh = refreshToken
 	atParams.Scope = "owner"
-	rti, err := manager.RefreshAccessToken(atParams)
+	rti, err := manager.RefreshAccessToken(ctx, atParams)
 	So(err, ShouldBeNil)
 
 	refreshAT := rti.GetAccess()
 	So(refreshAT, ShouldNotBeEmpty)
 
-	_, err = manager.LoadAccessToken(accessToken)
+	_, err = manager.LoadAccessToken(ctx, accessToken)
 	So(err, ShouldNotBeNil)
 
-	refreshAInfo, err := manager.LoadAccessToken(refreshAT)
+	refreshAInfo, err := manager.LoadAccessToken(ctx, refreshAT)
 	So(err, ShouldBeNil)
 	So(refreshAInfo.GetScope(), ShouldEqual, "owner")
 
-	err = manager.RemoveAccessToken(refreshAT)
+	err = manager.RemoveAccessToken(ctx, refreshAT)
 	So(err, ShouldBeNil)
 
-	_, err = manager.LoadAccessToken(refreshAT)
+	_, err = manager.LoadAccessToken(ctx, refreshAT)
 	So(err, ShouldNotBeNil)
 
-	err = manager.RemoveRefreshToken(refreshToken)
+	err = manager.RemoveRefreshToken(ctx, refreshToken)
 	So(err, ShouldBeNil)
 
-	_, err = manager.LoadRefreshToken(refreshToken)
+	_, err = manager.LoadRefreshToken(ctx, refreshToken)
 	So(err, ShouldNotBeNil)
 }
 
 func testZeroAccessExpirationManager(tgr *oauth2.TokenGenerateRequest, manager oauth2.Manager) {
+	ctx := context.Background()
 	config := manage.Config{
 		AccessTokenExp:    0, // Set explicitly as we're testing 0 (no) expiration
 		IsGenerateRefresh: true,
@@ -127,7 +131,7 @@ func testZeroAccessExpirationManager(tgr *oauth2.TokenGenerateRequest, manager o
 	So(ok, ShouldBeTrue)
 	m.SetAuthorizeCodeTokenCfg(&config)
 
-	cti, err := manager.GenerateAuthToken(oauth2.Code, tgr)
+	cti, err := manager.GenerateAuthToken(ctx, oauth2.Code, tgr)
 	So(err, ShouldBeNil)
 
 	code := cti.GetCode()
@@ -139,29 +143,30 @@ func testZeroAccessExpirationManager(tgr *oauth2.TokenGenerateRequest, manager o
 		RedirectURI:  tgr.RedirectURI,
 		Code:         code,
 	}
-	ati, err := manager.GenerateAccessToken(oauth2.AuthorizationCode, atParams)
+	ati, err := manager.GenerateAccessToken(ctx, oauth2.AuthorizationCode, atParams)
 	So(err, ShouldBeNil)
 
 	accessToken, refreshToken := ati.GetAccess(), ati.GetRefresh()
 	So(accessToken, ShouldNotBeEmpty)
 	So(refreshToken, ShouldNotBeEmpty)
 
-	tokenInfo, err := manager.LoadAccessToken(accessToken)
+	tokenInfo, err := manager.LoadAccessToken(ctx, accessToken)
 	So(err, ShouldBeNil)
 	So(tokenInfo, ShouldNotBeNil)
 	So(tokenInfo.GetAccess(), ShouldEqual, accessToken)
 	So(tokenInfo.GetAccessExpiresIn(), ShouldEqual, 0)
 }
 
 func testCannotRequestZeroExpirationAccessTokens(tgr *oauth2.TokenGenerateRequest, manager oauth2.Manager) {
+	ctx := context.Background()
 	config := manage.Config{
 		AccessTokenExp: time.Hour * 5,
 	}
 	m, ok := manager.(*manage.Manager)
 	So(ok, ShouldBeTrue)
 	m.SetAuthorizeCodeTokenCfg(&config)
 
-	cti, err := manager.GenerateAuthToken(oauth2.Code, tgr)
+	cti, err := manager.GenerateAuthToken(ctx, oauth2.Code, tgr)
 	So(err, ShouldBeNil)
 
 	code := cti.GetCode()
@@ -174,7 +179,7 @@ func testCannotRequestZeroExpirationAccessTokens(tgr *oauth2.TokenGenerateReques
 		AccessTokenExp: 0, // requesting token without expiration
 		Code:           code,
 	}
-	ati, err := manager.GenerateAccessToken(oauth2.AuthorizationCode, atParams)
+	ati, err := manager.GenerateAccessToken(ctx, oauth2.AuthorizationCode, atParams)
 	So(err, ShouldBeNil)
 
 	accessToken := ati.GetAccess()
@@ -183,6 +188,7 @@ func testCannotRequestZeroExpirationAccessTokens(tgr *oauth2.TokenGenerateReques
 }
 
 func testZeroRefreshExpirationManager(tgr *oauth2.TokenGenerateRequest, manager oauth2.Manager) {
+	ctx := context.Background()
 	config := manage.Config{
 		RefreshTokenExp:   0, // Set explicitly as we're testing 0 (no) expiration
 		IsGenerateRefresh: true,
@@ -191,7 +197,7 @@ func testZeroRefreshExpirationManager(tgr *oauth2.TokenGenerateRequest, manager
 	So(ok, ShouldBeTrue)
 	m.SetAuthorizeCodeTokenCfg(&config)
 
-	cti, err := manager.GenerateAuthToken(oauth2.Code, tgr)
+	cti, err := manager.GenerateAuthToken(ctx, oauth2.Code, tgr)
 	So(err, ShouldBeNil)
 
 	code := cti.GetCode()
@@ -204,21 +210,21 @@ func testZeroRefreshExpirationManager(tgr *oauth2.TokenGenerateRequest, manager
 		AccessTokenExp: time.Hour,
 		Code:           code,
 	}
-	ati, err := manager.GenerateAccessToken(oauth2.AuthorizationCode, atParams)
+	ati, err := manager.GenerateAccessToken(ctx, oauth2.AuthorizationCode, atParams)
 	So(err, ShouldBeNil)
 
 	accessToken, refreshToken := ati.GetAccess(), ati.GetRefresh()
 	So(accessToken, ShouldNotBeEmpty)
 	So(refreshToken, ShouldNotBeEmpty)
 
-	tokenInfo, err := manager.LoadRefreshToken(refreshToken)
+	tokenInfo, err := manager.LoadRefreshToken(ctx, refreshToken)
 	So(err, ShouldBeNil)
 	So(tokenInfo, ShouldNotBeNil)
 	So(tokenInfo.GetRefresh(), ShouldEqual, refreshToken)
 	So(tokenInfo.GetRefreshExpiresIn(), ShouldEqual, 0)
 
 	// LoadAccessToken also checks refresh expiry
-	tokenInfo, err = manager.LoadAccessToken(accessToken)
+	tokenInfo, err = manager.LoadAccessToken(ctx, accessToken)
 	So(err, ShouldBeNil)
 	So(tokenInfo, ShouldNotBeNil)
 	So(tokenInfo.GetRefresh(), ShouldEqual, refreshToken)