blob: 32009074e5a7e6293a02814a8df7ccb56e51261c (
plain)
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
|
package ratelimit
import (
"net/http"
"strings"
"sync"
"time"
"github.com/labstack/echo/v4"
"golang.org/x/time/rate"
)
type entry struct {
limiter *rate.Limiter
lastSeen time.Time
}
type IPRateLimiter struct {
entries sync.Map
rate rate.Limit
burst int
}
func NewIPRateLimiter(r rate.Limit, burst int) *IPRateLimiter {
rl := &IPRateLimiter{
rate: r,
burst: burst,
}
go rl.cleanup()
return rl
}
func (rl *IPRateLimiter) getLimiter(ip string) *rate.Limiter {
now := time.Now()
if v, ok := rl.entries.Load(ip); ok {
e := v.(*entry)
e.lastSeen = now
return e.limiter
}
limiter := rate.NewLimiter(rl.rate, rl.burst)
rl.entries.Store(ip, &entry{limiter: limiter, lastSeen: now})
return limiter
}
func (rl *IPRateLimiter) cleanup() {
for {
time.Sleep(10 * time.Minute)
rl.entries.Range(func(key, value any) bool {
e := value.(*entry)
if time.Since(e.lastSeen) > 10*time.Minute {
rl.entries.Delete(key)
}
return true
})
}
}
func LoginRateLimitMiddleware(rl *IPRateLimiter) echo.MiddlewareFunc {
return func(next echo.HandlerFunc) echo.HandlerFunc {
return func(c echo.Context) error {
if c.Request().Method != http.MethodPost || !strings.HasSuffix(c.Path(), "/login") {
return next(c)
}
ip := c.RealIP()
if !rl.getLimiter(ip).Allow() {
return c.JSON(http.StatusTooManyRequests, map[string]string{
"message": "Too many login attempts. Please try again later.",
})
}
return next(c)
}
}
}
|