ai-router/internal/server/proxy.go

package server

import (
	"ai-router/internal/protocol"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"sort"
	"strings"
	"time"
)

func resolve(cs []channel, k apiKey, model string) (channel, string, error) {
	name, alias, ok := strings.Cut(model, "/")
	if !ok || alias == "" {
		return channel{}, "", errors.New("model must be channel/alias")
	}
	for _, c := range cs {
		if c.Name != name {
			continue
		}
		if !c.Enabled || !k.allows(c.ID) {
			return channel{}, "", errors.New("channel unavailable")
		}
		if actual, ok := c.Models[alias]; ok {
			return c, actual, nil
		}
		return channel{}, "", errors.New("model unavailable")
	}
	return channel{}, "", errors.New("channel unavailable")
}
func (s *Server) proxy(w http.ResponseWriter, r *http.Request) {
	kinds := map[string]protocol.Kind{"/v1/chat/completions": protocol.Chat, "/v1/responses": protocol.Responses, "/v1/messages": protocol.Anthropic}
	down, ok := kinds[r.URL.Path]
	models := r.URL.Path == "/v1/models"
	if !ok && !models {
		errJSON(w, 404, "not found")
		return
	}
	if (models && r.Method != "GET") || (!models && r.Method != "POST") {
		errJSON(w, 405, "method not allowed")
		return
	}
	ip := s.clientIP(r)
	if !s.authAllowed(ip) {
		w.Header().Set("Retry-After", "60")
		errJSON(w, 429, "request limit exceeded")
		return
	}
	k, e := s.authenticate(r)
	if e != nil {
		s.authFail(ip)
		errJSON(w, 401, "invalid API key")
		return
	}
	release, retry := s.limiter.acquire(k.ID, time.Now())
	if release == nil {
		w.Header().Set("Retry-After", retry)
		errJSON(w, 429, "request limit exceeded")
		return
	}
	start := time.Now()
	l := requestLog{CreatedAt: start.UTC().Format(time.RFC3339Nano), KeyID: k.ID, KeyName: k.Name, Downstream: string(down), Status: 500}
	// Release the concurrency slot before the log write, which can block on a
	// full queue; rate-limited and unauthenticated requests stay unlogged.
	defer func() {
		release()
		if !models {
			l.DurationMS = time.Since(start).Milliseconds()
			s.record(l)
		}
	}()
	fail := func(status int, msg string) { l.Status = status; l.Error = msg; errJSON(w, status, msg) }
	cs, e := s.loadChannels()
	if e != nil {
		fail(500, "database error")
		return
	}
	if models {
		ids := []string{}
		for _, c := range cs {
			if c.Enabled && k.allows(c.ID) {
				for alias := range c.Models {
					ids = append(ids, c.Name+"/"+alias)
				}
			}
		}
		sort.Strings(ids)
		out := []any{}
		for _, id := range ids {
			out = append(out, map[string]string{"id": id, "object": "model", "owned_by": "ai-router"})
		}
		l.Status = 200
		writeJSON(w, map[string]any{"object": "list", "data": out})
		return
	}
	var b []byte
	e = readBody(w, r, maxBody, bodyReadTimeout, func() error {
		var err error
		b, err = io.ReadAll(r.Body)
		return err
	})
	if e != nil {
		var tooBig *http.MaxBytesError
		if errors.As(e, &tooBig) {
			fail(413, "request body too large")
		} else {
			fail(400, "request body unreadable")
		}
		return
	}
	var req struct {
		Model  string `json:"model"`
		Stream bool   `json:"stream"`
	}
	if json.Unmarshal(b, &req) != nil || req.Model == "" {
		fail(400, "invalid request")
		return
	}
	l.RequestModel = req.Model
	l.Stream = req.Stream
	c, actual, e := resolve(cs, k, req.Model)
	if e != nil {
		fail(403, "model or channel unavailable")
		return
	}
	l.Channel = c.Name
	l.ActualModel = actual
	l.Upstream = c.Protocol
	up := protocol.Kind(c.Protocol)
	payload, e := protocol.ConvertRequest(down, up, b, actual)
	if e != nil {
		fail(400, "unsupported request conversion")
		return
	}
	if req.Stream && up == protocol.Chat {
		var obj map[string]json.RawMessage
		if json.Unmarshal(payload, &obj) == nil {
			options := map[string]json.RawMessage{}
			if raw, ok := obj["stream_options"]; ok {
				_ = json.Unmarshal(raw, &options)
			}
			if options == nil {
				options = map[string]json.RawMessage{}
			}
			options["include_usage"] = json.RawMessage("true")
			obj["stream_options"], _ = json.Marshal(options)
			payload, _ = json.Marshal(obj)
		}
	}
	suffix := map[protocol.Kind]string{protocol.Chat: "/chat/completions", protocol.Responses: "/responses", protocol.Anthropic: "/v1/messages"}[up]
	ctx, cancel := context.WithTimeout(r.Context(), 10*time.Minute)
	defer cancel()
	request, e := http.NewRequestWithContext(ctx, "POST", c.BaseURL+suffix, bytes.NewReader(payload))
	if e != nil {
		fail(502, "invalid upstream configuration")
		return
	}
	request.Header.Set("Content-Type", "application/json")
	if up == protocol.Anthropic {
		request.Header.Set("x-api-key", c.APIKey)
		request.Header.Set("anthropic-version", "2023-06-01")
		if down == up {
			if v := r.Header.Get("anthropic-version"); safeForwardedHeader(v, 32) {
				request.Header.Set("anthropic-version", v)
			}
			if v := r.Header.Get("anthropic-beta"); safeForwardedHeader(v, 256) {
				request.Header.Set("anthropic-beta", v)
			}
		}
	} else {
		request.Header.Set("Authorization", "Bearer "+c.APIKey)
	}
	response, e := s.client.Do(request)
	if e != nil {
		fail(502, "upstream request failed")
		return
	}
	defer response.Body.Close()
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		status := response.StatusCode
		if status < 400 {
			status = 502
		}
		fail(status, "upstream rejected request")
		return
	}
	if req.Stream {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")
		w.Header().Set("X-Accel-Buffering", "no")
		w.WriteHeader(response.StatusCode)
		l.Status = response.StatusCode
		streamWriter := &deadlineWriter{ResponseWriter: w}
		u, e := protocol.Stream(ctx, streamWriter, response.Body, up, down, req.Model)
		l.Usage = u
		if e != nil {
			l.Status = 502
			l.Error = "upstream stream interrupted"
			l.Usage.Known = false
			if ctx.Err() == nil {
				event := "data: {\"error\":{\"type\":\"upstream_error\",\"message\":\"upstream stream interrupted\"}}\n\n"
				if down != protocol.Chat {
					event = "event: error\n" + event
				}
				_, _ = io.WriteString(streamWriter, event)
				if f, ok := w.(http.Flusher); ok {
					f.Flush()
				}
			}
		}
		return
	}
	b, e = io.ReadAll(io.LimitReader(response.Body, maxBody+1))
	if e != nil || len(b) > maxBody {
		fail(502, "invalid upstream response")
		return
	}
	l.Usage = protocol.ExtractUsage(up, b)
	out, e := protocol.ConvertResponse(up, down, b)
	if e != nil {
		fail(502, "unsupported response conversion")
		return
	}
	l.Status = response.StatusCode
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(response.StatusCode)
	if _, e = w.Write(out); e != nil {
		l.Status = 499
		l.Error = "client disconnected"
	}
}

type deadlineWriter struct{ http.ResponseWriter }

// safeForwardedHeader reports whether v is a short printable token list that
// is safe to forward upstream; anything else is dropped so clients cannot
// smuggle arbitrary features or exotic header content past the channel
// configuration.
func safeForwardedHeader(v string, max int) bool {
	if len(v) == 0 || len(v) > max {
		return false
	}
	for i := 0; i < len(v); i++ {
		c := v[i]
		switch {
		case c >= 'a' && c <= 'z', c >= 'A' && c <= 'Z', c >= '0' && c <= '9':
		case c == '-' || c == '_' || c == '.' || c == ',' || c == ' ':
		default:
			return false
		}
	}
	return true
}

func (w *deadlineWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
func (w *deadlineWriter) WriteHeader(status int) {
	_ = http.NewResponseController(w.ResponseWriter).SetWriteDeadline(time.Now().Add(30 * time.Second))
	w.ResponseWriter.WriteHeader(status)
}

func (w *deadlineWriter) Write(b []byte) (int, error) {
	_ = http.NewResponseController(w.ResponseWriter).SetWriteDeadline(time.Now().Add(30 * time.Second))
	return w.ResponseWriter.Write(b)
}
func (w *deadlineWriter) FlushError() error {
	_ = http.NewResponseController(w.ResponseWriter).SetWriteDeadline(time.Now().Add(30 * time.Second))
	return http.NewResponseController(w.ResponseWriter).Flush()
}
func (w *deadlineWriter) Flush() { _ = w.FlushError() }