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 728ef1e..22d2b9b 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 @@ -156,27 +157,28 @@ func New(modelDir string, backendType string, cloudCfg *cloud.Config) *Server { } return &Server{ - models: model.NewManager(modelDir), - backend: b, - startTime: time.Now(), - backendType: backendType, + models: model.NewManager(modelDir), + 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()