1 files changed, 38 insertions, 43 deletions

diff --git a/middleware/ratelimit.go b/middleware/ratelimit.go
index 065a0a2..fcb38cc 100644
--- a/middleware/ratelimit.go
+++ b/middleware/ratelimit.go
@@ -1,22 +1,15 @@
 package middleware
 
 import (
+	"auth-proxy/types"
 	"net"
 	"net/http"
 	"os"
 	"strconv"
-	"sync"
 	"time"
 )
 
-type RateLimiter struct {
-	requests map[string][]int64
-	mutex    sync.Mutex
-	max      int
-	window   time.Duration
-}
-
-func NewRateLimiter() *RateLimiter {
+func NewRateLimiter() *types.RateLimiter {
 	maxRequests, err := strconv.Atoi(os.Getenv("USER_CREATION_RATE_LIMIT_MAX_REQUESTS"))
 	if err != nil || maxRequests <= 0 {
 		maxRequests = 5 // Default value
@@ -27,48 +20,50 @@ func NewRateLimiter() *RateLimiter {
 		windowSeconds = 3600 // Default value (1 hour)
 	}
 
-	return &RateLimiter{
-		requests: make(map[string][]int64),
-		max:      maxRequests,
-		window:   time.Duration(windowSeconds) * time.Second,
+	return &types.RateLimiter{
+		Requests: make(map[string][]int64),
+		Max:      maxRequests,
+		Window:   time.Duration(windowSeconds) * time.Second,
 	}
 }
 
-func (l *RateLimiter) RateLimitMiddleware(next http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		l.mutex.Lock()
-		defer l.mutex.Unlock()
+func RateLimitMiddleware(l *types.RateLimiter) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			l.Mutex.Lock()
+			defer l.Mutex.Unlock()
 
-		ip, _, err := net.SplitHostPort(r.RemoteAddr)
-		if err != nil {
-			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
-			return
-		}
+			ip, _, err := net.SplitHostPort(r.RemoteAddr)
+			if err != nil {
+				http.Error(w, "Internal Server Error", http.StatusInternalServerError)
+				return
+			}
 
-		now := time.Now().Unix()
-		if _, ok := l.requests[ip]; !ok {
-			l.requests[ip] = []int64{}
-		}
+			now := time.Now().Unix()
+			if _, ok := l.Requests[ip]; !ok {
+				l.Requests[ip] = []int64{}
+			}
 
-		// Filter out old timestamps
-		var recentRequests []int64
-		for _, ts := range l.requests[ip] {
-			if now-ts < int64(l.window.Seconds()) {
-				recentRequests = append(recentRequests, ts)
+			// Filter out old timestamps
+			var recentRequests []int64
+			for _, ts := range l.Requests[ip] {
+				if now-ts < int64(l.Window.Seconds()) {
+					recentRequests = append(recentRequests, ts)
+				}
 			}
-		}
-		l.requests[ip] = recentRequests
+			l.Requests[ip] = recentRequests
 
-		if len(l.requests[ip]) >= l.max {
-			http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
-			return
-		}
+			if len(l.Requests[ip]) >= l.Max {
+				http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
+				return
+			}
 
-		// Only add a timestamp for POST requests (actual registration attempts)
-		if r.Method == http.MethodPost {
-			l.requests[ip] = append(l.requests[ip], now)
-		}
+			// Only add a timestamp for POST requests (actual registration attempts)
+			if r.Method == http.MethodPost {
+				l.Requests[ip] = append(l.Requests[ip], now)
+			}
 
-		next.ServeHTTP(w, r)
-	})
+			next.ServeHTTP(w, r)
+		})
+	}
 }