proxy.go

// Package proxy is the request path: resolve the caller, check policy, forward
// to the model API, stream the answer back, and account for what was spent.
//
// # The bug this package is built around
//
// A streaming LLM response is the one case where "the client went away" is not
// a free event. If the inbound request's context is not threaded into the
// upstream call, hanging up on the gateway stops nothing: the model keeps
// generating and the account keeps being billed for tokens that no longer have
// a reader. Nothing on the happy path notices, no error is logged anywhere, and
// the only artifact is a bill that does not match traffic.
//
// So two rules run through this file:
//
//  1. The upstream request is built with http.NewRequestWithContext(r.Context()),
//     never context.Background(). Client hangs up, upstream call dies with it.
//
//  2. The copy loop checks ctx.Done() before every read/write pair rather than
//     waiting to discover a write error. A departed client's writes can keep
//     succeeding into kernel and proxy buffers for a surprisingly long time, so
//     "the write failed" is a late and unreliable signal.
//
// Usage is recorded on abort too. Those tokens existed.
package proxy

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/comedianhhh/tsgate/internal/audit"
	"github.com/comedianhhh/tsgate/internal/identity"
	"github.com/comedianhhh/tsgate/internal/metrics"
	"github.com/comedianhhh/tsgate/internal/policy"
)

// maxRequestBytes caps the body we will buffer. Chat completion requests are
// small; anything larger is a mistake or an attack, and either way we would
// rather refuse it than hold it in memory.
const maxRequestBytes = 1 << 20 // 1 MiB

// statusClientClosed is nginx's 499. It is not an IANA status and is never
// written to the wire — by the time we use it there is nobody to write to. It
// exists so aborts are distinguishable in logs and metrics from real errors.
const statusClientClosed = 499

// Upstream is the model API being fronted.
type Upstream struct {
	// BaseURL is an origin, e.g. "https://api.openai.com". The inbound request
	// path is appended to it unchanged, so the gateway stays agnostic about
	// which endpoints the upstream offers.
	BaseURL string
	// APIKey is injected server-side. Callers never hold it — that is most of
	// the point of running this gateway.
	APIKey string
	Client *http.Client
}

// Options tunes the request path.
type Options struct {
	// ResponseHeaderTimeout bounds how long the gateway waits for the upstream
	// to begin responding.
	//
	// It deliberately does not bound the stream itself. A long generation is
	// not a failure, and an overall deadline would truncate exactly the
	// requests users care most about.
	ResponseHeaderTimeout time.Duration
}

// Proxy implements http.Handler.
type Proxy struct {
	up       Upstream
	resolver identity.Resolver
	policy   *policy.Engine
	audit    audit.Sink
	metrics  *metrics.Registry
	opts     Options
}

// New wires a Proxy. A nil audit sink or metrics registry is replaced with a
// no-op, so callers are never forced to construct observability they do not
// want in a test.
func New(up Upstream, r identity.Resolver, p *policy.Engine, a audit.Sink, m *metrics.Registry, opts Options) *Proxy {
	if up.Client == nil {
		up.Client = defaultClient(opts.ResponseHeaderTimeout)
	}
	if a == nil {
		a = audit.Discard{}
	}
	if m == nil {
		m = metrics.New()
	}
	return &Proxy{up: up, resolver: r, policy: p, audit: a, metrics: m, opts: opts}
}

func defaultClient(headerTimeout time.Duration) *http.Client {
	if headerTimeout <= 0 {
		headerTimeout = 30 * time.Second
	}
	tr := http.DefaultTransport.(*http.Transport).Clone()
	tr.ResponseHeaderTimeout = headerTimeout
	// No Client.Timeout: that would kill long streams. Cancellation is the
	// inbound context's job, and slow-to-start upstreams are the transport's.
	return &http.Client{Transport: tr}
}

// chatRequest is the sliver of the request body the gateway needs. Everything
// else is forwarded byte-for-byte, so new upstream fields need no change here.
type chatRequest struct {
	Model  string `json:"model"`
	Stream bool   `json:"stream"`
}

func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	start := time.Now()
	ctx := r.Context()

	var (
		ev       = audit.Event{Time: start}
		status   int
		tokens   int
		bytesOut int64
		aborted  bool
		failure  error
	)

	// One audit event and one metrics observation per request, on every exit
	// path including panics-free early returns. This is why the whole handler
	// uses named locals rather than writing the log at each return.
	defer func() {
		ev.Status = status
		ev.Tokens = tokens
		ev.Bytes = bytesOut
		ev.Aborted = aborted
		ev.DurationMS = float64(time.Since(start).Microseconds()) / 1000
		if failure != nil {
			ev.Error = failure.Error()
		}
		p.audit.Write(ev)
		p.metrics.Observe(status, ev.DurationMS, aborted, tokens)
	}()

	if r.Method != http.MethodPost {
		status, ev.Outcome = http.StatusMethodNotAllowed, "method_not_allowed"
		http.Error(w, "only POST is proxied", status)
		return
	}

	// 1. Identity, from the connection.
	id, err := p.resolver.Resolve(ctx, r.RemoteAddr)
	if err != nil {
		status, ev.Outcome, failure = http.StatusUnauthorized, "unidentified", err
		http.Error(w, "caller could not be identified on this tailnet", status)
		return
	}
	ev.Login, ev.Node = id.Login, id.Node

	// 2. Body. Buffered because we need the model name to make a policy
	// decision before a single byte reaches the upstream, and because the
	// upstream request has to be replayable from the start.
	body, err := io.ReadAll(io.LimitReader(r.Body, maxRequestBytes+1))
	if err != nil {
		status, ev.Outcome, failure = http.StatusBadRequest, "read_error", err
		http.Error(w, "could not read request body", status)
		return
	}
	if len(body) > maxRequestBytes {
		status, ev.Outcome = http.StatusRequestEntityTooLarge, "body_too_large"
		http.Error(w, "request body too large", status)
		return
	}

	var cr chatRequest
	if err := json.Unmarshal(body, &cr); err != nil || cr.Model == "" {
		status, ev.Outcome = http.StatusBadRequest, "bad_request"
		http.Error(w, `request body must be JSON containing a "model" field`, status)
		return
	}
	ev.Model = cr.Model

	// 3. Policy.
	decision := p.policy.Check(id.Login, cr.Model)
	ev.Outcome = decision.Outcome.String()
	if !decision.Allowed() {
		if decision.Outcome == policy.DenyQuota {
			status = http.StatusTooManyRequests
		} else {
			status = http.StatusForbidden
		}
		http.Error(w, decision.Reason, status)
		return
	}

	// 4. Upstream request. The inbound context is threaded through here; this
	// single argument is what makes a client hang-up stop the generation.
	upReq, err := http.NewRequestWithContext(ctx, http.MethodPost, p.upstreamURL(r), bytes.NewReader(body))
	if err != nil {
		status, ev.Outcome, failure = http.StatusInternalServerError, "request_build_error", err
		http.Error(w, "internal error", status)
		return
	}
	copyRequestHeaders(upReq, r)
	if p.up.APIKey != "" {
		upReq.Header.Set("Authorization", "Bearer "+p.up.APIKey)
	}

	resp, err := p.up.Client.Do(upReq)
	if err != nil {
		// Distinguish "the client left" from "the upstream broke". They look
		// identical here and mean opposite things operationally.
		if ctx.Err() != nil {
			aborted, status, ev.Outcome, failure = true, statusClientClosed, "client_gone_before_response", ctx.Err()
			return
		}
		if isTimeout(err) {
			status, ev.Outcome = http.StatusGatewayTimeout, "upstream_timeout"
		} else {
			status, ev.Outcome = http.StatusBadGateway, "upstream_error"
		}
		failure = err
		http.Error(w, "upstream unavailable", status)
		return
	}
	defer resp.Body.Close()

	copyResponseHeaders(w, resp)
	status = resp.StatusCode
	w.WriteHeader(status)

	if isEventStream(resp.Header.Get("Content-Type")) {
		bytesOut, tokens, failure = streamSSE(ctx, w, resp.Body)
	} else {
		bytesOut, failure = copyBody(ctx, w, resp.Body)
	}

	// 5. Account for what was actually generated, whether or not it was read.
	p.policy.Record(id.Login, tokens)

	if failure != nil && ctx.Err() != nil {
		aborted, ev.Outcome, failure = true, "client_gone_mid_stream", ctx.Err()
	}
}

func (p *Proxy) upstreamURL(r *http.Request) string {
	u := strings.TrimSuffix(p.up.BaseURL, "/") + r.URL.Path
	if r.URL.RawQuery != "" {
		u += "?" + r.URL.RawQuery
	}
	return u
}

// hopByHop headers are connection-scoped and must not be forwarded.
var hopByHop = map[string]bool{
	"connection":          true,
	"keep-alive":          true,
	"proxy-authenticate":  true,
	"proxy-authorization": true,
	"te":                  true,
	"trailer":             true,
	"transfer-encoding":   true,
	"upgrade":             true,
}

func copyRequestHeaders(dst *http.Request, src *http.Request) {
	for k, vs := range src.Header {
		lk := strings.ToLower(k)
		// The caller's own Authorization is dropped, not forwarded. Callers do
		// not hold upstream credentials, and letting one through would turn
		// the gateway into a way to bypass its own policy.
		if hopByHop[lk] || lk == "authorization" || lk == "host" {
			continue
		}
		for _, v := range vs {
			dst.Header.Add(k, v)
		}
	}
	if dst.Header.Get("Content-Type") == "" {
		dst.Header.Set("Content-Type", "application/json")
	}
}

func copyResponseHeaders(w http.ResponseWriter, resp *http.Response) {
	for k, vs := range resp.Header {
		if hopByHop[strings.ToLower(k)] {
			continue
		}
		for _, v := range vs {
			w.Header().Add(k, v)
		}
	}
}

func isEventStream(contentType string) bool {
	return strings.Contains(strings.ToLower(contentType), "text/event-stream")
}

var (
	ssePrefix = []byte("data:")
	sseDone   = []byte("[DONE]")
)

// streamSSE forwards a server-sent-event stream, flushing every line so the
// caller sees tokens as they are produced rather than when a buffer fills.
//
// The returned token count is approximate: it counts SSE data frames, and an
// OpenAI-style chat stream emits roughly one frame per token. Exact accounting
// would mean parsing every frame's JSON on the hot path to find a usage block
// that only arrives at the end — and never arrives at all on the aborted
// requests that most need to be charged for. An approximation available on
// every path beat an exact number available only on the happy one.
func streamSSE(ctx context.Context, w http.ResponseWriter, body io.Reader) (bytesOut int64, frames int, err error) {
	flusher, canFlush := w.(http.Flusher)
	br := bufio.NewReaderSize(body, 32<<10)

	for {
		// Checked before every read/write pair. See the package comment: a
		// write error is a late and unreliable way to learn the reader left.
		select {
		case <-ctx.Done():
			return bytesOut, frames, ctx.Err()
		default:
		}

		line, readErr := br.ReadBytes('\n')
		if len(line) > 0 {
			if bytes.HasPrefix(line, ssePrefix) && !bytes.Contains(line, sseDone) {
				frames++
			}
			n, writeErr := w.Write(line)
			bytesOut += int64(n)
			if canFlush {
				// Without this the stream sits in the response buffer and the
				// caller sees the whole answer at once, which defeats the
				// entire point of streaming it.
				flusher.Flush()
			}
			if writeErr != nil {
				return bytesOut, frames, writeErr
			}
		}
		if readErr != nil {
			if errors.Is(readErr, io.EOF) {
				return bytesOut, frames, nil
			}
			return bytesOut, frames, readErr
		}
	}
}

// copyBody forwards a non-streaming response. Same cancellation discipline,
// fixed-size buffer because there are no frame boundaries to respect.
func copyBody(ctx context.Context, w http.ResponseWriter, body io.Reader) (bytesOut int64, err error) {
	buf := make([]byte, 32<<10)
	for {
		select {
		case <-ctx.Done():
			return bytesOut, ctx.Err()
		default:
		}

		n, readErr := body.Read(buf)
		if n > 0 {
			written, writeErr := w.Write(buf[:n])
			bytesOut += int64(written)
			if writeErr != nil {
				return bytesOut, writeErr
			}
		}
		if readErr != nil {
			if errors.Is(readErr, io.EOF) {
				return bytesOut, nil
			}
			return bytesOut, readErr
		}
	}
}

func isTimeout(err error) bool {
	if errors.Is(err, context.DeadlineExceeded) {
		return true
	}
	var ne net.Error
	if errors.As(err, &ne) {
		return ne.Timeout()
	}
	return false
}