75 lines
2.1 KiB
Go
75 lines
2.1 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
)
|
|
|
|
// RedisClient is the subset of Redis operations needed for rate limiting.
|
|
type RedisClient interface {
|
|
Incr(ctx context.Context, key string) (int64, error)
|
|
Expire(ctx context.Context, key string, expiration time.Duration) error
|
|
Decr(ctx context.Context, key string) (int64, error)
|
|
}
|
|
|
|
// CheckMessageRate checks if a user has exceeded the message rate limit.
|
|
// Returns true if allowed, false if rate limited.
|
|
func CheckMessageRate(ctx context.Context, redis RedisClient, username string, max int, window time.Duration) (bool, error) {
|
|
if redis == nil {
|
|
return true, nil // No Redis = no rate limiting
|
|
}
|
|
key := fmt.Sprintf("rate:%s", username)
|
|
count, err := redis.Incr(ctx, key)
|
|
if err != nil {
|
|
return true, err // Fail open on Redis error
|
|
}
|
|
if count == 1 {
|
|
_ = redis.Expire(ctx, key, window)
|
|
}
|
|
return count <= int64(max), nil
|
|
}
|
|
|
|
// CheckLoginRate checks if an IP has exceeded the login attempt limit.
|
|
// Returns true if allowed, false if rate limited.
|
|
func CheckLoginRate(ctx context.Context, redis RedisClient, clientIP string) (bool, error) {
|
|
if redis == nil {
|
|
return true, nil
|
|
}
|
|
key := fmt.Sprintf("login_attempts:%s", clientIP)
|
|
count, err := redis.Incr(ctx, key)
|
|
if err != nil {
|
|
return true, err
|
|
}
|
|
if count == 1 {
|
|
_ = redis.Expire(ctx, key, 5*time.Minute)
|
|
}
|
|
return count <= 5, nil
|
|
}
|
|
|
|
// CheckWSConnRate checks if an IP has exceeded the WebSocket connection limit.
|
|
// Returns true if allowed, false if rate limited.
|
|
// Caller must call Decr on disconnect.
|
|
func CheckWSConnRate(ctx context.Context, redis RedisClient, clientIP string) (bool, error) {
|
|
if redis == nil {
|
|
return true, nil
|
|
}
|
|
key := fmt.Sprintf("ws_conn:%s", clientIP)
|
|
count, err := redis.Incr(ctx, key)
|
|
if err != nil {
|
|
return true, err
|
|
}
|
|
if count == 1 {
|
|
_ = redis.Expire(ctx, key, 60*time.Second)
|
|
}
|
|
return count <= 5, nil
|
|
}
|
|
|
|
// DecrWSConn decrements the WebSocket connection counter for an IP.
|
|
func DecrWSConn(ctx context.Context, redis RedisClient, clientIP string) {
|
|
if redis == nil {
|
|
return
|
|
}
|
|
key := fmt.Sprintf("ws_conn:%s", clientIP)
|
|
_, _ = redis.Decr(ctx, key)
|
|
} |