diff --git a/handler.go b/handler.go index 4f4515d..afd58a6 100644 --- a/handler.go +++ b/handler.go @@ -53,8 +53,9 @@ func handleGeminiRequest(conn net.Conn, sysConfig SysConfig, config UserConfig, if sysConfig.RateLimitEnable { noPort := logEntry.RemoteAddr.String() noPort = noPort[0:strings.LastIndex(noPort, ":")] - if !rl.Allowed(noPort) { - conn.Write([]byte("44 10 second cool down, please!\r\n")) + drips, allowed := rl.Allowed(noPort) + if !allowed { + conn.Write([]byte("44 " + strconv.Itoa(drips) + " second cool down, please!\r\n")) logEntry.Status = 44 return } diff --git a/ratelim.go b/ratelim.go index ff399f3..9c97ec9 100644 --- a/ratelim.go +++ b/ratelim.go @@ -35,18 +35,16 @@ func newRateLimiter(rate int, burst int) RateLimiter { return *rl } -func (rl *RateLimiter) Allowed(addr string) bool { +func (rl *RateLimiter) Allowed(addr string) (int, bool) { rl.mu.Lock() defer rl.mu.Unlock() drips, present := rl.bucket[addr] if !present { rl.bucket[addr] = 1 - return true + return 1, true } - if drips == rl.burst { - return false - } - rl.bucket[addr] = drips + 1 - return true + drips += 1 + rl.bucket[addr] = drips + return drips, drips < rl.burst }