traefik-llmgateway/routes_media.go

package traefikllmgateway

import (
	"bytes"
	"encoding/json"
	"errors"
	"io"
	"mime"
	"mime/multipart"
	"net/http"
)

// Route paths for the three unified media endpoints (spec §3, v0.2):
// image generation and both audio endpoints are openai-type-only native
// forwards, except images.generations, which gemini translates to Imagen
// (translate_gemini_images.go); anthropic answers 501 on all three
// (provider_anthropic.go). See llmgateway.go's ServeHTTP dispatch.
const (
	imagesGenerationsPath   = "/v1/images/generations"
	audioSpeechPath         = "/v1/audio/speech"
	audioTranscriptionsPath = "/v1/audio/transcriptions"
)

// maxMultipartModelFieldBytes caps how much of the "model" form field
// extractMultipartModel reads: a model id is always a short identifier,
// so this is generous while bounding memory against a form field a
// misbehaving or malicious client names "model" but fills with megabytes
// of data.
const maxMultipartModelFieldBytes = 4096

// decodeMediaJSONRequest claims a body-admission slot (acquireBodyAdmission),
// reads r.Body capped at maxRequestBytes (routes_unified.go's cap, shared
// here) and decodes it as a JSON object, then extracts its "model"
// field — releasing the slot before returning, on every path (security
// review finding 1b, round 3, 2026-08-22: the slot is scoped to this
// read+decode step alone, exactly like readAndDecodeUnifiedBody
// (routes_unified.go), never held across the adapter call that follows
// in the caller). It mirrors runUnified's own decode path exactly (spec
// §3's "reuse the unified decode path") — used by both the images and
// audio-speech routes, whose request bodies are both single JSON objects
// carrying a top-level "model" field, unlike audio-transcriptions'
// multipart body. A read failure or invalid/missing JSON is a 400,
// already written to sw; ok reports whether the caller may proceed.
func (g *Gateway) decodeMediaJSONRequest(sw *statusTrackingWriter, r *http.Request) (req map[string]any, model string, ok bool) {
	release, admitted := g.acquireBodyAdmission(sw)
	defer release()
	if !admitted {
		return nil, "", false
	}

	body, err := io.ReadAll(io.LimitReader(r.Body, maxRequestBytes))
	if err != nil {
		writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "cannot read request body")
		return nil, "", false
	}
	if err = json.Unmarshal(body, &req); err != nil {
		writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "invalid JSON body")
		return nil, "", false
	}
	model, _ = req["model"].(string)
	if model == "" {
		writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "model is required")
		return nil, "", false
	}
	return req, model, true
}

// resolveMediaModel runs the shared model-routing/authz step every media
// route needs before calling its adapter (spec §3's "auth → model routing
// (registry) → group authz" — group authz is folded into
// modelRegistry.resolve itself, the same as routes_unified.go's
// runUnified). On success it returns the resolved adapter and the
// upstream model id to route to; ok reports whether the caller may
// proceed — a resolve failure has already written its response to sw.
//
// This used to also run checkAndCount (a limit check), and was named
// resolveMediaRequest — security review finding 1a, 2026-08-22, split
// that out into admitRequest (routes_unified.go), called by every media
// handler BEFORE the request body is even read: checkAndCount needs only
// u and grp, both already known ahead of decode, so there was never a
// reason to make a rate-limited caller pay for reading and decoding a
// body that was always going to be discarded. See admitRequest's own doc
// comment for the full rationale, shared verbatim with runUnified's
// identical fix.
func (g *Gateway) resolveMediaModel(sw *statusTrackingWriter, grp *group, requestedModel string) (adapter providerAdapter, upstreamModel string, ok bool) {
	adapter, upstreamModel, _, err := g.registry.resolve(requestedModel, grp)
	if err != nil {
		writeModelResolveError(sw, err)
		return nil, "", false
	}
	return adapter, upstreamModel, true
}

// handleImagesGenerations implements POST /v1/images/generations (spec
// §3): an openai-type provider forwards natively; gemini translates to
// Imagen (translate_gemini_images.go); anthropic always answers 501
// (provider_anthropic.go's imagesGeneration). Images are never cached
// (spec §2) and never cost-accounted (spec §3) — only the request
// counters admitRequest's checkAndCount call already incremented move.
func (g *Gateway) handleImagesGenerations(w http.ResponseWriter, r *http.Request, u *user, grp *group) {
	sw := &statusTrackingWriter{ResponseWriter: w}

	if _, ok := g.admitRequest(sw, u, grp); !ok {
		return
	}

	req, requestedModel, ok := g.decodeMediaJSONRequest(sw, r)
	if !ok {
		return
	}
	adapter, upstreamModel, ok := g.resolveMediaModel(sw, grp, requestedModel)
	if !ok {
		return
	}
	req["model"] = upstreamModel

	// Feature A (v0.22): see routes_unified.go's runUnified for the
	// identical pattern this mirrors.
	ctx := withAttemptRecorder(r.Context(), func(resp *http.Response, attemptErr error) {
		g.limiter.recordProviderAttempt(adapter.name(), upstreamModel, resp, attemptErr)
	})
	if _, err := adapter.imagesGeneration(ctx, sw, req); err != nil {
		g.handleAdapterError(sw, err, adapter.name())
	}
}

// handleAudioSpeech implements POST /v1/audio/speech (spec §3): an
// openai-type provider forwards natively, streaming the binary audio
// response back via a flushWriter (provider_openai.go's audioSpeech);
// gemini and anthropic both always answer 501. req is re-marshaled to
// JSON after its "model" field is rewritten to the resolved upstream
// model id, mirroring how the unified chat/embeddings routes rewrite
// "model" before handing off to their adapter — audioSpeech's own
// interface method takes a raw []byte body rather than a map
// (providers.go), since its response is binary audio no translator ever
// needs to inspect as JSON again.
func (g *Gateway) handleAudioSpeech(w http.ResponseWriter, r *http.Request, u *user, grp *group) {
	sw := &statusTrackingWriter{ResponseWriter: w}

	if _, ok := g.admitRequest(sw, u, grp); !ok {
		return
	}

	req, requestedModel, ok := g.decodeMediaJSONRequest(sw, r)
	if !ok {
		return
	}
	adapter, upstreamModel, ok := g.resolveMediaModel(sw, grp, requestedModel)
	if !ok {
		return
	}
	req["model"] = upstreamModel
	// gatewayAliasKey must never reach the real provider, mirroring
	// openaiAdapter.imagesGeneration's identical delete
	// (provider_openai.go): routes_media.go never sets this key itself
	// for a media route, but a client independently sending a literal
	// "__alias" field of its own would otherwise survive the re-marshal
	// below and leak straight through as an unrecognized request field.
	// Unconditional, not conditional on whether the gateway set it — the
	// same defensive posture chatCompletion/embeddings/imagesGeneration
	// already take.
	delete(req, gatewayAliasKey)

	body, err := json.Marshal(req)
	if err != nil {
		// req decoded from the client's own JSON body plus one string-field
		// overwrite ("model") — a value json.Unmarshal has already accepted
		// once cannot fail to re-marshal; this is a defensive 500, not a
		// path any known input reaches.
		g.errorf("audio speech: re-marshal request: %v", err)
		writeOAIError(sw, http.StatusInternalServerError, "server_error", "internal error")
		return
	}

	// Feature A (v0.22): see routes_unified.go's runUnified for the
	// identical pattern this mirrors.
	ctx := withAttemptRecorder(r.Context(), func(resp *http.Response, attemptErr error) {
		g.limiter.recordProviderAttempt(adapter.name(), upstreamModel, resp, attemptErr)
	})
	if _, err := adapter.audioSpeech(ctx, sw, body, "application/json"); err != nil {
		g.handleAdapterError(sw, err, adapter.name())
	}
}

// handleAudioTranscriptions implements POST /v1/audio/transcriptions
// (spec §3): an openai-type provider forwards natively; gemini and
// anthropic both always answer 501. The multipart request body is
// replayed upstream unchanged only when the client's own "model" form
// field already equals the resolved upstream model id (the common case —
// a bare id); when it does not (a provider-prefixed or otherwise aliased
// id resolves to a different upstream model string), the body is
// rebuilt via rewriteMultipartModel so the alias never reaches the real
// provider, the same invariant every other unified route enforces
// (controller ruling, 2026-08-20, amends spec §3's original "REPLAY the
// full body upstream unchanged" — see rewriteMultipartModel's doc
// comment).
func (g *Gateway) handleAudioTranscriptions(w http.ResponseWriter, r *http.Request, u *user, grp *group) {
	sw := &statusTrackingWriter{ResponseWriter: w}

	if _, ok := g.admitRequest(sw, u, grp); !ok {
		return
	}

	body, oversize, admitted, err := g.readAdmittedCapped(sw, r.Body, maxRequestBytes)
	if !admitted {
		return
	}
	if err != nil {
		writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "cannot read request body")
		return
	}
	if oversize {
		writeOAIError(sw, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body too large")
		return
	}

	contentType := r.Header.Get("Content-Type")
	mediaType, params, err := mime.ParseMediaType(contentType)
	if err != nil {
		writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "invalid Content-Type")
		return
	}
	if mediaType != "multipart/form-data" {
		writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "Content-Type must be multipart/form-data")
		return
	}
	boundary := params["boundary"]
	if boundary == "" {
		writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "missing multipart boundary")
		return
	}

	// The three failure shapes below are kept distinct rather than folded
	// into one "model is required" 400 (review sweep, 2026-08-20): a
	// duplicate "model" part is ambiguous, a body extractMultipartModel
	// cannot even parse (a malformed or truncated tail) never reached a
	// point where "model" could be evaluated at all, and only a
	// successfully parsed body with no "model" part is actually missing
	// one — each names its own actual problem instead of collapsing every
	// case into the same misleading message.
	requestedModel, err := extractMultipartModel(body, boundary)
	switch {
	case errors.Is(err, errDuplicateModelField):
		writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "duplicate model field")
		return
	case err != nil:
		writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "malformed multipart body")
		return
	case requestedModel == "":
		writeOAIError(sw, http.StatusBadRequest, "invalid_request_error", "model is required")
		return
	}

	adapter, upstreamModel, ok := g.resolveMediaModel(sw, grp, requestedModel)
	if !ok {
		return
	}

	uploadBody := body
	uploadContentType := contentType
	if upstreamModel != requestedModel {
		rebuilt, rebuiltContentType, rerr := rewriteMultipartModel(body, boundary, upstreamModel)
		if rerr != nil {
			g.errorf("audio transcriptions: rewrite multipart model field: %v", rerr)
			writeOAIError(sw, http.StatusInternalServerError, "server_error", "internal error")
			return
		}
		uploadBody = rebuilt
		uploadContentType = rebuiltContentType
	}

	// Feature A (v0.22): see routes_unified.go's runUnified for the
	// identical pattern this mirrors.
	ctx := withAttemptRecorder(r.Context(), func(resp *http.Response, attemptErr error) {
		g.limiter.recordProviderAttempt(adapter.name(), upstreamModel, resp, attemptErr)
	})
	if _, err := adapter.audioTranscription(ctx, sw, uploadBody, uploadContentType); err != nil {
		g.handleAdapterError(sw, err, adapter.name())
	}
}

// readAdmittedCapped claims a body-admission slot (acquireBodyAdmission),
// calls readCapped, and releases the slot before returning — on every
// path, including a panic unwinding through readCapped itself — used by
// handleAudioTranscriptions (security review finding 1b, round 3,
// 2026-08-22). This route never json.Unmarshal's into a map[string]any
// (its own "decode", extractMultipartModel's bounded scan, runs over
// bytes already fully read, AFTER this returns and the slot is already
// free), so the slot only ever needs to be held for the read itself —
// narrower even than readAndDecodeUnifiedBody/decodeMediaJSONRequest,
// which both hold it through an actual json.Unmarshal too. admitted
// reports whether a slot was claimed at all (false means
// acquireBodyAdmission already wrote its own 503 to sw); body/oversize/
// err are readCapped's own results, valid only when admitted is true.
func (g *Gateway) readAdmittedCapped(sw *statusTrackingWriter, r io.Reader, limit int64) (body []byte, oversize, admitted bool, err error) {
	release, admitted := g.acquireBodyAdmission(sw)
	defer release()
	if !admitted {
		return nil, false, false, nil
	}
	body, oversize, err = readCapped(r, limit)
	return body, oversize, true, err
}

// readCapped reads up to limit+1 bytes from r, reporting oversize=true
// when more than limit bytes were available. Unlike
// decodeMediaJSONRequest's plain io.LimitReader-and-decode (which
// silently truncates an oversized JSON body into a parse failure),
// audio-transcriptions' multipart body cannot be truncated the same way —
// a cut multipart payload often still parses a valid prefix of it, which
// would forward truncated audio upstream instead of failing loudly — so
// this endpoint alone reports the cap explicitly (spec §3's "bodies above
// maxRequestBytes (10MB) → 413").
func readCapped(r io.Reader, limit int64) (body []byte, oversize bool, err error) {
	body, err = io.ReadAll(io.LimitReader(r, limit+1))
	if err != nil {
		return nil, false, err
	}
	if int64(len(body)) > limit {
		return body[:limit], true, nil
	}
	return body, false, nil
}

// errDuplicateModelField is extractMultipartModel's sentinel for a
// multipart body carrying more than one part named "model" — an
// ambiguous request a first-found parser and a last-found parser would
// resolve differently. handleAudioTranscriptions rejects it outright
// (400) rather than silently picking one.
var errDuplicateModelField = errors.New("multipart body has more than one model field")

// extractMultipartModel scans body — a full multipart/form-data payload
// already read into memory (bounded by handleAudioTranscriptions' own
// readCapped call before this is ever reached) — for its "model" form
// field, returning the field's value. It reads over a fresh
// bytes.NewReader(body), never consuming or mutating body itself, so the
// caller can replay or rewrite the exact same bytes afterward (spec §3).
// It always scans every part, never returning as soon as the first
// "model" part is found, so a second "model" part later in the body is
// caught as errDuplicateModelField instead of being silently ignored.
// Parts are read in whatever order the client sent them — a "model"
// field following a large audio-file part is found the same as one sent
// first. Returns "" with a nil error when no part named "model" is
// present; the caller treats an empty model the same as a missing one.
func extractMultipartModel(body []byte, boundary string) (string, error) {
	mr := multipart.NewReader(bytes.NewReader(body), boundary)
	found := false
	var model string
	for {
		part, err := mr.NextPart()
		if errors.Is(err, io.EOF) {
			return model, nil
		}
		if err != nil {
			return "", err
		}
		if part.FormName() != "model" {
			_ = part.Close() //nolint:errcheck // read-side close; nothing actionable on failure
			continue
		}
		if found {
			_ = part.Close() //nolint:errcheck // read-side close; nothing actionable on failure
			return "", errDuplicateModelField
		}
		val, err := io.ReadAll(io.LimitReader(part, maxMultipartModelFieldBytes))
		_ = part.Close() //nolint:errcheck // read-side close; nothing actionable on failure
		if err != nil {
			return "", err
		}
		model = string(val)
		found = true
	}
}

// rewriteMultipartModel rebuilds body — a multipart/form-data payload
// already validated to parse and to carry exactly one "model" field
// (extractMultipartModel) — with a fresh mime/multipart.Writer, copying
// every part verbatim (its original headers and content, in the
// client's original order) except the "model" part's value, which
// becomes upstreamModel.
//
// Controller ruling (2026-08-20, amends spec §3's original "REPLAY the
// full body upstream unchanged"): a client addressing a
// provider-prefixed or otherwise aliased model id must never have that
// alias reach the real upstream provider — the same invariant every
// other unified route already enforces by rewriting "model" before
// forwarding. handleAudioTranscriptions calls this only when the
// client's own model string differs from the resolved upstream model
// id; when they already match (the bare-id case), it replays the
// original bytes untouched instead (the fast path, no rebuild cost).
//
// The rebuilt body uses the Writer's own new boundary, not body's
// original one, and contentType is that Writer's own
// FormDataContentType() — byte-identical boundary preservation is not
// required, only a valid, equivalent multipart body.
func rewriteMultipartModel(body []byte, boundary, upstreamModel string) ([]byte, string, error) {
	mr := multipart.NewReader(bytes.NewReader(body), boundary)
	var buf bytes.Buffer
	// The rebuilt body is at least as large as the original (same parts,
	// same content, only the "model" value and the boundary differ) plus
	// the fresh Writer's own per-part header/boundary overhead — growing
	// buf to body's length up front, with headroom, avoids the repeated
	// reallocate-and-copy growth bytes.Buffer would otherwise do one
	// io.Copy at a time.
	buf.Grow(len(body) + 512)
	mw := multipart.NewWriter(&buf)

	for {
		part, err := mr.NextPart()
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			return nil, "", err
		}

		pw, err := mw.CreatePart(part.Header)
		if err != nil {
			_ = part.Close() //nolint:errcheck // read-side close; nothing actionable on failure
			return nil, "", err
		}

		if part.FormName() == "model" {
			_, err = pw.Write([]byte(upstreamModel))
		} else {
			_, err = io.Copy(pw, part)
		}
		_ = part.Close() //nolint:errcheck // read-side close; nothing actionable on failure
		if err != nil {
			return nil, "", err
		}
	}

	if err := mw.Close(); err != nil {
		return nil, "", err
	}
	return buf.Bytes(), mw.FormDataContentType(), nil
}