fix: ratelimit crash under concurrent requests

src/middleware/ratelimit.go

@@ -10,9 +10,39 @@ import (
 	"github.com/gin-gonic/gin"
 )
 
-var ipMap = sync.Map{}
+const (
+	ratelimitWindow   = time.Minute
+	ratelimitRequests = 30
+)
+
+type limiter struct {
+	mu          sync.Mutex
+	now         func() time.Time
+	windowStart time.Time
+	counts      map[string]int
+}
+
+func newLimiter(now func() time.Time) *limiter {
+	return &limiter{now: now, windowStart: now(), counts: map[string]int{}}
+}
+
+func (l *limiter) allow(ip string) bool {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	if now := l.now(); now.Sub(l.windowStart) >= ratelimitWindow {
+		l.windowStart = now
+		l.counts = map[string]int{}
+	}
+	l.counts[ip]++
+	return l.counts[ip] <= ratelimitRequests
+}
 
 func Ratelimit() gin.HandlerFunc {
+	return ratelimit(newLimiter(time.Now))
+}
+
+func ratelimit(l *limiter) gin.HandlerFunc {
 	return func(c *gin.Context) {
 
 		if os.Getenv("DISABLE_RATELIMIT") == "true" {
@@ -25,26 +55,7 @@ func Ratelimit() gin.HandlerFunc {
 			return
 		}
 
-		ip := c.ClientIP()
-
-		// log request count as the value, ip as key
-		// if the ip is not in the map, create a new entry with a value of 1
-		// if they exceed 30 requests in 1 minute, return a 429
-
-		// get the value from the map
-		val, ok := ipMap.Load(ip)
-		if !ok {
-			// if the ip is not in the map, create a new entry with a value of 1
-			ipMap.Store(ip, 1)
-			c.Next()
-			return
-		}
-
-		// if the ip is in the map, increment the value
-		ipMap.Store(ip, val.(int)+1)
-
-		// if they exceed 30 requests in 1 minute, return a 429
-		if val.(int) > 30 {
+		if !l.allow(c.ClientIP()) {
 			c.HTML(429, "home.html", gin.H{
 				"errorMessage": "You have exceeded the request limit. Please try again in a minute.",
 				"version":      config.Version,
@@ -53,17 +64,6 @@ func Ratelimit() gin.HandlerFunc {
 			return
 		}
 
-		// subtract 1 from the value after 1 minute if the value exists and is greater than 0
-		time.AfterFunc(time.Minute, func() {
-			val, ok := ipMap.Load(ip)
-			if ok && val.(int) > 0 {
-				ipMap.Store(ip, val.(int)-1)
-			}
-
-			// if the value is 0, delete the entry from the map
-			if val.(int) == 0 {
-				ipMap.Delete(ip)
-			}
-		})
+		c.Next()
 	}
 }

src/middleware/ratelimit_test.go

@@ -0,0 +1,111 @@
+package middleware
+
+import (
+	"fmt"
+	"html/template"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/gin-gonic/gin"
+)
+
+type fakeClock struct {
+	mu sync.Mutex
+	t  time.Time
+}
+
+func (f *fakeClock) now() time.Time {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return f.t
+}
+
+func (f *fakeClock) advance(d time.Duration) {
+	f.mu.Lock()
+	f.t = f.t.Add(d)
+	f.mu.Unlock()
+}
+
+func limitedEngine(l *limiter) *gin.Engine {
+	gin.SetMode(gin.TestMode)
+	r := gin.New()
+	r.SetHTMLTemplate(template.Must(template.New("home.html").Parse("limited")))
+	r.Use(ratelimit(l))
+	r.GET("/version", func(c *gin.Context) { c.String(http.StatusOK, "ok") })
+	return r
+}
+
+func request(r *gin.Engine, ip string) int {
+	req := httptest.NewRequest(http.MethodGet, "/version", nil)
+	req.RemoteAddr = ip + ":1234"
+	w := httptest.NewRecorder()
+	r.ServeHTTP(w, req)
+	return w.Code
+}
+
+func TestRatelimitLimitsAndRecovers(t *testing.T) {
+	clock := &fakeClock{t: time.Unix(1_700_000_000, 0)}
+	r := limitedEngine(newLimiter(clock.now))
+
+	for i := 1; i <= ratelimitRequests; i++ {
+		if code := request(r, "192.0.2.1"); code != http.StatusOK {
+			t.Fatalf("request %d: status %d, want 200", i, code)
+		}
+	}
+	// Overshoot well past the limit, as a script or a shared IP would.
+	for i := 0; i < 70; i++ {
+		if code := request(r, "192.0.2.1"); code != http.StatusTooManyRequests {
+			t.Fatalf("over-limit request: status %d, want 429", code)
+		}
+	}
+	if code := request(r, "192.0.2.2"); code != http.StatusOK {
+		t.Fatalf("other client limited too: status %d", code)
+	}
+
+	clock.advance(ratelimitWindow)
+	if code := request(r, "192.0.2.1"); code != http.StatusOK {
+		t.Fatalf("client still limited after the window: status %d", code)
+	}
+}
+
+func TestRatelimitForgetsOldClients(t *testing.T) {
+	clock := &fakeClock{t: time.Unix(1_700_000_000, 0)}
+	l := newLimiter(clock.now)
+	r := limitedEngine(l)
+
+	for i := 0; i < 5000; i++ {
+		request(r, fmt.Sprintf("198.51.%d.%d", i/256, i%256))
+	}
+	l.mu.Lock()
+	seen := len(l.counts)
+	l.mu.Unlock()
+	if seen != 5000 {
+		t.Fatalf("limiter tracked %d clients, want 5000", seen)
+	}
+
+	clock.advance(ratelimitWindow)
+	request(r, "192.0.2.1")
+
+	l.mu.Lock()
+	n := len(l.counts)
+	l.mu.Unlock()
+	if n != 1 {
+		t.Fatalf("limiter holds %d clients after the window, want 1", n)
+	}
+}
+
+func TestRatelimitConcurrent(t *testing.T) {
+	r := limitedEngine(newLimiter(time.Now))
+	var wg sync.WaitGroup
+	for i := 0; i < 200; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			request(r, "192.0.2.9")
+		}()
+	}
+	wg.Wait()
+}