1+ package util
2+
3+ import (
4+ "context"
5+ "sync"
6+ "time"
7+ )
8+
9+ // GetRemoteAddr extracts the remote IP address from context
10+ // This duplicates the logic from metrics package to avoid circular imports
11+ func GetRemoteAddr (ctx context.Context ) string {
12+ // Use the same context key pattern as metrics package
13+ type metricsContextKey int
14+ remoteAddrKey := metricsContextKey (1 )
15+
16+ if addr , ok := ctx .Value (remoteAddrKey ).(string ); ok {
17+ return addr
18+ }
19+ return ""
20+ }
21+
22+ // PasswordRateLimiter provides simple IP-based rate limiting for password attempts
23+ type PasswordRateLimiter struct {
24+ mu sync.RWMutex
25+ attempts map [string ][]time.Time
26+ maxAttempts int
27+ windowSize time.Duration
28+ cleanupTicker * time.Ticker
29+ done chan bool
30+ }
31+
32+ // NewPasswordRateLimiter creates a new rate limiter for password attempts
33+ // maxAttempts: maximum password attempts allowed per IP in the time window
34+ // windowSize: time window for counting attempts (e.g., 15 minutes)
35+ func NewPasswordRateLimiter (maxAttempts int , windowSize time.Duration ) * PasswordRateLimiter {
36+ rl := & PasswordRateLimiter {
37+ attempts : make (map [string ][]time.Time ),
38+ maxAttempts : maxAttempts ,
39+ windowSize : windowSize ,
40+ done : make (chan bool ),
41+ }
42+
43+ // Start cleanup routine to prevent memory leaks
44+ rl .cleanupTicker = time .NewTicker (windowSize )
45+ go rl .cleanup ()
46+
47+ return rl
48+ }
49+
50+ // IsAllowed checks if a password attempt from the given IP is allowed
51+ // Returns true if attempt is allowed, false if rate limited
52+ func (rl * PasswordRateLimiter ) IsAllowed (ctx context.Context , remoteAddr string ) bool {
53+ if remoteAddr == "" {
54+ return true // Allow if we can't determine IP
55+ }
56+
57+ rl .mu .RLock ()
58+ attempts , exists := rl .attempts [remoteAddr ]
59+ rl .mu .RUnlock ()
60+
61+ if ! exists {
62+ return true
63+ }
64+
65+ now := time .Now ()
66+ cutoff := now .Add (- rl .windowSize )
67+
68+ // Count valid attempts within window
69+ validAttempts := 0
70+ for _ , attemptTime := range attempts {
71+ if attemptTime .After (cutoff ) {
72+ validAttempts ++
73+ }
74+ }
75+
76+ return validAttempts < rl .maxAttempts
77+ }
78+
79+ // RecordFailedAttempt records a failed password attempt for the given IP
80+ func (rl * PasswordRateLimiter ) RecordFailedAttempt (ctx context.Context , remoteAddr string ) {
81+ if remoteAddr == "" {
82+ return // Nothing to record if we can't determine IP
83+ }
84+
85+ rl .mu .Lock ()
86+ defer rl .mu .Unlock ()
87+
88+ now := time .Now ()
89+ cutoff := now .Add (- rl .windowSize )
90+
91+ // Clean up old attempts for this IP and add new one
92+ attempts := rl .attempts [remoteAddr ]
93+ validAttempts := make ([]time.Time , 0 , len (attempts )+ 1 )
94+
95+ // Keep only recent attempts
96+ for _ , attemptTime := range attempts {
97+ if attemptTime .After (cutoff ) {
98+ validAttempts = append (validAttempts , attemptTime )
99+ }
100+ }
101+
102+ // Add current attempt
103+ validAttempts = append (validAttempts , now )
104+ rl .attempts [remoteAddr ] = validAttempts
105+ }
106+
107+ // cleanup periodically removes old entries to prevent memory leaks
108+ func (rl * PasswordRateLimiter ) cleanup () {
109+ for {
110+ select {
111+ case <- rl .cleanupTicker .C :
112+ rl .mu .Lock ()
113+ now := time .Now ()
114+ cutoff := now .Add (- rl .windowSize )
115+
116+ for ip , attempts := range rl .attempts {
117+ validAttempts := make ([]time.Time , 0 , len (attempts ))
118+ for _ , attemptTime := range attempts {
119+ if attemptTime .After (cutoff ) {
120+ validAttempts = append (validAttempts , attemptTime )
121+ }
122+ }
123+
124+ if len (validAttempts ) == 0 {
125+ delete (rl .attempts , ip )
126+ } else {
127+ rl .attempts [ip ] = validAttempts
128+ }
129+ }
130+ rl .mu .Unlock ()
131+ case <- rl .done :
132+ return
133+ }
134+ }
135+ }
136+
137+ // Close stops the cleanup routine
138+ func (rl * PasswordRateLimiter ) Close () {
139+ if rl .cleanupTicker != nil {
140+ rl .cleanupTicker .Stop ()
141+ }
142+ close (rl .done )
143+ }
0 commit comments