1
2
3
4
5
6
7
8
9
10
11
12
13
14
15 package auth
16
17
18
19
20 import (
21 "context"
22 "crypto/rand"
23 "fmt"
24 "math/big"
25 "strconv"
26 "strings"
27 "sync"
28 "time"
29
30 "go.uber.org/zap"
31 )
32
33 const (
34 letters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
35 defaultSimpleTokenLength = 16
36 )
37
38
39
40 var (
41 simpleTokenTTLDefault = 300 * time.Second
42 simpleTokenTTLResolution = 1 * time.Second
43 )
44
45 type simpleTokenTTLKeeper struct {
46 tokens map[string]time.Time
47 donec chan struct{}
48 stopc chan struct{}
49 deleteTokenFunc func(string)
50 mu *sync.Mutex
51 simpleTokenTTL time.Duration
52 }
53
54 func (tm *simpleTokenTTLKeeper) stop() {
55 select {
56 case tm.stopc <- struct{}{}:
57 case <-tm.donec:
58 }
59 <-tm.donec
60 }
61
62 func (tm *simpleTokenTTLKeeper) addSimpleToken(token string) {
63 tm.tokens[token] = time.Now().Add(tm.simpleTokenTTL)
64 }
65
66 func (tm *simpleTokenTTLKeeper) resetSimpleToken(token string) {
67 if _, ok := tm.tokens[token]; ok {
68 tm.tokens[token] = time.Now().Add(tm.simpleTokenTTL)
69 }
70 }
71
72 func (tm *simpleTokenTTLKeeper) deleteSimpleToken(token string) {
73 delete(tm.tokens, token)
74 }
75
76 func (tm *simpleTokenTTLKeeper) run() {
77 tokenTicker := time.NewTicker(simpleTokenTTLResolution)
78 defer func() {
79 tokenTicker.Stop()
80 close(tm.donec)
81 }()
82 for {
83 select {
84 case <-tokenTicker.C:
85 nowtime := time.Now()
86 tm.mu.Lock()
87 for t, tokenendtime := range tm.tokens {
88 if nowtime.After(tokenendtime) {
89 tm.deleteTokenFunc(t)
90 delete(tm.tokens, t)
91 }
92 }
93 tm.mu.Unlock()
94 case <-tm.stopc:
95 return
96 }
97 }
98 }
99
100 type tokenSimple struct {
101 lg *zap.Logger
102 indexWaiter func(uint64) <-chan struct{}
103 simpleTokenKeeper *simpleTokenTTLKeeper
104 simpleTokensMu sync.Mutex
105 simpleTokens map[string]string
106 simpleTokenTTL time.Duration
107 }
108
109 func (t *tokenSimple) genTokenPrefix() (string, error) {
110 ret := make([]byte, defaultSimpleTokenLength)
111
112 for i := 0; i < defaultSimpleTokenLength; i++ {
113 bInt, err := rand.Int(rand.Reader, big.NewInt(int64(len(letters))))
114 if err != nil {
115 return "", err
116 }
117
118 ret[i] = letters[bInt.Int64()]
119 }
120
121 return string(ret), nil
122 }
123
124 func (t *tokenSimple) assignSimpleTokenToUser(username, token string) {
125 t.simpleTokensMu.Lock()
126 defer t.simpleTokensMu.Unlock()
127 if t.simpleTokenKeeper == nil {
128 return
129 }
130
131 _, ok := t.simpleTokens[token]
132 if ok {
133 t.lg.Panic(
134 "failed to assign already-used simple token to a user",
135 zap.String("user-name", username),
136 zap.String("token", token),
137 )
138 }
139
140 t.simpleTokens[token] = username
141 t.simpleTokenKeeper.addSimpleToken(token)
142 }
143
144 func (t *tokenSimple) invalidateUser(username string) {
145 if t.simpleTokenKeeper == nil {
146 return
147 }
148 t.simpleTokensMu.Lock()
149 for token, name := range t.simpleTokens {
150 if name == username {
151 delete(t.simpleTokens, token)
152 t.simpleTokenKeeper.deleteSimpleToken(token)
153 }
154 }
155 t.simpleTokensMu.Unlock()
156 }
157
158 func (t *tokenSimple) enable() {
159 t.simpleTokensMu.Lock()
160 defer t.simpleTokensMu.Unlock()
161 if t.simpleTokenKeeper != nil {
162 return
163 }
164 if t.simpleTokenTTL <= 0 {
165 t.simpleTokenTTL = simpleTokenTTLDefault
166 }
167
168 delf := func(tk string) {
169 if username, ok := t.simpleTokens[tk]; ok {
170 t.lg.Info(
171 "deleted a simple token",
172 zap.String("user-name", username),
173 zap.String("token", tk),
174 )
175 delete(t.simpleTokens, tk)
176 }
177 }
178 t.simpleTokenKeeper = &simpleTokenTTLKeeper{
179 tokens: make(map[string]time.Time),
180 donec: make(chan struct{}),
181 stopc: make(chan struct{}),
182 deleteTokenFunc: delf,
183 mu: &t.simpleTokensMu,
184 simpleTokenTTL: t.simpleTokenTTL,
185 }
186 go t.simpleTokenKeeper.run()
187 }
188
189 func (t *tokenSimple) disable() {
190 t.simpleTokensMu.Lock()
191 tk := t.simpleTokenKeeper
192 t.simpleTokenKeeper = nil
193 t.simpleTokens = make(map[string]string)
194 t.simpleTokensMu.Unlock()
195 if tk != nil {
196 tk.stop()
197 }
198 }
199
200 func (t *tokenSimple) info(ctx context.Context, token string, revision uint64) (*AuthInfo, bool) {
201 if !t.isValidSimpleToken(ctx, token) {
202 return nil, false
203 }
204 t.simpleTokensMu.Lock()
205 username, ok := t.simpleTokens[token]
206 if ok && t.simpleTokenKeeper != nil {
207 t.simpleTokenKeeper.resetSimpleToken(token)
208 }
209 t.simpleTokensMu.Unlock()
210 return &AuthInfo{Username: username, Revision: revision}, ok
211 }
212
213 func (t *tokenSimple) assign(ctx context.Context, username string, rev uint64) (string, error) {
214
215 index := ctx.Value(AuthenticateParamIndex{}).(uint64)
216 simpleTokenPrefix := ctx.Value(AuthenticateParamSimpleTokenPrefix{}).(string)
217 token := fmt.Sprintf("%s.%d", simpleTokenPrefix, index)
218 t.assignSimpleTokenToUser(username, token)
219
220 return token, nil
221 }
222
223 func (t *tokenSimple) isValidSimpleToken(ctx context.Context, token string) bool {
224 splitted := strings.Split(token, ".")
225 if len(splitted) != 2 {
226 return false
227 }
228 index, err := strconv.ParseUint(splitted[1], 10, 0)
229 if err != nil {
230 return false
231 }
232
233 select {
234 case <-t.indexWaiter(uint64(index)):
235 return true
236 case <-ctx.Done():
237 }
238
239 return false
240 }
241
242 func newTokenProviderSimple(lg *zap.Logger, indexWaiter func(uint64) <-chan struct{}, TokenTTL time.Duration) *tokenSimple {
243 if lg == nil {
244 lg = zap.NewNop()
245 }
246 return &tokenSimple{
247 lg: lg,
248 simpleTokens: make(map[string]string),
249 indexWaiter: indexWaiter,
250 simpleTokenTTL: TokenTTL,
251 }
252 }
253
View as plain text