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
+)

internal/route.go

@@ -35,6 +35,10 @@ func init() {
 func (gatewayServer GatewayServer) Initialize() *mux.Router {
 	gateway := gatewayServer.gateway
 	middlewares := gatewayServer.middlewares
+	redisBased := false
+	if len(gateway.Redis.Addr) != 0 {
+		redisBased = true
+	}
 	//Routes background healthcheck
 	routesHealthCheck(gateway.Routes)
 	r := mux.NewRouter()
@@ -61,9 +65,9 @@ func (gatewayServer GatewayServer) Initialize() *mux.Router {
 		blockCommon := middleware.BlockCommon{}
 		r.Use(blockCommon.BlockExploitsMiddleware)
 	}
-	if gateway.RateLimit != 0 {
+	if gateway.RateLimit > 0 {
 		//rateLimiter := middleware.NewRateLimiter(gateway.RateLimit, time.Minute)
-		limiter := middleware.NewRateLimiterWindow(gateway.RateLimit, time.Minute, gateway.Cors.Origins) //  requests per minute
+		limiter := middleware.NewRateLimiterWindow(gateway.RateLimit, time.Minute, redisBased, gateway.Cors.Origins) //  requests per minute
 		// Add rate limit middleware to all routes, if defined
 		r.Use(limiter.RateLimitMiddleware())
 	}
@@ -229,7 +233,7 @@ func (gatewayServer GatewayServer) Initialize() *mux.Router {
 			// Apply route rate limit
 			if route.RateLimit > 0 {
 				//rateLimiter := middleware.NewRateLimiter(gateway.RateLimit, time.Minute)
-				limiter := middleware.NewRateLimiterWindow(route.RateLimit, time.Minute, route.Cors.Origins) //  requests per minute
+				limiter := middleware.NewRateLimiterWindow(route.RateLimit, time.Minute, redisBased, route.Cors.Origins) //  requests per minute
 				// Add rate limit middleware to all routes, if defined
 				router.Use(limiter.RateLimitMiddleware())
 			}

internal/server.go

@@ -19,7 +19,9 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
+	"github.com/jkaninda/goma-gateway/internal/middleware"
 	"github.com/jkaninda/goma-gateway/pkg/logger"
+	"github.com/redis/go-redis/v9"
 	"net/http"
 	"os"
 	"sync"
@@ -30,8 +32,19 @@ import (
 func (gatewayServer GatewayServer) Start(ctx context.Context) error {
 	logger.Info("Initializing routes...")
 	route := gatewayServer.Initialize()
+	gateway := gatewayServer.gateway
 	logger.Debug("Routes count=%d Middlewares count=%d", len(gatewayServer.gateway.Routes), len(gatewayServer.middlewares))
 	logger.Info("Initializing routes...done")
+	if len(gateway.Redis.Addr) != 0 {
+		middleware.InitRedis(gateway.Redis.Addr, gateway.Redis.Password)
+		defer func(Rdb *redis.Client) {
+			err := Rdb.Close()
+			if err != nil {
+				logger.Error("Redis connection closed with error: %v", err)
+			}
+		}(middleware.Rdb)
+	}
+
 	tlsConfig := &tls.Config{}
 	var listenWithTLS = false
 	if cert := gatewayServer.gateway.SSLCertFile; cert != "" && gatewayServer.gateway.SSLKeyFile != "" {

internal/types.go

@@ -178,6 +178,8 @@ type Gateway struct {
 	SSLCertFile string `yaml:"sslCertFile" env:"GOMA_SSL_CERT_FILE, overwrite"`
 	// SSLKeyFile SSL Private key  file
 	SSLKeyFile string `yaml:"sslKeyFile" env:"GOMA_SSL_KEY_FILE, overwrite"`
+	// Redis contains redis database details
+	Redis Redis `yaml:"redis"`
 	// WriteTimeout defines proxy write timeout
 	WriteTimeout int `yaml:"writeTimeout" env:"GOMA_WRITE_TIMEOUT, overwrite"`
 	// ReadTimeout defines proxy read timeout
@@ -204,6 +206,7 @@ type Gateway struct {
 	InterceptErrors []int `yaml:"interceptErrors"`
 	// Cors holds proxy global cors
 	Cors Cors `yaml:"cors"`
+
 	// Routes holds proxy routes
 	Routes []Route `yaml:"routes"`
 }
@@ -287,3 +290,8 @@ type Health struct {
 	Interval        string
 	HealthyStatuses []int
 }
+type Redis struct {
+	// Addr redis hostname and post number :
+	Addr     string `yaml:"addr"`
+	Password string `yaml:"password"`
+}