From d8c256ebd47dbfa473ae3c52cdee788611a79574 Mon Sep 17 00:00:00 2001
From: Jean Machuca <jeanmachuca@users.noreply.github.com>
Date: Tue, 10 Nov 2026 21:48:23 -0300
Subject: [PATCH] fix: add request validation and rate limiting middleware
 (#39)

---
 internal/server/middleware.go | 90 +++++++++++++++++++++++++++++++++++
 internal/server/server.go     | 24 +++++-----
 2 files changed, 103 insertions(+), 11 deletions(-)
 create mode 100644 internal/server/middleware.go

diff --git a/internal/server/middleware.go b/internal/server/middleware.go
new file mode 100644
index 0000000..fe890b1
--- /dev/null
+++ b/internal/server/middleware.go
@@ -0,0 +1,90 @@
+package server
+
+import (
+	"encoding/json"
+	"net/http"
+	"sync"
+	"time"
+)
+
+// rateLimiter implements a simple fixed-window rate limiter per client IP.
+type rateLimiter struct {
+	mu          sync.Mutex
+	clients     map[string]*clientWindow
+	window      time.Duration
+	maxRequests int
+}
+
+type clientWindow struct {
+	count      int
+	windowStart time.Time
+}
+
+func newRateLimiter(window time.Duration, maxRequests int) *rateLimiter {
+	return &rateLimiter{
+		clients:     make(map[string]*clientWindow),
+		window:      window,
+		maxRequests: maxRequests,
+	}
+}
+
+func (rl *rateLimiter) allow(ip string) bool {
+	rl.mu.Lock()
+	defer rl.mu.Unlock()
+
+	now := time.Now()
+	cw, ok := rl.clients[ip]
+	if !ok || now.Sub(cw.windowStart) >= rl.window {
+		rl.clients[ip] = &clientWindow{count: 1, windowStart: now}
+		return true
+	}
+
+	if cw.count >= rl.maxRequests {
+		return false
+	}
+
+	cw.count++
+	return true
+}
+
+// validateJSONBody checks that the request body is valid JSON for POST requests.
+func validateJSONBody(next http.HandlerFunc) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		if r.Method == http.MethodPost && r.Body != nil {
+			var v interface{}
+			if err := json.NewDecoder(r.Body).Decode(&v); err != nil {
+				http.Error(w, `{"error":"invalid JSON body","code":"E_INVALID_JSON"}`, http.StatusBadRequest)
+				return
+			}
+			_ = v
+		}
+		next.ServeHTTP(w, r)
+	}
+}
+
+// rateLimitMiddleware wraps a handler with per-IP rate limiting.
+func (s *Server) rateLimitMiddleware(next http.HandlerFunc) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		ip := r.RemoteAddr
+		if colon := len(r.RemoteAddr) - 1; colon >= 0 {
+			for i := len(r.RemoteAddr) - 1; i >= 0; i-- {
+				if r.RemoteAddr[i] == ':' {
+					colon = i
+					break
+				}
+			}
+			ip = r.RemoteAddr[:colon]
+		}
+
+		if !s.rateLimiter.allow(ip) {
+			http.Error(w, `{"error":"rate limit exceeded","code":"E_RATE_LIMITED"}`, http.StatusTooManyRequests)
+			return
+		}
+		next.ServeHTTP(w, r)
+	}
+}
+
+// middleware wraps handlers with validation and rate limiting.
+func (s *Server) middleware(next http.HandlerFunc) http.HandlerFunc {
+	return s.rateLimitMiddleware(validateJSONBody(next))
+}
diff --git a/internal/server/server.go b/internal/server/server.go
index 152614f..4fe6595 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -27,6 +27,7 @@ type Server struct {
 	mu        sync.Mutex
 
 	backendType string
+	rateLimiter *rateLimiter
 
 	lastActivity time.Time
 	lastError    string
@@ -160,23 +161,24 @@ func New(modelDir string, backendType string, cloudCfg *cloud.Config) *Server {
 		backend:     b,
 		startTime:   time.Now(),
 		backendType: backendType,
+		rateLimiter: newRateLimiter(1*time.Minute, 60),
 		lastActivity: time.Now(),
 	}
 }
 
 func (s *Server) Listen(addr string) (*http.Server, error) {
 	mux := http.NewServeMux()
-	mux.HandleFunc("/api/generate", s.handleGenerate)
-	mux.HandleFunc("/api/chat", s.handleChat)
-	mux.HandleFunc("/api/tags", s.handleTags)
-	mux.HandleFunc("/api/pull", s.handlePull)
-	mux.HandleFunc("/api/ps", s.handlePs)
-	mux.HandleFunc("/api/delete", s.handleDelete)
-	mux.HandleFunc("/api/adapter", s.handleAdapter)
-	mux.HandleFunc("/api/negotiate", s.handleNegotiate)
-	mux.HandleFunc("/cognitiveos/status", s.handleStatus)
-	mux.HandleFunc("/cognitiveos/capabilities", s.handleCapabilities)
-	mux.HandleFunc("/health", s.handleHealth)
+	mux.HandleFunc("/api/generate", s.middleware(s.handleGenerate))
+	mux.HandleFunc("/api/chat", s.middleware(s.handleChat))
+	mux.HandleFunc("/api/tags", s.middleware(s.handleTags))
+	mux.HandleFunc("/api/pull", s.middleware(s.handlePull))
+	mux.HandleFunc("/api/ps", s.middleware(s.handlePs))
+	mux.HandleFunc("/api/delete", s.middleware(s.handleDelete))
+	mux.HandleFunc("/api/adapter", s.middleware(s.handleAdapter))
+	mux.HandleFunc("/api/negotiate", s.middleware(s.handleNegotiate))
+	mux.HandleFunc("/cognitiveos/status", s.middleware(s.handleStatus))
+	mux.HandleFunc("/cognitiveos/capabilities", s.middleware(s.handleCapabilities))
+	mux.HandleFunc("/health", s.middleware(s.handleHealth))
 
 	go s.idleTimeoutLoop()