init push

backend/pkg/ratelimit/rate_limiter.go

@@ -0,0 +1,102 @@
+package ratelimit
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+
+	"github.com/chaitin/panda-wiki/log"
+	"github.com/chaitin/panda-wiki/store/cache"
+)
+
+type RateLimiter struct {
+	logger *log.Logger
+	cache  *cache.Cache
+}
+
+func NewRateLimiter(logger *log.Logger, cache *cache.Cache) *RateLimiter {
+	return &RateLimiter{
+		logger: logger,
+		cache:  cache,
+	}
+}
+
+const (
+	LockThreshold1    = 5  // 第一次锁定阈值
+	LockThreshold2    = 10 // 第二次锁定阈值
+	LockThreshold3    = 15 // 第三次锁定阈值
+	AttemptsKeyExpiry = 24 * time.Hour
+)
+
+// CheckIPLocked checks if the IP is currently locked
+// Returns:
+// - bool: whether the IP is locked
+// - time.Duration: remaining lockout duration
+func (r *RateLimiter) CheckIPLocked(ctx context.Context, ip string) (bool, time.Duration) {
+	lockKey := fmt.Sprintf("login_lock:%s", ip)
+
+	ttl, err := r.cache.TTL(ctx, lockKey).Result()
+	if err != nil {
+		r.logger.Error("failed to check lock status", "error", err, "ip", ip)
+		return false, 0
+	}
+
+	if ttl > 0 {
+		return true, ttl
+	}
+
+	return false, 0
+}
+
+func (r *RateLimiter) LockAttempt(ctx context.Context, ip string) {
+	attemptsKey := fmt.Sprintf("login_attempts:%s", ip)
+	lockKey := fmt.Sprintf("login_lock:%s", ip)
+
+	attempts, err := r.cache.Incr(ctx, attemptsKey).Result()
+	if err != nil {
+		r.logger.Error("failed to increment attempts", "error", err, "ip", ip)
+		return
+	}
+
+	if err := r.cache.Expire(ctx, attemptsKey, AttemptsKeyExpiry).Err(); err != nil {
+		r.logger.Error("failed to set expiry on attempts key", "error", err, "ip", ip)
+	}
+
+	var lockDuration time.Duration
+
+	if attempts%5 == 0 {
+		switch {
+		case attempts == LockThreshold1:
+			lockDuration = time.Minute
+		case attempts == LockThreshold2:
+			lockDuration = 15 * time.Minute
+		case attempts >= LockThreshold3:
+			lockDuration = time.Hour
+		}
+		if lockDuration > 0 {
+			if err := r.cache.Set(ctx, lockKey, 1, lockDuration).Err(); err != nil {
+				r.logger.Error("failed to set lock key", "error", err, "ip", ip)
+				return
+			}
+			r.logger.Info("IP has been locked", "ip", ip, "lockDuration", lockDuration)
+		}
+	}
+}
+
+// ResetLoginAttempts resets the login attempt counter and lock for an IP
+func (r *RateLimiter) ResetLoginAttempts(ctx context.Context, ip string) error {
+	attemptsKey := fmt.Sprintf("login_attempts:%s", ip)
+	lockKey := fmt.Sprintf("login_lock:%s", ip)
+
+	pipe := r.cache.Pipeline()
+	pipe.Del(ctx, attemptsKey)
+	pipe.Del(ctx, lockKey)
+	_, err := pipe.Exec(ctx)
+	if err != nil && !errors.Is(err, redis.Nil) {
+		return fmt.Errorf("failed to reset login attempts: %w", err)
+	}
+	return nil
+}