forked from gofiber/fiber
-
Notifications
You must be signed in to change notification settings - Fork 0
/
limiter_sliding.go
137 lines (110 loc) · 3.81 KB
/
limiter_sliding.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
package limiter
import (
"strconv"
"sync"
"sync/atomic"
"time"
"github.com/boomhut/fiber/v2"
"github.com/boomhut/fiber/v2/utils"
)
type SlidingWindow struct{}
// New creates a new sliding window middleware handler
func (SlidingWindow) New(cfg Config) fiber.Handler {
var (
// Limiter variables
mux = &sync.RWMutex{}
max = strconv.Itoa(cfg.Max)
expiration = uint64(cfg.Expiration.Seconds())
)
// Create manager to simplify storage operations ( see manager.go )
manager := newManager(cfg.Storage)
// Update timestamp every second
utils.StartTimeStampUpdater()
// Return new handler
return func(c *fiber.Ctx) error {
// Don't execute middleware if Next returns true
if cfg.Next != nil && cfg.Next(c) {
return c.Next()
}
// Get key from request
key := cfg.KeyGenerator(c)
// Lock entry
mux.Lock()
// Get entry from pool and release when finished
e := manager.get(key)
// Get timestamp
ts := uint64(atomic.LoadUint32(&utils.Timestamp))
// Set expiration if entry does not exist
if e.exp == 0 {
e.exp = ts + expiration
} else if ts >= e.exp {
// The entry has expired, handle the expiration.
// Set the prevHits to the current hits and reset the hits to 0.
e.prevHits = e.currHits
// Reset the current hits to 0.
e.currHits = 0
// Check how much into the current window it currently is and sets the
// expiry based on that, otherwise this would only reset on
// the next request and not show the correct expiry.
elapsed := ts - e.exp
if elapsed >= expiration {
e.exp = ts + expiration
} else {
e.exp = ts + expiration - elapsed
}
}
// Increment hits
e.currHits++
// Calculate when it resets in seconds
resetInSec := e.exp - ts
// weight = time until current window reset / total window length
weight := float64(resetInSec) / float64(expiration)
// rate = request count in previous window - weight + request count in current window
rate := int(float64(e.prevHits)*weight) + e.currHits
// Calculate how many hits can be made based on the current rate
remaining := cfg.Max - rate
// Update storage. Garbage collect when the next window ends.
// |--------------------------|--------------------------|
// ^ ^ ^ ^
// ts e.exp End sample window End next window
// <------------>
// resetInSec
// resetInSec = e.exp - ts - time until end of current window.
// duration + expiration = end of next window.
// Because we don't want to garbage collect in the middle of a window
// we add the expiration to the duration.
// Otherwise after the end of "sample window", attackers could launch
// a new request with the full window length.
manager.set(key, e, time.Duration(resetInSec+expiration)*time.Second)
// Unlock entry
mux.Unlock()
// Check if hits exceed the cfg.Max
if remaining < 0 {
// Return response with Retry-After header
// https://tools.ietf.org/html/rfc6584
c.Set(fiber.HeaderRetryAfter, strconv.FormatUint(resetInSec, 10))
// Call LimitReached handler
return cfg.LimitReached(c)
}
// Continue stack for reaching c.Response().StatusCode()
// Store err for returning
err := c.Next()
// Check for SkipFailedRequests and SkipSuccessfulRequests
if (cfg.SkipSuccessfulRequests && c.Response().StatusCode() < fiber.StatusBadRequest) ||
(cfg.SkipFailedRequests && c.Response().StatusCode() >= fiber.StatusBadRequest) {
// Lock entry
mux.Lock()
e = manager.get(key)
e.currHits--
remaining++
manager.set(key, e, cfg.Expiration)
// Unlock entry
mux.Unlock()
}
// We can continue, update RateLimit headers
c.Set(xRateLimitLimit, max)
c.Set(xRateLimitRemaining, strconv.Itoa(remaining))
c.Set(xRateLimitReset, strconv.FormatUint(resetInSec, 10))
return err
}
}