limiter.go (1377B)
1 package ssh 2 3 import ( 4 "sync" 5 "time" 6 ) 7 8 type ipLimiter struct { 9 mu sync.Mutex 10 window time.Duration 11 max int 12 blockFor time.Duration 13 maxEntries int 14 entries map[string]ipEntry 15 } 16 17 type ipEntry struct { 18 count int 19 windowStart time.Time 20 blockedTill time.Time 21 lastSeen time.Time 22 } 23 24 func newIPLimiter(window time.Duration, max int, blockFor time.Duration, maxEntries int) *ipLimiter { 25 return &ipLimiter{ 26 window: window, 27 max: max, 28 blockFor: blockFor, 29 maxEntries: maxEntries, 30 entries: make(map[string]ipEntry), 31 } 32 } 33 34 func (l *ipLimiter) Allow(ip string) (bool, time.Duration) { 35 now := time.Now() 36 37 l.mu.Lock() 38 defer l.mu.Unlock() 39 40 if len(l.entries) > l.maxEntries { 41 l.compact(now) 42 } 43 44 entry := l.entries[ip] 45 entry.lastSeen = now 46 47 if entry.blockedTill.After(now) { 48 l.entries[ip] = entry 49 return false, entry.blockedTill.Sub(now) 50 } 51 52 if entry.windowStart.IsZero() || now.Sub(entry.windowStart) >= l.window { 53 entry.windowStart = now 54 entry.count = 0 55 } 56 57 entry.count++ 58 if entry.count > l.max { 59 entry.blockedTill = now.Add(l.blockFor) 60 l.entries[ip] = entry 61 return false, l.blockFor 62 } 63 64 l.entries[ip] = entry 65 return true, 0 66 } 67 68 func (l *ipLimiter) compact(now time.Time) { 69 for ip, entry := range l.entries { 70 if now.Sub(entry.lastSeen) > (l.blockFor + l.window) { 71 delete(l.entries, ip) 72 } 73 } 74 }