Files
traefik-llmgateway/users_file.go
T
lukaszraczylo 4c01c75276 fix(gateway): close billing, limit, routing and translation gaps
Billing:
- Passthrough SSE streams now bill the usage the upstream reports.
  The gateway falls back to a request-size estimate only on a 2xx
  stream or a truncated JSON body.
- A unified stream that fails after it starts is estimated, not
  billed at zero.
- Billing uses the built-in LiteLLM price table and honours
  modelMeta free: true.
- Pricing entries are validated. Cost arithmetic rounds the total and
  saturates on overflow.
- Gemini thinking tokens count as completion tokens.

Limits and storage:
- A rejected request rolls back its counts on the other scopes, so one
  user cannot lock out a group.
- Redis EXPIRE replies are checked. Keys recreated without a TTL are
  repaired at once.
- A mutating pipeline is not retried after reply bytes arrive.
- History, model totals and target usage report a store outage.

Routing and translation:
- Discovered ids with a provider-like prefix are listed only in
  owner/id form, so listing and routing agree.
- Anthropic and Gemini discovery follows pagination.
- The developer role maps to system.
- stream_options.include_usage is honoured on translated streams.
- tool_choice none keeps tools when history holds tool blocks.
- refusal maps to content_filter. Data URI parameters are stripped.

Streaming, retry and failover:
- SSE accepts CR, LF and CRLF line ends. This closes a CR injection
  path.
- Retry-After is never jittered. A cancelled wait returns the
  context error.
- The watchdog timer stops at end of body.
- Oversized bodies return 413.
- A failed stream ends with an error event.
- Failover restores pre-existing headers. Client cancels do not count
  as provider successes.

Also: "admin" is a reserved provider name. metrics.path must start
with "/". Pricing warnings go to the right gateway instance. The admin
targets view sends access as null or [] so the two states differ.
Upstream organisation headers are stripped. The users-file reload
error is logged once per error. Dead code is removed, and test-only
helpers move to test files.

Claude-Session: https://claude.ai/code/session_01V5jYJNxpDFdJkyyExpADrv
2026-09-23 09:44:31 +01:00

166 lines
6.7 KiB
Go

package traefikllmgateway
import (
"encoding/json"
"fmt"
"os"
"time"
)
// reloadEvery is the minimum interval between users-file mtime checks on the
// request path. maybeReload never stats the file more often than this, even
// under heavy concurrent traffic.
const reloadEvery = 5 * time.Second
// usersFile is a JSON-encoded file of API-key users at path, in the shape
// {"users":[{"name":...,"group":...,"apiKey":...,"limits":{...}}]}.
type usersFile struct {
path string
}
// newUsersFile returns a usersFile reading from path. It does not touch the
// filesystem — call load to read and parse the file.
func newUsersFile(path string) *usersFile {
return &usersFile{path: path}
}
// usersFileDoc is the on-disk shape of a users file.
type usersFileDoc struct {
Users []*UserConfig `json:"users"`
}
// load reads and parses uf's file. path is operator-supplied middleware
// configuration, not untrusted request input.
func (uf *usersFile) load() ([]*UserConfig, error) {
b, err := os.ReadFile(uf.path) // #nosec G304 -- operator-supplied users file path from middleware config
if err != nil {
return nil, fmt.Errorf("llmgateway: cannot read users file %q: %w", uf.path, err)
}
var doc usersFileDoc
if err := json.Unmarshal(b, &doc); err != nil {
return nil, fmt.Errorf("llmgateway: cannot parse users file %q: %w", uf.path, err)
}
return doc.Users, nil
}
// gatewayLogger is the logging surface maybeReload needs. *Gateway
// implements it (see logger.go); authStore depends on this small interface
// rather than *Gateway directly, so the reload path can be tested without an
// HTTP handler chain.
type gatewayLogger interface {
logf(format string, args ...any)
errorf(format string, args ...any)
}
// maybeReload reloads a's file-sourced users when a.usersFile is configured,
// reloadEvery has elapsed since the last check, and the file's mtime differs
// from the last observed mtime (forward or backward). It runs on the request
// path (ServeHTTP calls it at entry), so it must stay cheap when nothing
// changed and must never panic: on any error — stat, load, or
// replaceFileUsers rejecting the new set — it logs via logReloadError (rate-
// limited to once per distinct error/mtime pair, finding F11) and keeps the
// last good set. A successful reload that changes the file-sourced user
// count logs via a.log.logf, unconditionally — that line is inherently rare
// (only on an actual reload, itself throttled), unlike the error path above.
//
// Returns true only when a reload actually replaced the user set (review
// fix, adversarial verification 2026-08-23) — every early-return path
// above (no file configured, throttled, unchanged mtime, a stat/load
// error, or replaceFileUsers itself rejecting the new set) returns
// false. ServeHTTP uses this to gate limiter.pruneRejections: that call
// walks every tracked rejection scope, which is cheap only relative to a
// real reload (itself throttled to at most once per reloadEvery and
// gated on the file's mtime actually changing) — running it on every
// request regardless, the way maybeReload itself must, would add real
// per-request cost to ServeHTTP's hottest path for a reason (deleted
// users hot-reloaded out of the file) that is rare by construction.
func (a *authStore) maybeReload() bool {
if a.usersFile == nil {
return false
}
// TryLock, not Lock: ServeHTTP calls maybeReload on every request, and a
// reload's build phase can block on file I/O (resolveSecret's "file:"
// reads inside replaceFileUsers). Blocking every concurrent request
// behind one full reload would undo the auth.go buildMu/mu lock split's
// benefit one layer up. A request that finds a reload already in flight
// just skips this cycle — the next request's check (or the next
// throttle window) covers it. Reload is deliberately best-effort and
// throttled, not "every change lands on the very next request."
if !a.reloadMu.TryLock() {
return false
}
defer a.reloadMu.Unlock()
now := a.nowFn()
if now.Sub(a.lastCheck) < reloadEvery {
return false
}
a.lastCheck = now
info, statErr := os.Stat(a.usersFile.path)
if statErr != nil {
// A stat failure has no fresh mtime to dedup against — use the
// zero time.Time as that half of the key, so this branch's own
// dedup is driven by the error TEXT changing alone (see
// logReloadError's own doc comment).
a.logReloadError(time.Time{}, "llmgateway: cannot stat users file %q: %v", a.usersFile.path, statErr)
return false
}
// Equal, not After: an After-only check misses a backward mtime move —
// a restore from backup, or a `cp -p` from an older file. Any
// difference from the last observed mtime, forward or backward, must
// trigger a reload; only an unchanged mtime skips it.
if info.ModTime().Equal(a.lastModTime) {
return false
}
users, err := a.usersFile.load()
if err != nil {
a.logReloadError(info.ModTime(), "llmgateway: users file %q reload failed: %v", a.usersFile.path, err)
return false
}
a.mu.RLock()
prevCount := a.fileUserCount
a.mu.RUnlock()
if err := a.replaceFileUsers(users); err != nil {
a.logReloadError(info.ModTime(), "llmgateway: users file %q reload rejected: %v", a.usersFile.path, err)
return false
}
a.lastModTime = info.ModTime()
// The file is good again — clear the dedup key so a LATER failure
// (after this success) logs immediately rather than staying
// suppressed by a stale key from a much earlier failure episode.
a.lastReloadErrMsg = ""
a.lastReloadErrModTime = time.Time{}
if len(users) != prevCount {
a.log.logf("llmgateway: users file %q reload: user count changed from %d to %d", a.usersFile.path, prevCount, len(users))
}
return true
}
// logReloadError logs a maybeReload failure via a.log.errorf, but only
// when it is NEW: a distinct message from the last one logged, or the
// same message attempted against a distinct file mtime (review-auth
// finding F11, 2026-09 audit — "once per distinct error / mtime change").
// Without this, a users file left broken gets re-attempted AND re-logged
// every reloadEvery for as long as it stays broken — retrying is
// intentional (the operator may fix it any second), but the identical log
// line every 5s is not: roughly 17k lines/day at the default throttle,
// unbounded like every other log line in this package already is not
// (logStoreError, limits.go; the failure/throttle logGates, auth.go).
// modTime is compared with Equal, matching maybeReload's own mtime-change
// convention immediately above.
func (a *authStore) logReloadError(modTime time.Time, format string, args ...any) {
msg := fmt.Sprintf(format, args...)
if msg == a.lastReloadErrMsg && modTime.Equal(a.lastReloadErrModTime) {
return
}
a.lastReloadErrMsg = msg
a.lastReloadErrModTime = modTime
a.log.errorf("%s", msg)
}