/
token.go
168 lines (152 loc) · 3.99 KB
/
token.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
package store
import (
"encoding/json"
"time"
oauth2 "github.com/armonia-tech/at.oauth2"
models "github.com/armonia-tech/at.oauth2/models"
uuid "github.com/armonia-tech/at.oauth2/utils/uuid"
"github.com/tidwall/buntdb"
)
// NewMemoryTokenStore create a token store instance based on memory
func NewMemoryTokenStore() (store oauth2.TokenStore, err error) {
store, err = NewFileTokenStore(":memory:")
return
}
// NewFileTokenStore create a token store instance based on file
func NewFileTokenStore(filename string) (store oauth2.TokenStore, err error) {
db, err := buntdb.Open(filename)
if err != nil {
return
}
store = &TokenStore{db: db}
return
}
// TokenStore token storage based on buntdb(https://github.com/tidwall/buntdb)
type TokenStore struct {
db *buntdb.DB
}
// Create create and store the new token information
func (ts *TokenStore) Create(info oauth2.TokenInfo) (err error) {
ct := time.Now()
jv, err := json.Marshal(info)
if err != nil {
return
}
err = ts.db.Update(func(tx *buntdb.Tx) (err error) {
if code := info.GetCode(); code != "" {
_, _, err = tx.Set(code, string(jv), &buntdb.SetOptions{Expires: true, TTL: info.GetCodeExpiresIn()})
return
}
basicID := uuid.Must(uuid.NewRandom()).String()
aexp := info.GetAccessExpiresIn()
rexp := aexp
expires := true
if refresh := info.GetRefresh(); refresh != "" {
rexp = info.GetRefreshCreateAt().Add(info.GetRefreshExpiresIn()).Sub(ct)
if aexp.Seconds() > rexp.Seconds() {
aexp = rexp
}
expires = info.GetRefreshExpiresIn() != 0
_, _, err = tx.Set(refresh, basicID, &buntdb.SetOptions{Expires: expires, TTL: rexp})
if err != nil {
return
}
}
_, _, err = tx.Set(basicID, string(jv), &buntdb.SetOptions{Expires: expires, TTL: rexp})
if err != nil {
return
}
_, _, err = tx.Set(info.GetAccess(), basicID, &buntdb.SetOptions{Expires: expires, TTL: aexp})
return
})
return
}
// remove key
func (ts *TokenStore) remove(key string) (err error) {
verr := ts.db.Update(func(tx *buntdb.Tx) (err error) {
_, err = tx.Delete(key)
return
})
if verr == buntdb.ErrNotFound {
return
}
err = verr
return
}
// RemoveByCode use the authorization code to delete the token information
func (ts *TokenStore) RemoveByCode(code string) (err error) {
err = ts.remove(code)
return
}
// RemoveByAccess use the access token to delete the token information
func (ts *TokenStore) RemoveByAccess(access string) (err error) {
err = ts.remove(access)
return
}
// RemoveByRefresh use the refresh token to delete the token information
func (ts *TokenStore) RemoveByRefresh(refresh string) (err error) {
err = ts.remove(refresh)
return
}
func (ts *TokenStore) getData(key string) (ti oauth2.TokenInfo, err error) {
verr := ts.db.View(func(tx *buntdb.Tx) (err error) {
jv, err := tx.Get(key)
if err != nil {
return
}
var tm models.Token
err = json.Unmarshal([]byte(jv), &tm)
if err != nil {
return
}
ti = &tm
return
})
if verr != nil {
if verr == buntdb.ErrNotFound {
return
}
err = verr
}
return
}
func (ts *TokenStore) getBasicID(key string) (basicID string, err error) {
verr := ts.db.View(func(tx *buntdb.Tx) (err error) {
v, err := tx.Get(key)
if err != nil {
return
}
basicID = v
return
})
if verr != nil {
if verr == buntdb.ErrNotFound {
return
}
err = verr
}
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) {
basicID, err := ts.getBasicID(access)
if err != nil {
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) {
basicID, err := ts.getBasicID(refresh)
if err != nil {
return
}
ti, err = ts.getData(basicID)
return
}