-
Notifications
You must be signed in to change notification settings - Fork 1.8k
/
authratelimiter.go
109 lines (89 loc) · 2.58 KB
/
authratelimiter.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
package home
import (
"sync"
"time"
)
// failedAuthTTL is the period of time for which the failed attempt will stay in
// cache.
const failedAuthTTL = 1 * time.Minute
// failedAuth is an entry of authRateLimiter's cache.
type failedAuth struct {
until time.Time
num uint
}
// authRateLimiter used to cache failed authentication attempts.
type authRateLimiter struct {
failedAuths map[string]failedAuth
// failedAuthsLock protects failedAuths.
failedAuthsLock sync.Mutex
blockDur time.Duration
maxAttempts uint
}
// newAuthRateLimiter returns properly initialized *authRateLimiter.
func newAuthRateLimiter(blockDur time.Duration, maxAttempts uint) (ab *authRateLimiter) {
return &authRateLimiter{
failedAuths: make(map[string]failedAuth),
blockDur: blockDur,
maxAttempts: maxAttempts,
}
}
// cleanupLocked checks each blocked users removing ones with expired TTL. For
// internal use only.
func (ab *authRateLimiter) cleanupLocked(now time.Time) {
for k, v := range ab.failedAuths {
if now.After(v.until) {
delete(ab.failedAuths, k)
}
}
}
// checkLocked checks the attempter for it's state. For internal use only.
func (ab *authRateLimiter) checkLocked(usrID string, now time.Time) (left time.Duration) {
a, ok := ab.failedAuths[usrID]
if !ok {
return 0
}
if a.num < ab.maxAttempts {
return 0
}
return a.until.Sub(now)
}
// check returns the time left until unblocking. The nonpositive result should
// be interpreted as not blocked attempter.
func (ab *authRateLimiter) check(usrID string) (left time.Duration) {
now := time.Now()
ab.failedAuthsLock.Lock()
defer ab.failedAuthsLock.Unlock()
ab.cleanupLocked(now)
return ab.checkLocked(usrID, now)
}
// incLocked increments the number of unsuccessful attempts for attempter with
// ip and updates it's blocking moment if needed. For internal use only.
func (ab *authRateLimiter) incLocked(usrID string, now time.Time) {
var until time.Time = now.Add(failedAuthTTL)
var attNum uint = 1
a, ok := ab.failedAuths[usrID]
if ok {
until = a.until
attNum = a.num + 1
}
if attNum >= ab.maxAttempts {
until = now.Add(ab.blockDur)
}
ab.failedAuths[usrID] = failedAuth{
num: attNum,
until: until,
}
}
// inc updates the failed attempt in cache.
func (ab *authRateLimiter) inc(usrID string) {
now := time.Now()
ab.failedAuthsLock.Lock()
defer ab.failedAuthsLock.Unlock()
ab.incLocked(usrID, now)
}
// remove stops any tracking and any blocking of the user.
func (ab *authRateLimiter) remove(usrID string) {
ab.failedAuthsLock.Lock()
defer ab.failedAuthsLock.Unlock()
delete(ab.failedAuths, usrID)
}