feat: add Redis based rate limiting for multiple instances

internal/middleware/rate-limit.go

@@ -16,9 +16,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 import (
+	"errors"
 	"fmt"
+	"github.com/go-redis/redis_rate/v10"
 	"github.com/gorilla/mux"
 	"github.com/jkaninda/goma-gateway/pkg/logger"
+	"github.com/redis/go-redis/v9"
+	"golang.org/x/net/context"
 	"net/http"
 	"time"
 )
@@ -49,30 +53,62 @@ func (rl *TokenRateLimiter) RateLimitMiddleware() mux.MiddlewareFunc {
 func (rl *RateLimiter) RateLimitMiddleware() mux.MiddlewareFunc {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			clientID := getRealIP(r)
-			rl.mu.Lock()
-			client, exists := rl.ClientMap[clientID]
-			if !exists || time.Now().After(client.ExpiresAt) {
-				client = &Client{
-					RequestCount: 0,
-					ExpiresAt:    time.Now().Add(rl.Window),
+			clientIP := getRealIP(r)
+			logger.Debug("rate limiter: clientID: %s, redisBased: %s", clientIP, rl.RedisBased)
+			if rl.RedisBased {
+				err := redisRateLimiter(clientIP, rl.Requests)
+				if err != nil {
+					logger.Error("Redis Rate limiter error: %s", err.Error())
+					RespondWithError(w, http.StatusTooManyRequests, fmt.Sprintf("%d Too many requests, API rate limit exceeded. Please try again later", http.StatusTooManyRequests), rl.ErrorInterceptor)
+					return
 				}
-				rl.ClientMap[clientID] = client
-			}
-			client.RequestCount++
-			rl.mu.Unlock()
+				// Proceed to the next handler if rate limit is not exceeded
+				next.ServeHTTP(w, r)
+			} else {
+				rl.mu.Lock()
+				client, exists := rl.ClientMap[clientIP]
+				if !exists || time.Now().After(client.ExpiresAt) {
+					client = &Client{
+						RequestCount: 0,
+						ExpiresAt:    time.Now().Add(rl.Window),
+					}
+					rl.ClientMap[clientIP] = client
+				}
+				client.RequestCount++
+				rl.mu.Unlock()
 
-			if client.RequestCount > rl.Requests {
-				logger.Error("Too many requests from IP: %s %s %s", clientID, r.URL, r.UserAgent())
-				//Update Origin Cors Headers
-				if allowedOrigin(rl.Origins, r.Header.Get("Origin")) {
-					w.Header().Set("Access-Control-Allow-Origin", r.Header.Get("Origin"))
+				if client.RequestCount > rl.Requests {
+					logger.Error("Too many requests from IP: %s %s %s", clientIP, r.URL, r.UserAgent())
+					//Update Origin Cors Headers
+					if allowedOrigin(rl.Origins, r.Header.Get("Origin")) {
+						w.Header().Set("Access-Control-Allow-Origin", r.Header.Get("Origin"))
+					}
+					RespondWithError(w, http.StatusTooManyRequests, fmt.Sprintf("%d Too many requests, API rate limit exceeded. Please try again later", http.StatusTooManyRequests), rl.ErrorInterceptor)
+					return
 				}
-				RespondWithError(w, http.StatusTooManyRequests, fmt.Sprintf("%d Too many requests, API rate limit exceeded. Please try again later", http.StatusTooManyRequests), rl.ErrorInterceptor)
-				return
 			}
 			// Proceed to the next handler if rate limit is not exceeded
 			next.ServeHTTP(w, r)
 		})
 	}
 }
+func redisRateLimiter(clientIP string, rate int) error {
+	ctx := context.Background()
+
+	res, err := limiter.Allow(ctx, clientIP, redis_rate.PerMinute(rate))
+	if err != nil {
+		return err
+	}
+	if res.Remaining == 0 {
+		return errors.New("rate limit exceeded")
+	}
+
+	return nil
+}
+func InitRedis(addr, password string) {
+	Rdb = redis.NewClient(&redis.Options{
+		Addr:     addr,
+		Password: password,
+	})
+	limiter = redis_rate.NewLimiter(Rdb)
+}

internal/middleware/types.go

@@ -33,6 +33,7 @@ type RateLimiter struct {
 	mu               sync.Mutex
 	Origins          []string
 	ErrorInterceptor errorinterceptor.ErrorInterceptor
+	RedisBased       bool
 }
 
 // Client stores request count and window expiration for each client.
@@ -42,12 +43,13 @@ type Client struct {
 }
 
 // NewRateLimiterWindow creates a new RateLimiter.
-func NewRateLimiterWindow(requests int, window time.Duration, origin []string) *RateLimiter {
+func NewRateLimiterWindow(requests int, window time.Duration, redisBased bool, origin []string) *RateLimiter {
 	return &RateLimiter{
-		Requests:  requests,
-		Window:    window,
-		ClientMap: make(map[string]*Client),
-		Origins:   origin,
+		Requests:   requests,
+		Window:     window,
+		ClientMap:  make(map[string]*Client),
+		Origins:    origin,
+		RedisBased: redisBased,
 	}
 }
 

internal/middleware/var.go

@@ -17,7 +17,17 @@
 
 package middleware
 
+import (
+	"github.com/go-redis/redis_rate/v10"
+	"github.com/redis/go-redis/v9"
+)
+
 // sqlPatterns contains SQL injections patters
 const sqlPatterns = `(?i)(union|select|drop|insert|delete|update|create|alter|exec|;|--)`
 const traversalPatterns = `\.\./`
 const xssPatterns = `(?i)<script|onerror|onload`
+
+var (
+	Rdb     *redis.Client
+	limiter *redis_rate.Limiter
+)