Initial commit: Open sourcing all of the Maple Open Technologies code.

cloud/maplepress-backend/pkg/ratelimit/ratelimiter.go

@@ -0,0 +1,172 @@
+package ratelimit
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+	"go.uber.org/zap"
+)
+
+// RateLimiter provides rate limiting functionality using Redis
+type RateLimiter interface {
+	// Allow checks if a request should be allowed based on the key
+	// Returns true if allowed, false if rate limit exceeded
+	Allow(ctx context.Context, key string) (bool, error)
+
+	// AllowN checks if N requests should be allowed
+	AllowN(ctx context.Context, key string, n int) (bool, error)
+
+	// Reset resets the rate limit for a key
+	Reset(ctx context.Context, key string) error
+
+	// GetRemaining returns the number of remaining requests
+	GetRemaining(ctx context.Context, key string) (int, error)
+}
+
+// Config holds rate limiter configuration
+type Config struct {
+	// MaxRequests is the maximum number of requests allowed
+	MaxRequests int
+	// Window is the time window for rate limiting
+	Window time.Duration
+	// KeyPrefix is the prefix for Redis keys
+	KeyPrefix string
+}
+
+type rateLimiter struct {
+	client *redis.Client
+	config Config
+	logger *zap.Logger
+}
+
+// NewRateLimiter creates a new rate limiter
+func NewRateLimiter(client *redis.Client, config Config, logger *zap.Logger) RateLimiter {
+	return &rateLimiter{
+		client: client,
+		config: config,
+		logger: logger.Named("rate-limiter"),
+	}
+}
+
+// Allow checks if a request should be allowed
+func (r *rateLimiter) Allow(ctx context.Context, key string) (bool, error) {
+	return r.AllowN(ctx, key, 1)
+}
+
+// AllowN checks if N requests should be allowed using sliding window counter
+func (r *rateLimiter) AllowN(ctx context.Context, key string, n int) (bool, error) {
+	redisKey := r.getRedisKey(key)
+	now := time.Now()
+	windowStart := now.Add(-r.config.Window)
+
+	// Use Redis transaction to ensure atomicity
+	pipe := r.client.Pipeline()
+
+	// Remove old entries outside the window
+	pipe.ZRemRangeByScore(ctx, redisKey, "0", fmt.Sprintf("%d", windowStart.UnixNano()))
+
+	// Count current requests in window
+	countCmd := pipe.ZCount(ctx, redisKey, fmt.Sprintf("%d", windowStart.UnixNano()), "+inf")
+
+	// Execute pipeline
+	_, err := pipe.Exec(ctx)
+	if err != nil && err != redis.Nil {
+		r.logger.Error("failed to check rate limit",
+			zap.String("key", key),
+			zap.Error(err))
+		// Fail open: allow request if Redis is down
+		return true, err
+	}
+
+	currentCount := countCmd.Val()
+
+	// Check if adding N requests would exceed limit
+	if currentCount+int64(n) > int64(r.config.MaxRequests) {
+		r.logger.Warn("rate limit exceeded",
+			zap.String("key", key),
+			zap.Int64("current_count", currentCount),
+			zap.Int("max_requests", r.config.MaxRequests))
+		return false, nil
+	}
+
+	// Add the new request(s) to the sorted set
+	pipe2 := r.client.Pipeline()
+	for i := 0; i < n; i++ {
+		// Use nanosecond timestamp with incremental offset to ensure uniqueness
+		timestamp := now.Add(time.Duration(i) * time.Nanosecond).UnixNano()
+		pipe2.ZAdd(ctx, redisKey, redis.Z{
+			Score:  float64(timestamp),
+			Member: fmt.Sprintf("%d-%d", timestamp, i),
+		})
+	}
+
+	// Set expiration on the key (window + buffer)
+	pipe2.Expire(ctx, redisKey, r.config.Window+time.Minute)
+
+	// Execute pipeline
+	_, err = pipe2.Exec(ctx)
+	if err != nil && err != redis.Nil {
+		r.logger.Error("failed to record request",
+			zap.String("key", key),
+			zap.Error(err))
+		// Already counted, so return true
+		return true, err
+	}
+
+	r.logger.Debug("rate limit check passed",
+		zap.String("key", key),
+		zap.Int64("current_count", currentCount),
+		zap.Int("max_requests", r.config.MaxRequests))
+
+	return true, nil
+}
+
+// Reset resets the rate limit for a key
+func (r *rateLimiter) Reset(ctx context.Context, key string) error {
+	redisKey := r.getRedisKey(key)
+	err := r.client.Del(ctx, redisKey).Err()
+	if err != nil {
+		r.logger.Error("failed to reset rate limit",
+			zap.String("key", key),
+			zap.Error(err))
+		return err
+	}
+
+	r.logger.Info("rate limit reset",
+		zap.String("key", key))
+
+	return nil
+}
+
+// GetRemaining returns the number of remaining requests in the current window
+func (r *rateLimiter) GetRemaining(ctx context.Context, key string) (int, error) {
+	redisKey := r.getRedisKey(key)
+	now := time.Now()
+	windowStart := now.Add(-r.config.Window)
+
+	// Count current requests in window
+	count, err := r.client.ZCount(ctx, redisKey,
+		fmt.Sprintf("%d", windowStart.UnixNano()),
+		"+inf").Result()
+
+	if err != nil && err != redis.Nil {
+		r.logger.Error("failed to get remaining requests",
+			zap.String("key", key),
+			zap.Error(err))
+		return 0, err
+	}
+
+	remaining := r.config.MaxRequests - int(count)
+	if remaining < 0 {
+		remaining = 0
+	}
+
+	return remaining, nil
+}
+
+// getRedisKey constructs the Redis key with prefix
+func (r *rateLimiter) getRedisKey(key string) string {
+	return fmt.Sprintf("%s:%s", r.config.KeyPrefix, key)
+}