package server import ( "context" "log/slog" "net" "net/http" "strconv" "strings" "sync" "time" "github.com/VictorVargas/rony-chat-bot/internal/config" ) type ctxKey string const ctxKeyRequestID ctxKey = "requestID" // RequestID assigns a short id per request and exposes it via context + header. func RequestID(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { id := strconv.FormatInt(time.Now().UnixNano(), 36) ctx := context.WithValue(r.Context(), ctxKeyRequestID, id) w.Header().Set("X-Request-ID", id) next.ServeHTTP(w, r.WithContext(ctx)) }) } // Logging emits one structured log line per request after it completes. func Logging(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { start := time.Now() ww := &statusRecorder{ResponseWriter: w, status: http.StatusOK} next.ServeHTTP(ww, r) id, _ := r.Context().Value(ctxKeyRequestID).(string) slog.Info("http", "id", id, "method", r.Method, "path", r.URL.Path, "status", ww.status, "bytes", ww.bytes, "duration_ms", time.Since(start).Milliseconds(), "remote", clientIP(r), ) }) } type statusRecorder struct { http.ResponseWriter status int bytes int } func (r *statusRecorder) WriteHeader(code int) { r.status = code r.ResponseWriter.WriteHeader(code) } func (r *statusRecorder) Write(b []byte) (int, error) { n, err := r.ResponseWriter.Write(b) r.bytes += n return n, err } func (r *statusRecorder) Flush() { if f, ok := r.ResponseWriter.(http.Flusher); ok { f.Flush() } } // CORS rejects requests whose Origin isn't on the allowlist. Requests with // no Origin header (curl, server-to-server) are allowed through. func CORS(allowed []string) func(http.Handler) http.Handler { set := make(map[string]struct{}, len(allowed)) allowAll := false for _, o := range allowed { o = strings.TrimSpace(o) if o == "*" { allowAll = true continue } set[o] = struct{}{} } return func(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { origin := r.Header.Get("Origin") w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS") w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization") if origin != "" { if !allowAll { if _, ok := set[origin]; !ok { http.Error(w, "origin not allowed", http.StatusForbidden) return } } w.Header().Set("Access-Control-Allow-Origin", origin) w.Header().Set("Vary", "Origin") w.Header().Set("Access-Control-Allow-Credentials", "true") } if r.Method == http.MethodOptions { w.WriteHeader(http.StatusNoContent) return } next.ServeHTTP(w, r) }) } } // RateLimit is a per-IP token bucket. cfg.RateLimit.RequestsPerMinute sets // the refill rate; cfg.RateLimit.Burst is the bucket size. func RateLimit(cfg config.RateLimit) func(http.Handler) http.Handler { type bucket struct { tokens float64 lastFill time.Time } var mu sync.Mutex buckets := make(map[string]*bucket) rate := float64(cfg.RequestsPerMinute) / 60.0 burst := float64(cfg.Burst) if rate <= 0 { rate = 0.5 } if burst <= 0 { burst = 5 } return func(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { ip := clientIP(r) now := time.Now() mu.Lock() b, ok := buckets[ip] if !ok { b = &bucket{tokens: burst, lastFill: now} buckets[ip] = b } elapsed := now.Sub(b.lastFill).Seconds() b.tokens += elapsed * rate if b.tokens > burst { b.tokens = burst } b.lastFill = now if b.tokens < 1 { mu.Unlock() w.Header().Set("Retry-After", "60") http.Error(w, "rate limit exceeded", http.StatusTooManyRequests) return } b.tokens-- mu.Unlock() next.ServeHTTP(w, r) }) } } func clientIP(r *http.Request) string { if xff := r.Header.Get("X-Forwarded-For"); xff != "" { if i := strings.IndexByte(xff, ','); i > 0 { return strings.TrimSpace(xff[:i]) } return strings.TrimSpace(xff) } host, _, err := net.SplitHostPort(r.RemoteAddr) if err != nil { return r.RemoteAddr } return host }