ai-router/internal/server/security.go

package server

import (
	"fmt"
	"net"
	"net/http"
	"net/netip"
	"strconv"
	"strings"
	"sync"
	"time"
)

const bodyReadTimeout = 30 * time.Second
const maxLimitEntries = 10000

// Login rate limits: failures only, per IP and globally; successful logins
// never consume or hit these buckets.
const (
	loginFailuresPerMinute       = 10
	loginGlobalFailuresPerMinute = 300
)

// authFailuresPerMinute bounds how many failed /v1 authentications one client
// IP may attempt per minute before further attempts are rejected before the
// database lookup.
const authFailuresPerMinute = 60

func pruneFailures(m map[string][]time.Time, now time.Time) {
	for k, ts := range m {
		if len(ts) == 0 || now.Sub(ts[len(ts)-1]) > time.Minute {
			delete(m, k)
		}
	}
}

// trimRecent drops entries older than one minute while reusing the backing
// array, so a long-lived failure window cannot grow without bound.
func trimRecent(ts []time.Time, now time.Time) []time.Time {
	i := 0
	for i < len(ts) && now.Sub(ts[i]) >= time.Minute {
		i++
	}
	if i > 0 {
		ts = append(ts[:0], ts[i:]...)
	}
	return ts
}

// authAllowed reports whether the client IP may attempt authentication. Only
// failures are counted, so legitimate traffic is unaffected.
func (s *Server) authAllowed(ip string) bool {
	now := time.Now()
	s.mu.Lock()
	defer s.mu.Unlock()
	if now.Sub(s.authFailuresSeen) >= time.Minute {
		pruneFailures(s.authFailures, now)
		s.authFailuresSeen = now
	}
	recent := 0
	for _, t := range s.authFailures[ip] {
		if now.Sub(t) < time.Minute {
			recent++
		}
	}
	return recent < authFailuresPerMinute && (recent > 0 || len(s.authFailures) < maxLimitEntries)
}

// authFail records one failed authentication from ip.
func (s *Server) authFail(ip string) {
	now := time.Now()
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.authFailures[ip]; !ok && len(s.authFailures) >= maxLimitEntries {
		return
	}
	s.authFailures[ip] = append(s.authFailures[ip], now)
}

func readBody(w http.ResponseWriter, r *http.Request, limit int64, timeout time.Duration, read func() error) error {
	controller := http.NewResponseController(w)
	_ = controller.SetReadDeadline(time.Now().Add(timeout))
	defer controller.SetReadDeadline(time.Time{})
	r.Body = http.MaxBytesReader(w, r.Body, limit)
	defer r.Body.Close()
	return read()
}

// SetTrustedProxies must be called before serving requests.
func (s *Server) SetTrustedProxies(cidrs []string) error {
	prefixes := make([]netip.Prefix, 0, len(cidrs))
	for _, cidr := range cidrs {
		p, err := netip.ParsePrefix(cidr)
		if err != nil {
			return fmt.Errorf("invalid trusted proxy CIDR %q: %w", cidr, err)
		}
		if p.Bits() == 0 {
			return fmt.Errorf("trusted proxy CIDR must not trust all addresses")
		}
		prefixes = append(prefixes, p)
	}
	s.trustedProxies = prefixes
	return nil
}

func (s *Server) trusted(ip netip.Addr) bool {
	for _, p := range s.trustedProxies {
		if p.Contains(ip) {
			return true
		}
	}
	return false
}

func (s *Server) clientIP(r *http.Request) string {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		host = r.RemoteAddr
	}
	peer, err := netip.ParseAddr(host)
	if err != nil {
		return host
	}
	peer = peer.Unmap()
	if !s.trusted(peer) {
		return peer.String()
	}
	parts := strings.Split(strings.Join(r.Header.Values("X-Forwarded-For"), ","), ",")
	ips := make([]netip.Addr, len(parts))
	for i, part := range parts {
		ip, err := netip.ParseAddr(strings.TrimSpace(part))
		if err != nil || ip.Zone() != "" {
			return peer.String()
		}
		ips[i] = ip.Unmap()
	}
	for i := len(ips) - 1; i >= 0; i-- {
		if !s.trusted(ips[i]) {
			return ips[i].String()
		}
	}
	return peer.String()
}

type limitEntry struct {
	window           time.Time
	requests, active int
}
type requestLimiter struct {
	mu                          sync.Mutex
	rpm, perKey, global, active int
	entries                     map[string]*limitEntry
	nextCleanup                 time.Time
}

func newRequestLimiter(rpm, perKey, global int) *requestLimiter {
	return &requestLimiter{rpm: rpm, perKey: perKey, global: global, entries: make(map[string]*limitEntry)}
}

// SetLimits must be called before serving requests.
func (s *Server) SetLimits(rpm, perKey, global int) error {
	if rpm <= 0 || perKey <= 0 || global <= 0 {
		return fmt.Errorf("limits must be positive integers")
	}
	s.limiter = newRequestLimiter(rpm, perKey, global)
	return nil
}
func (l *requestLimiter) acquire(key string, now time.Time) (func(), string) {
	l.mu.Lock()
	defer l.mu.Unlock()
	if !now.Before(l.nextCleanup) {
		for k, e := range l.entries {
			if e.active == 0 && now.Sub(e.window) >= time.Minute {
				delete(l.entries, k)
			}
		}
		l.nextCleanup = now.Add(time.Minute)
	}
	e := l.entries[key]
	if e == nil {
		if len(l.entries) >= maxLimitEntries {
			return nil, "60"
		}
		e = &limitEntry{window: now}
		l.entries[key] = e
	}
	if now.Sub(e.window) >= time.Minute {
		e.window = now
		e.requests = 0
	}
	if e.requests >= l.rpm {
		seconds := int((time.Minute - now.Sub(e.window) + time.Second - 1) / time.Second)
		return nil, strconv.Itoa(seconds)
	}
	if e.active >= l.perKey || l.active >= l.global {
		return nil, "1"
	}
	e.requests++
	e.active++
	l.active++
	var once sync.Once
	return func() { once.Do(func() { l.mu.Lock(); defer l.mu.Unlock(); e.active--; l.active-- }) }, ""
}