Add the authorization code storage support

Author: LyricTian

token.go

@@ -81,6 +81,20 @@ func (ts *TokenStore) c(name string) *mgo.Collection {
 
 // Create Create and store the new token information
 func (ts *TokenStore) Create(info oauth2.TokenInfo) (err error) {
+	jv, err := json.Marshal(info)
+	if err != nil {
+		return
+	}
+
+	if code := info.GetCode(); code != "" {
+		err = ts.c(ts.tcfg.BasicCName).Insert(basicData{
+			ID:        code,
+			Data:      jv,
+			ExpiredAt: info.GetCodeCreateAt().Add(info.GetCodeExpiresIn()),
+		})
+		return
+	}
+
 	aexp := info.GetAccessCreateAt().Add(info.GetAccessExpiresIn())
 	rexp := aexp
 	if refresh := info.GetRefresh(); refresh != "" {
@@ -90,10 +104,6 @@ func (ts *TokenStore) Create(info oauth2.TokenInfo) (err error) {
 		}
 	}
 	id := bson.NewObjectId().Hex()
-	jv, err := json.Marshal(info)
-	if err != nil {
-		return
-	}
 	runner := txn.NewRunner(ts.c(ts.tcfg.TxnCName))
 	ops := []txn.Op{{
 		C:      ts.tcfg.BasicCName,
@@ -127,9 +137,9 @@ func (ts *TokenStore) Create(info oauth2.TokenInfo) (err error) {
 	return
 }
 
-// RemoveByAccess Use the access token to delete the token information
-func (ts *TokenStore) RemoveByAccess(access string) (err error) {
-	verr := ts.c(ts.tcfg.AccessCName).RemoveId(access)
+// RemoveByCode Use the authorization code to delete the token information
+func (ts *TokenStore) RemoveByCode(code string) (err error) {
+	verr := ts.c(ts.tcfg.BasicCName).RemoveId(code)
 	if verr != nil {
 		if verr == mgo.ErrNotFound {
 			return
@@ -139,9 +149,9 @@ func (ts *TokenStore) RemoveByAccess(access string) (err error) {
 	return
 }
 
-// RemoveByRefresh Use the refresh token to delete the token information
-func (ts *TokenStore) RemoveByRefresh(refresh string) (err error) {
-	verr := ts.c(ts.tcfg.RefreshCName).RemoveId(refresh)
+// RemoveByAccess Use the access token to delete the token information
+func (ts *TokenStore) RemoveByAccess(access string) (err error) {
+	verr := ts.c(ts.tcfg.AccessCName).RemoveId(access)
 	if verr != nil {
 		if verr == mgo.ErrNotFound {
 			return
@@ -151,19 +161,21 @@ func (ts *TokenStore) RemoveByRefresh(refresh string) (err error) {
 	return
 }
 
-// get
-func (ts *TokenStore) get(cname, token string) (ti oauth2.TokenInfo, err error) {
-	var td tokenData
-	verr := ts.c(cname).FindId(token).One(&td)
+// RemoveByRefresh Use the refresh token to delete the token information
+func (ts *TokenStore) RemoveByRefresh(refresh string) (err error) {
+	verr := ts.c(ts.tcfg.RefreshCName).RemoveId(refresh)
 	if verr != nil {
 		if verr == mgo.ErrNotFound {
 			return
 		}
 		err = verr
-		return
 	}
+	return
+}
+
+func (ts *TokenStore) getData(basicID string) (ti oauth2.TokenInfo, err error) {
 	var bd basicData
-	verr = ts.c(ts.tcfg.BasicCName).FindId(td.BasicID).One(&bd)
+	verr := ts.c(ts.tcfg.BasicCName).FindId(basicID).One(&bd)
 	if verr != nil {
 		if verr == mgo.ErrNotFound {
 			return
@@ -180,15 +192,43 @@ func (ts *TokenStore) get(cname, token string) (ti oauth2.TokenInfo, err error)
 	return
 }
 
+func (ts *TokenStore) getBasicID(cname, token string) (basicID string, err error) {
+	var td tokenData
+	verr := ts.c(cname).FindId(token).One(&td)
+	if verr != nil {
+		if verr == mgo.ErrNotFound {
+			return
+		}
+		err = verr
+		return
+	}
+	basicID = td.BasicID
+	return
+}
+
+// GetByCode Use the authorization code for token information data
+func (ts *TokenStore) GetByCode(code string) (ti oauth2.TokenInfo, err error) {
+	ti, err = ts.getData(code)
+	return
+}
+
 // GetByAccess Use the access token for token information data
 func (ts *TokenStore) GetByAccess(access string) (ti oauth2.TokenInfo, err error) {
-	ti, err = ts.get(ts.tcfg.AccessCName, access)
+	basicID, err := ts.getBasicID(ts.tcfg.AccessCName, access)
+	if err != nil && basicID == "" {
+		return
+	}
+	ti, err = ts.getData(basicID)
 	return
 }
 
 // GetByRefresh Use the refresh token for token information data
 func (ts *TokenStore) GetByRefresh(refresh string) (ti oauth2.TokenInfo, err error) {
-	ti, err = ts.get(ts.tcfg.RefreshCName, refresh)
+	basicID, err := ts.getBasicID(ts.tcfg.RefreshCName, refresh)
+	if err != nil && basicID == "" {
+		return
+	}
+	ti, err = ts.getData(basicID)
 	return
 }
 

token_test.go

@@ -16,6 +16,31 @@ func TestTokenStore(t *testing.T) {
 		store, err := mongo.NewTokenStore(mcfg)
 		So(err, ShouldBeNil)
 
+		Convey("Test authorization code store", func() {
+			info := &models.Token{
+				ClientID:      "1",
+				UserID:        "1_1",
+				RedirectURI:   "http://localhost/",
+				Scope:         "all",
+				Code:          "11_11_11",
+				CodeCreateAt:  time.Now(),
+				CodeExpiresIn: time.Second * 5,
+			}
+			err := store.Create(info)
+			So(err, ShouldBeNil)
+
+			cinfo, err := store.GetByCode(info.Code)
+			So(err, ShouldBeNil)
+			So(cinfo.GetUserID(), ShouldEqual, info.UserID)
+
+			err = store.RemoveByCode(info.Code)
+			So(err, ShouldBeNil)
+
+			cinfo, err = store.GetByCode(info.Code)
+			So(err, ShouldBeNil)
+			So(cinfo, ShouldBeNil)
+		})
+
 		Convey("Test access token store", func() {
 			info := &models.Token{
 				ClientID:        "1",
@@ -52,7 +77,7 @@ func TestTokenStore(t *testing.T) {
 				AccessExpiresIn:  time.Second * 5,
 				Refresh:          "1_2_2",
 				RefreshCreateAt:  time.Now(),
-				RefreshExpiresIn: time.Minute * 1,
+				RefreshExpiresIn: time.Second * 15,
 			}
 			err := store.Create(info)
 			So(err, ShouldBeNil)