diff --git a/config.example.yaml b/config.example.yaml index 3c087ca98..f5b9c8857 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -9,6 +9,17 @@ # The compatibility notes at the end show old spellings for existing deployments. config-version: 8 +# Optional model catalog sources: http(s) URLs or absolute local file paths. +# Each explicit source overrides --local-model; otherwise --local-model uses +# embedded catalogs, and normal mode uses the official remote sources. +# Failed refreshes retain the last valid catalog; custom sources never fall back +# to official URLs. Files and URLs are refreshed every 3 hours and on source changes. +# Home mode never loads the general catalog, even when explicitly configured. +models: + catalog: "" + codex-catalog: "" + devin-catalog: "" + # Listener, TLS, network discovery, and server operation. server: # Server host/interface to bind to. Default is empty ("") to bind all interfaces (IPv4 + IPv6). @@ -18,6 +29,25 @@ server: # Server port port: 8317 + # Maximum allowed size (in bytes) for incoming HTTP request bodies. + # Requests exceeding this limit are rejected with HTTP 413 Payload Too Large. + # Default is 0 (disabled / no limit). + max-request-body-bytes: 0 + + # Negotiated compression of downstream HTTP responses (including SSE streams). + # Responses are compressed only when the client advertises zstd or gzip in + # Accept-Encoding; WebSocket upgrades are never wrapped. + response-compression: + # Default: true. + enabled: true + # Non-streaming responses smaller than this many bytes are sent uncompressed. + min-bytes: 1024 + + # Global token for GitHub requests (for example, release and asset downloads). + # A non-empty config token takes precedence over the GITHUB_TOKEN environment variable. + # Leave empty to use GITHUB_TOKEN, or unauthenticated requests if neither is set. + github-token: "" + # IP addresses or CIDRs of proxies allowed to supply forwarded client IP headers. # An empty list (the default) does not trust forwarded headers. Restart after changing this setting. trusted-proxies: [] @@ -37,6 +67,9 @@ server: # Enable local network mDNS / DNS-SD service discovery (RFC 6762 / RFC 6763). # When enabled, broadcasts _ai-gateway._tcp with API protocol subtypes on the local network. + # TXT records follow draft-su-ai-gateway-discovery-00: + # txtvers, path, tls, auth, protocols, auth_type, product, id, role + # The TXT protocols list mirrors the advertised subtypes. Management endpoints are never advertised. # NOTE for Docker users: Link-local multicast (224.0.0.251) cannot cross default NAT bridge networks. # To broadcast across the physical LAN, run the container with host networking (--net=host / network_mode: host). discovery: @@ -49,12 +82,12 @@ server: # - "_messages" # - "_generate-content" # - "_interactions" + # # Optional extended subtypes: _completions, _responses-ws, _realtime, _images, + # # _videos, _audio, _count-tokens, _search # interfaces: # include: [] # White-listed interfaces (empty: auto-detect physical interfaces) # exclude: [] # Additional interface patterns to exclude - # auth-required: true # Advertise that credentials are required - - # advertise-management: false # Advertise management availability + # auth-required: true # Advertise auth=1 (credentials required) # Management API settings management: @@ -150,11 +183,22 @@ routing: session-affinity-ttl: "1h" # When true (default), subagents (child sessions with parent references) inherit and bind # to the parent's upstream credential across all providers (Claude, Codex, Antigravity, Gemini), - # maximizing Prompt/KV Cache reuse and minimizing TTFT. + # for the same model, maximizing Prompt/KV Cache reuse and minimizing TTFT. + # Guardian reviews (X-Openai-Subagent: guardian) may inherit across models from the + # parent's most recent selection on the same real provider, if that credential is eligible. + # This includes mixed-provider pools; result-driven TTL refreshes do not reorder selections. + # Existing child bindings take precedence. Parent lookups do not extend the parent's TTL. # When false, subagents are distributed across the credential pool via the fallback selector # for maximum parallel concurrency. # Ignored when routing.session-affinity is false. session-affinity-subagents: true + # Controls which credential failures rebind an established session to the failover credential: + # - "any" (default): any credential-attributed failure unbinds and rebinds on failover. + # - "quota": transient errors (408, 5xx, empty stream, buffer limit) fail over for that + # request only while keeping the session bound to its original credential; quota and auth + # errors (429, 401/402/403, 404/model support, invalid_grant) rebind as normal. + # Ignored when routing.session-affinity is false. + session-affinity-rebind-on: "any" # When true, unprefixed model requests only use credentials without a prefix (except when prefix == model name). force-model-prefix: false @@ -205,6 +249,17 @@ requests: # Per-entry proxy-url also supports "direct" or "none" to bypass both the global proxy-url and environment proxies explicitly. proxy-url: "" + # Maximum time, in seconds, for upstream TLS handshakes made by the shared built-in + # uTLS client. It applies to both the Claude Code Anthropic transport and the Chrome + # ChatGPT/Codex transport; the standard fallback transport is unaffected. The timer + # starts after TCP or proxy connection establishment returns; it does not bound dialing + # or proxy CONNECT. When the upstream path accepts the connection but never answers the + # TLS ClientHello, the request would otherwise wait indefinitely and request-retry could + # never engage. The bound applies only during the TLS handshake and is lifted once it + # succeeds, so inference and long-running streams are never time-limited. + # Set to 0 to disable (default). + tls-handshake-timeout-seconds: 0 + # When true, forward filtered upstream response headers to downstream clients. # Default is false (disabled). passthrough-headers: false @@ -214,8 +269,9 @@ requests: # Streaming behavior (SSE keep-alives + safe bootstrap retries). # streaming: - # keepalive-seconds: 15 # Default: 0 (disabled). <= 0 disables keep-alives. - # bootstrap-retries: 1 # Default: 0 (disabled). Retries before first byte is sent. + # keepalive-seconds: 15 # Default: 0 (disabled). <= 0 disables keep-alives. + # bootstrap-retries: 1 # Default: 0 (disabled). Retries before first byte is sent. + # send-keepalive-before-ft: false # Default: false. When true and keepalive-seconds > 0, send SSE heartbeats before the first token. # Optional payload configuration # payload: @@ -295,6 +351,12 @@ client: # Default false. Applies to both OAuth and API-key credentials across providers. # Legacy provider paths migrate here; this path wins conflicts, including explicit false. optimize-multi-agent-v2: false + # Cache upstream web search results per session so replayed Codex web_search_call items + # that omit results still retain search results in later turns. + web-search-result-cache: + enabled: false + ttl: "24h" + limit: 1024 # Shared upstream behavior applies to OAuth and API-key credentials. upstream: @@ -302,6 +364,17 @@ upstream: # Experimental full-duplex Codex Responses WebSocket steering (default: false). # Keeps one upstream account/model per socket; accepted input is never replayed. response-steering: false + # Compress request bodies sent to the ChatGPT Codex backend with zstd, + # mirroring the official Codex client. Applies only to OAuth credentials + # targeting https://chatgpt.com/backend-api/codex (API-key credentials, + # custom base URLs, and SIWC requests are never compressed). A 400/415 + # rejection is retried once with an uncompressed body. Default: false. + upstream-request-compression: false + # Request bodies smaller than this many bytes are sent uncompressed. + upstream-request-compression-min-bytes: 1024 + # Enable permessage-deflate compression on Codex WebSocket connections + # (upstream and downstream) after negotiation. Default: true. + websocket-compression: true # Disable forcing the official Codex User-Agent and Originator headers on HTTP/SSE and WebSocket requests. disable-codex-cloaking: false # Hold back the frames that carry nothing the client has seen - the handshake (response.created, @@ -383,6 +456,20 @@ upstream: # the entire credential across all sibling models. Default is false. model-level-cooling: false + # When true, enable Codex client tool_search on Claude Responses routes. Credentials on + # api.anthropic.com (or keys with native-tool-search: true) receive it as Anthropic native + # tool search (defer_loading + tool_reference), which keeps the prompt-cache prefix stable + # across tool discoveries; other credentials fall back to the plain translation. The Codex + # model catalog advertises supports_search_tool for a Claude model only when every + # credential serving it supports native tool search. + # Default false keeps the translation and catalog behavior that predate tool_search support. + tool-search: false + + # Days to retain the cloaked currentDate pin per session, so an idle gap or a credential + # switch cannot re-anchor the date and invalidate earlier thinking signatures and the + # prompt cache. Default 7; an explicit 0 keeps the continuity-scoped pin (1h lifetime). + session-date-pin-days: 7 + # When true, disable Claude request cloaking on OAuth and API-key CLI profiles (the Claude Code CLI disguise and # system prompt replacement), so the original system prompt is passed through to Claude as-is. # Credential attributes and explicit per-key cloak modes override this shared default. @@ -405,13 +492,14 @@ upstream: # # Shared Claude CLI fingerprint defaults. Update only after measuring a new Claude Code release. # Unconfirmed clients use this CLI baseline. Verified native Claude Code CLI, sdk-cli, - # and VSCode requests preserve their measured entrypoint when the Claude Code version is a + # sdk-ts, and VSCode requests preserve their measured entrypoint when the Claude Code version is a # patch release in the configured major/minor line at or above this baseline. Older versions # and newer major/minor lines fall back to the baseline until measured. In legacy mode, # timeout is a fallback and # verified native OS/arch values remain client-supplied. When stabilize-device-profile is # enabled, OS/arch are pinned to the values below and cached profiles remain constrained to - # the same exact software baseline rather than learning newer client versions. + # the same exact software baseline rather than learning newer client versions (unless the + # credential resolves to cloak_mode: never, which preserves the confirmed caller's headers). # header-defaults: # user-agent: "claude-cli/2.1.280 (external, cli)" # package-version: "0.112.1" @@ -421,6 +509,20 @@ upstream: # timeout: "600" # timezone: "Asia/Singapore" # fallback IANA timezone for cloaked currentDate; a credential JSON "timezone" takes priority # stabilize-device-profile: false # optional, default false; set true to enable per-auth fingerprint pinning + # # Advanced operator overrides for exact-parity or gateway lanes. They apply to every Claude + # # request, after the per-request derivation above, and can produce a beta set real Claude Code + # # never sends, which api.anthropic.com may reject. Names are matched exactly; strip runs first. + # # A credential's own "headers" entries are more specific and keep precedence where they already + # # did: Anthropic-Beta on custom gateways and x-client-request-id everywhere. + # anthropic-beta-strip: [] # remove these names from the derived or caller-supplied set, keeping order + # anthropic-beta-append: [] # append these names at the end unless already present + # # Cloaked requests to api.anthropic.com get a "diagnostics" body field like Claude Code, which in + # # turn adds cache-diagnosis-2026-04-07. Set false to stop injecting it (default true). + # inject-diagnostics: true + # # x-client-request-id on first-party CLI-profile requests: "on" (default) synthesizes a fresh + # # UUID per request (a confirmed Claude Code client keeps its own), "passthrough" forwards only + # # the caller's value and never synthesizes one, "off" never sends the header. + # client-request-id: "on" xai: # When true, inject the native x_search tool when the request does not declare it. @@ -553,6 +655,7 @@ upstream: # base-url: "https://example.com/api" # optional, e.g. https://zenmux.ai/api; falls back to Google Vertex when omitted # proxy-url: "socks5://proxy.example.com:1080" # optional per-key proxy override # # proxy-url: "direct" # optional: explicit direct connect for this credential +# # interactions: true # optional: opt into native Vertex Interactions API routing (/v1beta1/interactions). For service account JSON files under auths/, add "interactions": true at the root of the file. # headers: # X-Custom-Header: "custom-value" # # Values starting with "$" dynamically copy the header value from downstream client requests. @@ -659,10 +762,11 @@ upstream: # - api-key: "sk-atSM..." # weight: 5 # optional: weighted-round-robin share; omitted defaults to 1; maximum 1,000,000 # rebuild-mid-system-message: false # optional: default is false; when true, move messages with role "system" into the top-level Claude system field +# native-tool-search: false # optional: default is false; when true, trust this key's custom base-url with Anthropic native tool search (defer_loading + tool_reference) for Codex tool_search; requires upstream.claude.tool-search: true # cloak: # optional: explicitly enable request cloaking for non-Claude-Code clients # mode: "auto" # "auto" (default inside this block): cloak only when client is not Claude Code # # "always": cloak every unconfirmed client; confirmed native Claude Code still passes through -# # "never": never apply cloaking +# # "never" / "uncloak": never apply cloaking; on non-CLI API keys, also strip incoming Claude Code identity/billing attribution # # This "cloak" block applies to this api-keys.claude key entry only. For Claude OAuth # # credentials, set the same options in the auth/token JSON file via "cloak_mode" / # # "cloak_strict_mode" / "cloak_sensitive_words" / "cloak_cache_user_id". @@ -924,6 +1028,11 @@ oauth: # max-context-length: 1048576 # OAuth provider excluded models + # Applies to OAuth/file-backed credentials of each provider. + # A credential's own non-empty "excluded_models" list (set in its auth JSON file) takes priority + # over these rules for that credential and is not merged with them. Credentials without their own + # list (missing, empty, or null) inherit the current rules here; changes to these rules + # take effect for them on config reload. # excluded-models: # vertex: # - "gemini-3-pro-preview" @@ -1025,6 +1134,9 @@ oauth: # username: "user" # credential: "secret" + # Identifier (agent_name_hint and User-Agent prefix) for Sign in with ChatGPT (SIWC) OAuth and API requests. + signin-with-chatgpt-identifier: "CLIProxyAPI" + # Default headers for Codex OAuth model requests. # These are used only for file-backed/OAuth Codex requests when the client # does not send the header. `user-agent` applies to HTTP and websocket requests; diff --git a/examples/custom-provider/main.go b/examples/custom-provider/main.go index 984175f60..98184e281 100644 --- a/examples/custom-provider/main.go +++ b/examples/custom-provider/main.go @@ -50,7 +50,7 @@ const ( // transformation logic between OpenAI format and your provider's format. func init() { sdktr.Register(fOpenAI, fMyProv, - func(model string, raw []byte, stream bool) []byte { return raw }, + func(model string, raw []byte, stream bool) ([]byte, error) { return raw, nil }, sdktr.ResponseTransform{ Stream: func(ctx context.Context, model string, originalReq, translatedReq, raw []byte, param *any) [][]byte { return [][]byte{raw} diff --git a/examples/plugin/claude-web-search-router/go/go.mod b/examples/plugin/claude-web-search-router/go/go.mod index aff999128..659a184f2 100644 --- a/examples/plugin/claude-web-search-router/go/go.mod +++ b/examples/plugin/claude-web-search-router/go/go.mod @@ -1,6 +1,6 @@ module github.com/router-for-me/CLIProxyAPI/v7/examples/plugin/claude-web-search-router/go -go 1.26.0 +go 1.26.10 require ( github.com/router-for-me/CLIProxyAPI/v7 v7.0.0 @@ -12,7 +12,7 @@ require ( github.com/sirupsen/logrus v1.9.3 // indirect github.com/tidwall/match v1.1.1 // indirect github.com/tidwall/pretty v1.2.0 // indirect - golang.org/x/sys v0.47.0 // indirect + golang.org/x/sys v0.48.0 // indirect ) replace github.com/router-for-me/CLIProxyAPI/v7 => ../../../.. diff --git a/examples/plugin/claude-web-search-router/go/go.sum b/examples/plugin/claude-web-search-router/go/go.sum index 79cf47e97..08317b1ac 100644 --- a/examples/plugin/claude-web-search-router/go/go.sum +++ b/examples/plugin/claude-web-search-router/go/go.sum @@ -18,6 +18,7 @@ golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBc golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc= golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/sys v0.48.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/examples/plugin/codex-service-tier/go/go.mod b/examples/plugin/codex-service-tier/go/go.mod index 599588ee1..28422bc1a 100644 --- a/examples/plugin/codex-service-tier/go/go.mod +++ b/examples/plugin/codex-service-tier/go/go.mod @@ -1,6 +1,6 @@ module github.com/router-for-me/CLIProxyAPI/v7/examples/plugin/codex-service-tier/go -go 1.26.0 +go 1.26.10 require ( github.com/router-for-me/CLIProxyAPI/v7 v7.0.0 diff --git a/examples/plugin/frontend-auth-exclusive/go/go.mod b/examples/plugin/frontend-auth-exclusive/go/go.mod index c5f0e70a4..95c25967c 100644 --- a/examples/plugin/frontend-auth-exclusive/go/go.mod +++ b/examples/plugin/frontend-auth-exclusive/go/go.mod @@ -1,6 +1,6 @@ module github.com/router-for-me/CLIProxyAPI/v7/examples/plugin/frontend-auth-exclusive/go -go 1.26.0 +go 1.26.10 require github.com/router-for-me/CLIProxyAPI/v7 v7.0.0 diff --git a/examples/plugin/host-callback-auth-files/go/go.mod b/examples/plugin/host-callback-auth-files/go/go.mod index c67dbc66f..6b449ddd9 100644 --- a/examples/plugin/host-callback-auth-files/go/go.mod +++ b/examples/plugin/host-callback-auth-files/go/go.mod @@ -1,6 +1,6 @@ module github.com/router-for-me/CLIProxyAPI/v7/examples/plugin/host-callback-auth-files/go -go 1.26.0 +go 1.26.10 require github.com/router-for-me/CLIProxyAPI/v7 v7.0.0 diff --git a/examples/plugin/host-model-callback/go/go.mod b/examples/plugin/host-model-callback/go/go.mod index 95672b7e6..52e799aa6 100644 --- a/examples/plugin/host-model-callback/go/go.mod +++ b/examples/plugin/host-model-callback/go/go.mod @@ -1,6 +1,6 @@ module github.com/router-for-me/CLIProxyAPI/v7/examples/plugin/host-model-callback/go -go 1.26.0 +go 1.26.10 require github.com/router-for-me/CLIProxyAPI/v7 v7.0.0 diff --git a/examples/plugin/request-lifecycle/go/go.mod b/examples/plugin/request-lifecycle/go/go.mod index 420d628e8..fc152e93d 100644 --- a/examples/plugin/request-lifecycle/go/go.mod +++ b/examples/plugin/request-lifecycle/go/go.mod @@ -1,6 +1,6 @@ module github.com/router-for-me/CLIProxyAPI/v7/examples/plugin/request-lifecycle/go -go 1.26.0 +go 1.26.10 require ( github.com/router-for-me/CLIProxyAPI/v7 v7.0.0 diff --git a/examples/plugin/scheduler/go/go.mod b/examples/plugin/scheduler/go/go.mod index 99ead9836..34f7afc2d 100644 --- a/examples/plugin/scheduler/go/go.mod +++ b/examples/plugin/scheduler/go/go.mod @@ -1,6 +1,6 @@ module github.com/router-for-me/CLIProxyAPI/v7/examples/plugin/scheduler/go -go 1.26.0 +go 1.26.10 require ( github.com/router-for-me/CLIProxyAPI/v7 v7.0.0 diff --git a/examples/plugin/simple/go/go.mod b/examples/plugin/simple/go/go.mod index 7dd60e3f4..bcb1dc416 100644 --- a/examples/plugin/simple/go/go.mod +++ b/examples/plugin/simple/go/go.mod @@ -1,6 +1,6 @@ module github.com/router-for-me/CLIProxyAPI/v7/examples/plugin/simple/go -go 1.26.0 +go 1.26.10 require github.com/router-for-me/CLIProxyAPI/v7 v7.0.0 diff --git a/go.mod b/go.mod index 97c9a70c2..afe03d836 100644 --- a/go.mod +++ b/go.mod @@ -1,6 +1,6 @@ module github.com/router-for-me/CLIProxyAPI/v8 -go 1.26.0 +go 1.26.10 require ( github.com/andybalholm/brotli v1.0.6 @@ -17,6 +17,7 @@ require ( github.com/joho/godotenv v1.5.1 github.com/klauspost/compress v1.18.7 github.com/libp2p/zeroconf/v2 v2.2.0 + github.com/miekg/dns v1.1.43 github.com/minio/minio-go/v7 v7.0.66 github.com/pion/ice/v4 v4.3.0 github.com/pion/interceptor v0.1.45 @@ -44,7 +45,6 @@ require ( require ( github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/dlclark/regexp2/v2 v2.5.1 // indirect - github.com/miekg/dns v1.1.43 // indirect github.com/pion/datachannel v1.6.2 // indirect github.com/pion/dtls/v3 v3.1.5 // indirect github.com/pion/logging v0.2.4 // indirect @@ -107,7 +107,7 @@ require ( github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect github.com/muesli/cancelreader v0.2.2 // indirect github.com/muesli/termenv v0.16.0 // indirect - github.com/pelletier/go-toml/v2 v2.2.2 // indirect + github.com/pelletier/go-toml/v2 v2.4.3-0.20260624044523-35f78d5908f1 // indirect github.com/pierrec/xxHash v0.1.5 github.com/pjbgf/sha1cd v0.6.0 // indirect github.com/rivo/uniseg v0.4.7 // indirect diff --git a/internal/auth/antigravity/auth.go b/internal/auth/antigravity/auth.go index 316ae583e..e4450ea39 100644 --- a/internal/auth/antigravity/auth.go +++ b/internal/auth/antigravity/auth.go @@ -257,13 +257,12 @@ func (o *AntigravityAuth) FetchProjectID(ctx context.Context, accessToken string return "", fmt.Errorf("marshal request body: %w", errMarshal) } - endpointURL := fmt.Sprintf("%s/%s:loadCodeAssist", APIEndpoint, APIVersion) + endpointURL := fmt.Sprintf("%s/%s:loadCodeAssist", DailyAPIEndpoint, APIVersion) req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpointURL, strings.NewReader(string(rawBody))) if err != nil { return "", fmt.Errorf("create request: %w", err) } req.Header.Set("Authorization", "Bearer "+accessToken) - req.Header.Set("Accept", "*/*") req.Header.Set("Content-Type", "application/json") req.Header.Set("User-Agent", userAgent) diff --git a/internal/auth/codex/filelock_unix.go b/internal/auth/codex/filelock_unix.go new file mode 100644 index 000000000..d9ccce8b2 --- /dev/null +++ b/internal/auth/codex/filelock_unix.go @@ -0,0 +1,25 @@ +//go:build !windows + +package codex + +import ( + "fmt" + "os" + "syscall" +) + +func acquireInterProcessLockFile(authFilePath string) (func(), error) { + lockPath := canonicalLockPath(authFilePath) + ".lock" + f, errOpen := os.OpenFile(lockPath, os.O_CREATE|os.O_RDWR, 0o600) + if errOpen != nil { + return nil, fmt.Errorf("failed to open credential lock file %s: %w", lockPath, errOpen) + } + if errFlock := syscall.Flock(int(f.Fd()), syscall.LOCK_EX); errFlock != nil { + _ = f.Close() + return nil, fmt.Errorf("failed to acquire flock on %s: %w", lockPath, errFlock) + } + return func() { + _ = syscall.Flock(int(f.Fd()), syscall.LOCK_UN) + _ = f.Close() + }, nil +} diff --git a/internal/auth/codex/filelock_windows.go b/internal/auth/codex/filelock_windows.go new file mode 100644 index 000000000..bd9ed8f86 --- /dev/null +++ b/internal/auth/codex/filelock_windows.go @@ -0,0 +1,44 @@ +//go:build windows + +package codex + +import ( + "fmt" + "os" + "syscall" + "unsafe" +) + +var ( + modkernel32 = syscall.NewLazyDLL("kernel32.dll") + procLockFileEx = modkernel32.NewProc("LockFileEx") + procUnlockFileEx = modkernel32.NewProc("UnlockFileEx") +) + +const lockfileExclusiveLock = 0x00000002 + +func acquireInterProcessLockFile(authFilePath string) (func(), error) { + lockPath := canonicalLockPath(authFilePath) + ".lock" + f, errOpen := os.OpenFile(lockPath, os.O_CREATE|os.O_RDWR, 0o600) + if errOpen != nil { + return nil, fmt.Errorf("failed to open credential lock file %s: %w", lockPath, errOpen) + } + var overlapped syscall.Overlapped + r1, _, errLock := procLockFileEx.Call( + f.Fd(), + uintptr(lockfileExclusiveLock), + 0, + 1, + 0, + uintptr(unsafe.Pointer(&overlapped)), + ) + if r1 == 0 { + _ = f.Close() + return nil, fmt.Errorf("failed to lock credential file %s: %w", lockPath, errLock) + } + return func() { + var ov syscall.Overlapped + _, _, _ = procUnlockFileEx.Call(f.Fd(), 0, 1, 0, uintptr(unsafe.Pointer(&ov))) + _ = f.Close() + }, nil +} diff --git a/internal/auth/codex/oauth_server.go b/internal/auth/codex/oauth_server.go index 9c6a6c5b7..a71994a86 100644 --- a/internal/auth/codex/oauth_server.go +++ b/internal/auth/codex/oauth_server.go @@ -39,6 +39,8 @@ type OAuthResult struct { Code string // State is the state parameter used to prevent CSRF attacks State string + // ClientID is the dynamically issued client ID returned by Sign in with ChatGPT + ClientID string // Error contains any error message if the OAuth flow failed Error string } @@ -175,6 +177,7 @@ func (s *OAuthServer) handleCallback(w http.ResponseWriter, r *http.Request) { query := r.URL.Query() code := query.Get("code") state := query.Get("state") + clientID := strings.TrimSpace(query.Get("client_id")) errorParam := query.Get("error") // Validate required parameters @@ -210,8 +213,9 @@ func (s *OAuthServer) handleCallback(w http.ResponseWriter, r *http.Request) { // Send successful result result := &OAuthResult{ - Code: code, - State: state, + Code: code, + State: state, + ClientID: clientID, } s.sendResult(result) diff --git a/internal/auth/codex/openai.go b/internal/auth/codex/openai.go index 7aeffe9f3..a942fd0d3 100644 --- a/internal/auth/codex/openai.go +++ b/internal/auth/codex/openai.go @@ -29,6 +29,52 @@ type CodexTokenData struct { PlanType string `json:"plan_type,omitempty"` } +// CodexSIWCTokenData holds the OAuth token information obtained from OpenAI's Sign in with ChatGPT (SIWC) flow. +type CodexSIWCTokenData struct { + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token"` + IDToken string `json:"id_token,omitempty"` + ClientID string `json:"client_id"` + Scope string `json:"scope,omitempty"` + Scopes []string `json:"scopes,omitempty"` + Expire string `json:"expired"` + ExpiresAt string `json:"expires_at,omitempty"` + LastRefresh string `json:"last_refresh"` + DeviceID string `json:"device_id,omitempty"` +} + +// ToMap converts CodexSIWCTokenData into a map suitable for credential metadata persistence. +func (d *CodexSIWCTokenData) ToMap() map[string]any { + if d == nil { + return nil + } + out := map[string]any{ + "access_token": d.AccessToken, + "refresh_token": d.RefreshToken, + "client_id": d.ClientID, + "expired": d.Expire, + "last_refresh": d.LastRefresh, + } + if d.ExpiresAt != "" { + out["expires_at"] = d.ExpiresAt + } + if d.IDToken != "" { + out["id_token"] = d.IDToken + } + if d.Scope != "" { + out["scope"] = d.Scope + } + if len(d.Scopes) > 0 { + scopesCopy := make([]string, len(d.Scopes)) + copy(scopesCopy, d.Scopes) + out["scopes"] = scopesCopy + } + if d.DeviceID != "" { + out["device_id"] = d.DeviceID + } + return out +} + // CodexAuthBundle aggregates all authentication-related data after the OAuth flow is complete. // This includes the API key, token data, and the timestamp of the last refresh. type CodexAuthBundle struct { diff --git a/internal/auth/codex/openai_auth.go b/internal/auth/codex/openai_auth.go index f6761169b..bc9ce24a0 100644 --- a/internal/auth/codex/openai_auth.go +++ b/internal/auth/codex/openai_auth.go @@ -33,7 +33,8 @@ const ( // It manages the HTTP client and provides methods for generating authorization URLs, // exchanging authorization codes for tokens, and refreshing access tokens. type CodexAuth struct { - httpClient *http.Client + httpClient *http.Client + siwcAgentNameHint string } var codexRefreshGroup singleflight.Group @@ -49,15 +50,18 @@ func NewCodexAuth(cfg *config.Config) *CodexAuth { func NewCodexAuthWithProxyURL(cfg *config.Config, proxyURL string) *CodexAuth { effectiveProxyURL := strings.TrimSpace(proxyURL) var sdkCfg config.SDKConfig + agentNameHint := SIWCAgentNameHint if cfg != nil { sdkCfg = cfg.SDKConfig if effectiveProxyURL == "" { effectiveProxyURL = strings.TrimSpace(cfg.ProxyURL) } + agentNameHint = cfg.Codex.EffectiveSignInWithChatGPTIdentifier() } sdkCfg.ProxyURL = effectiveProxyURL return &CodexAuth{ - httpClient: util.SetProxy(&sdkCfg, &http.Client{}), + httpClient: util.SetProxy(&sdkCfg, &http.Client{}), + siwcAgentNameHint: agentNameHint, } } diff --git a/internal/auth/codex/refresh_patch.go b/internal/auth/codex/refresh_patch.go new file mode 100644 index 000000000..e3ce17f52 --- /dev/null +++ b/internal/auth/codex/refresh_patch.go @@ -0,0 +1,165 @@ +package codex + +import "reflect" + +var codexSIWCCredentialKeys = []string{ + "access_token", "refresh_token", "id_token", "client_id", "device_id", + "expired", "expires_at", "last_refresh", "scope", "scopes", +} + +// CredentialRefreshPatch is an immutable, partition-specific refresh result. +// Its original token identities (both in-memory and on-disk baselines) act as versions, +// independent of timestamp precision. +type CredentialRefreshPatch struct { + outerBefore map[string]any + outerDiskBefore map[string]any + outerAfter map[string]any + siwcBefore map[string]any + siwcDiskBefore map[string]any + siwcAfter map[string]any +} + +// NewCredentialRefreshPatch records only credential partitions actually refreshed. +func NewCredentialRefreshPatch(before, after map[string]any, outer, siwc bool) *CredentialRefreshPatch { + return NewCredentialRefreshPatchWithDiskBaseline(before, nil, after, outer, siwc) +} + +// NewCredentialRefreshPatchWithDiskBaseline records both the in-memory baseline (before) +// and the on-disk baseline (diskBefore) observed under the file lock before refreshing to after. +func NewCredentialRefreshPatchWithDiskBaseline(before, diskBefore, after map[string]any, outer, siwc bool) *CredentialRefreshPatch { + patch := &CredentialRefreshPatch{} + if outer { + patch.outerBefore = credentialFields(before, codexOuterCredentialKeys) + if len(diskBefore) > 0 { + patch.outerDiskBefore = credentialFields(diskBefore, codexOuterCredentialKeys) + } + patch.outerAfter = credentialFields(after, codexOuterCredentialKeys) + } + if siwc { + beforeSIWC, _ := before["siwc"].(map[string]any) + diskBeforeSIWC, _ := diskBefore["siwc"].(map[string]any) + afterSIWC, _ := after["siwc"].(map[string]any) + patch.siwcBefore = credentialFields(beforeSIWC, codexSIWCCredentialKeys) + if len(diskBeforeSIWC) > 0 { + patch.siwcDiskBefore = credentialFields(diskBeforeSIWC, codexSIWCCredentialKeys) + } + patch.siwcAfter = credentialFields(afterSIWC, codexSIWCCredentialKeys) + } + return patch +} + +func credentialFields(metadata map[string]any, keys []string) map[string]any { + fields := make(map[string]any) + for _, key := range keys { + if value, ok := metadata[key]; ok { + fields[key] = value + } + } + return cloneStringAnyMap(fields) +} + +func sameCredentialVersion(a, b map[string]any) bool { + for _, key := range []string{"access_token", "refresh_token", "id_token", "account_id", "client_id", "device_id"} { + if !reflect.DeepEqual(a[key], b[key]) { + return false + } + } + return true +} + +func matchesBaselineOrAfter(target, memBefore, diskBefore, after map[string]any) (matched bool, effectiveBefore map[string]any) { + if len(diskBefore) > 0 && sameCredentialVersion(target, diskBefore) { + return true, diskBefore + } + if sameCredentialVersion(target, memBefore) { + return true, memBefore + } + if len(after) > 0 && sameCredentialVersion(target, after) { + return true, after + } + return false, nil +} + +func applyCredentialFields(disk, before, after map[string]any) { + for key, value := range after { + // Preserve fields changed on disk that the refresh did not modify. + oldValue, existed := before[key] + if !existed || !reflect.DeepEqual(oldValue, value) { + disk[key] = value + } + } + if _, hadExpiresAt := before["expires_at"]; hadExpiresAt { + if _, hasExpiresAt := after["expires_at"]; !hasExpiresAt { + delete(disk, "expires_at") + } + } +} + +// IsCodexCredentialKey reports whether key belongs to a Codex credential partition or switch +// managed by CredentialRefreshPatch. +func IsCodexCredentialKey(key string) bool { + if key == "siwc" || key == "use_siwc" { + return true + } + for _, k := range codexOuterCredentialKeys { + if k == key { + return true + } + } + return false +} + +// RefreshedOuter reports whether p carries a refreshed outer Codex credential partition. +func (p *CredentialRefreshPatch) RefreshedOuter() bool { + return p != nil && p.outerAfter != nil +} + +// RefreshedSIWC reports whether p carries a refreshed SIWC credential partition. +func (p *CredentialRefreshPatch) RefreshedSIWC() bool { + return p != nil && p.siwcAfter != nil +} + +// CanApplyTo reports whether at least one refreshed partition in p still matches disk's version. +func (p *CredentialRefreshPatch) CanApplyTo(disk map[string]any) bool { + if p == nil { + return false + } + if p.outerAfter != nil { + if ok, _ := matchesBaselineOrAfter(disk, p.outerBefore, p.outerDiskBefore, p.outerAfter); ok { + return true + } + } + diskSIWC, _ := disk["siwc"].(map[string]any) + if p.siwcAfter != nil { + if ok, _ := matchesBaselineOrAfter(diskSIWC, p.siwcBefore, p.siwcDiskBefore, p.siwcAfter); ok { + return true + } + } + return false +} + +// Apply overlays successful refreshes on the locked disk snapshot. A partition +// rotated since refresh began is retained, without blocking the other partition. +func (p *CredentialRefreshPatch) Apply(disk map[string]any) map[string]any { + merged := cloneStringAnyMap(disk) + if merged == nil { + merged = make(map[string]any) + } + if p.outerAfter != nil { + if ok, base := matchesBaselineOrAfter(disk, p.outerBefore, p.outerDiskBefore, p.outerAfter); ok { + applyCredentialFields(merged, base, p.outerAfter) + } + } + diskSIWC, _ := disk["siwc"].(map[string]any) + if p.siwcAfter != nil { + if ok, base := matchesBaselineOrAfter(diskSIWC, p.siwcBefore, p.siwcDiskBefore, p.siwcAfter); ok { + mergedSIWC := cloneStringAnyMap(diskSIWC) + if mergedSIWC == nil { + mergedSIWC = make(map[string]any) + } + applyCredentialFields(mergedSIWC, base, p.siwcAfter) + merged["siwc"] = mergedSIWC + } + } + return merged +} diff --git a/internal/auth/codex/siwc_auth.go b/internal/auth/codex/siwc_auth.go new file mode 100644 index 000000000..975251652 --- /dev/null +++ b/internal/auth/codex/siwc_auth.go @@ -0,0 +1,267 @@ +package codex + +import ( + "context" + "crypto/rand" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "regexp" + "slices" + "strings" + "time" + + log "github.com/sirupsen/logrus" + "golang.org/x/sync/singleflight" +) + +// OAuth configuration constants for OpenAI Sign in with ChatGPT (SIWC). +const ( + SIWCDynamicClientID = "dynamic_agent_client" + SIWCAgentNameHint = "CLIProxyAPI" + SIWCAuthorizeURL = "https://auth.openai.com/api/accounts/authorize" + SIWCTokenURL = "https://auth.openai.com/api/accounts/oauth/token" + SIWCResource = "https://api.openai.com/v1" + SIWCCallbackHost = "127.0.0.1" + SIWCCallbackPort = 1455 + SIWCCallbackPath = "/auth/callback" + SIWCRedirectURI = "http://127.0.0.1:1455/auth/callback" + SIWCDirectTokenScope = "chatgpt.tokens.use.direct" + SIWCScope = "openid profile email offline_access resource.invoke chatgpt.tokens.use.direct" + SIWCExpiryMargin = 3 * time.Minute + SIWCRefreshLead = 5 * time.Minute +) + +var ( + siwcUUIDPattern = regexp.MustCompile(`(?i)^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$`) + siwcTokenEndpointURL = SIWCTokenURL + siwcRefreshGroup singleflight.Group +) + +// IsValidUUID reports whether raw is a canonical 36-char UUID string. +func IsValidUUID(raw string) bool { + return siwcUUIDPattern.MatchString(strings.TrimSpace(raw)) +} + +// GenerateRandomBase64URL generates a 32-byte cryptographically random base64url string. +func GenerateRandomBase64URL() (string, error) { + buf := make([]byte, 32) + if _, errRead := rand.Read(buf); errRead != nil { + return "", fmt.Errorf("failed to generate random bytes: %w", errRead) + } + return base64.RawURLEncoding.EncodeToString(buf), nil +} + +// GenerateSIWCAuthURL constructs the OpenAI Sign in with ChatGPT authorization URL. +func (o *CodexAuth) GenerateSIWCAuthURL(state, nonce, deviceID, redirectURI string, pkceCodes *PKCECodes) (string, error) { + if pkceCodes == nil { + return "", fmt.Errorf("PKCE codes are required") + } + deviceID = strings.TrimSpace(deviceID) + if !IsValidUUID(deviceID) { + return "", fmt.Errorf("Sign in with ChatGPT requires a valid device UUID") + } + if strings.TrimSpace(state) == "" { + return "", fmt.Errorf("state is required") + } + if strings.TrimSpace(nonce) == "" { + return "", fmt.Errorf("nonce is required") + } + if strings.TrimSpace(redirectURI) == "" { + redirectURI = SIWCRedirectURI + } + + agentNameHint := SIWCAgentNameHint + if o != nil && strings.TrimSpace(o.siwcAgentNameHint) != "" { + agentNameHint = strings.TrimSpace(o.siwcAgentNameHint) + } + params := url.Values{ + "client_id": {SIWCDynamicClientID}, + "agent_name_hint": {agentNameHint}, + "ext_agent_host_id": {"urn:uuid:" + strings.ToLower(deviceID)}, + "response_type": {"code"}, + "redirect_uri": {strings.TrimSpace(redirectURI)}, + "resource": {SIWCResource}, + "scope": {SIWCScope}, + "state": {strings.TrimSpace(state)}, + "code_challenge": {pkceCodes.CodeChallenge}, + "code_challenge_method": {"S256"}, + "nonce": {strings.TrimSpace(nonce)}, + } + return fmt.Sprintf("%s?%s", SIWCAuthorizeURL, params.Encode()), nil +} + +type siwcTokenResponse struct { + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token"` + IDToken string `json:"id_token"` + Scope string `json:"scope"` + ExpiresIn int `json:"expires_in"` +} + +func (o *CodexAuth) requestSIWCToken(ctx context.Context, form url.Values, clientID, deviceID string, requireIDToken bool) (*CodexSIWCTokenData, error) { + req, errReq := http.NewRequestWithContext(ctx, http.MethodPost, siwcTokenEndpointURL, strings.NewReader(form.Encode())) + if errReq != nil { + return nil, fmt.Errorf("failed to create SIWC token request: %w", errReq) + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + req.Header.Set("Accept", "application/json") + + resp, errDo := o.httpClient.Do(req) + if errDo != nil { + return nil, fmt.Errorf("SIWC token request failed: %w", errDo) + } + defer func() { + if errClose := resp.Body.Close(); errClose != nil { + log.Errorf("SIWC token response body close error: %v", errClose) + } + }() + + body, errRead := io.ReadAll(resp.Body) + if errRead != nil { + return nil, fmt.Errorf("failed to read SIWC token response: %w", errRead) + } + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("OpenAI OAuth token request failed (%d): %s", resp.StatusCode, strings.TrimSpace(string(body))) + } + + var tokenResp siwcTokenResponse + if errUnmarshal := json.Unmarshal(body, &tokenResp); errUnmarshal != nil { + return nil, fmt.Errorf("failed to parse SIWC token response: %w", errUnmarshal) + } + if strings.TrimSpace(tokenResp.AccessToken) == "" { + return nil, fmt.Errorf("OpenAI OAuth token response has invalid access_token") + } + if strings.TrimSpace(tokenResp.RefreshToken) == "" { + return nil, fmt.Errorf("OpenAI OAuth token response has invalid refresh_token") + } + if requireIDToken && strings.TrimSpace(tokenResp.IDToken) == "" { + return nil, fmt.Errorf("OpenAI OAuth token response did not contain an ID token") + } + if strings.TrimSpace(tokenResp.Scope) == "" { + return nil, fmt.Errorf("OpenAI OAuth token response has invalid scope") + } + if tokenResp.ExpiresIn <= 0 { + return nil, fmt.Errorf("OpenAI OAuth token response has invalid expires_in") + } + + scopes := strings.Fields(tokenResp.Scope) + if !slices.Contains(scopes, SIWCDirectTokenScope) { + return nil, fmt.Errorf("OpenAI OAuth grant did not include %s", SIWCDirectTokenScope) + } + + now := time.Now() + rawTTL := time.Duration(tokenResp.ExpiresIn) * time.Second + ttl := rawTTL - SIWCExpiryMargin + if ttl <= 0 { + ttl = rawTTL + } + + return &CodexSIWCTokenData{ + AccessToken: tokenResp.AccessToken, + RefreshToken: tokenResp.RefreshToken, + IDToken: tokenResp.IDToken, + ClientID: strings.TrimSpace(clientID), + Scope: strings.TrimSpace(tokenResp.Scope), + Scopes: scopes, + Expire: now.Add(ttl).Format(time.RFC3339), + ExpiresAt: now.Add(rawTTL).Format(time.RFC3339), + LastRefresh: now.Format(time.RFC3339), + DeviceID: strings.TrimSpace(deviceID), + }, nil +} + +// ExchangeSIWCCodeForTokens exchanges a Sign in with ChatGPT authorization code for tokens. +func (o *CodexAuth) ExchangeSIWCCodeForTokens(ctx context.Context, code, issuedClientID, redirectURI, deviceID string, pkceCodes *PKCECodes) (*CodexSIWCTokenData, error) { + if pkceCodes == nil { + return nil, fmt.Errorf("PKCE codes are required for SIWC token exchange") + } + issuedClientID = strings.TrimSpace(issuedClientID) + if issuedClientID == "" { + return nil, fmt.Errorf("OpenAI OAuth registration callback did not contain an issued client ID") + } + code = strings.TrimSpace(code) + if code == "" { + return nil, fmt.Errorf("authorization code is required") + } + if strings.TrimSpace(redirectURI) == "" { + redirectURI = SIWCRedirectURI + } + + form := url.Values{ + "grant_type": {"authorization_code"}, + "client_id": {issuedClientID}, + "code": {code}, + "code_verifier": {pkceCodes.CodeVerifier}, + "redirect_uri": {strings.TrimSpace(redirectURI)}, + "resource": {SIWCResource}, + } + return o.requestSIWCToken(ctx, form, issuedClientID, deviceID, true) +} + +// RefreshSIWCTokens refreshes a Sign in with ChatGPT access token using its refresh token and issued client ID. +func (o *CodexAuth) RefreshSIWCTokens(ctx context.Context, refreshToken, clientID, deviceID string) (*CodexSIWCTokenData, error) { + refreshToken = strings.TrimSpace(refreshToken) + if refreshToken == "" { + return nil, fmt.Errorf("SIWC refresh token is required") + } + clientID = strings.TrimSpace(clientID) + if clientID == "" { + return nil, fmt.Errorf("Stored OpenAI OAuth credential does not contain an issued client ID; reconnect ChatGPT") + } + if ctx == nil { + ctx = context.Background() + } + + sfKey := clientID + ":" + refreshToken + result, errSF, _ := siwcRefreshGroup.Do(sfKey, func() (any, error) { + refreshCtx, cancelRefresh := context.WithTimeout(context.WithoutCancel(ctx), codexRefreshTimeout) + defer cancelRefresh() + form := url.Values{ + "grant_type": {"refresh_token"}, + "client_id": {clientID}, + "refresh_token": {refreshToken}, + "resource": {SIWCResource}, + } + return o.requestSIWCToken(refreshCtx, form, clientID, deviceID, false) + }) + if errSF != nil { + return nil, errSF + } + tokenData, ok := result.(*CodexSIWCTokenData) + if !ok || tokenData == nil { + return nil, fmt.Errorf("SIWC token refresh failed: invalid single-flight result") + } + return tokenData, nil +} + +// RefreshSIWCTokensWithRetry refreshes SIWC tokens with exponential backoff. +func (o *CodexAuth) RefreshSIWCTokensWithRetry(ctx context.Context, refreshToken, clientID, deviceID string, maxRetries int) (*CodexSIWCTokenData, error) { + if maxRetries <= 0 { + maxRetries = 1 + } + var lastErr error + for attempt := 0; attempt < maxRetries; attempt++ { + if attempt > 0 { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(time.Duration(attempt) * time.Second): + } + } + tokenData, errRefresh := o.RefreshSIWCTokens(ctx, refreshToken, clientID, deviceID) + if errRefresh == nil { + return tokenData, nil + } + if isNonRetryableRefreshErr(errRefresh) { + log.Warnf("SIWC token refresh attempt %d failed with non-retryable error: %v", attempt+1, errRefresh) + return nil, errRefresh + } + lastErr = errRefresh + log.Warnf("SIWC token refresh attempt %d failed: %v", attempt+1, errRefresh) + } + return nil, fmt.Errorf("SIWC token refresh failed after %d attempts: %w", maxRetries, lastErr) +} diff --git a/internal/auth/codex/token.go b/internal/auth/codex/token.go index 05d95d9b4..6baffb944 100644 --- a/internal/auth/codex/token.go +++ b/internal/auth/codex/token.go @@ -8,11 +8,222 @@ import ( "fmt" "os" "path/filepath" + "strconv" + "strings" + "sync" "github.com/router-for-me/CLIProxyAPI/v8/internal/misc" log "github.com/sirupsen/logrus" ) +var codexFileLocks sync.Map + +func lockForFilePath(authFilePath string) *sync.Mutex { + cleaned := canonicalLockPath(authFilePath) + if actual, ok := codexFileLocks.Load(cleaned); ok { + return actual.(*sync.Mutex) + } + mu := &sync.Mutex{} + actual, _ := codexFileLocks.LoadOrStore(cleaned, mu) + return actual.(*sync.Mutex) +} + +func canonicalLockPath(authFilePath string) string { + cleaned := filepath.Clean(strings.TrimSpace(authFilePath)) + if abs, errAbs := filepath.Abs(cleaned); errAbs == nil && strings.TrimSpace(abs) != "" { + cleaned = filepath.Clean(abs) + } + if real, errEval := filepath.EvalSymlinks(cleaned); errEval == nil && strings.TrimSpace(real) != "" { + cleaned = filepath.Clean(real) + } else if parent, errParent := filepath.EvalSymlinks(filepath.Dir(cleaned)); errParent == nil && strings.TrimSpace(parent) != "" { + cleaned = filepath.Clean(filepath.Join(parent, filepath.Base(cleaned))) + } + return cleaned +} + +// WithFileLock executes fn while holding both an in-process mutex and an inter-process lockfile for authFilePath. +func WithFileLock(authFilePath string, fn func() error) error { + mu := lockForFilePath(authFilePath) + mu.Lock() + defer mu.Unlock() + + if errDir := os.MkdirAll(filepath.Dir(authFilePath), 0o700); errDir != nil { + return fmt.Errorf("failed to create credential directory: %w", errDir) + } + releaseLockFile, errLock := acquireInterProcessLockFile(authFilePath) + if errLock != nil { + return errLock + } + defer releaseLockFile() + + return fn() +} + +// WriteJSONFileAtomic writes data as JSON to authFilePath atomically via a temporary file in the same directory. +func WriteJSONFileAtomic(authFilePath string, data map[string]any) error { + dir := filepath.Dir(authFilePath) + tmpFile, errTmp := os.CreateTemp(dir, filepath.Base(authFilePath)+".tmp.*") + if errTmp != nil { + return fmt.Errorf("failed to create temp credential file: %w", errTmp) + } + tmpPath := tmpFile.Name() + cleanupTmp := true + defer func() { + if cleanupTmp { + _ = os.Remove(tmpPath) + } + }() + + if errChmod := tmpFile.Chmod(0o600); errChmod != nil { + _ = tmpFile.Close() + return fmt.Errorf("failed to chmod temp credential file: %w", errChmod) + } + raw, errMarshal := json.Marshal(data) + if errMarshal != nil { + _ = tmpFile.Close() + return fmt.Errorf("failed to marshal token data: %w", errMarshal) + } + if _, errWrite := tmpFile.Write(raw); errWrite != nil { + _ = tmpFile.Close() + return fmt.Errorf("failed to write token to temp file: %w", errWrite) + } + if errClose := tmpFile.Close(); errClose != nil { + log.Errorf("codex token storage: close temp token file error: %v", errClose) + return fmt.Errorf("failed to close temp credential file: %w", errClose) + } + if errRename := os.Rename(tmpPath, authFilePath); errRename != nil { + return fmt.Errorf("failed to replace credential file: %w", errRename) + } + cleanupTmp = false + return nil +} + +var codexOuterCredentialKeys = []string{ + "id_token", + "access_token", + "refresh_token", + "account_id", + "last_refresh", + "email", + "expired", + "plan_type", +} + +func cloneStringAnyMap(src map[string]any) map[string]any { + if src == nil { + return nil + } + dst := make(map[string]any, len(src)) + for k, v := range src { + switch typed := v.(type) { + case map[string]any: + dst[k] = cloneStringAnyMap(typed) + case []string: + s := make([]string, len(typed)) + copy(s, typed) + dst[k] = s + case []any: + s := make([]any, len(typed)) + copy(s, typed) + dst[k] = s + default: + dst[k] = v + } + } + return dst +} + +// MergeWithDiskMap preserves disk credentials when saving a non-refresh snapshot. +// Explicit login replacements and refresh patches are handled separately. +func MergeWithDiskMap(diskMap, incomingMap map[string]any) map[string]any { + merged := cloneStringAnyMap(diskMap) + if merged == nil { + merged = make(map[string]any) + } + for key, value := range incomingMap { + if value == nil { + continue + } + if key == "siwc" || key == "use_siwc" { + if _, exists := diskMap[key]; exists { + continue + } + } + if text, ok := value.(string); ok && strings.TrimSpace(text) == "" { + if _, exists := diskMap[key]; exists { + continue + } + } + merged[key] = value + } + for _, key := range codexOuterCredentialKeys { + if value, exists := diskMap[key]; exists { + merged[key] = value + } else if len(diskMap) > 0 { + delete(merged, key) + } + } + if _, exists := diskMap["siwc"]; !exists && len(diskMap) > 0 { + delete(merged, "siwc") + } + return merged +} + +func hasDiskOAuthOuter(diskMap map[string]any) bool { + if len(diskMap) == 0 { + return false + } + if rt, ok := diskMap["refresh_token"].(string); ok && strings.TrimSpace(rt) != "" { + return true + } + if lr, ok := diskMap["last_refresh"].(string); ok && strings.TrimSpace(lr) != "" { + return true + } + return false +} + +// MergeCodexCredentialsFromDisk protects disk credentials during configuration +// saves while allowing intentional changes or deletions of non-credential settings. +func MergeCodexCredentialsFromDisk(diskMap, incomingMetadata map[string]any) map[string]any { + merged := cloneStringAnyMap(incomingMetadata) + if merged == nil { + merged = make(map[string]any) + } + if hasDiskOAuthOuter(diskMap) { + for _, key := range codexOuterCredentialKeys { + if value, exists := diskMap[key]; exists { + merged[key] = value + } else { + delete(merged, key) + } + } + } else { + for _, key := range codexOuterCredentialKeys { + if value, exists := diskMap[key]; exists { + if inVal, hadIn := merged[key]; !hadIn || inVal == nil || (func() bool { + s, ok := inVal.(string) + return ok && strings.TrimSpace(s) == "" + }()) { + merged[key] = value + } + } + } + } + if value, exists := diskMap["siwc"]; exists { + if siwcMap, ok := value.(map[string]any); ok { + merged["siwc"] = cloneStringAnyMap(siwcMap) + } else { + merged["siwc"] = value + } + } + if value, exists := diskMap["use_siwc"]; exists { + if _, specified := incomingMetadata["use_siwc"]; !specified { + merged["use_siwc"] = value + } + } + return merged +} + // CodexTokenStorage stores OAuth2 token information for OpenAI Codex API authentication. // It maintains compatibility with the existing auth system while adding Codex-specific fields // for managing access tokens, refresh tokens, and user account information. @@ -35,6 +246,15 @@ type CodexTokenStorage struct { Expire string `json:"expired"` // PlanType indicates the ChatGPT subscription plan type, defaults to "free" if not present. PlanType string `json:"plan_type,omitempty"` + // UseSIWC controls whether API requests use the Sign in with ChatGPT (SIWC) credential section. + UseSIWC *bool `json:"use_siwc,omitempty"` + // SIWC holds the Sign in with ChatGPT OAuth credential fields. + SIWC map[string]any `json:"siwc,omitempty"` + + // ReplaceCredentials permits an explicit login to replace existing outer tokens. + ReplaceCredentials bool `json:"-"` + // RefreshPatch updates only refreshed partitions during a direct storage save. + RefreshPatch *CredentialRefreshPatch `json:"-"` // Metadata holds arbitrary key-value pairs injected via hooks. // It is not exported to JSON directly to allow flattening during serialization. @@ -46,41 +266,42 @@ func (ts *CodexTokenStorage) SetMetadata(meta map[string]any) { ts.Metadata = meta } -// SaveTokenToFile serializes the Codex token storage to a JSON file. -// This method creates the necessary directory structure and writes the token -// data in JSON format to the specified file path for persistent storage. -// It merges any injected metadata into the top-level JSON object. -// -// Parameters: -// - authFilePath: The full path where the token file should be saved -// -// Returns: -// - error: An error if the operation fails, nil otherwise +// SaveTokenToFile serializes the Codex token storage to a JSON file under a per-file and inter-process lock, +// reading and merging any existing JSON on disk before atomically replacing the file. func (ts *CodexTokenStorage) SaveTokenToFile(authFilePath string) error { misc.LogSavingCredentials(authFilePath) ts.Type = "codex" - if err := os.MkdirAll(filepath.Dir(authFilePath), 0700); err != nil { - return fmt.Errorf("failed to create directory: %v", err) - } + _ = strconv.IntSize - // Merge metadata using helper - data, errMerge := misc.MergeMetadata(ts, ts.Metadata) - if errMerge != nil { - return fmt.Errorf("failed to merge metadata: %w", errMerge) - } + return WithFileLock(authFilePath, func() error { + data, errMerge := misc.MergeMetadata(ts, ts.Metadata) + if errMerge != nil { + return fmt.Errorf("failed to merge metadata: %w", errMerge) + } - f, err := os.Create(authFilePath) - if err != nil { - return fmt.Errorf("failed to create token file: %w", err) - } - defer func() { - if errClose := f.Close(); errClose != nil { - log.Errorf("codex token storage: close token file error: %v", errClose) + existingBytes, errRead := os.ReadFile(authFilePath) + if errRead != nil && !os.IsNotExist(errRead) { + return fmt.Errorf("failed to read existing credentials: %w", errRead) + } + if errRead == nil { + var diskMap map[string]any + if errUnmarshal := json.Unmarshal(existingBytes, &diskMap); errUnmarshal != nil || diskMap == nil { + return fmt.Errorf("invalid existing credential JSON") + } + if ts.RefreshPatch != nil { + return WriteJSONFileAtomic(authFilePath, ts.RefreshPatch.Apply(diskMap)) + } + merged := MergeWithDiskMap(diskMap, data) + if ts.ReplaceCredentials { + for _, key := range codexOuterCredentialKeys { + if value, exists := data[key]; exists { + merged[key] = value + } + } + } + data = merged } - }() - if err = json.NewEncoder(f).Encode(data); err != nil { - return fmt.Errorf("failed to write token to file: %w", err) - } - return nil + return WriteJSONFileAtomic(authFilePath, data) + }) } diff --git a/internal/cache/antigravity_interactions_session.go b/internal/cache/antigravity_interactions_session.go new file mode 100644 index 000000000..97ad5f8e1 --- /dev/null +++ b/internal/cache/antigravity_interactions_session.go @@ -0,0 +1,76 @@ +package cache + +import ( + "sort" + "strings" + "sync" + "time" +) + +// InteractionsContinuation contains only upstream continuation identifiers. +type InteractionsContinuation struct { + ID string + Environment string +} + +type interactionsContinuationEntry struct { + state InteractionsContinuation + expires time.Time +} + +// InteractionsSessionCache is a bounded, process-local cache. Keys must already +// isolate the caller, credential, model, conversation, and pending call IDs. +type InteractionsSessionCache struct { + mu sync.Mutex + entries map[string]interactionsContinuationEntry + now func() time.Time +} + +func NewInteractionsSessionCache() *InteractionsSessionCache { + return &InteractionsSessionCache{entries: make(map[string]interactionsContinuationEntry), now: time.Now} +} + +// InteractionsCallKey rejects ambiguous or incomplete call sets. +func InteractionsCallKey(ids []string) string { + ids = append([]string(nil), ids...) + sort.Strings(ids) + for i, id := range ids { + if id == "" || strings.ContainsRune(id, '\x00') || (i > 0 && ids[i-1] == id) { + return "" + } + } + return strings.Join(ids, "\x00") +} + +func (c *InteractionsSessionCache) Get(key string) (InteractionsContinuation, bool) { + c.mu.Lock() + defer c.mu.Unlock() + entry, ok := c.entries[key] + if !ok || !c.now().Before(entry.expires) { + delete(c.entries, key) + return InteractionsContinuation{}, false + } + return entry.state, true +} + +func (c *InteractionsSessionCache) Put(key string, state InteractionsContinuation) { + c.mu.Lock() + defer c.mu.Unlock() + now := c.now() + for k, entry := range c.entries { + if !now.Before(entry.expires) { + delete(c.entries, k) + } + } + if _, exists := c.entries[key]; !exists && len(c.entries) >= 1024 { + var oldestKey string + var oldest time.Time + for k, entry := range c.entries { + if oldest.IsZero() || entry.expires.Before(oldest) { + oldestKey, oldest = k, entry.expires + } + } + delete(c.entries, oldestKey) + } + c.entries[key] = interactionsContinuationEntry{state: state, expires: now.Add(30 * time.Minute)} +} diff --git a/internal/cache/claude_web_search_result_cache.go b/internal/cache/claude_web_search_result_cache.go new file mode 100644 index 000000000..6573881e6 --- /dev/null +++ b/internal/cache/claude_web_search_result_cache.go @@ -0,0 +1,221 @@ +package cache + +import ( + "bytes" + "container/list" + "context" + "crypto/sha256" + "encoding/hex" + "strings" + "sync" + "time" + + homekv "github.com/router-for-me/CLIProxyAPI/v8/internal/home" + log "github.com/sirupsen/logrus" +) + +const ( + defaultClaudeWebSearchResultTTL = 24 * time.Hour + defaultClaudeWebSearchResultLimit = 1024 +) + +type claudeWebSearchResultEntry struct { + key string + content []byte + ttl time.Duration + expiresAt time.Time +} + +var ( + claudeWebSearchResultMu sync.Mutex + claudeWebSearchResultIndex = make(map[string]*list.Element) + claudeWebSearchResultOrder = list.New() + claudeWebSearchResultNow = time.Now +) + +type claudeWebSearchResultKVClient interface { + KVGet(ctx context.Context, key string) ([]byte, bool, error) + KVSet(ctx context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error) + KVExpire(ctx context.Context, key string, ttl time.Duration) (bool, error) +} + +var currentClaudeWebSearchResultKVClient = func() (claudeWebSearchResultKVClient, bool, error) { + return homekv.CurrentKVClient() +} + +func claudeWebSearchResultKey(sessionKey, toolUseID string) (string, bool) { + sessionKey = strings.TrimSpace(sessionKey) + toolUseID = strings.TrimSpace(toolUseID) + if sessionKey == "" || toolUseID == "" { + return "", false + } + sum := sha256.Sum256([]byte(sessionKey + "\x00" + toolUseID)) + return hex.EncodeToString(sum[:]), true +} + +func claudeWebSearchResultKVKey(sessionKey, toolUseID string) string { + key, ok := claudeWebSearchResultKey(sessionKey, toolUseID) + if !ok { + return "" + } + return "cpa:claude:web-search-result:" + key +} + +// CacheClaudeWebSearchResult stores the raw content of one web_search_tool_result +// block for a session. Memory is bounded by the maximum entry count (limit); +// least recently used entries are evicted when the entry limit is exceeded. +func CacheClaudeWebSearchResult(sessionKey, toolUseID string, content []byte, ttl time.Duration, limit int) { + CacheClaudeWebSearchResultBestEffort(context.Background(), sessionKey, toolUseID, content, ttl, limit) +} + +// CacheClaudeWebSearchResultBestEffort stores the raw content of one +// web_search_tool_result block for completed response paths. +func CacheClaudeWebSearchResultBestEffort(ctx context.Context, sessionKey, toolUseID string, content []byte, ttl time.Duration, limit int) bool { + key, ok := claudeWebSearchResultKey(sessionKey, toolUseID) + if !ok || len(content) == 0 { + return false + } + if ttl <= 0 { + ttl = defaultClaudeWebSearchResultTTL + } + if ctx == nil { + ctx = context.Background() + } + if client, homeMode, errClient := currentClaudeWebSearchResultKVClient(); homeMode { + if errClient != nil { + log.Errorf("home kv best-effort claude web search result set failed prefix=cpa:claude:*: %v", errClient) + return false + } + written, errSet := client.KVSet(ctx, "cpa:claude:web-search-result:"+key, bytes.Clone(content), homekv.KVSetOptions{EX: ttl}) + if errSet != nil { + log.Errorf("home kv best-effort claude web search result set failed prefix=cpa:claude:*: %v", errSet) + return false + } + return written + } + if limit <= 0 { + limit = defaultClaudeWebSearchResultLimit + } + cacheCleanupOnce.Do(startCacheCleanup) + + claudeWebSearchResultMu.Lock() + defer claudeWebSearchResultMu.Unlock() + + now := claudeWebSearchResultNow() + if element, exists := claudeWebSearchResultIndex[key]; exists { + removeClaudeWebSearchResultLocked(element) + } + if claudeWebSearchResultOrder.Len() >= limit { + purgeExpiredClaudeWebSearchResultLocked(now) + } + entry := &claudeWebSearchResultEntry{ + key: key, + content: bytes.Clone(content), + ttl: ttl, + expiresAt: now.Add(ttl), + } + claudeWebSearchResultIndex[key] = claudeWebSearchResultOrder.PushFront(entry) + for claudeWebSearchResultOrder.Len() > limit { + oldest := claudeWebSearchResultOrder.Back() + if oldest == nil { + break + } + removeClaudeWebSearchResultLocked(oldest) + } + return true +} + +// GetClaudeWebSearchResult returns a copy of the stored content and extends its +// lifetime. Expired entries are removed and reported as missing. +func GetClaudeWebSearchResult(sessionKey, toolUseID string, ttl ...time.Duration) ([]byte, bool) { + content, ok, errGet := GetClaudeWebSearchResultRequired(context.Background(), sessionKey, toolUseID, ttl...) + if errGet != nil { + return nil, false + } + return content, ok +} + +// GetClaudeWebSearchResultRequired retrieves a copy of the stored content for +// request-time paths and extends its lifetime. +func GetClaudeWebSearchResultRequired(ctx context.Context, sessionKey, toolUseID string, ttl ...time.Duration) ([]byte, bool, error) { + key, ok := claudeWebSearchResultKey(sessionKey, toolUseID) + if !ok { + return nil, false, nil + } + if ctx == nil { + ctx = context.Background() + } + if client, homeMode, errClient := currentClaudeWebSearchResultKVClient(); homeMode { + if errClient != nil { + return nil, false, errClient + } + kvKey := "cpa:claude:web-search-result:" + key + raw, found, errGet := client.KVGet(ctx, kvKey) + if errGet != nil || !found || len(raw) == 0 { + return nil, false, errGet + } + effectiveTTL := defaultClaudeWebSearchResultTTL + if len(ttl) > 0 && ttl[0] > 0 { + effectiveTTL = ttl[0] + } + if _, errExpire := client.KVExpire(ctx, kvKey, effectiveTTL); errExpire != nil { + log.Warnf("home kv claude web search result expire failed prefix=cpa:claude:*: %v", errExpire) + } + return bytes.Clone(raw), true, nil + } + claudeWebSearchResultMu.Lock() + defer claudeWebSearchResultMu.Unlock() + + element, exists := claudeWebSearchResultIndex[key] + if !exists { + return nil, false, nil + } + entry := element.Value.(*claudeWebSearchResultEntry) + now := claudeWebSearchResultNow() + if !now.Before(entry.expiresAt) { + removeClaudeWebSearchResultLocked(element) + return nil, false, nil + } + effectiveTTL := entry.ttl + if len(ttl) > 0 && ttl[0] > 0 { + effectiveTTL = ttl[0] + entry.ttl = effectiveTTL + } + if effectiveTTL <= 0 { + effectiveTTL = defaultClaudeWebSearchResultTTL + } + entry.expiresAt = now.Add(effectiveTTL) + claudeWebSearchResultOrder.MoveToFront(element) + return bytes.Clone(entry.content), true, nil +} + +// ClearClaudeWebSearchResultCache removes every stored entry. +func ClearClaudeWebSearchResultCache() { + claudeWebSearchResultMu.Lock() + defer claudeWebSearchResultMu.Unlock() + claudeWebSearchResultIndex = make(map[string]*list.Element) + claudeWebSearchResultOrder.Init() +} + +func purgeExpiredClaudeWebSearchResultCache(now time.Time) { + claudeWebSearchResultMu.Lock() + defer claudeWebSearchResultMu.Unlock() + purgeExpiredClaudeWebSearchResultLocked(now) +} + +func purgeExpiredClaudeWebSearchResultLocked(now time.Time) { + for element := claudeWebSearchResultOrder.Back(); element != nil; { + prev := element.Prev() + entry := element.Value.(*claudeWebSearchResultEntry) + if !now.Before(entry.expiresAt) { + removeClaudeWebSearchResultLocked(element) + } + element = prev + } +} + +func removeClaudeWebSearchResultLocked(element *list.Element) { + entry := element.Value.(*claudeWebSearchResultEntry) + delete(claudeWebSearchResultIndex, entry.key) + claudeWebSearchResultOrder.Remove(element) +} diff --git a/internal/cache/codex_reasoning_replay_cache.go b/internal/cache/codex_reasoning_replay_cache.go index 7c63b96d3..d89921882 100644 --- a/internal/cache/codex_reasoning_replay_cache.go +++ b/internal/cache/codex_reasoning_replay_cache.go @@ -405,7 +405,7 @@ func normalizeCodexReasoningReplayReasoningItem(itemResult gjson.Result) ([]byte if encryptedContent != strings.TrimSpace(encryptedContent) { return nil, false } - if _, err := signature.InspectGPTReasoningSignature(encryptedContent); err != nil { + if !signature.IsValidGPTReasoningSignature(encryptedContent) && !signature.IsValidMetaReasoningSignature(encryptedContent) { return nil, false } diff --git a/internal/cache/signature_cache.go b/internal/cache/signature_cache.go index 2a6c0a484..da7186c0a 100644 --- a/internal/cache/signature_cache.go +++ b/internal/cache/signature_cache.go @@ -113,6 +113,7 @@ func purgeExpiredCaches() { purgeExpiredAntigravityReasoningReplayCache(now) purgeExpiredKimiThinkingReplayCache(now) purgeExpiredClaudeThinkingReplayCache(now) + purgeExpiredClaudeWebSearchResultCache(now) } // CacheSignature stores a thinking signature for a given model group and text. diff --git a/internal/client/codex/live/capabilities.go b/internal/client/codex/live/capabilities.go index ac7af04cb..8c393c681 100644 --- a/internal/client/codex/live/capabilities.go +++ b/internal/client/codex/live/capabilities.go @@ -135,7 +135,7 @@ func (h *Handler) HandleHangup(c *gin.Context) { body, errRead := readBody(c.Request.Body) if errRead != nil { - writeRealtimeError(c, http.StatusBadRequest, errRead.Error(), "invalid_request_error", "invalid_request") + writeRealtimeError(c, clienterror.HTTPStatusFromErrorOr(errRead, http.StatusBadRequest), errRead.Error(), "invalid_request_error", "invalid_request") return } upstreamURL := h.realtimeHTTPBaseURL() + "/realtime/calls/" + url.PathEscape(callID) + "/hangup" diff --git a/internal/client/codex/live/client_secret.go b/internal/client/codex/live/client_secret.go index 5f44300bd..ef9811042 100644 --- a/internal/client/codex/live/client_secret.go +++ b/internal/client/codex/live/client_secret.go @@ -13,6 +13,7 @@ import ( "time" "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v8/internal/clienterror" ) const ( @@ -212,7 +213,7 @@ func (h *Handler) CreateClientSecret(c *gin.Context) { } body, errRead := readClientSecretBody(c.Request.Body) if errRead != nil { - status := http.StatusBadRequest + status := clienterror.HTTPStatusFromErrorOr(errRead, http.StatusBadRequest) if errors.Is(errRead, errBodyTooLarge) { status = http.StatusRequestEntityTooLarge } @@ -237,7 +238,7 @@ func (h *Handler) CreateLegacySession(c *gin.Context) { } body, errRead := readClientSecretBody(c.Request.Body) if errRead != nil { - status := http.StatusBadRequest + status := clienterror.HTTPStatusFromErrorOr(errRead, http.StatusBadRequest) if errors.Is(errRead, errBodyTooLarge) { status = http.StatusRequestEntityTooLarge } diff --git a/internal/client/codex/models/apply_patch.go b/internal/client/codex/models/apply_patch.go index 0883b7fd0..c2934ad11 100644 --- a/internal/client/codex/models/apply_patch.go +++ b/internal/client/codex/models/apply_patch.go @@ -9,7 +9,7 @@ type ApplyPatchCapabilityForModelFunc func(string) bool func applyCodexClientApplyPatchCapability(entry map[string]any, id string, capability ApplyPatchCapabilityForModelFunc) { templateSupported := entry["apply_patch_tool_type"] == "freeform" entry["apply_patch_tool_type"] = nil - // Built-in image/video IDs can inherit a generic text template in Home. + // Built-in image, video, and speech IDs can inherit a generic text template in Home. // This classification is a model restriction, not routing capability evidence. baseID := strings.ToLower(strings.TrimSpace(id)) if index := strings.LastIndex(baseID, "/"); index != -1 { diff --git a/internal/client/codex/models/models.go b/internal/client/codex/models/models.go index 97001de2e..d15faca40 100644 --- a/internal/client/codex/models/models.go +++ b/internal/client/codex/models/models.go @@ -23,6 +23,10 @@ type ProvidersForModelFunc func(string) []string // metadata for an exact public model ID. nil means unknown. type WebSearchCapabilityForModelFunc func(string) *bool +// ToolSearchCapabilityForModelFunc reports whether every credential serving a +// public model ID can receive Codex tool_search as Anthropic native tool search. +type ToolSearchCapabilityForModelFunc func(id string) bool + var ( codexClientModelTemplatesMu sync.Mutex codexClientModelTemplatesLoaded bool @@ -66,14 +70,22 @@ func BuildResponseForClient(availableModels []map[string]any, providersForModel // BuildResponseForClientWithCPACapabilities builds a client response while // allowing Home to supply capability metadata independent of the local registry. func BuildResponseForClientWithCPACapabilities(availableModels []map[string]any, providersForModel ProvidersForModelFunc, webSearchCapabilityForModel WebSearchCapabilityForModelFunc, optimizeMultiAgentV2 bool, clientVersion string) map[string]any { - return BuildResponseForClientWithToolCapabilities(availableModels, providersForModel, webSearchCapabilityForModel, nil, optimizeMultiAgentV2, clientVersion) + return BuildResponseForClientWithToolCapabilities(availableModels, providersForModel, webSearchCapabilityForModel, nil, nil, optimizeMultiAgentV2, clientVersion) } // BuildResponseForClientWithToolCapabilities adds executor-backed tool support // independently of template metadata and legacy provider restrictions. -func BuildResponseForClientWithToolCapabilities(availableModels []map[string]any, providersForModel ProvidersForModelFunc, webSearchCapabilityForModel WebSearchCapabilityForModelFunc, applyPatchCapabilityForModel ApplyPatchCapabilityForModelFunc, optimizeMultiAgentV2 bool, clientVersion string) map[string]any { +func BuildResponseForClientWithToolCapabilities(availableModels []map[string]any, providersForModel ProvidersForModelFunc, webSearchCapabilityForModel WebSearchCapabilityForModelFunc, applyPatchCapabilityForModel ApplyPatchCapabilityForModelFunc, toolSearchCapabilityForModel ToolSearchCapabilityForModelFunc, optimizeMultiAgentV2 bool, clientVersion string) map[string]any { + return BuildResponseForClientWithSearchToolRoutes(availableModels, providersForModel, nil, webSearchCapabilityForModel, applyPatchCapabilityForModel, toolSearchCapabilityForModel, optimizeMultiAgentV2, clientVersion) +} + +// BuildResponseForClientWithSearchToolRoutes lets Home supply routing evidence +// for supports_search_tool without the other provider-dependent metadata. +// A nil searchToolProvidersForModel uses providersForModel. A nil +// toolSearchCapabilityForModel keeps Claude routes from advertising search. +func BuildResponseForClientWithSearchToolRoutes(availableModels []map[string]any, providersForModel, searchToolProvidersForModel ProvidersForModelFunc, webSearchCapabilityForModel WebSearchCapabilityForModelFunc, applyPatchCapabilityForModel ApplyPatchCapabilityForModelFunc, toolSearchCapabilityForModel ToolSearchCapabilityForModelFunc, optimizeMultiAgentV2 bool, clientVersion string) map[string]any { return map[string]any{ - "models": buildCodexClientModelsWithToolCapabilities(availableModels, providersForModel, webSearchCapabilityForModel, applyPatchCapabilityForModel, optimizeMultiAgentV2, clientVersion), + "models": buildCodexClientModelsWithToolCapabilities(availableModels, providersForModel, searchToolProvidersForModel, webSearchCapabilityForModel, applyPatchCapabilityForModel, toolSearchCapabilityForModel, optimizeMultiAgentV2, clientVersion), } } @@ -90,14 +102,17 @@ func MarshalCompact(payload any) ([]byte, error) { } func buildCodexClientModels(models []map[string]any, providersForModel ProvidersForModelFunc, webSearchCapabilityForModel WebSearchCapabilityForModelFunc, optimizeMultiAgentV2 bool, clientVersion string) []map[string]any { - return buildCodexClientModelsWithToolCapabilities(models, providersForModel, webSearchCapabilityForModel, nil, optimizeMultiAgentV2, clientVersion) + return buildCodexClientModelsWithToolCapabilities(models, providersForModel, nil, webSearchCapabilityForModel, nil, nil, optimizeMultiAgentV2, clientVersion) } -func buildCodexClientModelsWithToolCapabilities(models []map[string]any, providersForModel ProvidersForModelFunc, webSearchCapabilityForModel WebSearchCapabilityForModelFunc, applyPatchCapabilityForModel ApplyPatchCapabilityForModelFunc, optimizeMultiAgentV2 bool, clientVersion string) []map[string]any { +func buildCodexClientModelsWithToolCapabilities(models []map[string]any, providersForModel, searchToolProvidersForModel ProvidersForModelFunc, webSearchCapabilityForModel WebSearchCapabilityForModelFunc, applyPatchCapabilityForModel ApplyPatchCapabilityForModelFunc, toolSearchCapabilityForModel ToolSearchCapabilityForModelFunc, optimizeMultiAgentV2 bool, clientVersion string) []map[string]any { templates, defaultTemplate, errLoadCodexClientModelTemplates := loadCodexClientModelTemplates() if errLoadCodexClientModelTemplates != nil || defaultTemplate == nil { return nil } + if searchToolProvidersForModel == nil { + searchToolProvidersForModel = providersForModel + } result := make([]map[string]any, 0, len(models)) for _, model := range models { @@ -121,7 +136,8 @@ func buildCodexClientModelsWithToolCapabilities(models []map[string]any, provide if thinkingSupport := codexClientThinkingSupport(model); thinkingSupport != nil { applyCodexClientThinkingMetadata(entry, thinkingSupport, clientVersion) } - applyCodexClientProviderCapabilities(entry, id, true, providersForModel) + applyCodexClientProviderCapabilities(entry, id, providersForModel) + applyCodexClientSearchToolSupport(entry, id, true, searchToolProvidersForModel, toolSearchCapabilityForModel) applyCPAWebSearchCapability(entry, id, webSearchCapabilityForModel, clientVersion) sanitizeCodexClientReasoningMetadata(entry, clientVersion) applyCodexClientVisibilityOverride(entry, id) @@ -137,7 +153,7 @@ func buildCodexClientModelsWithToolCapabilities(models []map[string]any, provide entry := cloneCodexClientModelMap(defaultTemplate) applyCodexClientModelMetadata(entry, id, model, optimizeMultiAgentV2, clientVersion) applyCodexClientMaxTokens(entry, model) - applyCodexClientProviderCapabilities(entry, id, false, providersForModel) + applyCodexClientSearchToolSupport(entry, id, false, searchToolProvidersForModel, toolSearchCapabilityForModel) applyCPAWebSearchCapability(entry, id, webSearchCapabilityForModel, clientVersion) sanitizeCodexClientReasoningMetadata(entry, clientVersion) applyCodexClientVisibilityOverride(entry, id) @@ -233,7 +249,7 @@ func loadCodexClientModelTemplatesSnapshot(raw []byte, revision uint64) (map[str continue } templates[slug] = cloneCodexClientModelMap(model) - if slug == "gpt-5.5" { + if defaultTemplate == nil || slug == "gpt-5.5" { defaultTemplate = cloneCodexClientModelMap(model) } } @@ -542,25 +558,21 @@ func applyCPAWebSearchCapability(entry map[string]any, id string, capabilityForM } } -func applyCodexClientProviderCapabilities(entry map[string]any, id string, isTemplate bool, providersForModel ProvidersForModelFunc) { - if !isTemplate { - applyCodexClientSearchToolSupport(entry, id, false, providersForModel) - return - } +func applyCodexClientProviderCapabilities(entry map[string]any, id string, providersForModel ProvidersForModelFunc) { if providersForModel != nil && !isPureCodexProvider(id, providersForModel) { - entry["supports_search_tool"] = false entry["prefer_websockets"] = false entry["service_tiers"] = []any{} nullCodexClientRequiredOptions(entry) - return } - applyCodexClientSearchToolSupport(entry, id, isTemplate, providersForModel) } // nullCodexClientRequiredOptions clears capabilities that non-Codex models must not advertise. // Codex requires these keys to be present; omitting them rejects the entire catalog. func nullCodexClientRequiredOptions(entry map[string]any) { entry["apply_patch_tool_type"] = nil + entry["tool_mode"] = "direct" + entry["use_responses_lite"] = false + entry["supports_reasoning_effort_updates"] = false entry["upgrade"] = nil entry["availability_nux"] = nil } @@ -586,53 +598,58 @@ func isPureCodexProvider(id string, providersForModel ProvidersForModelFunc) boo if providersForModel == nil { return true } + return allCodexClientProvidersAre(codexClientModelProviders(id, providersForModel), "codex") +} + +// codexClientModelProviders resolves the providers serving a model, falling +// back to the base ID for prefixed aliases. +func codexClientModelProviders(id string, providersForModel ProvidersForModelFunc) []string { providers := providersForModel(id) if len(providers) == 0 && strings.Contains(id, "/") { idx := strings.Index(id, "/") base := strings.TrimSpace(id[idx+1:]) providers = providersForModel(base) } + return providers +} + +// allCodexClientProvidersAre reports whether providers is non-empty and every +// entry is the named provider. +func allCodexClientProvidersAre(providers []string, name string) bool { if len(providers) == 0 { return false } for _, provider := range providers { - if !strings.EqualFold(strings.TrimSpace(provider), "codex") { + if !strings.EqualFold(strings.TrimSpace(provider), name) { return false } } return true } -func applyCodexClientSearchToolSupport(entry map[string]any, id string, templateModel bool, providersForModel ProvidersForModelFunc) { +// applyCodexClientSearchToolSupport keeps supports_search_tool only when every +// route serving the model can close Codex's client tool_search loop: pure Codex +// template models, or pure Claude models whose every credential supports +// Anthropic native tool search. Mixed and unknown routes fail closed. +func applyCodexClientSearchToolSupport(entry map[string]any, id string, templateModel bool, providersForModel ProvidersForModelFunc, toolSearchCapabilityForModel ToolSearchCapabilityForModelFunc) { supportsSearch, _ := entry["supports_search_tool"].(bool) if !supportsSearch { return } - if !templateModel { - entry["supports_search_tool"] = false - return - } - if providersForModel == nil { + if !templateModel { + entry["supports_search_tool"] = false + } return } - providers := providersForModel(id) - if len(providers) == 0 && strings.Contains(id, "/") { - idx := strings.Index(id, "/") - base := strings.TrimSpace(id[idx+1:]) - providers = providersForModel(base) - } - if len(providers) == 0 { - entry["supports_search_tool"] = false + providers := codexClientModelProviders(id, providersForModel) + if allCodexClientProvidersAre(providers, "claude") && toolSearchCapabilityForModel != nil && toolSearchCapabilityForModel(id) { return } - for _, provider := range providers { - if !strings.EqualFold(strings.TrimSpace(provider), "codex") { - entry["supports_search_tool"] = false - return - } + if !templateModel || !allCodexClientProvidersAre(providers, "codex") { + entry["supports_search_tool"] = false } } @@ -730,13 +747,15 @@ func applyCodexClientVisibilityOverride(entry map[string]any, id string) { } } +// isCodexClientImageOrVideoModel reports catalog IDs that are not chat models. +// Speech IDs are included so the Codex client does not offer them as text models. func isCodexClientImageOrVideoModel(id string) bool { target := strings.TrimSpace(id) if idx := strings.Index(target, "/"); idx != -1 { target = strings.TrimSpace(target[idx+1:]) } switch target { - case "grok-imagine-image-quality", "gpt-image-1.5", "gpt-image-2", "gpt-image-2.5-flare", "gpt-image-2.5-sunburst", "gpt-image-2.5", "grok-imagine-image", "grok-imagine-image-2.0", "grok-imagine-video", "grok-imagine-video-1.5", "grok-imagine-video-1.5-preview": + case "grok-imagine-image-quality", "gpt-image-1.5", "gpt-image-2", "gpt-image-2.5-flare", "gpt-image-2.5-sunburst", "gpt-image-2.5", "grok-imagine-image", "grok-imagine-image-2.0", "grok-imagine-video", "grok-imagine-video-1.5", "grok-imagine-video-1.5-preview", "grok-tts", "grok-voice-tts-1.0": return true default: return false diff --git a/internal/client/codex/optimize-multi-agent-v2/optimize_multi_agent_v2.go b/internal/client/codex/optimize-multi-agent-v2/optimize_multi_agent_v2.go index 86bcc30fb..2fda7f71d 100644 --- a/internal/client/codex/optimize-multi-agent-v2/optimize_multi_agent_v2.go +++ b/internal/client/codex/optimize-multi-agent-v2/optimize_multi_agent_v2.go @@ -270,7 +270,7 @@ func loadCodexCatalogTemplates() (map[string]map[string]any, map[string]any, uin continue } templates[modelID] = model - if modelID == "gpt-5.5" { + if defaultTemplate == nil || modelID == "gpt-5.5" { defaultTemplate = model } } @@ -417,7 +417,7 @@ func codexSpawnAgentModelsFromSources(availableModels []map[string]any, catalogJ continue } templates[modelID] = model - if modelID == "gpt-5.5" { + if defaultTemplate == nil || modelID == "gpt-5.5" { defaultTemplate = model } } diff --git a/internal/clienterror/client_error.go b/internal/clienterror/client_error.go index e5af82b76..2c78b3a72 100644 --- a/internal/clienterror/client_error.go +++ b/internal/clienterror/client_error.go @@ -54,6 +54,10 @@ func HTTPStatusFromError(err error) int { return code } } + var maxBytesErr *http.MaxBytesError + if errors.As(err, &maxBytesErr) { + return http.StatusRequestEntityTooLarge + } if errors.Is(err, context.Canceled) { return StatusClientClosedRequest } @@ -96,6 +100,11 @@ func IsRequestFault(status int, err error) bool { if status == http.StatusUnauthorized && hasAuthenticationErrorBody(err) { return false } + // Claude's missing thread state is tied to the stale continuation ID rather + // than the selected credential. Let the client replay the full conversation. + if IsClaudeThreadNotFound(status, err) { + return true + } // Model not found indicates a credential-model capability mismatch rather than // a caller request error. Preserve rotation and cooldown for the model. if hasModelNotFoundErrorBody(err) { @@ -133,6 +142,33 @@ func IsItemNotPersisted(message string) bool { strings.Contains(lower, "items are not persisted when `store` is set to false") } +// IsClaudeThreadNotFound reports the Claude 404 response for a stale +// previous_message_id continuation from either a structured JSON response +// or a plain-text error message. +func IsClaudeThreadNotFound(status int, err error) bool { + if status != http.StatusNotFound || err == nil { + return false + } + body := strings.TrimSpace(err.Error()) + type responseBodyProvider interface { + ResponseBody() []byte + } + var responseBody responseBodyProvider + if errors.As(err, &responseBody) && responseBody != nil && len(responseBody.ResponseBody()) > 0 { + body = strings.TrimSpace(string(responseBody.ResponseBody())) + } + if body == "" { + return false + } + if json.Valid([]byte(body)) { + return strings.EqualFold(strings.TrimSpace(gjson.Get(body, "error.type").String()), "not_found_error") && + strings.Contains(strings.ToLower(gjson.Get(body, "error.message").String()), "thread state") && + strings.Contains(strings.ToLower(gjson.Get(body, "error.message").String()), "previous_message_id") + } + lower := strings.ToLower(body) + return strings.Contains(lower, "thread state") && strings.Contains(lower, "previous_message_id") +} + func hasModelNotFoundErrorBody(err error) bool { if err == nil { return false diff --git a/internal/discovery/browse_queries.go b/internal/discovery/browse_queries.go new file mode 100644 index 000000000..b36128efd --- /dev/null +++ b/internal/discovery/browse_queries.go @@ -0,0 +1,75 @@ +package discovery + +import ( + "time" + + "github.com/miekg/dns" +) + +// browseQueryRetry keeps discovery backoff independent from cached PTR lifetime. +type browseQueryRetry struct { + next, last time.Time + interval time.Duration +} + +func (r *browseQueryRetry) sent(now time.Time) { + if r.interval == 0 { + r.interval = time.Second + } else if r.interval < time.Hour { + r.interval *= 2 + if r.interval > time.Hour { + r.interval = time.Hour + } + } + r.last = now + r.next = now.Add(r.interval) +} + +func (c *browseCache) queryDue(index int, q dns.Question, state browseQueryRetry, now time.Time) bool { + if !now.Before(state.next) { + return true + } + if q.Qtype != dns.TypePTR { + return false + } + // RFC 6762 section 5.2: active browse PTRs must be refreshed even after the + // discovery backoff has grown longer than their TTL. Retry about every 5% + // of TTL from 80% onward; never generate sub-second retries. + for _, r := range c.rrset(index, q.Name, dns.TypePTR) { + ttl := time.Duration(r.rr.Header().Ttl) * time.Second + retry := max(time.Second, ttl/20) + if now.Sub(r.received) >= ttl*4/5 && now.Sub(state.last) >= retry { + return true + } + } + return false +} + +// queryMessages implements known-answer suppression (RFC 6762 sections 7.1 and +// 7.2). Oversized answer lists continue in TC-marked query packets rather than +// forcing IP fragmentation or silently discarding the remaining known answers. +func (c *browseCache) queryMessages(index int, q dns.Question, now time.Time) []*dns.Msg { + const maxPayload = 1200 + current := &dns.Msg{Question: []dns.Question{q}} + var messages []*dns.Msg + for _, r := range c.rrset(index, q.Name, q.Qtype) { + remaining := r.expires.Sub(now) + original := time.Duration(r.rr.Header().Ttl) * time.Second + if remaining <= 0 || remaining < original/2 { + continue + } + rr := dns.Copy(r.rr) + rr.Header().Class &^= 0x8000 + rr.Header().Ttl = uint32(remaining / time.Second) + current.Answer = append(current.Answer, rr) + wire, err := current.Pack() + if err == nil && len(wire) <= maxPayload { + continue + } + current.Answer = current.Answer[:len(current.Answer)-1] + current.Truncated = true + messages = append(messages, current) + current = &dns.Msg{Answer: []dns.RR{rr}} + } + return append(messages, current) +} diff --git a/internal/discovery/browse_resolver.go b/internal/discovery/browse_resolver.go new file mode 100644 index 000000000..4f19cf539 --- /dev/null +++ b/internal/discovery/browse_resolver.go @@ -0,0 +1,663 @@ +package discovery + +import ( + "context" + "fmt" + "net" + "sort" + "strings" + "sync" + "time" + + "github.com/libp2p/zeroconf/v2" + "github.com/miekg/dns" + log "github.com/sirupsen/logrus" + "golang.org/x/net/ipv4" + "golang.org/x/net/ipv6" +) + +const ( + browseCacheRRLimit = maxDiscoveredServices * (maxDiscoveredAddresses + 4) + browseCacheByteLimit = 4 << 20 + browseIPLimit = maxDiscoveredAddresses +) + +type browseLink struct { + iface net.Interface + networks []*net.IPNet + local []net.IP +} +type browseRR struct { + rr dns.RR + expires, received time.Time + size int + serial uint64 +} +type browseRRKey struct { + link int + name string + kind uint16 + data string +} +type browseCache struct { + browse, base, service, domain string + links map[int]browseLink + records map[browseRRKey]browseRR + bytes int + serial uint64 +} + +func browseName(s string) string { return strings.ToLower(dns.Fqdn(s)) } + +func newBrowseCache(ifaces []net.Interface, service, domain string) *browseCache { + if service == "" { + service = DefaultServiceType + } + if domain == "" { + domain = DefaultDomain + } + domain = dns.Fqdn(domain) + service = strings.TrimSuffix(service, ".") + // Accept both a service type and its fully qualified browse name. + service = strings.TrimSuffix(service, "."+strings.TrimSuffix(domain, ".")) + base := service + if i := strings.Index(strings.ToLower(base), "._sub."); i >= 0 { + base = base[i+6:] + } + c := &browseCache{browse: browseName(service + "." + domain), base: browseName(base + "." + domain), service: base, domain: domain, links: make(map[int]browseLink), records: make(map[browseRRKey]browseRR)} + for _, iface := range ifaces { + l := browseLink{iface: iface} + addrs, _ := iface.Addrs() + for _, a := range addrs { + if n, ok := a.(*net.IPNet); ok { + l.networks = append(l.networks, n) + l.local = append(l.local, n.IP) + } + } + if iface.Index > 0 { + c.links[iface.Index] = l + } + } + return c +} + +// RFC 6762 section 11 permits off-subnet sources for multicast responses. +// Unicast responses must originate on-link and be addressed to this interface. +func (c *browseCache) validPeer(index int, src *net.UDPAddr, dst net.IP) bool { + l, ok := c.links[index] + if !ok || src == nil || src.Port != mdnsPort || src.IP == nil || src.IP.IsUnspecified() || src.IP.IsMulticast() || src.IP.IsLoopback() || dst == nil { + return false + } + if (src.IP.To4() != nil) != (dst.To4() != nil) { + return false + } + if dst.IsMulticast() { + return dst.Equal(mdnsGroupAddrIPv4.IP) || dst.Equal(mdnsGroupAddrIPv6.IP) + } + local := false + for _, ip := range l.local { + if ip.Equal(dst) { + local = true + break + } + } + if !local { + return false + } + if src.IP.IsLinkLocalUnicast() { + return true + } + for _, n := range l.networks { + if n.Contains(src.IP) { + return true + } + } + return false +} +func (c *browseCache) prune(now time.Time) { + for k, r := range c.records { + if !now.Before(r.expires) { + delete(c.records, k) + c.bytes -= r.size + } + } +} +func (c *browseCache) instance(name string) bool { + labels := dns.SplitDomainName(name) + base := dns.SplitDomainName(c.base) + return len(labels) == len(base)+1 && strings.EqualFold(strings.Join(labels[1:], "."), strings.Join(base, ".")) +} +func browseTXTValid(txt []string) bool { + if len(txt) > maxBrowseTXTRecords { + return false + } + size := 0 + for _, s := range txt { + if len(s) > 255 { + return false + } + size += 1 + len(s) + } + limit := 16 * 1024 + if _, ok := ParseTXTRecords(txt)["txtvers"]; ok { + limit = 400 + } + return size <= limit +} +func (c *browseCache) admissible(rr dns.RR) bool { + h := rr.Header() + name := browseName(h.Name) + if h.Class&0x7fff != dns.ClassINET { + return false + } + switch r := rr.(type) { + case *dns.PTR: + return name == c.browse && c.instance(r.Ptr) + case *dns.SRV: + return c.instance(name) && (r.Target == "." || validBrowseHost(r.Target)) + case *dns.TXT: + return c.instance(name) && browseTXTValid(r.Txt) + case *dns.A: + return validBrowseHost(name) && r.A.To4() != nil + case *dns.AAAA: + return validBrowseHost(name) && r.AAAA.To16() != nil && r.AAAA.To4() == nil + } + return false +} + +// SRV targets are hostnames, not URLs. They need not share the service's +// browsing domain; the PTR owner still confines service membership. +func validBrowseHost(host string) bool { + host = strings.TrimSuffix(host, ".") + if host == "" || len(host) > 253 { + return false + } + for _, label := range strings.Split(host, ".") { + if label == "" || len(label) > 63 || label[0] == '-' || label[len(label)-1] == '-' { + return false + } + for _, ch := range label { + if !((ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') || (ch >= '0' && ch <= '9') || ch == '-') { + return false + } + } + } + return true +} + +func browseRecordKey(index int, rr dns.RR) browseRRKey { + copyRR := dns.Copy(rr) + h := copyRR.Header() + h.Name = browseName(h.Name) + h.Ttl = 0 + h.Class = dns.ClassINET + switch r := copyRR.(type) { + case *dns.PTR: + r.Ptr = browseName(r.Ptr) + case *dns.SRV: + r.Target = browseName(r.Target) + } + return browseRRKey{index, h.Name, h.Rrtype, copyRR.String()} +} + +// Only answers and additional records in responses enter the cache. Query known +// answers and probe authorities must never manufacture a discovered service. +func (c *browseCache) ingest(msg *dns.Msg, index int, src *net.UDPAddr, dst net.IP, now time.Time) { + c.prune(now) + if msg == nil || !msg.Response || msg.Opcode != dns.OpcodeQuery || msg.Rcode != dns.RcodeSuccess || msg.Id != 0 || !c.validPeer(index, src, dst) { + return + } + records := append(append([]dns.RR(nil), msg.Answer...), msg.Extra...) + incoming := make(map[browseRRKey]bool, len(records)) + for _, rr := range records { + if c.admissible(rr) && rr.Header().Ttl > 0 { + incoming[browseRecordKey(index, rr)] = true + } + } + flushVictims := make(map[browseRRKey]bool) + evictFlushed := func(matches func(browseRRKey) bool) bool { + var oldest browseRRKey + sequence := ^uint64(0) + for key := range flushVictims { + if record, ok := c.records[key]; ok && matches(key) && record.serial < sequence { + oldest, sequence = key, record.serial + } + } + if sequence == ^uint64(0) { + return false + } + c.bytes -= c.records[oldest].size + delete(c.records, oldest) + delete(flushVictims, oldest) + return true + } + // Apply flush before inserting anything, preserving all members of this packet + // and the one-second immunity for RRsets split across consecutive packets. + for _, rr := range records { + if !c.admissible(rr) || rr.Header().Ttl == 0 || rr.Header().Class&0x8000 == 0 || rr.Header().Rrtype == dns.TypePTR { + continue + } + k := browseRecordKey(index, rr) + for old, r := range c.records { + if old.link == index && old.name == k.name && old.kind == k.kind && now.Sub(r.received) >= time.Second && !incoming[old] { + flushVictims[old] = true + if r.expires.After(now.Add(time.Second)) { + r.expires = now.Add(time.Second) + c.records[old] = r + } + } + } + } + for _, rr := range records { + if !c.admissible(rr) { + continue + } + k := browseRecordKey(index, rr) + h := rr.Header() + if h.Ttl == 0 { + if r, ok := c.records[k]; ok && r.expires.After(now.Add(time.Second)) { + r.expires = now.Add(time.Second) + c.records[k] = r + } + continue + } + size := len(k.data) + len(k.name) + 64 + old, exists := c.records[k] + if !exists { + for len(c.records) >= browseCacheRRLimit || c.bytes+size > browseCacheByteLimit { + if !evictFlushed(func(browseRRKey) bool { return true }) { + break + } + } + if len(c.records) >= browseCacheRRLimit || c.bytes+size > browseCacheByteLimit { + continue + } + count := 0 + instances := make(map[string]struct{}) + for key, r := range c.records { + if key.link == index && key.name == k.name && (key.kind == k.kind || ((k.kind == dns.TypeA || k.kind == dns.TypeAAAA) && (key.kind == dns.TypeA || key.kind == dns.TypeAAAA))) { + count++ + } + if c.instance(key.name) { + instances[fmt.Sprintf("%d/%s", key.link, key.name)] = struct{}{} + } + if p, ok := r.rr.(*dns.PTR); ok { + instances[fmt.Sprintf("%d/%s", key.link, browseName(p.Ptr))] = struct{}{} + } + } + instance := "" + if c.instance(k.name) { + instance = k.name + } + if p, ok := rr.(*dns.PTR); ok { + instance = browseName(p.Ptr) + } + if instance != "" { + if _, ok := instances[fmt.Sprintf("%d/%s", index, instance)]; !ok && len(instances) >= maxDiscoveredServices { + continue + } + } + limit := browseIPLimit + if k.kind == dns.TypePTR { + limit = maxDiscoveredServices + } + if k.kind == dns.TypeSRV || k.kind == dns.TypeTXT { + limit = 2 + } + if k.kind == dns.TypeA || k.kind == dns.TypeAAAA { + for count >= limit && evictFlushed(func(key browseRRKey) bool { + return key.link == k.link && key.name == k.name && (key.kind == dns.TypeA || key.kind == dns.TypeAAAA) + }) { + count-- + } + } + if count >= limit { + if k.kind != dns.TypeSRV && k.kind != dns.TypeTXT { + continue + } + // Unique data changes must not be blocked by the bounded grace + // history when a peer publishes several rapid updates. + var oldest browseRRKey + sequence := ^uint64(0) + for key, value := range c.records { + if key.link == index && key.name == k.name && key.kind == k.kind && value.serial < sequence { + oldest = key + sequence = value.serial + } + } + c.bytes -= c.records[oldest].size + delete(c.records, oldest) + } + } + c.serial++ + c.records[k] = browseRR{dns.Copy(rr), now.Add(time.Duration(h.Ttl) * time.Second), now, size, c.serial} + c.bytes += size - old.size + } +} + +func (c *browseCache) rrset(index int, name string, kind uint16) []browseRR { + var out []browseRR + for k, r := range c.records { + if k.link == index && k.name == browseName(name) && k.kind == kind { + out = append(out, r) + } + } + sort.Slice(out, func(i, j int) bool { return out[i].serial > out[j].serial }) + return out +} +func (c *browseCache) results(now time.Time) []DiscoveredService { + c.prune(now) + var out []DiscoveredService + for k, r := range c.records { + ptr, ok := r.rr.(*dns.PTR) + if !ok { + continue + } + srvs := c.rrset(k.link, ptr.Ptr, dns.TypeSRV) + txts := c.rrset(k.link, ptr.Ptr, dns.TypeTXT) + if len(srvs) == 0 || len(txts) == 0 { + continue + } + srv := srvs[0].rr.(*dns.SRV) + if srv.Target == "." || srv.Port == 0 { + continue + } + labels := dns.SplitDomainName(ptr.Ptr) + entry := &zeroconf.ServiceEntry{ServiceRecord: zeroconf.ServiceRecord{Instance: labels[0], Service: c.service, Domain: c.domain}, HostName: srv.Target, Port: int(srv.Port), Text: txts[0].rr.(*dns.TXT).Txt} + for _, kind := range []uint16{dns.TypeA, dns.TypeAAAA} { + for _, addr := range c.rrset(k.link, srv.Target, kind) { + var ip net.IP + switch v := addr.rr.(type) { + case *dns.A: + ip = v.A + case *dns.AAAA: + ip = v.AAAA + } + if ip.IsMulticast() || ip.IsUnspecified() || ip.IsLoopback() { + continue + } + if kind == dns.TypeA { + entry.AddrIPv4 = append(entry.AddrIPv4, ip) + } else { + entry.AddrIPv6 = append(entry.AddrIPv6, ip) + } + } + } + svc, valid := entryToDiscovered(entry) + if !valid || len(svc.IPv4)+len(svc.IPv6) == 0 { + continue + } + svc.InterfaceIndex = k.link + svc.IPv6Zones = make(map[string]string) + for _, ip := range svc.IPv6 { + if ip.IsLinkLocalUnicast() { + svc.IPv6Zones[ip.String()] = c.links[k.link].iface.Name + } + } + out = append(out, svc) + } + sort.Slice(out, func(i, j int) bool { + if out[i].InstanceName == out[j].InstanceName { + return out[i].InterfaceIndex < out[j].InterfaceIndex + } + return out[i].InstanceName < out[j].InstanceName + }) + return out +} + +// queries includes missing dependencies and refreshes records after 80% of TTL. +func (c *browseCache) queries(index int, now time.Time) []dns.Question { + c.prune(now) + seen := make(map[string]bool) + var out []dns.Question + add := func(name string, kind uint16) { + key := fmt.Sprintf("%s/%d", browseName(name), kind) + if !seen[key] { + seen[key] = true + out = append(out, dns.Question{Name: name, Qtype: kind, Qclass: dns.ClassINET}) + } + } + add(c.browse, dns.TypePTR) + need := func(name string, kind uint16) { + set := c.rrset(index, name, kind) + if len(set) == 0 { + add(name, kind) + return + } + for _, r := range set { + if now.Sub(r.received) >= (r.expires.Sub(r.received) * 4 / 5) { + add(name, kind) + break + } + } + } + for k, r := range c.records { + if k.link != index { + continue + } + if ptr, ok := r.rr.(*dns.PTR); ok { + need(ptr.Ptr, dns.TypeSRV) + need(ptr.Ptr, dns.TypeTXT) + for _, s := range c.rrset(index, ptr.Ptr, dns.TypeSRV) { + target := s.rr.(*dns.SRV).Target + if target != "." { + need(target, dns.TypeA) + need(target, dns.TypeAAAA) + } + } + } + } + return out +} + +type browsePacket struct { + data []byte + index int + src *net.UDPAddr + dst net.IP +} + +// browseLAN has no private deadline. Closing sockets interrupts blocked readers; +// the sole timer schedules queries, never ends the caller's browse window. +func browseLAN(ctx context.Context, ifaces []net.Interface, serviceType, domain string) ([]DiscoveredService, error) { + if ctx == nil { + ctx = context.TODO() + } + if ctx.Err() != nil { + return nil, nil + } + if len(ifaces) == 0 { + return nil, fmt.Errorf("discovery: no allowed browse interfaces") + } + cache := newBrowseCache(ifaces, serviceType, domain) + conn4, err4 := openMDNSConn4(ifaces) + conn6, err6 := openMDNSConn6(ifaces) + // Interface and destination metadata are mandatory for receive isolation. + if conn4 != nil { + if err := conn4.SetControlMessage(ipv4.FlagInterface|ipv4.FlagDst, true); err != nil { + err4 = err + if errClose := conn4.Close(); errClose != nil { + log.WithError(errClose).Debug("discovery: close unusable IPv4 browse socket") + } + conn4 = nil + } + } + if conn6 != nil { + if err := conn6.SetControlMessage(ipv6.FlagInterface|ipv6.FlagDst, true); err != nil { + err6 = err + if errClose := conn6.Close(); errClose != nil { + log.WithError(errClose).Debug("discovery: close unusable IPv6 browse socket") + } + conn6 = nil + } + } + if conn4 == nil && conn6 == nil { + return nil, fmt.Errorf("discovery: open mDNS sockets: IPv4: %v; IPv6: %v", err4, err6) + } + packets := make(chan browsePacket, 32) + readerErrors := make(chan error, 2) + readers := 0 + stop := make(chan struct{}) + var wg sync.WaitGroup + var closeOnce sync.Once + closeSockets := func() { + closeOnce.Do(func() { + if conn4 != nil { + if errClose := conn4.Close(); errClose != nil { + log.WithError(errClose).Debug("discovery: close IPv4 browse socket") + } + } + if conn6 != nil { + if errClose := conn6.Close(); errClose != nil { + log.WithError(errClose).Debug("discovery: close IPv6 browse socket") + } + } + }) + } + // Cancellation interrupts reads and writes without installing socket deadlines. + stopClose := context.AfterFunc(ctx, closeSockets) + defer func() { + stopClose() + close(stop) + closeSockets() + wg.Wait() + }() + for _, v6 := range []bool{false, true} { + if (!v6 && conn4 == nil) || (v6 && conn6 == nil) { + continue + } + wg.Add(1) + readers++ + go func(v6 bool) { + defer wg.Done() + buf := make([]byte, 65536) + for { + var p browsePacket + var from net.Addr + if v6 { + n, cm, src, e := conn6.ReadFrom(buf) + if e != nil { + readerErrors <- e + return + } + if cm != nil { + p.index = cm.IfIndex + p.dst = cm.Dst + } + from = src + p.data = append([]byte(nil), buf[:n]...) + } else { + n, cm, src, e := conn4.ReadFrom(buf) + if e != nil { + readerErrors <- e + return + } + if cm != nil { + p.index = cm.IfIndex + p.dst = cm.Dst + } + from = src + p.data = append([]byte(nil), buf[:n]...) + } + p.src, _ = from.(*net.UDPAddr) + select { + case packets <- p: + case <-stop: + return + case <-ctx.Done(): + return + } + } + }(v6) + } + pending := make(map[string]browseQueryRetry) + send := func(now time.Time) error { + active := make(map[string]bool) + attempted, sent := 0, 0 + var lastSendErr error + for _, iface := range ifaces { + qs := cache.queries(iface.Index, now) + for _, q := range qs { + key := fmt.Sprintf("%d/%s/%d", iface.Index, browseName(q.Name), q.Qtype) + active[key] = true + state := pending[key] + if !cache.queryDue(iface.Index, q, state, now) { + continue + } + state.sent(now) + pending[key] = state + for _, m := range cache.queryMessages(iface.Index, q, now) { + wire, err := m.Pack() + if err != nil { + continue + } + if conn4 != nil { + attempted++ + errSend := conn4.SetMulticastInterface(&iface) + if errSend == nil { + _, errSend = conn4.WriteTo(wire, nil, mdnsGroupAddrIPv4) + } + if errSend != nil { + lastSendErr = errSend + } else { + sent++ + } + } + if conn6 != nil { + attempted++ + errSend := conn6.SetMulticastInterface(&iface) + if errSend == nil { + _, errSend = conn6.WriteTo(wire, nil, mdnsGroupAddrIPv6) + } + if errSend != nil { + lastSendErr = errSend + } else { + sent++ + } + } + } + } + } + for key := range pending { + if !active[key] { + delete(pending, key) + } + } + if attempted > 0 && sent == 0 && ctx.Err() == nil { + return fmt.Errorf("discovery: send browse queries: %w", lastSendErr) + } + return nil + } + if errSend := send(time.Now()); errSend != nil { + return nil, errSend + } + timer := time.NewTimer(time.Second) + defer timer.Stop() + for { + select { + case <-ctx.Done(): + return cache.results(time.Now()), nil + case errRead := <-readerErrors: + if ctx.Err() != nil { + return cache.results(time.Now()), nil + } + readers-- + if readers == 0 { + return cache.results(time.Now()), fmt.Errorf("discovery: mDNS readers stopped: %w", errRead) + } + case p := <-packets: + msg := new(dns.Msg) + if msg.Unpack(p.data) == nil { + cache.ingest(msg, p.index, p.src, p.dst, time.Now()) + } + if result := cache.results(time.Now()); len(result) >= maxDiscoveredServices { + return result, nil + } + case now := <-timer.C: + if errSend := send(now); errSend != nil { + return cache.results(time.Now()), errSend + } + timer.Reset(time.Second) + } + } +} diff --git a/internal/discovery/service.go b/internal/discovery/service.go index 6b4180318..7f909903e 100644 --- a/internal/discovery/service.go +++ b/internal/discovery/service.go @@ -49,6 +49,9 @@ func validateServiceType(st string) error { if len(name) == 0 || len(name) > 15 { return fmt.Errorf("service name %q must be between 1 and 15 characters (RFC 6335)", name) } + if !strings.ContainsAny(name, "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ") || strings.Contains(name, "--") { + return fmt.Errorf("service name %q must contain a letter and cannot contain consecutive hyphens (RFC 6335)", name) + } for i, r := range name { if !((r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '-') { return fmt.Errorf("service type contains invalid character %q (RFC 6335)", r) @@ -90,6 +93,47 @@ func sanitizeSubtype(sub string) string { return sub } +// normalizeBrowseQuery accepts the RFC 6763 subtype owner form as well as the +// comma form historically accepted by zeroconf. It returns lowercase DNS names. +func normalizeBrowseQuery(service, domain string) (string, string, error) { + domain = strings.ToLower(strings.TrimSuffix(strings.TrimSpace(domain), ".")) + if domain == "" { + domain = strings.TrimSuffix(DefaultDomain, ".") + } + if len(domain) > 253 { + return "", "", fmt.Errorf("discovery: browse domain too long") + } + for _, label := range strings.Split(domain, ".") { + if label == "" || len(label) > 63 || label[0] == '-' || label[len(label)-1] == '-' { + return "", "", fmt.Errorf("discovery: invalid browse domain") + } + for _, ch := range label { + if !((ch >= 'a' && ch <= 'z') || (ch >= '0' && ch <= '9') || ch == '-') { + return "", "", fmt.Errorf("discovery: invalid browse domain") + } + } + } + service = strings.ToLower(strings.TrimSuffix(strings.TrimSpace(service), ".")) + if service == "" { + service = DefaultServiceType + } + service = strings.TrimSuffix(service, "."+domain) + if base, subtype, comma := strings.Cut(service, ","); comma { + service = subtype + "._sub." + base + } + base := service + if subtype, rest, sub := strings.Cut(service, "._sub."); sub { + if !strings.HasPrefix(subtype, "_") || sanitizeSubtype(subtype) != subtype { + return "", "", fmt.Errorf("discovery: invalid browse subtype") + } + base = rest + } + if err := validateServiceType(base); err != nil { + return "", "", fmt.Errorf("discovery: invalid browse service: %w", err) + } + return service, domain + ".", nil +} + func truncateRunesTo(s string, maxBytes int) string { if maxBytes <= 0 { return "" @@ -232,24 +276,28 @@ func BuildServiceSpec(cfg *config.Config, port int, tlsEnabled bool) (ServiceSpe filtered = append(filtered, advertisedIP) } } + if len(filtered) == 0 { + filtered = append(filtered, bindIP.String()) + } advertisedIPs = filtered } if len(advertisedIPs) == 0 { return ServiceSpec{}, fmt.Errorf("discovery: no advertised address matches bind host %q", bindHost) } - // 6. Build TXT records + // 6. Build TXT records; protocols mirror the advertised subtypes. txtOpts := DefaultTXTOptions() txtOpts.InstanceID = instanceID txtOpts.TLS = tlsEnabled - txtOpts.AdvertiseManagement = discCfg.AdvertiseManagement if discCfg.AuthRequired != nil { txtOpts.AuthRequired = *discCfg.AuthRequired - } else { - txtOpts.AuthRequired = true } + txtOpts.Protocols = protocolsFromSubtypes(subtypes) - txtRecords := BuildTXTRecords(txtOpts) + txtRecords, errTXT := BuildTXTRecords(txtOpts) + if errTXT != nil { + return ServiceSpec{}, errTXT + } return ServiceSpec{ InstanceName: instanceName, @@ -262,3 +310,13 @@ func BuildServiceSpec(cfg *config.Config, port int, tlsEnabled bool) (ServiceSpe AdvertisedIPs: advertisedIPs, }, nil } + +// protocolsFromSubtypes derives TXT protocol tokens from advertised subtype labels +// so the protocols attribute and DNS-SD subtypes stay consistent. +func protocolsFromSubtypes(subtypes []string) []string { + protocols := make([]string, 0, len(subtypes)) + for _, subtype := range subtypes { + protocols = append(protocols, ProtocolFromSubtype(subtype)) + } + return protocols +} diff --git a/internal/discovery/spec.go b/internal/discovery/spec.go index 532fe8582..8d4e86c6b 100644 --- a/internal/discovery/spec.go +++ b/internal/discovery/spec.go @@ -6,161 +6,399 @@ import ( ) const ( - maxTXTRecordBytes = 255 // RFC 6763 §6.3 single string limit - maxTXTBytes = 400 // Conservative LAN UDP packet safety boundary + maxTXTRecordBytes = 255 // RFC 6763 §6.1 single character-string limit + maxTXTBytes = 400 // draft-su-ai-gateway-discovery-00 TXT budget + + txtSchemaVersion = "1" + defaultMountPath = "/" +) + +// TXT keys defined by draft-su-ai-gateway-discovery-00. +// Every key is at most nine characters long (RFC 6763 §6.4). +const ( + txtKeyVersion = "txtvers" + txtKeyPath = "path" + txtKeyTLS = "tls" + txtKeyAuth = "auth" + txtKeyProtocols = "protocols" + txtKeyAuthType = "auth_type" + txtKeyProduct = "product" + txtKeyID = "id" + txtKeyRole = "role" ) -func sanitizeTXTValue(v string) string { - var b strings.Builder - for _, r := range v { - if r >= 32 && r < 127 { - b.WriteRune(r) +// Legacy TXT keys advertised by CPA releases before the draft schema. +// Browsers still accept them so older gateways remain discoverable. +const ( + legacyKeyVersion = "version" + legacyKeyAuth = "auth_required" + legacyKeyAuthMethods = "auth_methods" + legacyKeyInstanceID = "instance_id" + legacyKeyNodeRole = "node_role" + legacyKeyOpenAI = "api_openai" + legacyKeyAnthropic = "api_anthropic" + legacyKeyGemini = "api_gemini" +) + +// protocolRoute maps a protocol token to its standard route below the gateway mount path. +type protocolRoute struct { + version string // Versioned prefix below path, e.g. /v1 or /v1beta + suffix string // Operation route below the versioned prefix + legacyField string // Legacy api_* key that carried the versioned base path +} + +var protocolRoutes = map[string]protocolRoute{ + "chat-completions": {version: "/v1", suffix: "/chat/completions", legacyField: legacyKeyOpenAI}, + "completions": {version: "/v1", suffix: "/completions", legacyField: legacyKeyOpenAI}, + "responses": {version: "/v1", suffix: "/responses", legacyField: legacyKeyOpenAI}, + "responses-ws": {version: "/v1", suffix: "/responses", legacyField: legacyKeyOpenAI}, + "realtime": {version: "/v1", suffix: "/realtime", legacyField: legacyKeyOpenAI}, + "images": {version: "/v1", suffix: "/images/generations", legacyField: legacyKeyOpenAI}, + "videos": {version: "/v1", suffix: "/videos", legacyField: legacyKeyOpenAI}, + "audio": {version: "/v1", suffix: "/audio/speech", legacyField: legacyKeyOpenAI}, + "search": {version: "/v1", suffix: "/alpha/search", legacyField: legacyKeyOpenAI}, + "messages": {version: "/v1", suffix: "/messages", legacyField: legacyKeyAnthropic}, + "count-tokens": {version: "/v1", suffix: "/messages/count_tokens", legacyField: legacyKeyAnthropic}, + "generate-content": {version: "/v1beta", suffix: "/models/{model}:generateContent", legacyField: legacyKeyGemini}, + "interactions": {version: "/v1beta", suffix: "/interactions", legacyField: legacyKeyGemini}, +} + +// printableTXTASCII checks bytes without Unicode folding or whitespace normalization. +func printableTXTASCII(value string) bool { + for i := 0; i < len(value); i++ { + if value[i] < 0x20 || value[i] > 0x7e { + return false } } - return b.String() + return true } -// TXTOptions contains input fields for generating safe DNS-SD TXT records. +// TXTOptions contains input fields for generating DNS-SD TXT records. type TXTOptions struct { - Version string - Product string - Protocols []string - Features []string - APIPathOpenAI string - APIPathAnthropic string - APIPathGemini string - TLS bool - AuthRequired bool - AuthMethods []string - InstanceID string - NodeRole string - AdvertiseManagement bool + Path string + Product string + Protocols []string + TLS bool + AuthRequired bool + AuthMethods []string + InstanceID string + NodeRole string } // DefaultTXTOptions provides safe defaults for CPA TXT records. func DefaultTXTOptions() TXTOptions { return TXTOptions{ - Version: "1", - Product: ProductCPA, - Protocols: []string{"chat-completions", "responses", "messages", "generate-content", "interactions"}, - Features: []string{"chat", "responses", "messages", "generate_content", "interactions"}, - APIPathOpenAI: "/v1", - APIPathAnthropic: "/v1", - APIPathGemini: "/v1beta", - TLS: false, - AuthRequired: true, - AuthMethods: []string{"api_key"}, - NodeRole: "standalone", + Path: defaultMountPath, + Product: ProductCPA, + Protocols: []string{"chat-completions", "responses", "messages", "generate-content", "interactions"}, + TLS: false, + AuthRequired: true, + AuthMethods: []string{"api_key"}, + NodeRole: "standalone", } } -// BuildTXTRecords constructs key=value TXT records strictly within 400 bytes total -// and 255 bytes per record (RFC 6763 §6.3), prioritizing essential routing metadata. -func BuildTXTRecords(opts TXTOptions) []string { - if opts.Version == "" { - opts.Version = "1" +// BuildTXTRecords constructs the draft-su-ai-gateway-discovery-00 TXT record. +// Required attributes are always emitted first; optional attributes are skipped +// when they would exceed the 255-byte string or 400-byte record budget. +func BuildTXTRecords(opts TXTOptions) ([]string, error) { + path := opts.Path + if path == "" { + path = defaultMountPath } - if opts.Product == "" { - opts.Product = ProductCPA + if validateMountPath(path) == "" { + return nil, fmt.Errorf("discovery: invalid TXT path %q", path) } - if opts.APIPathOpenAI == "" { - opts.APIPathOpenAI = "/v1" + + records := []string{ + txtKeyVersion + "=" + txtSchemaVersion, + txtKeyPath + "=" + path, + txtKeyTLS + "=" + txtFlag(opts.TLS), + txtKeyAuth + "=" + txtFlag(opts.AuthRequired), } - if opts.APIPathAnthropic == "" { - opts.APIPathAnthropic = "/v1" + totalBytes := 0 + for _, record := range records { + if len(record) > maxTXTRecordBytes { + return nil, fmt.Errorf("discovery: TXT record %q exceeds %d bytes", record, maxTXTRecordBytes) + } + totalBytes += len(record) + 1 // +1 for the DNS character-string length prefix } - if opts.APIPathGemini == "" { - opts.APIPathGemini = "/v1beta" + if totalBytes > maxTXTBytes { + return nil, fmt.Errorf("discovery: required TXT records exceed %d bytes", maxTXTBytes) } - type kv struct { - key string - val string - priority int // 1: core, 2: standard, 3: optional + for _, tokens := range [][]string{opts.Protocols, opts.AuthMethods} { + for _, token := range tokens { + if !validTXTToken(token) { + return nil, fmt.Errorf("discovery: invalid TXT token %q", token) + } + } } - - var candidates []kv - addCandidate := func(key, val string, priority int) { - val = sanitizeTXTValue(strings.TrimSpace(val)) - if val != "" { - candidates = append(candidates, kv{key: key, val: val, priority: priority}) + for _, value := range []string{opts.Product, opts.InstanceID, opts.NodeRole} { + if !printableTXTASCII(value) { + return nil, fmt.Errorf("discovery: optional TXT value must be printable ASCII") } } - // Priority 1: Core metadata - addCandidate("version", opts.Version, 1) - addCandidate("product", opts.Product, 1) - addCandidate("instance_id", opts.InstanceID, 1) - if opts.TLS { - addCandidate("tls", "1", 1) - } else { - addCandidate("tls", "0", 1) + optional := []struct { + key string + value string + }{ + {txtKeyProtocols, joinTXTTokens(opts.Protocols)}, + {txtKeyAuthType, joinTXTTokens(opts.AuthMethods)}, + {txtKeyProduct, opts.Product}, + {txtKeyID, opts.InstanceID}, + {txtKeyRole, opts.NodeRole}, } - if opts.AuthRequired { - addCandidate("auth_required", "true", 1) - } else { - addCandidate("auth_required", "false", 1) + for _, item := range optional { + if item.value == "" { + continue + } + record := item.key + "=" + item.value + cost := len(record) + 1 + if len(record) > maxTXTRecordBytes || totalBytes+cost > maxTXTBytes { + continue + } + records = append(records, record) + totalBytes += cost } - addCandidate("api_openai", opts.APIPathOpenAI, 1) + return records, nil +} - // Priority 2: Standard endpoints and protocols - addCandidate("api_anthropic", opts.APIPathAnthropic, 2) - addCandidate("api_gemini", opts.APIPathGemini, 2) - if opts.AdvertiseManagement { - addCandidate("management", "true", 2) - } else { - addCandidate("management", "false", 2) - } - if len(opts.AuthMethods) > 0 { - addCandidate("auth_methods", strings.Join(opts.AuthMethods, ","), 2) - } - if len(opts.Protocols) > 0 { - addCandidate("protocols", strings.Join(opts.Protocols, ","), 2) +func txtFlag(enabled bool) string { + if enabled { + return "1" } + return "0" +} - // Priority 3: Optional discovery hints - if opts.NodeRole != "" { - addCandidate("node_role", opts.NodeRole, 3) +// parseTXTFlag accepts the draft boolean encoding ("0" or "1") only. +func parseTXTFlag(value string) (bool, bool) { + switch value { + case "1": + return true, true + case "0": + return false, true + default: + return false, false } - if len(opts.Features) > 0 { - addCandidate("features", strings.Join(opts.Features, ","), 3) - } - - // Assemble records respecting single-string (255) and total length (400) limits - var records []string - totalBytes := 0 +} - for _, item := range candidates { - entry := fmt.Sprintf("%s=%s", item.key, item.val) - if len(entry) > maxTXTRecordBytes { +func joinTXTTokens(values []string) string { + tokens := make([]string, 0, len(values)) + seen := make(map[string]struct{}, len(values)) + for _, value := range values { + token := value + if !validTXTToken(token) { continue } - entryCost := len(entry) + 1 // +1 for DNS TXT length prefix byte - if totalBytes+entryCost > maxTXTBytes { + if _, ok := seen[token]; ok { continue } - - records = append(records, entry) - totalBytes += entryCost + seen[token] = struct{}{} + tokens = append(tokens, token) } + return strings.Join(tokens, ",") +} - return records +func validTXTToken(token string) bool { + if token == "" { + return false + } + for i := 0; i < len(token); i++ { + c := token[i] + if !((c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || c == '-' || c == '_') { + return false + } + } + return true } -// ParseTXTRecords parses a slice of "key=value" strings into a key-value map. +// ParseTXTRecords parses "key=value" strings into a map with lowercase keys. +// Per RFC 6763 §6.4, only the first occurrence of a key is used. func ParseTXTRecords(txt []string) map[string]string { result := make(map[string]string, len(txt)) for _, entry := range txt { - parts := strings.SplitN(entry, "=", 2) - key := strings.ToLower(strings.TrimSpace(parts[0])) - if key == "" { + key, value, _ := strings.Cut(entry, "=") + if key == "" || !printableTXTASCII(key) { continue } - if len(parts) == 2 { - result[key] = parts[1] - } else if len(parts) == 1 { - result[key] = "" + key = strings.ToLower(key) + if _, exists := result[key]; exists { + continue } + result[key] = value } return result } + +// validateMountPath returns p when it is a safe absolute path for the TXT path attribute. +// It rejects authorities, queries, fragments, escapes, traversal segments, and trailing slashes. +func validateMountPath(p string) string { + if p == defaultMountPath { + return p + } + if len(p) < 2 || p[0] != '/' || p[len(p)-1] == '/' { + return "" + } + for i := 0; i < len(p); i++ { + c := p[i] + if !((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || + (c >= '0' && c <= '9') || strings.ContainsRune("/-._~!$&'()*+,;=:@", rune(c))) { + return "" + } + } + for _, segment := range strings.Split(p[1:], "/") { + if segment == "" || segment == "." || segment == ".." { + return "" + } + } + return p +} + +// joinURLPath appends route below base without producing a double slash. +func joinURLPath(base, route string) string { + if base == defaultMountPath { + return route + } + return base + route +} + +// applyTXTMetadata fills svc from parsed TXT attributes and reports whether the +// record is usable. Draft records (txtvers) are validated strictly; legacy CPA +// records (version + api_*) remain accepted for mixed-version networks. +func applyTXTMetadata(svc *DiscoveredService, parsed map[string]string) bool { + // Clear schema-specific metadata if the destination was previously populated. + svc.RawTXT = parsed + svc.Product = parsed[txtKeyProduct] + svc.InstanceID, svc.NodeRole, svc.Path = "", "", "" + svc.Protocols, svc.AuthMethods = nil, nil + svc.AuthRequired, svc.TLS = false, false + svc.Endpoints = make(map[string]string) + if version, draft := parsed[txtKeyVersion]; draft { + svc.Version = version + return applyDraftTXTMetadata(svc, parsed) + } + return applyLegacyTXTMetadata(svc, parsed) +} + +func applyDraftTXTMetadata(svc *DiscoveredService, parsed map[string]string) bool { + for _, key := range []string{txtKeyVersion, txtKeyPath, txtKeyTLS, txtKeyAuth, + txtKeyProtocols, txtKeyAuthType, txtKeyProduct, txtKeyID, txtKeyRole} { + if value, present := parsed[key]; present && (value == "" || !printableTXTASCII(value)) { + return false + } + } + tls, okTLS := parseTXTFlag(parsed[txtKeyTLS]) + auth, okAuth := parseTXTFlag(parsed[txtKeyAuth]) + path := validateMountPath(parsed[txtKeyPath]) + if svc.Version != txtSchemaVersion || !okTLS || !okAuth || path == "" { + return false + } + for _, key := range []string{txtKeyProtocols, txtKeyAuthType} { + if value, present := parsed[key]; present { + for _, token := range strings.Split(value, ",") { + if !validTXTToken(token) { + return false + } + } + } + } + svc.TLS, svc.AuthRequired, svc.Path = tls, auth, path + svc.InstanceID, svc.NodeRole = parsed[txtKeyID], parsed[txtKeyRole] + if value, present := parsed[txtKeyProtocols]; present { + for _, token := range draftTXTList(value) { + if _, known := protocolRoutes[token]; known { + svc.Protocols = append(svc.Protocols, token) + } + } + } + if value, present := parsed[txtKeyAuthType]; present { + for _, token := range draftTXTList(value) { + if token == "api_key" || token == "oauth" { + svc.AuthMethods = append(svc.AuthMethods, token) + } + } + } + for _, protocol := range svc.Protocols { + if route, known := protocolRoutes[protocol]; known { + svc.Endpoints[protocol] = joinURLPath(path, route.version+route.suffix) + } + } + return true +} + +// draftTXTList bounds retained hints without changing token spelling or syntax. +func draftTXTList(value string) []string { + var result []string + seen := make(map[string]bool) + for _, token := range strings.Split(value, ",") { + if !seen[token] { + result = append(result, token) + seen[token] = true + if len(result) == maxDiscoveredMetadataItems { + break + } + } + } + return result +} + +func applyLegacyTXTMetadata(svc *DiscoveredService, parsed map[string]string) bool { + svc.Version = parsed[legacyKeyVersion] + tls, okTLS := parseTXTFlag(parsed[txtKeyTLS]) + if svc.Version != txtSchemaVersion || !okTLS { + return false + } + svc.TLS = tls + switch parsed[legacyKeyAuth] { + case "true": + svc.AuthRequired = true + case "false": + svc.AuthRequired = false + default: + return false + } + svc.InstanceID, svc.NodeRole = parsed[legacyKeyInstanceID], parsed[legacyKeyNodeRole] + svc.Protocols = parseTXTList(parsed[txtKeyProtocols]) + svc.AuthMethods = parseTXTList(parsed[legacyKeyAuthMethods]) + if _, present := parsed[txtKeyProtocols]; !present { + // A base path does not imply support for any particular operation. + for _, key := range []string{legacyKeyOpenAI, legacyKeyAnthropic, legacyKeyGemini} { + if base := validateMountPath(parsed[key]); base != "" { + svc.Endpoints[key] = base + } + } + } + for _, protocol := range svc.Protocols { + route, known := protocolRoutes[protocol] + if !known { + continue + } + if base := validateMountPath(parsed[route.legacyField]); base != "" { + svc.Endpoints[protocol] = joinURLPath(base, route.suffix) + } + } + return len(svc.Endpoints) > 0 +} + +// validateDraftTXTRecords enforces wire-size limits before parsing loses duplicate +// strings and their length prefixes. Call this for draft entries before conversion. +func validateDraftTXTRecords(records []string) bool { + total := 0 + for _, record := range records { + if len(record) > maxTXTRecordBytes { + return false + } + total += len(record) + 1 + if total > maxTXTBytes { + return false + } + } + return true +} + +// ProtocolFromSubtype converts a DNS-SD subtype label (e.g. "_responses") to its protocol token. +func ProtocolFromSubtype(subtype string) string { + return strings.ToLower(strings.TrimPrefix(strings.TrimSpace(subtype), "_")) +} diff --git a/internal/discovery/types.go b/internal/discovery/types.go index 811f3a06c..68db9a17e 100644 --- a/internal/discovery/types.go +++ b/internal/discovery/types.go @@ -22,6 +22,22 @@ const ( SubtypeGenerateContent = "_generate-content" // SubtypeInteractions identifies the Interactions API protocol. SubtypeInteractions = "_interactions" + // SubtypeCompletions identifies the legacy text Completions API protocol. + SubtypeCompletions = "_completions" + // SubtypeResponsesWS identifies the persistent WebSocket Responses stream. + SubtypeResponsesWS = "_responses-ws" + // SubtypeRealtime identifies realtime WebSocket / WebRTC sessions. + SubtypeRealtime = "_realtime" + // SubtypeImages identifies image generation and editing. + SubtypeImages = "_images" + // SubtypeVideos identifies video generation, editing, and retrieval. + SubtypeVideos = "_videos" + // SubtypeAudio identifies speech synthesis. + SubtypeAudio = "_audio" + // SubtypeCountTokens identifies standalone token counting. + SubtypeCountTokens = "_count-tokens" + // SubtypeSearch identifies agentic web/code search. + SubtypeSearch = "_search" // ProductCPA is the product identifier for CLIProxyAPI. ProductCPA = "cliproxyapi" ) @@ -40,22 +56,31 @@ type ServiceSpec struct { // DiscoveredService describes a discovered AI gateway service on the LAN. type DiscoveredService struct { - InstanceName string `json:"instance_name"` - ServiceType string `json:"service_type"` - Domain string `json:"domain"` - Host string `json:"host"` - Port int `json:"port"` - IPv4 []net.IP `json:"ipv4"` - IPv6 []net.IP `json:"ipv6"` - Protocols []string `json:"protocols,omitempty"` - Features []string `json:"features,omitempty"` - Product string `json:"product,omitempty"` - AuthRequired bool `json:"auth_required"` - AuthMethods []string `json:"auth_methods,omitempty"` - Endpoints map[string]string `json:"endpoints,omitempty"` - NodeRole string `json:"node_role,omitempty"` - Version string `json:"version,omitempty"` - RawTXT map[string]string `json:"raw_txt,omitempty"` + InstanceName string `json:"instance_name"` + ServiceType string `json:"service_type"` + Domain string `json:"domain"` + Host string `json:"host"` + Port int `json:"port"` + IPv4 []net.IP `json:"ipv4"` + IPv6 []net.IP `json:"ipv6"` + // InterfaceIndex separates identical service names observed on different links. + InterfaceIndex int `json:"interface_index,omitempty"` + // IPv6Zones maps link-local address strings to their receiving interface names. + IPv6Zones map[string]string `json:"ipv6_zones,omitempty"` + // Path is the gateway mount path from the TXT path key; empty for legacy advertisers. + Path string `json:"path,omitempty"` + TLS bool `json:"tls"` + Protocols []string `json:"protocols,omitempty"` + Product string `json:"product,omitempty"` + InstanceID string `json:"instance_id,omitempty"` + AuthRequired bool `json:"auth_required"` + AuthMethods []string `json:"auth_methods,omitempty"` + // Endpoints maps each advertised protocol token to its route on the gateway origin. + Endpoints map[string]string `json:"endpoints,omitempty"` + NodeRole string `json:"node_role,omitempty"` + // Version is the discovery TXT schema version (txtvers, or legacy version). + Version string `json:"version,omitempty"` + RawTXT map[string]string `json:"raw_txt,omitempty"` } // Advertiser defines the server-side mDNS advertisement lifecycle. diff --git a/internal/discovery/zeroconf.go b/internal/discovery/zeroconf.go index a08a72860..c132ff173 100644 --- a/internal/discovery/zeroconf.go +++ b/internal/discovery/zeroconf.go @@ -1,45 +1,2445 @@ package discovery import ( + "bytes" "context" + "crypto/rand" + "encoding/binary" "fmt" "net" "os" + "runtime" + "sort" "strings" "sync" + "sync/atomic" + "time" "github.com/libp2p/zeroconf/v2" log "github.com/sirupsen/logrus" + "golang.org/x/net/dns/dnsmessage" + "golang.org/x/net/ipv4" + "golang.org/x/net/ipv6" ) +const ( + maxDNSLabelBytes = 63 + maxHostCollisionAttempts = 8 + mdnsCacheFlushBit = 1 << 15 + mdnsServiceRecordTTL = 3200 + mdnsHostRecordTTL = 120 + mdnsLegacyRecordTTL = 10 + mdnsPort = 5353 + mdnsHostProbeCount = 3 + mdnsHostProbeInterval = 250 * time.Millisecond + mdnsAnnounceRepetitions = 2 + mdnsAnnounceInterval = time.Second + dnsTypeNSEC = dnsmessage.Type(47) +) + +var ( + mdnsGroupAddrIPv4 = &net.UDPAddr{IP: net.IPv4(224, 0, 0, 251), Port: mdnsPort} + mdnsGroupAddrIPv6 = &net.UDPAddr{IP: net.ParseIP("ff02::fb"), Port: mdnsPort} + mdnsWildcardAddrIPv4 = &net.UDPAddr{IP: net.IPv4(224, 0, 0, 0), Port: mdnsPort} + mdnsWildcardAddrIPv6 = &net.UDPAddr{IP: net.ParseIP("ff02::"), Port: mdnsPort} +) + +type zeroconfServer interface { + Shutdown() +} + +type registerProxyFunc func(instance, service, domain string, port int, host string, ips []string, text []string, ifaces []net.Interface, opts ...zeroconf.ServerOption) (zeroconfServer, error) +type hostConflictProbeFunc func(ctx context.Context, hostFQDN string, ifaces []net.Interface, localIPs []string) (bool, error) + +func validateRegisterProxyArgs(instance, service, domain string, port int, host string, ips []string) error { + if strings.TrimSpace(instance) == "" { + return fmt.Errorf("missing service instance name") + } + if strings.TrimSpace(service) == "" { + return fmt.Errorf("missing service name") + } + if strings.TrimSpace(host) == "" { + return fmt.Errorf("missing host name") + } + if port <= 0 || port > 65535 { + return fmt.Errorf("invalid port %d", port) + } + for _, ip := range ips { + if net.ParseIP(strings.TrimSpace(ip)) == nil { + return fmt.Errorf("failed to parse given IP: %v", ip) + } + } + return nil +} + +func defaultRegisterProxy(instance, service, domain string, port int, host string, ips []string, text []string, ifaces []net.Interface, _ ...zeroconf.ServerOption) (zeroconfServer, error) { + if errValid := validateRegisterProxyArgs(instance, service, domain, port, host, ips); errValid != nil { + return nil, errValid + } + resp, err := newHostRecordResponderWithRebind(nil, host, domain, ips, ifaces, nil) + if err != nil { + return nil, err + } + if hr, ok := resp.(*hostRecordResponder); ok && hr != nil { + hr.configureService(instance, service, domain, port, text) + hr.announceService() + } + return resp, nil +} + +func formatHostFQDN(host, domain string) string { + h := strings.TrimSuffix(strings.TrimSpace(host), ".") + d := strings.Trim(strings.TrimSpace(domain), ".") + if d == "" { + d = "local" + } + if !strings.HasSuffix(strings.ToLower(h), "."+strings.ToLower(d)) { + h = h + "." + d + } + return h + "." +} + +func splitIPv4AndIPv6(ips []string) (v4 [][4]byte, v6 [][16]byte) { + for _, raw := range ips { + ip := net.ParseIP(strings.TrimSpace(raw)) + if ip == nil || ip.IsLoopback() || ip.IsUnspecified() { + continue + } + if ip4 := ip.To4(); ip4 != nil { + v4 = append(v4, [4]byte(ip4)) + continue + } + if ip16 := ip.To16(); ip16 != nil { + v6 = append(v6, [16]byte(ip16)) + } + } + return v4, v6 +} + +func packWireDomainName(fqdn string) []byte { + trimmed := strings.Trim(strings.TrimSpace(fqdn), ".") + if trimmed == "" { + return []byte{0} + } + parts := strings.Split(trimmed, ".") + var out []byte + for _, part := range parts { + if len(part) == 0 || len(part) > maxDNSLabelBytes { + return nil + } + out = append(out, byte(len(part))) + out = append(out, part...) + } + out = append(out, 0) + return out +} + +func buildHostNSECRData(hostFQDN string, hasA, hasAAAA bool) []byte { + wireName := packWireDomainName(hostFQDN) + if len(wireName) == 0 || (!hasA && !hasAAAA) { + return nil + } + var bitmap [4]byte + bitmapLen := 0 + if hasA { + bitmap[0] |= 1 << (7 - 1) // TypeA = 1 + bitmapLen = 1 + } + if hasAAAA { + bitmap[3] |= 1 << (7 - (28 - 24)) // TypeAAAA = 28 + bitmapLen = 4 + } + out := append(wireName, 0x00, byte(bitmapLen)) + return append(out, bitmap[:bitmapLen]...) +} + +func buildHostAddressReply(query []byte, hostFQDN string, ips []string, legacyUnicast bool) (reply []byte, unicast bool) { + return buildHostAddressReplyFiltered(query, hostFQDN, ips, legacyUnicast, nil) +} + +func buildHostAddressReplyFiltered(query []byte, hostFQDN string, ips []string, legacyUnicast bool, onlyQU *bool) (reply []byte, unicast bool) { + normQuery := rewriteWireLabelChar(query, '.', dotInLabelSentinel) + var p dnsmessage.Parser + hdr, err := p.Start(normQuery) + if err != nil || hdr.Response { + return nil, false + } + questions, errQ := p.AllQuestions() + if errQ != nil || len(questions) == 0 { + return nil, false + } + if errSkipAns := p.SkipAllAnswers(); errSkipAns != nil { + return nil, false + } + if errSkipAuth := p.SkipAllAuthorities(); errSkipAuth != nil { + return nil, false + } + + var wantA, wantAAAA, extraA, extraAAAA bool + var matchedQuestions []dnsmessage.Question + seenQ := make(map[string]struct{}, len(questions)) + for _, q := range questions { + isQU := uint16(q.Class)&mdnsCacheFlushBit != 0 + if onlyQU != nil && isQU != *onlyQU { + continue + } + if !strings.EqualFold(q.Name.String(), hostFQDN) { + continue + } + baseClass := uint16(q.Class) &^ mdnsCacheFlushBit + if baseClass != uint16(dnsmessage.ClassINET) && baseClass != uint16(dnsmessage.ClassANY) { + continue + } + switch q.Type { + case dnsmessage.TypeA: + wantA = true + extraAAAA = true + case dnsmessage.TypeAAAA: + wantAAAA = true + extraA = true + case dnsmessage.TypeALL: + wantA = true + wantAAAA = true + default: + continue + } + qKey := fmt.Sprintf("%s\x00%d\x00%d", strings.ToLower(q.Name.String()), q.Type, baseClass) + if _, dup := seenQ[qKey]; !dup { + seenQ[qKey] = struct{}{} + matchedQuestions = append(matchedQuestions, dnsmessage.Question{ + Name: q.Name, + Type: q.Type, + Class: dnsmessage.Class(baseClass), + }) + } + if isQU { + unicast = true + } + } + if !wantA && !wantAAAA { + return nil, false + } + if legacyUnicast { + unicast = true + } + + v4, v6 := splitIPv4AndIPv6(ips) + if len(v4) == 0 && len(v6) == 0 { + return nil, false + } + hasAnswerA := wantA && len(v4) > 0 + hasAnswerAAAA := wantAAAA && len(v6) > 0 + needNSEC := (wantA && len(v4) == 0) || (wantAAAA && len(v6) == 0) + + name, errName := dnsmessage.NewName(hostFQDN) + if errName != nil { + return nil, false + } + rrClass := dnsmessage.ClassINET | dnsmessage.Class(mdnsCacheFlushBit) + ttl := uint32(mdnsHostRecordTTL) + if legacyUnicast { + rrClass = dnsmessage.ClassINET + ttl = mdnsLegacyRecordTTL + } + builder := dnsmessage.NewBuilder(nil, dnsmessage.Header{ + ID: hdr.ID, + Response: true, + Authoritative: true, + }) + builder.EnableCompression() + if legacyUnicast && len(matchedQuestions) > 0 { + if errStartQ := builder.StartQuestions(); errStartQ != nil { + return nil, false + } + for _, mq := range matchedQuestions { + if errQAdd := builder.Question(mq); errQAdd != nil { + return nil, false + } + } + } + if errAns := builder.StartAnswers(); errAns != nil { + return nil, false + } + if hasAnswerA { + for _, addr := range v4 { + if errAdd := builder.AResource(dnsmessage.ResourceHeader{ + Name: name, + Type: dnsmessage.TypeA, + Class: rrClass, + TTL: ttl, + }, dnsmessage.AResource{A: addr}); errAdd != nil { + return nil, false + } + } + } + if hasAnswerAAAA { + for _, addr := range v6 { + if errAdd := builder.AAAAResource(dnsmessage.ResourceHeader{ + Name: name, + Type: dnsmessage.TypeAAAA, + Class: rrClass, + TTL: ttl, + }, dnsmessage.AAAAResource{AAAA: addr}); errAdd != nil { + return nil, false + } + } + } + if needNSEC { + nsecData := buildHostNSECRData(hostFQDN, len(v4) > 0, len(v6) > 0) + if len(nsecData) == 0 { + return nil, false + } + if errNSEC := builder.UnknownResource(dnsmessage.ResourceHeader{ + Name: name, + Type: dnsTypeNSEC, + Class: rrClass, + TTL: ttl, + }, dnsmessage.UnknownResource{ + Type: dnsTypeNSEC, + Data: nsecData, + }); errNSEC != nil { + return nil, false + } + } + + addExtraA := !hasAnswerA && extraA && len(v4) > 0 + addExtraAAAA := !hasAnswerAAAA && extraAAAA && len(v6) > 0 + if addExtraA || addExtraAAAA { + if errExt := builder.StartAdditionals(); errExt != nil { + return nil, false + } + if addExtraA { + for _, addr := range v4 { + _ = builder.AResource(dnsmessage.ResourceHeader{ + Name: name, + Type: dnsmessage.TypeA, + Class: rrClass, + TTL: ttl, + }, dnsmessage.AResource{A: addr}) + } + } + if addExtraAAAA { + for _, addr := range v6 { + _ = builder.AAAAResource(dnsmessage.ResourceHeader{ + Name: name, + Type: dnsmessage.TypeAAAA, + Class: rrClass, + TTL: ttl, + }, dnsmessage.AAAAResource{AAAA: addr}) + } + } + } + + out, errFinish := builder.Finish() + if errFinish != nil { + return nil, false + } + return rewriteWireLabelChar(out, dotInLabelSentinel, '.'), unicast +} + +func buildHostGoodbyePacket(hostFQDN string, ips []string) []byte { + v4, v6 := splitIPv4AndIPv6(ips) + if len(v4) == 0 && len(v6) == 0 { + return nil + } + name, errName := dnsmessage.NewName(hostFQDN) + if errName != nil { + return nil + } + rrClass := dnsmessage.ClassINET | dnsmessage.Class(mdnsCacheFlushBit) + builder := dnsmessage.NewBuilder(nil, dnsmessage.Header{ + Response: true, + Authoritative: true, + }) + builder.EnableCompression() + if errAns := builder.StartAnswers(); errAns != nil { + return nil + } + for _, addr := range v4 { + if errAdd := builder.AResource(dnsmessage.ResourceHeader{ + Name: name, + Type: dnsmessage.TypeA, + Class: rrClass, + TTL: 0, + }, dnsmessage.AResource{A: addr}); errAdd != nil { + return nil + } + } + for _, addr := range v6 { + if errAdd := builder.AAAAResource(dnsmessage.ResourceHeader{ + Name: name, + Type: dnsmessage.TypeAAAA, + Class: rrClass, + TTL: 0, + }, dnsmessage.AAAAResource{AAAA: addr}); errAdd != nil { + return nil + } + } + out, errFinish := builder.Finish() + if errFinish != nil { + return nil + } + return out +} + +const dotInLabelSentinel byte = '\x01' + +type mdnsServiceConfig struct { + serviceTypeName string + serviceName string + instanceName string + subtypes []string + port uint16 + text []string +} + +func rewriteWireName(buf []byte, off int, fromByte, toByte byte) int { + for off < len(buf) { + b := buf[off] + if b == 0 { + return off + 1 + } + if b&0xC0 == 0xC0 { + if off+2 > len(buf) { + return -1 + } + return off + 2 + } + if b >= 64 { + return -1 + } + end := off + 1 + int(b) + if end > len(buf) { + return -1 + } + for i := off + 1; i < end; i++ { + if buf[i] == fromByte { + buf[i] = toByte + } + } + off = end + } + return -1 +} + +func rewriteWireLabelChar(pkt []byte, fromByte, toByte byte) []byte { + if len(pkt) < 12 { + return pkt + } + buf := bytes.Clone(pkt) + qd := int(binary.BigEndian.Uint16(buf[4:6])) + an := int(binary.BigEndian.Uint16(buf[6:8])) + ns := int(binary.BigEndian.Uint16(buf[8:10])) + ar := int(binary.BigEndian.Uint16(buf[10:12])) + off := 12 + for i := 0; i < qd; i++ { + off = rewriteWireName(buf, off, fromByte, toByte) + if off < 0 || off+4 > len(buf) { + return buf + } + off += 4 + } + totalRR := an + ns + ar + for i := 0; i < totalRR; i++ { + off = rewriteWireName(buf, off, fromByte, toByte) + if off < 0 || off+10 > len(buf) { + return buf + } + rrType := dnsmessage.Type(binary.BigEndian.Uint16(buf[off : off+2])) + rdLen := int(binary.BigEndian.Uint16(buf[off+8 : off+10])) + rdataStart := off + 10 + rdataEnd := rdataStart + rdLen + if rdataEnd > len(buf) { + return buf + } + switch rrType { + case dnsmessage.TypePTR, dnsmessage.TypeCNAME, dnsmessage.TypeNS: + _ = rewriteWireName(buf[:rdataEnd], rdataStart, fromByte, toByte) + case dnsmessage.TypeSRV: + if rdLen >= 6 { + _ = rewriteWireName(buf[:rdataEnd], rdataStart+6, fromByte, toByte) + } + } + off = rdataEnd + } + return buf +} + +func parseMDNSServiceConfig(instance, service, domain string, port int, text []string) mdnsServiceConfig { + cleanDom := strings.Trim(strings.TrimSpace(domain), ".") + if cleanDom == "" { + cleanDom = "local" + } + parts := strings.Split(service, ",") + baseSvc := strings.Trim(strings.TrimSpace(parts[0]), ".") + svcName := baseSvc + "." + cleanDom + "." + cleanInst := strings.ReplaceAll(strings.Trim(strings.TrimSpace(instance), "."), ".", string(dotInLabelSentinel)) + var subs []string + for _, sub := range parts[1:] { + s := strings.Trim(strings.TrimSpace(sub), ".") + if s != "" { + subs = append(subs, s+"._sub."+svcName) + } + } + return mdnsServiceConfig{ + serviceTypeName: "_services._dns-sd._udp." + cleanDom + ".", + serviceName: svcName, + instanceName: cleanInst + "." + svcName, + subtypes: subs, + port: uint16(port), + text: append([]string(nil), text...), + } +} + +func appendHostAddrResources(builder *dnsmessage.Builder, hostName dnsmessage.Name, rrClass dnsmessage.Class, ttl uint32, v4 [][4]byte, v6 [][16]byte) error { + for _, addr := range v4 { + if err := builder.AResource(dnsmessage.ResourceHeader{ + Name: hostName, + Type: dnsmessage.TypeA, + Class: rrClass, + TTL: ttl, + }, dnsmessage.AResource{A: addr}); err != nil { + return err + } + } + for _, addr := range v6 { + if err := builder.AAAAResource(dnsmessage.ResourceHeader{ + Name: hostName, + Type: dnsmessage.TypeAAAA, + Class: rrClass, + TTL: ttl, + }, dnsmessage.AAAAResource{AAAA: addr}); err != nil { + return err + } + } + return nil +} + +func buildServicePacket(svc mdnsServiceConfig, hostFQDN string, ips []string, ttl uint32) []byte { + v4, v6 := splitIPv4AndIPv6(ips) + if svc.serviceName == "" { + if ttl == 0 { + return buildHostGoodbyePacket(hostFQDN, ips) + } + return nil + } + svcTypeName, err1 := dnsmessage.NewName(svc.serviceTypeName) + svcName, err2 := dnsmessage.NewName(svc.serviceName) + instName, err3 := dnsmessage.NewName(svc.instanceName) + hostName, err4 := dnsmessage.NewName(hostFQDN) + if err1 != nil || err2 != nil || err3 != nil || err4 != nil { + return nil + } + var subNames []dnsmessage.Name + for _, sub := range svc.subtypes { + if sn, errS := dnsmessage.NewName(sub); errS == nil { + subNames = append(subNames, sn) + } + } + + svcTTL := ttl + hostTTL := ttl + if ttl > 0 { + svcTTL = mdnsServiceRecordTTL + hostTTL = mdnsHostRecordTTL + } + flushClass := dnsmessage.ClassINET | dnsmessage.Class(mdnsCacheFlushBit) + builder := dnsmessage.NewBuilder(nil, dnsmessage.Header{ + Response: true, + Authoritative: true, + }) + builder.EnableCompression() + if errAns := builder.StartAnswers(); errAns != nil { + return nil + } + if err := builder.SRVResource(dnsmessage.ResourceHeader{ + Name: instName, + Type: dnsmessage.TypeSRV, + Class: flushClass, + TTL: svcTTL, + }, dnsmessage.SRVResource{Port: svc.port, Target: hostName}); err != nil { + return nil + } + if err := builder.TXTResource(dnsmessage.ResourceHeader{ + Name: instName, + Type: dnsmessage.TypeTXT, + Class: flushClass, + TTL: svcTTL, + }, dnsmessage.TXTResource{TXT: svc.text}); err != nil { + return nil + } + if err := builder.PTRResource(dnsmessage.ResourceHeader{ + Name: svcName, + Type: dnsmessage.TypePTR, + Class: dnsmessage.ClassINET, + TTL: svcTTL, + }, dnsmessage.PTRResource{PTR: instName}); err != nil { + return nil + } + if err := builder.PTRResource(dnsmessage.ResourceHeader{ + Name: svcTypeName, + Type: dnsmessage.TypePTR, + Class: dnsmessage.ClassINET, + TTL: svcTTL, + }, dnsmessage.PTRResource{PTR: svcName}); err != nil { + return nil + } + for _, sn := range subNames { + if err := builder.PTRResource(dnsmessage.ResourceHeader{ + Name: sn, + Type: dnsmessage.TypePTR, + Class: dnsmessage.ClassINET, + TTL: svcTTL, + }, dnsmessage.PTRResource{PTR: instName}); err != nil { + return nil + } + } + if err := appendHostAddrResources(&builder, hostName, flushClass, hostTTL, v4, v6); err != nil { + return nil + } + out, errFinish := builder.Finish() + if errFinish != nil { + return nil + } + return rewriteWireLabelChar(out, dotInLabelSentinel, '.') +} + +func ptrKnownKey(ownerName, targetName string) string { + return strings.ToLower(ownerName) + "\x00" + strings.ToLower(targetName) +} + +func collectKnownPTRAnswers(p *dnsmessage.Parser, replyTTL uint32) (map[string]struct{}, error) { + var known map[string]struct{} + for { + rh, errH := p.AnswerHeader() + if errH != nil { + if errH == dnsmessage.ErrSectionDone { + return known, nil + } + return nil, errH + } + if rh.Type != dnsmessage.TypePTR || rh.TTL < replyTTL/2 { + if errSkip := p.SkipAnswer(); errSkip != nil { + return nil, errSkip + } + continue + } + ptr, errPTR := p.PTRResource() + if errPTR != nil { + return nil, errPTR + } + if known == nil { + known = make(map[string]struct{}) + } + known[ptrKnownKey(rh.Name.String(), ptr.PTR.String())] = struct{}{} + } +} + +func isKnownPTRAnswer(known map[string]struct{}, ownerName, targetName string) bool { + if len(known) == 0 { + return false + } + _, ok := known[ptrKnownKey(ownerName, targetName)] + return ok +} + +func buildServiceQueryReply(query []byte, svc mdnsServiceConfig, hostFQDN string, ips []string, legacyUnicast bool) (reply []byte, unicast bool) { + return buildServiceQueryReplyFiltered(query, svc, hostFQDN, ips, legacyUnicast, nil) +} + +func buildServiceQueryReplyFiltered(query []byte, svc mdnsServiceConfig, hostFQDN string, ips []string, legacyUnicast bool, onlyQU *bool) (reply []byte, unicast bool) { + if svc.serviceName == "" { + return nil, false + } + normQuery := rewriteWireLabelChar(query, '.', dotInLabelSentinel) + var p dnsmessage.Parser + hdr, err := p.Start(normQuery) + if err != nil || hdr.Response { + return nil, false + } + questions, errQ := p.AllQuestions() + if errQ != nil || len(questions) == 0 { + return nil, false + } + + svcTTL := uint32(mdnsServiceRecordTTL) + hostTTL := uint32(mdnsHostRecordTTL) + rrClass := dnsmessage.ClassINET | dnsmessage.Class(mdnsCacheFlushBit) + if legacyUnicast { + svcTTL = mdnsLegacyRecordTTL + hostTTL = mdnsLegacyRecordTTL + rrClass = dnsmessage.ClassINET + unicast = true + } + + knownPTRs, errKnown := collectKnownPTRAnswers(&p, svcTTL) + if errKnown != nil { + return nil, false + } + + svcTypeName, err1 := dnsmessage.NewName(svc.serviceTypeName) + svcName, err2 := dnsmessage.NewName(svc.serviceName) + instName, err3 := dnsmessage.NewName(svc.instanceName) + hostName, err4 := dnsmessage.NewName(hostFQDN) + if err1 != nil || err2 != nil || err3 != nil || err4 != nil { + return nil, false + } + + var matchedQuestions []dnsmessage.Question + var answerTypeEnum, answerBrowse, answerLookup bool + var browseSubtypes []dnsmessage.Name + seenSubtypes := make(map[string]struct{}, len(svc.subtypes)) + seenQ := make(map[string]struct{}, len(questions)) + + for _, q := range questions { + isQU := uint16(q.Class)&mdnsCacheFlushBit != 0 + if onlyQU != nil && isQU != *onlyQU { + continue + } + baseClass := uint16(q.Class) &^ mdnsCacheFlushBit + if baseClass != uint16(dnsmessage.ClassINET) && baseClass != uint16(dnsmessage.ClassANY) { + continue + } + qNameStr := q.Name.String() + switch { + case strings.EqualFold(qNameStr, svc.serviceTypeName): + if q.Type != dnsmessage.TypePTR && q.Type != dnsmessage.TypeALL { + continue + } + if isKnownPTRAnswer(knownPTRs, svc.serviceTypeName, svc.serviceName) { + continue + } + answerTypeEnum = true + case strings.EqualFold(qNameStr, svc.serviceName): + if q.Type != dnsmessage.TypePTR && q.Type != dnsmessage.TypeALL { + continue + } + if isKnownPTRAnswer(knownPTRs, svc.serviceName, svc.instanceName) { + continue + } + answerBrowse = true + case strings.EqualFold(qNameStr, svc.instanceName): + if q.Type != dnsmessage.TypeSRV && q.Type != dnsmessage.TypeTXT && q.Type != dnsmessage.TypePTR && q.Type != dnsmessage.TypeALL { + continue + } + answerLookup = true + default: + matchedSub := "" + for _, sub := range svc.subtypes { + if strings.EqualFold(qNameStr, sub) { + matchedSub = sub + break + } + } + if matchedSub == "" || (q.Type != dnsmessage.TypePTR && q.Type != dnsmessage.TypeALL) { + continue + } + if isKnownPTRAnswer(knownPTRs, matchedSub, svc.instanceName) { + continue + } + if _, dupSub := seenSubtypes[matchedSub]; !dupSub { + seenSubtypes[matchedSub] = struct{}{} + if sn, errSN := dnsmessage.NewName(matchedSub); errSN == nil { + browseSubtypes = append(browseSubtypes, sn) + } + } + } + qKey := fmt.Sprintf("%s\x00%d\x00%d", strings.ToLower(qNameStr), q.Type, baseClass) + if _, dup := seenQ[qKey]; !dup { + seenQ[qKey] = struct{}{} + matchedQuestions = append(matchedQuestions, dnsmessage.Question{ + Name: q.Name, + Type: q.Type, + Class: dnsmessage.Class(baseClass), + }) + } + if isQU { + unicast = true + } + } + + if !answerTypeEnum && !answerBrowse && !answerLookup && len(browseSubtypes) == 0 { + return nil, false + } + + v4, v6 := splitIPv4AndIPv6(ips) + builder := dnsmessage.NewBuilder(nil, dnsmessage.Header{ + ID: hdr.ID, + Response: true, + Authoritative: true, + }) + builder.EnableCompression() + if legacyUnicast && len(matchedQuestions) > 0 { + if errQ := builder.StartQuestions(); errQ != nil { + return nil, false + } + for _, mq := range matchedQuestions { + if errAddQ := builder.Question(mq); errAddQ != nil { + return nil, false + } + } + } + if errAns := builder.StartAnswers(); errAns != nil { + return nil, false + } + + if answerTypeEnum { + if err := builder.PTRResource(dnsmessage.ResourceHeader{ + Name: svcTypeName, + Type: dnsmessage.TypePTR, + Class: dnsmessage.ClassINET, + TTL: svcTTL, + }, dnsmessage.PTRResource{PTR: svcName}); err != nil { + return nil, false + } + } + if answerBrowse { + if err := builder.PTRResource(dnsmessage.ResourceHeader{ + Name: svcName, + Type: dnsmessage.TypePTR, + Class: dnsmessage.ClassINET, + TTL: svcTTL, + }, dnsmessage.PTRResource{PTR: instName}); err != nil { + return nil, false + } + } + for _, sn := range browseSubtypes { + if err := builder.PTRResource(dnsmessage.ResourceHeader{ + Name: sn, + Type: dnsmessage.TypePTR, + Class: dnsmessage.ClassINET, + TTL: svcTTL, + }, dnsmessage.PTRResource{PTR: instName}); err != nil { + return nil, false + } + } + if answerLookup { + if err := builder.SRVResource(dnsmessage.ResourceHeader{ + Name: instName, + Type: dnsmessage.TypeSRV, + Class: rrClass, + TTL: svcTTL, + }, dnsmessage.SRVResource{Port: svc.port, Target: hostName}); err != nil { + return nil, false + } + if err := builder.TXTResource(dnsmessage.ResourceHeader{ + Name: instName, + Type: dnsmessage.TypeTXT, + Class: rrClass, + TTL: svcTTL, + }, dnsmessage.TXTResource{TXT: svc.text}); err != nil { + return nil, false + } + if !answerBrowse { + _ = builder.PTRResource(dnsmessage.ResourceHeader{ + Name: svcName, + Type: dnsmessage.TypePTR, + Class: dnsmessage.ClassINET, + TTL: svcTTL, + }, dnsmessage.PTRResource{PTR: instName}) + } + if err := appendHostAddrResources(&builder, hostName, rrClass, hostTTL, v4, v6); err != nil { + return nil, false + } + } else if answerBrowse || len(browseSubtypes) > 0 { + if errExt := builder.StartAdditionals(); errExt != nil { + return nil, false + } + _ = builder.SRVResource(dnsmessage.ResourceHeader{ + Name: instName, + Type: dnsmessage.TypeSRV, + Class: dnsmessage.ClassINET, + TTL: svcTTL, + }, dnsmessage.SRVResource{Port: svc.port, Target: hostName}) + _ = builder.TXTResource(dnsmessage.ResourceHeader{ + Name: instName, + Type: dnsmessage.TypeTXT, + Class: dnsmessage.ClassINET, + TTL: svcTTL, + }, dnsmessage.TXTResource{TXT: svc.text}) + _ = appendHostAddrResources(&builder, hostName, dnsmessage.ClassINET, hostTTL, v4, v6) + } + + out, errFinish := builder.Finish() + if errFinish != nil { + return nil, false + } + return rewriteWireLabelChar(out, dotInLabelSentinel, '.'), unicast +} + +type hostRecordResponder struct { + mu sync.Mutex + inner zeroconfServer + svc mdnsServiceConfig + baseHost string + osHost string + domain string + attempt int + hostFQDN string + fixedIPs []string + ifaces []net.Interface + rebindInner func(string) (zeroconfServer, error) + probeConflict hostConflictProbeFunc + onSend func(hostFQDN string, unicast bool, pkt []byte) + inConflict bool + stopped bool + ipv4conn *ipv4.PacketConn + ipv6conn *ipv6.PacketConn + stopCh chan struct{} + wg sync.WaitGroup + shutdownOnce sync.Once +} + +func (r *hostRecordResponder) configureService(instance, service, domain string, port int, text []string) { + r.mu.Lock() + r.svc = parseMDNSServiceConfig(instance, service, domain, port, text) + r.mu.Unlock() +} + +func (r *hostRecordResponder) sendIfHostCurrent(expectedHostFQDN string, unicast bool, buf []byte, ifIndex int, from net.Addr) bool { + if len(buf) == 0 || expectedHostFQDN == "" { + return false + } + r.mu.Lock() + defer r.mu.Unlock() + if r.stopped || r.hostFQDN == "" || r.hostFQDN != expectedHostFQDN { + return false + } + if r.onSend != nil { + r.onSend(expectedHostFQDN, unicast, buf) + } + if unicast && from != nil { + r.sendUnicast(buf, ifIndex, from) + return true + } + r.sendMulticast(buf, ifIndex) + return true +} + +func (r *hostRecordResponder) announceServiceOnce() { + r.mu.Lock() + if r.stopped || r.hostFQDN == "" || r.svc.serviceName == "" { + r.mu.Unlock() + return + } + svc := r.svc + hostFQDN := r.hostFQDN + r.mu.Unlock() + + for _, iface := range r.ifaces { + ips := r.ipsForInterface(iface.Index) + if len(ips) == 0 { + continue + } + if pkt := buildServicePacket(svc, hostFQDN, ips, mdnsServiceRecordTTL); len(pkt) > 0 { + _ = r.sendIfHostCurrent(hostFQDN, false, pkt, iface.Index, nil) + } + } +} + +func (r *hostRecordResponder) announceService() { + r.scheduleServiceAnnouncements(mdnsAnnounceRepetitions, mdnsAnnounceInterval, r.announceServiceOnce) +} + +func (r *hostRecordResponder) scheduleServiceAnnouncements(repetitions int, initialDelay time.Duration, sendOnce func()) { + if repetitions < 1 || sendOnce == nil { + return + } + r.mu.Lock() + if r.stopped || r.hostFQDN == "" || r.svc.serviceName == "" { + r.mu.Unlock() + return + } + if repetitions > 1 { + r.wg.Add(1) + } + r.mu.Unlock() + + sendOnce() + if repetitions <= 1 { + return + } + go func() { + defer r.wg.Done() + delay := initialDelay + for i := 1; i < repetitions; i++ { + timer := time.NewTimer(delay) + select { + case <-r.stopCh: + timer.Stop() + return + case <-timer.C: + sendOnce() + delay *= 2 + } + } + }() +} + +func newHostRecordResponder(inner zeroconfServer, host, domain string, fixedIPs []string, ifaces []net.Interface) (zeroconfServer, error) { + return newHostRecordResponderWithRebind(inner, host, domain, fixedIPs, ifaces, nil) +} + +func newHostRecordResponderWithRebind(inner zeroconfServer, host, domain string, fixedIPs []string, ifaces []net.Interface, rebind func(string) (zeroconfServer, error)) (zeroconfServer, error) { + ipv4conn, err4 := openMDNSConn4(ifaces) + ipv6conn, err6 := openMDNSConn6(ifaces) + if ipv4conn == nil && ipv6conn == nil { + if inner != nil { + inner.Shutdown() + } + return nil, fmt.Errorf("discovery: failed to start host record responder (ipv4: %v, ipv6: %v)", err4, err6) + } + r := &hostRecordResponder{ + inner: inner, + baseHost: host, + osHost: extractCleanHost(), + domain: domain, + attempt: 1, + hostFQDN: formatHostFQDN(host, domain), + fixedIPs: append([]string(nil), fixedIPs...), + ifaces: append([]net.Interface(nil), ifaces...), + rebindInner: rebind, + probeConflict: probeHostConflictOnLAN, + ipv4conn: ipv4conn, + ipv6conn: ipv6conn, + stopCh: make(chan struct{}), + } + if r.ipv4conn != nil { + r.wg.Add(1) + go r.recv4() + } + if r.ipv6conn != nil { + r.wg.Add(1) + go r.recv6() + } + return r, nil +} + +func openMDNSConn4(ifaces []net.Interface) (*ipv4.PacketConn, error) { + udpConn, err := net.ListenUDP("udp4", mdnsWildcardAddrIPv4) + if err != nil { + return nil, err + } + pkConn := ipv4.NewPacketConn(udpConn) + _ = pkConn.SetControlMessage(ipv4.FlagInterface|ipv4.FlagDst, true) + joined := 0 + var lastErr error + for i := range ifaces { + if errJoin := pkConn.JoinGroup(&ifaces[i], mdnsGroupAddrIPv4); errJoin == nil { + joined++ + } else { + lastErr = errJoin + } + } + if joined == 0 { + _ = pkConn.Close() + if lastErr == nil { + lastErr = fmt.Errorf("no multicast interfaces joined") + } + return nil, lastErr + } + _ = pkConn.SetMulticastTTL(255) + return pkConn, nil +} + +func openMDNSConn6(ifaces []net.Interface) (*ipv6.PacketConn, error) { + udpConn, err := net.ListenUDP("udp6", mdnsWildcardAddrIPv6) + if err != nil { + return nil, err + } + pkConn := ipv6.NewPacketConn(udpConn) + _ = pkConn.SetControlMessage(ipv6.FlagInterface|ipv6.FlagDst, true) + joined := 0 + var lastErr error + for i := range ifaces { + if errJoin := pkConn.JoinGroup(&ifaces[i], mdnsGroupAddrIPv6); errJoin == nil { + joined++ + } else { + lastErr = errJoin + } + } + if joined == 0 { + _ = pkConn.Close() + if lastErr == nil { + lastErr = fmt.Errorf("no multicast interfaces joined") + } + return nil, lastErr + } + _ = pkConn.SetMulticastHopLimit(255) + return pkConn, nil +} + +func (r *hostRecordResponder) ifaceAllowed(ifIndex int) bool { + if ifIndex == 0 { + return true + } + for _, iface := range r.ifaces { + if iface.Index == ifIndex { + return true + } + } + return false +} + +func (r *hostRecordResponder) ipsForInterface(ifIndex int) []string { + if len(r.fixedIPs) > 0 { + return r.fixedIPs + } + if ifIndex != 0 { + for _, iface := range r.ifaces { + if iface.Index == ifIndex { + addrs, err := iface.Addrs() + if err != nil { + return nil + } + return extractInterfaceAddrs(addrs) + } + } + if iface, err := net.InterfaceByIndex(ifIndex); err == nil && iface != nil { + if addrs, errAddrs := iface.Addrs(); errAddrs == nil { + return extractInterfaceAddrs(addrs) + } + } + return nil + } + return extractInterfaceIPs(r.ifaces) +} + +func (r *hostRecordResponder) recv4() { + defer r.wg.Done() + buf := make([]byte, 65536) + for { + select { + case <-r.stopCh: + return + default: + n, cm, from, err := r.ipv4conn.ReadFrom(buf) + if err != nil { + continue + } + ifIndex := 0 + isMulticastDst := false + if cm != nil { + ifIndex = cm.IfIndex + if len(cm.Dst) > 0 { + isMulticastDst = cm.Dst.IsMulticast() + } + } + r.handlePacketWithMeta(buf[:n], ifIndex, from, isMulticastDst) + } + } +} + +func (r *hostRecordResponder) recv6() { + defer r.wg.Done() + buf := make([]byte, 65536) + for { + select { + case <-r.stopCh: + return + default: + n, cm, from, err := r.ipv6conn.ReadFrom(buf) + if err != nil { + continue + } + ifIndex := 0 + isMulticastDst := false + if cm != nil { + ifIndex = cm.IfIndex + if len(cm.Dst) > 0 { + isMulticastDst = cm.Dst.IsMulticast() + } + } + r.handlePacketWithMeta(buf[:n], ifIndex, from, isMulticastDst) + } + } +} + +func (r *hostRecordResponder) currentHostFQDN() string { + r.mu.Lock() + defer r.mu.Unlock() + return r.hostFQDN +} + +func (r *hostRecordResponder) goodbyePacketsLocked(hostFQDN string) map[int][]byte { + if hostFQDN == "" { + return nil + } + if len(r.fixedIPs) > 0 && r.inner != nil { + return nil + } + out := make(map[int][]byte, len(r.ifaces)) + for _, iface := range r.ifaces { + ips := r.ipsForInterface(iface.Index) + if len(ips) == 0 { + continue + } + var pkt []byte + if r.svc.serviceName != "" { + pkt = buildServicePacket(r.svc, hostFQDN, ips, 0) + } else { + pkt = buildHostGoodbyePacket(hostFQDN, ips) + } + if len(pkt) > 0 { + out[iface.Index] = pkt + } + } + return out +} + +func (r *hostRecordResponder) goodbyePacketsForHost(hostFQDN string) map[int][]byte { + r.mu.Lock() + defer r.mu.Unlock() + return r.goodbyePacketsLocked(hostFQDN) +} + +func (r *hostRecordResponder) sendGoodbyeLocked(hostFQDN string) { + for ifIndex, pkt := range r.goodbyePacketsLocked(hostFQDN) { + r.sendMulticast(pkt, ifIndex) + } +} + +func (r *hostRecordResponder) sendGoodbyeForHost(hostFQDN string) { + r.mu.Lock() + defer r.mu.Unlock() + r.sendGoodbyeLocked(hostFQDN) +} + +func (r *hostRecordResponder) withdrawConflictedHost(observedFQDN string) { + r.mu.Lock() + if r.stopped || r.hostFQDN == "" || r.hostFQDN != observedFQDN { + r.mu.Unlock() + return + } + oldFQDN := r.hostFQDN + oldInner := r.inner + r.hostFQDN = "" + r.inner = nil + r.sendGoodbyeLocked(oldFQDN) + r.mu.Unlock() + + if oldInner != nil { + oldInner.Shutdown() + } +} + +func (r *hostRecordResponder) handleConflict(observedFQDN string) { + r.mu.Lock() + if r.stopped || r.inConflict || r.hostFQDN == "" || r.hostFQDN != observedFQDN { + r.mu.Unlock() + return + } + if r.attempt >= maxHostCollisionAttempts { + r.mu.Unlock() + log.Warnf("discovery: conflict attempts exhausted for %s; withdrawing host advertisement", observedFQDN) + r.withdrawConflictedHost(observedFQDN) + return + } + r.inConflict = true + startAttempt := r.attempt + 1 + baseHost := r.baseHost + osHost := r.osHost + domain := r.domain + rebind := r.rebindInner + probe := r.probeConflict + r.mu.Unlock() + + defer func() { + r.mu.Lock() + r.inConflict = false + r.mu.Unlock() + }() + + localIPs := r.allLocalIPs() + nextHost, nextAttempt, errResolve := resolveAdvertisedHostCandidate(context.Background(), baseHost, startAttempt, osHost, domain, r.ifaces, localIPs, probe) + if errResolve != nil { + log.Warnf("discovery: unable to resolve replacement host after conflict on %s: %v; withdrawing host advertisement", observedFQDN, errResolve) + r.withdrawConflictedHost(observedFQDN) + return + } + + var nextInner zeroconfServer + if rebind != nil { + var errRebind error + nextInner, errRebind = rebind(nextHost) + if errRebind != nil || nextInner == nil { + log.Warnf("discovery: failed to rebind service under replacement host %s: %v; withdrawing host advertisement", nextHost, errRebind) + r.withdrawConflictedHost(observedFQDN) + return + } + } + + r.mu.Lock() + if r.stopped || r.hostFQDN != observedFQDN { + r.mu.Unlock() + if nextInner != nil { + nextInner.Shutdown() + } + return + } + oldFQDN := r.hostFQDN + oldInner := r.inner + r.attempt = nextAttempt + r.hostFQDN = formatHostFQDN(nextHost, domain) + if nextInner != nil { + r.inner = nextInner + } + r.sendGoodbyeLocked(oldFQDN) + r.mu.Unlock() + + if oldInner != nil && nextInner != nil && oldInner != nextInner { + oldInner.Shutdown() + } + r.announceService() +} + +func (r *hostRecordResponder) allLocalIPs() []string { + if len(r.fixedIPs) > 0 { + return r.fixedIPs + } + return extractInterfaceIPs(r.ifaces) +} + +func hostAnswerHasForeignIP(answerIPs, localIPs []string) bool { + if len(answerIPs) == 0 { + return false + } + if len(localIPs) == 0 { + return true + } + localSet := make(map[string]struct{}, len(localIPs)) + for _, raw := range localIPs { + if ip := net.ParseIP(strings.TrimSpace(raw)); ip != nil { + localSet[ip.String()] = struct{}{} + } + } + for _, raw := range answerIPs { + ip := net.ParseIP(strings.TrimSpace(raw)) + if ip == nil { + continue + } + if _, ok := localSet[ip.String()]; !ok { + return true + } + } + return false +} + +func isOnLinkSource(from net.Addr, ifIndex int, ifaces []net.Interface, configuredIPs []string) bool { + uaddr, ok := from.(*net.UDPAddr) + if !ok || uaddr == nil || uaddr.IP == nil { + return false + } + if uaddr.IP.IsLoopback() || uaddr.IP.IsLinkLocalUnicast() { + return true + } + checkedKernelAddrs := false + for _, iface := range ifaces { + if ifIndex != 0 && iface.Index != ifIndex { + continue + } + addrs, err := iface.Addrs() + if err != nil || len(addrs) == 0 { + continue + } + checkedKernelAddrs = true + for _, addr := range addrs { + switch v := addr.(type) { + case *net.IPNet: + if v.Contains(uaddr.IP) { + return true + } + case *net.IPAddr: + if v.IP.Equal(uaddr.IP) { + return true + } + } + } + } + if !checkedKernelAddrs { + for _, iface := range ifaces { + addrs, err := iface.Addrs() + if err != nil || len(addrs) == 0 { + continue + } + checkedKernelAddrs = true + for _, addr := range addrs { + switch v := addr.(type) { + case *net.IPNet: + if v.Contains(uaddr.IP) { + return true + } + case *net.IPAddr: + if v.IP.Equal(uaddr.IP) { + return true + } + } + } + } + } + for _, raw := range configuredIPs { + ip := net.ParseIP(strings.TrimSpace(raw)) + if ip == nil || isLocalInterfaceIP(ip, ifaces) { + continue + } + if ip4 := ip.To4(); ip4 != nil { + mask := net.CIDRMask(24, 32) + subnet := net.IPNet{IP: ip4.Mask(mask), Mask: mask} + if subnet.Contains(uaddr.IP) { + return true + } + continue + } + if ip16 := ip.To16(); ip16 != nil { + mask := net.CIDRMask(64, 128) + subnet := net.IPNet{IP: ip16.Mask(mask), Mask: mask} + if subnet.Contains(uaddr.IP) { + return true + } + } + } + return false +} + +func isValidMDNSResponseSource(from net.Addr, ifIndex int, ifaces []net.Interface, configuredIPs []string) bool { + uaddr, ok := from.(*net.UDPAddr) + if !ok || uaddr == nil || uaddr.Port != mdnsPort { + return false + } + return isOnLinkSource(from, ifIndex, ifaces, configuredIPs) +} + +func isValidMDNSResponsePeer(from net.Addr, ifIndex int, ifaces []net.Interface, configuredIPs []string, isMulticastDst bool) bool { + uaddr, ok := from.(*net.UDPAddr) + if !ok || uaddr == nil || uaddr.Port != mdnsPort || uaddr.IP == nil { + return false + } + if isMulticastDst { + return true + } + return isOnLinkSource(from, ifIndex, ifaces, configuredIPs) +} + +func (r *hostRecordResponder) handlePacket(pkt []byte, ifIndex int, from net.Addr) { + r.handlePacketWithMeta(pkt, ifIndex, from, false) +} + +type mdnsQueryReply struct { + payload []byte + unicast bool +} + +func buildCombinedQueryReplies(pkt []byte, svc mdnsServiceConfig, hostFQDN string, ips []string, legacyUnicast bool) []mdnsQueryReply { + var out []mdnsQueryReply + if legacyUnicast { + if hostReply, hostUnicast := buildHostAddressReplyFiltered(pkt, hostFQDN, ips, true, nil); len(hostReply) > 0 { + out = append(out, mdnsQueryReply{payload: hostReply, unicast: hostUnicast}) + } + if svc.serviceName != "" { + if svcReply, svcUnicast := buildServiceQueryReplyFiltered(pkt, svc, hostFQDN, ips, true, nil); len(svcReply) > 0 { + out = append(out, mdnsQueryReply{payload: svcReply, unicast: svcUnicast}) + } + } + return out + } + for _, wantQU := range []bool{false, true} { + qu := wantQU + if hostReply, hostUnicast := buildHostAddressReplyFiltered(pkt, hostFQDN, ips, false, &qu); len(hostReply) > 0 { + out = append(out, mdnsQueryReply{payload: hostReply, unicast: hostUnicast}) + } + if svc.serviceName != "" { + if svcReply, svcUnicast := buildServiceQueryReplyFiltered(pkt, svc, hostFQDN, ips, false, &qu); len(svcReply) > 0 { + out = append(out, mdnsQueryReply{payload: svcReply, unicast: svcUnicast}) + } + } + } + return out +} + +func (r *hostRecordResponder) handlePacketWithMeta(pkt []byte, ifIndex int, from net.Addr, isMulticastDst bool) { + if !r.ifaceAllowed(ifIndex) { + return + } + r.mu.Lock() + hostFQDN := r.hostFQDN + svc := r.svc + r.mu.Unlock() + if hostFQDN == "" { + return + } + ips := r.ipsForInterface(ifIndex) + if len(ips) == 0 { + return + } + allLocal := r.allLocalIPs() + if !isMulticastDst && from != nil && !isOnLinkSource(from, ifIndex, r.ifaces, allLocal) { + return + } + if answerIPs := extractHostAnswerIPs(pkt, hostFQDN); len(answerIPs) > 0 { + if isValidMDNSResponsePeer(from, ifIndex, r.ifaces, allLocal, isMulticastDst) && hostAnswerHasForeignIP(answerIPs, allLocal) { + r.handleConflict(hostFQDN) + } + return + } + if hasHostResponseRecord(pkt, hostFQDN) { + if isValidMDNSResponsePeer(from, ifIndex, r.ifaces, allLocal, isMulticastDst) { + r.handleConflict(hostFQDN) + } + return + } + legacyUnicast := false + if uaddr, ok := from.(*net.UDPAddr); ok && uaddr != nil && uaddr.Port != mdnsPort { + legacyUnicast = true + } + for _, rep := range buildCombinedQueryReplies(pkt, svc, hostFQDN, ips, legacyUnicast) { + _ = r.sendIfHostCurrent(hostFQDN, rep.unicast && from != nil, rep.payload, ifIndex, from) + } +} + +func (r *hostRecordResponder) sendUnicast(buf []byte, ifIndex int, from net.Addr) { + uaddr, ok := from.(*net.UDPAddr) + if !ok || uaddr == nil { + return + } + if uaddr.IP.To4() != nil { + if r.ipv4conn == nil { + return + } + if ifIndex != 0 && supportsPerInterfaceAddrs(runtime.GOOS) { + wcm := ipv4.ControlMessage{IfIndex: ifIndex} + _, _ = r.ipv4conn.WriteTo(buf, &wcm, uaddr) + return + } + _, _ = r.ipv4conn.WriteTo(buf, nil, uaddr) + return + } + if r.ipv6conn == nil { + return + } + if ifIndex != 0 && supportsPerInterfaceAddrs(runtime.GOOS) { + wcm := ipv6.ControlMessage{IfIndex: ifIndex} + _, _ = r.ipv6conn.WriteTo(buf, &wcm, uaddr) + return + } + _, _ = r.ipv6conn.WriteTo(buf, nil, uaddr) +} + +func (r *hostRecordResponder) sendMulticast(buf []byte, ifIndex int) { + if r.ipv4conn != nil { + if ifIndex != 0 { + r.writeMulticast4(buf, ifIndex) + } else { + for _, iface := range r.ifaces { + r.writeMulticast4(buf, iface.Index) + } + } + } + if r.ipv6conn != nil { + if ifIndex != 0 { + r.writeMulticast6(buf, ifIndex) + } else { + for _, iface := range r.ifaces { + r.writeMulticast6(buf, iface.Index) + } + } + } +} + +func (r *hostRecordResponder) writeMulticast4(buf []byte, ifIndex int) { + if supportsPerInterfaceAddrs(runtime.GOOS) { + wcm := ipv4.ControlMessage{IfIndex: ifIndex} + _, _ = r.ipv4conn.WriteTo(buf, &wcm, mdnsGroupAddrIPv4) + return + } + if iface, err := net.InterfaceByIndex(ifIndex); err == nil && iface != nil { + _ = r.ipv4conn.SetMulticastInterface(iface) + } + _, _ = r.ipv4conn.WriteTo(buf, nil, mdnsGroupAddrIPv4) +} + +func (r *hostRecordResponder) writeMulticast6(buf []byte, ifIndex int) { + if supportsPerInterfaceAddrs(runtime.GOOS) { + wcm := ipv6.ControlMessage{IfIndex: ifIndex} + _, _ = r.ipv6conn.WriteTo(buf, &wcm, mdnsGroupAddrIPv6) + return + } + if iface, err := net.InterfaceByIndex(ifIndex); err == nil && iface != nil { + _ = r.ipv6conn.SetMulticastInterface(iface) + } + _, _ = r.ipv6conn.WriteTo(buf, nil, mdnsGroupAddrIPv6) +} + +func (r *hostRecordResponder) sendGoodbye() { + r.sendGoodbyeForHost(r.currentHostFQDN()) +} + +func (r *hostRecordResponder) Shutdown() { + r.shutdownOnce.Do(func() { + r.mu.Lock() + r.stopped = true + inner := r.inner + r.inner = nil + r.mu.Unlock() + + r.sendGoodbye() + close(r.stopCh) + if r.ipv4conn != nil { + _ = r.ipv4conn.Close() + } + if r.ipv6conn != nil { + _ = r.ipv6conn.Close() + } + r.wg.Wait() + if inner != nil { + inner.Shutdown() + } + }) +} + // ZeroconfAdvertiser wraps libp2p/zeroconf/v2 Server with defensive lifecycle management. type ZeroconfAdvertiser struct { - mu sync.Mutex - server *zeroconf.Server - inFlight bool - closed bool + mu sync.Mutex + server zeroconfServer + inFlight bool + closed bool + registerProxy registerProxyFunc + interfaceIPs func([]net.Interface) []string + osHostname func() (string, error) + probeHostConflict hostConflictProbeFunc +} + +// NewZeroconfAdvertiser returns a new ZeroconfAdvertiser. +func NewZeroconfAdvertiser() *ZeroconfAdvertiser { + return &ZeroconfAdvertiser{} +} + +func cleanHostname(h string) string { + h = strings.TrimSpace(h) + h = strings.TrimSuffix(h, ".") + if len(h) >= 6 && strings.EqualFold(h[len(h)-6:], ".local") { + h = h[:len(h)-6] + } + h = strings.TrimSuffix(h, ".") + if h == "" { + return "localhost" + } + return h +} + +// extractCleanHost returns bare hostname without any .local suffix +// to prevent libp2p/zeroconf/v2 from appending a duplicate .local. (e.g. on macOS). +func extractCleanHost() string { + h, err := os.Hostname() + if err != nil || h == "" { + return "localhost" + } + return cleanHostname(h) +} + +func sanitizeDNSLabel(s string) string { + var b strings.Builder + prevHyphen := false + for i := 0; i < len(s); i++ { + c := s[i] + switch { + case c >= 'a' && c <= 'z', c >= '0' && c <= '9': + b.WriteByte(c) + prevHyphen = false + case c >= 'A' && c <= 'Z': + b.WriteByte(c + ('a' - 'A')) + prevHyphen = false + default: + if b.Len() > 0 && !prevHyphen { + b.WriteByte('-') + prevHyphen = true + } + } + } + return strings.Trim(b.String(), "-") +} + +// buildAdvertisedHost derives a proxy-specific RFC 1123 host label from the +// instance name so mDNS advertisement never claims or flushes the OS-owned +// .local. record. +func buildAdvertisedHost(instanceName, osHost string) string { + clean := sanitizeDNSLabel(instanceName) + var rest string + switch { + case clean == "cpa": + rest = "0001" + case strings.HasPrefix(clean, "cpa-"): + rest = strings.TrimPrefix(clean, "cpa-") + default: + rest = clean + } + rest = strings.Trim(rest, "-") + if rest == "" { + rest = "0001" + } + + const prefix = "cpa-" + maxRest := maxDNSLabelBytes - len(prefix) + if len(rest) > maxRest { + if len(rest) > 5 && rest[len(rest)-5] == '-' && isValidHex4(rest[len(rest)-4:]) { + suffix := rest[len(rest)-5:] + head := strings.TrimRight(rest[:maxRest-len(suffix)], "-") + if head == "" { + rest = strings.TrimPrefix(suffix, "-") + } else { + rest = head + suffix + } + } else { + rest = strings.TrimRight(rest[:maxRest], "-") + } + if rest == "" { + rest = "0001" + } + } + + label := prefix + rest + cleanOS := strings.ToLower(cleanHostname(osHost)) + if cleanOS != "" && cleanOS != "localhost" && strings.EqualFold(label, cleanOS) { + const altSuffix = "-svc" + if len(label)+len(altSuffix) <= maxDNSLabelBytes { + label += altSuffix + } else { + trimmed := strings.TrimRight(label[:maxDNSLabelBytes-len(altSuffix)], "-") + if trimmed == "" { + label = "cpa-0001" + altSuffix + } else { + label = trimmed + altSuffix + } + } + } + return label +} + +func appendHostCollisionSuffix(baseLabel string, attempt int) string { + if attempt <= 1 { + return baseLabel + } + suffix := fmt.Sprintf("-%d", attempt) + if len(baseLabel)+len(suffix) <= maxDNSLabelBytes { + return baseLabel + suffix + } + trimmed := strings.TrimRight(baseLabel[:maxDNSLabelBytes-len(suffix)], "-") + if trimmed == "" { + return "cpa" + suffix + } + return trimmed + suffix +} + +func extractHostAnswerIPs(pkt []byte, hostFQDN string) []string { + normPkt := rewriteWireLabelChar(pkt, '.', dotInLabelSentinel) + var p dnsmessage.Parser + hdr, err := p.Start(normPkt) + if err != nil || !hdr.Response { + return nil + } + if errSkipQ := p.SkipAllQuestions(); errSkipQ != nil { + return nil + } + var ips []string + collect := func(nextHeader func() (dnsmessage.ResourceHeader, error), skip func() error) bool { + for { + rh, errH := nextHeader() + if errH != nil { + return errH == dnsmessage.ErrSectionDone + } + if !strings.EqualFold(rh.Name.String(), hostFQDN) || rh.TTL == 0 { + if errSkip := skip(); errSkip != nil { + return false + } + continue + } + switch rh.Type { + case dnsmessage.TypeA: + r, errA := p.AResource() + if errA != nil { + return false + } + ips = append(ips, net.IP(r.A[:]).String()) + case dnsmessage.TypeAAAA: + r, errAAAA := p.AAAAResource() + if errAAAA != nil { + return false + } + ips = append(ips, net.IP(r.AAAA[:]).String()) + default: + if errSkip := skip(); errSkip != nil { + return false + } + } + } + } + if !collect(p.AnswerHeader, p.SkipAnswer) { + return ips + } + if errSkipAuth := p.SkipAllAuthorities(); errSkipAuth != nil { + return ips + } + _ = collect(p.AdditionalHeader, p.SkipAdditional) + return ips +} + +func hasHostResponseRecord(pkt []byte, hostFQDN string) bool { + normPkt := rewriteWireLabelChar(pkt, '.', dotInLabelSentinel) + var p dnsmessage.Parser + hdr, err := p.Start(normPkt) + if err != nil || !hdr.Response { + return false + } + if errSkipQ := p.SkipAllQuestions(); errSkipQ != nil { + return false + } + found := false + scan := func(nextHeader func() (dnsmessage.ResourceHeader, error), skip func() error) bool { + for { + rh, errH := nextHeader() + if errH != nil { + return errH == dnsmessage.ErrSectionDone + } + if errSkip := skip(); errSkip != nil { + return false + } + if rh.TTL > 0 && strings.EqualFold(rh.Name.String(), hostFQDN) { + found = true + return false + } + } + } + if !scan(p.AnswerHeader, p.SkipAnswer) || found { + return found + } + if errSkipAuth := p.SkipAllAuthorities(); errSkipAuth != nil { + return false + } + _ = scan(p.AdditionalHeader, p.SkipAdditional) + return found +} + +func extractHostAuthorityIPs(pkt []byte, hostFQDN string) (uint16, []string, bool) { + normPkt := rewriteWireLabelChar(pkt, '.', dotInLabelSentinel) + var p dnsmessage.Parser + hdr, err := p.Start(normPkt) + if err != nil || hdr.Response { + return 0, nil, false + } + qs, errQ := p.AllQuestions() + if errQ != nil || len(qs) == 0 { + return 0, nil, false + } + matchedQ := false + for _, q := range qs { + if strings.EqualFold(q.Name.String(), hostFQDN) { + matchedQ = true + break + } + } + if !matchedQ { + return 0, nil, false + } + if errSkipAns := p.SkipAllAnswers(); errSkipAns != nil { + return 0, nil, false + } + var ips []string + hasOtherType := false + for { + rh, errH := p.AuthorityHeader() + if errH != nil { + break + } + if !strings.EqualFold(rh.Name.String(), hostFQDN) || rh.TTL == 0 { + if errSkip := p.SkipAuthority(); errSkip != nil { + break + } + continue + } + switch rh.Type { + case dnsmessage.TypeA: + r, errA := p.AResource() + if errA != nil { + return hdr.ID, ips, hasOtherType + } + ips = append(ips, net.IP(r.A[:]).String()) + case dnsmessage.TypeAAAA: + r, errAAAA := p.AAAAResource() + if errAAAA != nil { + return hdr.ID, ips, hasOtherType + } + ips = append(ips, net.IP(r.AAAA[:]).String()) + default: + if errSkip := p.SkipAuthority(); errSkip != nil { + return hdr.ID, ips, hasOtherType + } + hasOtherType = true + } + } + return hdr.ID, ips, hasOtherType +} + +func canonicalHostRecordList(ips []string) [][]byte { + v4, v6 := splitIPv4AndIPv6(ips) + records := make([][]byte, 0, len(v4)+len(v6)) + for _, addr := range v4 { + rec := append([]byte{0, byte(dnsmessage.TypeA)}, addr[:]...) + records = append(records, rec) + } + for _, addr := range v6 { + rec := append([]byte{0, byte(dnsmessage.TypeAAAA)}, addr[:]...) + records = append(records, rec) + } + sort.Slice(records, func(i, j int) bool { + return bytes.Compare(records[i], records[j]) < 0 + }) + dedup := records[:0] + for _, rec := range records { + if len(dedup) == 0 || !bytes.Equal(dedup[len(dedup)-1], rec) { + dedup = append(dedup, rec) + } + } + return dedup +} + +// compareHostRecordSets compares two proposed record sets lexicographically per RFC 6762 §8.2. +// Returns -1 if localIPs < peerIPs (local loses), +1 if localIPs > peerIPs (local wins), 0 if equal. +func compareHostRecordSets(localIPs, peerIPs []string) int { + ours := canonicalHostRecordList(localIPs) + theirs := canonicalHostRecordList(peerIPs) + limit := min(len(ours), len(theirs)) + for i := 0; i < limit; i++ { + if cmp := bytes.Compare(ours[i], theirs[i]); cmp != 0 { + return cmp + } + } + switch { + case len(ours) < len(theirs): + return -1 + case len(ours) > len(theirs): + return 1 + default: + return 0 + } +} + +func isLocalInterfaceIP(target net.IP, ifaces []net.Interface) bool { + if target == nil { + return false + } + for _, iface := range ifaces { + if interfaceHasIP(iface, target) { + return true + } + } + return false +} + +func isOwnLoopbackProbe(pkt, ownQuery []byte, from net.Addr, ownPorts map[int]struct{}, ifaces []net.Interface) bool { + if len(ownQuery) == 0 || !bytes.Equal(pkt, ownQuery) { + return false + } + uaddr, ok := from.(*net.UDPAddr) + if !ok || uaddr == nil { + return false + } + if _, portMatch := ownPorts[uaddr.Port]; !portMatch { + return false + } + return isLocalInterfaceIP(uaddr.IP, ifaces) +} + +func isHostProbeConflict(pkt []byte, hostFQDN string, localIPs []string, ownPort, peerPort int) bool { + if hasHostResponseRecord(pkt, hostFQDN) { + return true + } + theirID, authIPs, hasOtherType := extractHostAuthorityIPs(pkt, hostFQDN) + if hasOtherType { + return true + } + if len(authIPs) == 0 { + return false + } + cmp := compareHostRecordSets(localIPs, authIPs) + if cmp < 0 { + return true + } + if cmp > 0 { + return false + } + if ownPort > 0 && peerPort > 0 && ownPort != peerPort { + return ownPort < peerPort + } + return theirID != 0 +} + +func ifaceInList(ifIndex int, ifaces []net.Interface) bool { + if ifIndex <= 0 { + return false + } + for _, iface := range ifaces { + if iface.Index == ifIndex { + return true + } + } + return false +} + +func isValidProbePeer(from net.Addr, ifIndex int, ifaces []net.Interface, configuredIPs []string, isMulticast, isResponse bool) bool { + uaddr, ok := from.(*net.UDPAddr) + if !ok || uaddr == nil || uaddr.IP == nil { + return false + } + if isResponse && uaddr.Port != mdnsPort { + return false + } + if isMulticast || ifaceInList(ifIndex, ifaces) { + return true + } + return isOnLinkSource(from, ifIndex, ifaces, configuredIPs) +} + +func buildHostProbeQuery(hostFQDN string, probeID uint16, localIPs []string) ([]byte, error) { + qName, errName := dnsmessage.NewName(hostFQDN) + if errName != nil { + return nil, errName + } + builder := dnsmessage.NewBuilder(nil, dnsmessage.Header{ID: probeID}) + builder.EnableCompression() + if errStartQ := builder.StartQuestions(); errStartQ != nil { + return nil, errStartQ + } + if errQ := builder.Question(dnsmessage.Question{ + Name: qName, + Type: dnsmessage.TypeALL, + Class: dnsmessage.ClassINET, + }); errQ != nil { + return nil, errQ + } + v4, v6 := splitIPv4AndIPv6(localIPs) + if len(v4) > 0 || len(v6) > 0 { + if errAuth := builder.StartAuthorities(); errAuth != nil { + return nil, errAuth + } + for _, addr := range v4 { + if errAdd := builder.AResource(dnsmessage.ResourceHeader{ + Name: qName, + Type: dnsmessage.TypeA, + Class: dnsmessage.ClassINET, + TTL: mdnsHostRecordTTL, + }, dnsmessage.AResource{A: addr}); errAdd != nil { + return nil, errAdd + } + } + for _, addr := range v6 { + if errAdd := builder.AAAAResource(dnsmessage.ResourceHeader{ + Name: qName, + Type: dnsmessage.TypeAAAA, + Class: dnsmessage.ClassINET, + TTL: mdnsHostRecordTTL, + }, dnsmessage.AAAAResource{AAAA: addr}); errAdd != nil { + return nil, errAdd + } + } + } + return builder.Finish() +} + +func randomProbeID() uint16 { + var b [2]byte + if _, err := rand.Read(b[:]); err == nil { + if id := binary.BigEndian.Uint16(b[:]); id != 0 { + return id + } + } + return uint16(os.Getpid()&0x7FFF) | 1 +} + +func awaitProbeConflictResult(conflictCh <-chan bool, doneCh <-chan struct{}, detected *atomic.Bool, ctx context.Context, stopReaders func()) (bool, error) { + select { + case <-conflictCh: + if stopReaders != nil { + stopReaders() + } + <-doneCh + return true, nil + case <-doneCh: + if detected != nil && detected.Load() { + return true, nil + } + select { + case <-conflictCh: + return true, nil + default: + } + if ctx != nil && ctx.Err() != nil { + return false, ctx.Err() + } + return false, nil + } } -// NewZeroconfAdvertiser returns a new ZeroconfAdvertiser. -func NewZeroconfAdvertiser() *ZeroconfAdvertiser { - return &ZeroconfAdvertiser{} +func runHostProbeSchedule(ctx context.Context, probeCount int, probeInterval time.Duration, sendProbe func() error, conflictCh <-chan bool, doneCh <-chan struct{}, detected *atomic.Bool, stopReaders func()) (bool, error) { + if probeCount < 1 { + probeCount = 1 + } + var ctxDone <-chan struct{} + if ctx != nil { + ctxDone = ctx.Done() + } + for i := 0; i < probeCount; i++ { + if ctx != nil && ctx.Err() != nil { + if stopReaders != nil { + stopReaders() + } + <-doneCh + if detected != nil && detected.Load() { + return true, nil + } + return false, ctx.Err() + } + if errSend := sendProbe(); errSend != nil { + if stopReaders != nil { + stopReaders() + } + <-doneCh + return false, errSend + } + timer := time.NewTimer(probeInterval) + select { + case <-conflictCh: + timer.Stop() + if stopReaders != nil { + stopReaders() + } + <-doneCh + return true, nil + case <-ctxDone: + timer.Stop() + if stopReaders != nil { + stopReaders() + } + <-doneCh + if detected != nil && detected.Load() { + return true, nil + } + return false, ctx.Err() + case <-doneCh: + timer.Stop() + if detected != nil && detected.Load() { + return true, nil + } + select { + case <-conflictCh: + return true, nil + default: + } + if ctx != nil && ctx.Err() != nil { + return false, ctx.Err() + } + return false, fmt.Errorf("discovery: host conflict probe readers terminated before probe window completed") + case <-timer.C: + if detected != nil && detected.Load() { + if stopReaders != nil { + stopReaders() + } + <-doneCh + return true, nil + } + } + } + if stopReaders != nil { + stopReaders() + } + return awaitProbeConflictResult(conflictCh, doneCh, detected, ctx, stopReaders) } -// extractCleanHost returns bare hostname without any .local suffix -// to prevent libp2p/zeroconf/v2 from appending a duplicate .local. (e.g. on macOS). -func extractCleanHost() string { - h, err := os.Hostname() - if err != nil || h == "" { - return "localhost" +func probeHostConflictOnLAN(ctx context.Context, hostFQDN string, ifaces []net.Interface, localIPs []string) (bool, error) { + if len(ifaces) == 0 { + return false, fmt.Errorf("discovery: cannot probe host conflict with empty interface list") } - h = strings.TrimSpace(h) - h = strings.TrimSuffix(h, ".") - h = strings.TrimSuffix(h, ".local") - h = strings.TrimSuffix(h, ".") - if h == "" { - return "localhost" + if ctx != nil { + if errCtx := ctx.Err(); errCtx != nil { + return false, errCtx + } + if deadline, ok := ctx.Deadline(); ok && time.Until(deadline) < time.Duration(mdnsHostProbeCount)*mdnsHostProbeInterval { + return false, context.DeadlineExceeded + } } - return h + probeID := randomProbeID() + queryBytes, errQuery := buildHostProbeQuery(hostFQDN, probeID, localIPs) + if errQuery != nil { + return false, errQuery + } + + type probeConn struct { + conn *net.UDPConn + pconn4 *ipv4.PacketConn + pconn6 *ipv6.PacketConn + isMulticastDst bool + } + var conns []probeConn + defer func() { + for _, c := range conns { + _ = c.conn.Close() + } + }() + + ownPorts := make(map[int]struct{}, 2) + ownPort4 := 0 + ownPort6 := 0 + if mConn4, errM4 := net.ListenUDP("udp4", mdnsWildcardAddrIPv4); errM4 == nil { + mp4 := ipv4.NewPacketConn(mConn4) + _ = mp4.SetControlMessage(ipv4.FlagInterface|ipv4.FlagDst, true) + joined := 0 + for i := range ifaces { + if mp4.JoinGroup(&ifaces[i], mdnsGroupAddrIPv4) == nil { + joined++ + } + } + if joined > 0 { + _ = mp4.SetMulticastTTL(255) + conns = append(conns, probeConn{conn: mConn4, pconn4: mp4, isMulticastDst: true}) + } else { + _ = mConn4.Close() + } + } + if mConn6, errM6 := net.ListenUDP("udp6", mdnsWildcardAddrIPv6); errM6 == nil { + mp6 := ipv6.NewPacketConn(mConn6) + _ = mp6.SetControlMessage(ipv6.FlagInterface|ipv6.FlagDst, true) + joined := 0 + for i := range ifaces { + if mp6.JoinGroup(&ifaces[i], mdnsGroupAddrIPv6) == nil { + joined++ + } + } + if joined > 0 { + _ = mp6.SetMulticastHopLimit(255) + conns = append(conns, probeConn{conn: mConn6, pconn6: mp6, isMulticastDst: true}) + } else { + _ = mConn6.Close() + } + } + + var pConn4 *ipv4.PacketConn + var pConn6 *ipv6.PacketConn + var lastInitErr error + if uConn4, errListen4 := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IPv4zero, Port: 0}); errListen4 != nil { + lastInitErr = errListen4 + } else { + pConn4 = ipv4.NewPacketConn(uConn4) + _ = pConn4.SetControlMessage(ipv4.FlagInterface|ipv4.FlagDst, true) + conns = append(conns, probeConn{conn: uConn4, pconn4: pConn4, isMulticastDst: false}) + if laddr, ok := uConn4.LocalAddr().(*net.UDPAddr); ok && laddr != nil { + ownPorts[laddr.Port] = struct{}{} + ownPort4 = laddr.Port + } + } + if uConn6, errListen6 := net.ListenUDP("udp6", &net.UDPAddr{IP: net.IPv6zero, Port: 0}); errListen6 != nil { + lastInitErr = errListen6 + } else { + pConn6 = ipv6.NewPacketConn(uConn6) + _ = pConn6.SetControlMessage(ipv6.FlagInterface|ipv6.FlagDst, true) + conns = append(conns, probeConn{conn: uConn6, pconn6: pConn6, isMulticastDst: false}) + if laddr, ok := uConn6.LocalAddr().(*net.UDPAddr); ok && laddr != nil { + ownPorts[laddr.Port] = struct{}{} + ownPort6 = laddr.Port + } + } + if pConn4 == nil && pConn6 == nil { + if lastInitErr == nil { + lastInitErr = fmt.Errorf("no probe sockets available") + } + return false, fmt.Errorf("discovery: failed to initialize mDNS host conflict probe for %s: %w", hostFQDN, lastInitErr) + } + + sendOnce := func() error { + var lastErr error + sentPerIface := make([]int, len(ifaces)) + if pConn4 != nil { + for i := range ifaces { + if errSet := pConn4.SetMulticastInterface(&ifaces[i]); errSet != nil { + lastErr = errSet + continue + } + if _, errWrite := pConn4.WriteTo(queryBytes, nil, mdnsGroupAddrIPv4); errWrite == nil { + sentPerIface[i]++ + } else { + lastErr = errWrite + } + } + } + if pConn6 != nil { + for i := range ifaces { + if errSet := pConn6.SetMulticastInterface(&ifaces[i]); errSet != nil { + lastErr = errSet + continue + } + if _, errWrite := pConn6.WriteTo(queryBytes, nil, mdnsGroupAddrIPv6); errWrite == nil { + sentPerIface[i]++ + } else { + lastErr = errWrite + } + } + } + for i := range ifaces { + if sentPerIface[i] == 0 { + if lastErr == nil { + lastErr = fmt.Errorf("no probe packets sent") + } + return fmt.Errorf("discovery: failed to send mDNS host conflict probe for %s on interface %s: %w", hostFQDN, ifaces[i].Name, lastErr) + } + } + return nil + } + + var detected atomic.Bool + conflictCh := make(chan bool, len(conns)) + var wg sync.WaitGroup + for _, pc := range conns { + wg.Add(1) + go func(entry probeConn) { + defer wg.Done() + buf := make([]byte, 4096) + for { + if ctx != nil && ctx.Err() != nil { + return + } + var ( + n int + ifIndex int + fromAddr net.Addr + errRead error + isMcast bool + ) + switch { + case entry.pconn4 != nil: + var cm *ipv4.ControlMessage + n, cm, fromAddr, errRead = entry.pconn4.ReadFrom(buf) + if cm != nil { + ifIndex = cm.IfIndex + if len(cm.Dst) > 0 { + isMcast = cm.Dst.IsMulticast() + } + } + case entry.pconn6 != nil: + var cm *ipv6.ControlMessage + n, cm, fromAddr, errRead = entry.pconn6.ReadFrom(buf) + if cm != nil { + ifIndex = cm.IfIndex + if len(cm.Dst) > 0 { + isMcast = cm.Dst.IsMulticast() + } + } + default: + n, fromAddr, errRead = entry.conn.ReadFromUDP(buf) + } + if errRead != nil { + return + } + pkt := buf[:n] + if isOwnLoopbackProbe(pkt, queryBytes, fromAddr, ownPorts, ifaces) { + continue + } + isResp := hasHostResponseRecord(pkt, hostFQDN) + if !isValidProbePeer(fromAddr, ifIndex, ifaces, localIPs, isMcast, isResp) { + continue + } + peerPort := 0 + ownPort := ownPort4 + if uaddr, ok := fromAddr.(*net.UDPAddr); ok && uaddr != nil { + peerPort = uaddr.Port + if uaddr.IP != nil && uaddr.IP.To4() == nil && ownPort6 != 0 { + ownPort = ownPort6 + } + } + if isHostProbeConflict(pkt, hostFQDN, localIPs, ownPort, peerPort) { + detected.Store(true) + conflictCh <- true + return + } + } + }(pc) + } + + doneCh := make(chan struct{}) + go func() { + wg.Wait() + close(doneCh) + }() + + var stopOnce sync.Once + stopReaders := func() { + stopOnce.Do(func() { + for _, c := range conns { + _ = c.conn.Close() + } + }) + } + return runHostProbeSchedule(ctx, mdnsHostProbeCount, mdnsHostProbeInterval, sendOnce, conflictCh, doneCh, &detected, stopReaders) +} + +func resolveAdvertisedHostCandidate(ctx context.Context, baseHost string, startAttempt int, osHost, domain string, ifaces []net.Interface, localIPs []string, probeConflict hostConflictProbeFunc) (string, int, error) { + cleanOS := strings.ToLower(cleanHostname(osHost)) + if probeConflict == nil { + probeConflict = probeHostConflictOnLAN + } + if startAttempt < 1 { + startAttempt = 1 + } + for attempt := startAttempt; attempt <= maxHostCollisionAttempts; attempt++ { + if ctx != nil && ctx.Err() != nil { + return "", 0, ctx.Err() + } + candidate := appendHostCollisionSuffix(baseHost, attempt) + if cleanOS != "" && cleanOS != "localhost" && strings.EqualFold(candidate, cleanOS) { + continue + } + conflict, errProbe := probeConflict(ctx, formatHostFQDN(candidate, domain), ifaces, localIPs) + if errProbe != nil { + return "", 0, errProbe + } + if !conflict { + return candidate, attempt, nil + } + } + return "", 0, fmt.Errorf("discovery: unable to resolve non-conflicting mDNS host label for %q after %d attempts", baseHost, maxHostCollisionAttempts) +} + +func resolveAdvertisedHost(ctx context.Context, instanceName, osHost, domain string, ifaces []net.Interface, localIPs []string, probeConflict hostConflictProbeFunc) (string, error) { + baseHost := buildAdvertisedHost(instanceName, osHost) + candidate, _, err := resolveAdvertisedHostCandidate(ctx, baseHost, 1, osHost, domain, ifaces, localIPs, probeConflict) + return candidate, err +} + +func extractInterfaceAddrs(addrs []net.Addr) []string { + var v4, v6Global, v6LinkLocal []string + for _, addr := range addrs { + var ip net.IP + switch v := addr.(type) { + case *net.IPNet: + ip = v.IP + case *net.IPAddr: + ip = v.IP + } + if ip == nil || ip.IsLoopback() || ip.IsUnspecified() { + continue + } + if ip4 := ip.To4(); ip4 != nil { + v4 = append(v4, ip4.String()) + continue + } + if ip16 := ip.To16(); ip16 != nil { + if ip16.IsGlobalUnicast() { + v6Global = append(v6Global, ip16.String()) + } else if ip16.IsLinkLocalUnicast() { + v6LinkLocal = append(v6LinkLocal, ip16.String()) + } + } + } + if len(v6Global) > 0 { + return append(v4, v6Global...) + } + return append(v4, v6LinkLocal...) } // extractInterfaceIPs collects non-loopback IP addresses from the selected interfaces. @@ -50,23 +2450,71 @@ func extractInterfaceIPs(ifaces []net.Interface) []string { if err != nil { continue } - for _, addr := range addrs { - var ip net.IP - switch v := addr.(type) { - case *net.IPNet: - ip = v.IP - case *net.IPAddr: - ip = v.IP - } - if ip == nil || ip.IsLoopback() || ip.IsUnspecified() { - continue - } - ips = append(ips, ip.String()) - } + ips = append(ips, extractInterfaceAddrs(addrs)...) } return ips } +func supportsPerInterfaceAddrs(goos string) bool { + switch goos { + case "darwin", "ios", "linux": + return true + default: + return false + } +} + +func sameIPSet(a, b []string) bool { + if len(a) == 0 || len(b) == 0 { + return len(a) == len(b) + } + setA := make(map[string]struct{}, len(a)) + for _, raw := range a { + trimmed := strings.TrimSpace(raw) + if ip := net.ParseIP(trimmed); ip != nil { + setA[ip.String()] = struct{}{} + } else if trimmed != "" { + setA[trimmed] = struct{}{} + } + } + setB := make(map[string]struct{}, len(b)) + for _, raw := range b { + trimmed := strings.TrimSpace(raw) + if ip := net.ParseIP(trimmed); ip != nil { + setB[ip.String()] = struct{}{} + } else if trimmed != "" { + setB[trimmed] = struct{}{} + } + } + if len(setA) != len(setB) { + return false + } + for k := range setA { + if _, ok := setB[k]; !ok { + return false + } + } + return true +} + +func resolveRegisterProxyIPs(ifaces []net.Interface, advertisedIPs []string, goos string, lookupIPs func([]net.Interface) []string) ([]string, error) { + if lookupIPs == nil { + lookupIPs = extractInterfaceIPs + } + interfaceIPs := lookupIPs(ifaces) + ips := advertisedIPs + if len(ips) == 0 { + ips = interfaceIPs + } + if len(ips) == 0 { + return nil, fmt.Errorf("discovery: no usable IP addresses found on specified interfaces") + } + if supportsPerInterfaceAddrs(goos) && len(interfaceIPs) > 0 && (len(advertisedIPs) == 0 || (len(ifaces) > 1 && sameIPSet(advertisedIPs, interfaceIPs))) { + return nil, nil + } + return ips, nil +} + // Start registers and starts mDNS advertisement for the primary service type // and any configured API protocol subtypes. func (a *ZeroconfAdvertiser) Start(ctx context.Context, spec ServiceSpec) (err error) { @@ -89,7 +2537,7 @@ func (a *ZeroconfAdvertiser) Start(ctx context.Context, spec ServiceSpec) (err e a.closed = false a.mu.Unlock() - var server *zeroconf.Server + var server zeroconfServer defer func() { if r := recover(); r != nil { err = fmt.Errorf("discovery: panic starting advertiser: %v", r) @@ -121,19 +2569,40 @@ func (a *ZeroconfAdvertiser) Start(ctx context.Context, spec ServiceSpec) (err e serviceType = DefaultServiceType } - cleanHost := extractCleanHost() - ips := spec.AdvertisedIPs - if len(ips) == 0 { - ips = extractInterfaceIPs(spec.Interfaces) - } - if len(ips) == 0 { - return fmt.Errorf("discovery: no usable IP addresses found on specified interfaces") + proxyIPs, errIPs := resolveRegisterProxyIPs(spec.Interfaces, spec.AdvertisedIPs, runtime.GOOS, a.interfaceIPs) + if errIPs != nil { + return errIPs } spec.InstanceName = sanitizeInstanceName(spec.InstanceName) if spec.InstanceName == "" { spec.InstanceName = DefaultInstancePrefix + "0001" } + cleanHost := extractCleanHost() + if a.osHostname != nil { + if rawHost, errHost := a.osHostname(); errHost == nil && rawHost != "" { + cleanHost = cleanHostname(rawHost) + } else { + cleanHost = "localhost" + } + } + localIPs := spec.AdvertisedIPs + if len(localIPs) == 0 { + if a.interfaceIPs != nil { + localIPs = a.interfaceIPs(spec.Interfaces) + } else { + localIPs = extractInterfaceIPs(spec.Interfaces) + } + } + probeConflict := a.probeHostConflict + if probeConflict == nil && a.registerProxy != nil { + probeConflict = func(context.Context, string, []net.Interface, []string) (bool, error) { return false, nil } + } + baseHost := buildAdvertisedHost(spec.InstanceName, cleanHost) + advertisedHost, initialAttempt, errHost := resolveAdvertisedHostCandidate(ctx, baseHost, 1, cleanHost, domain, spec.Interfaces, localIPs, probeConflict) + if errHost != nil { + return errHost + } if ctx != nil && ctx.Err() != nil { return ctx.Err() } @@ -145,20 +2614,35 @@ func (a *ZeroconfAdvertiser) Start(ctx context.Context, spec ServiceSpec) (err e } } + registerProxy := a.registerProxy + if registerProxy == nil { + registerProxy = defaultRegisterProxy + } + var errRegister error - server, errRegister = zeroconf.RegisterProxy( + server, errRegister = registerProxy( spec.InstanceName, primaryService, domain, spec.Port, - cleanHost, - ips, + advertisedHost, + proxyIPs, spec.TextRecords, spec.Interfaces, ) if errRegister != nil { return fmt.Errorf("discovery: failed to register primary service %s: %w", primaryService, errRegister) } + if hr, ok := server.(*hostRecordResponder); ok && hr != nil { + hr.mu.Lock() + hr.baseHost = baseHost + hr.osHost = cleanHost + hr.attempt = initialAttempt + if probeConflict != nil { + hr.probeConflict = probeConflict + } + hr.mu.Unlock() + } if ctx != nil && ctx.Err() != nil { return ctx.Err() } @@ -194,24 +2678,19 @@ func (a *ZeroconfAdvertiser) Stop() error { // ZeroconfBrowser provides DNS-SD browsing with fallback support. type ZeroconfBrowser struct { - options []zeroconf.ClientOption + ifaces []net.Interface } // NewZeroconfBrowser creates a new ZeroconfBrowser. func NewZeroconfBrowser(ifaces ...net.Interface) *ZeroconfBrowser { - var opts []zeroconf.ClientOption - if len(ifaces) > 0 { - opts = append(opts, zeroconf.SelectIfaces(ifaces)) - } - return &ZeroconfBrowser{options: opts} + return &ZeroconfBrowser{ifaces: append([]net.Interface(nil), ifaces...)} } const ( maxDiscoveredServices = 256 - maxBrowseEntries = 256 maxBrowseTXTRecords = 64 maxBrowseTXTBytes = 16 * 1024 - maxDiscoveredAddresses = 32 + maxDiscoveredAddresses = 64 maxDiscoveredMetadataItems = 32 maxMetadataItemBytes = 64 ) @@ -245,162 +2724,20 @@ func (b *ZeroconfBrowser) Browse(ctx context.Context, serviceType, domain string serviceType = DefaultServiceType } - entries := make(chan *zeroconf.ServiceEntry, 32) - browseCtx, cancelBrowse := context.WithCancel(ctx) - defer cancelBrowse() - - var discovered []DiscoveredService - var mu sync.Mutex - seen := make(map[string]int) - - // Stop the underlying zeroconf client after a bounded number of entries. This - // is separate from the caller's context so a LAN flood cannot grow the - // dependency's internal sentEntries cache for the full browse lifetime. - doneCh := make(chan struct{}) - go func() { - defer close(doneCh) - entriesSeen := 0 - for entry := range entries { - if entriesSeen >= maxBrowseEntries { - continue - } - entriesSeen++ - - if browseEntryWithinLimits(entry) { - svc := entryToDiscovered(entry) - if svc.Port != 0 && (len(svc.IPv4) != 0 || len(svc.IPv6) != 0) { - key := discoveredServiceKey(svc) - mu.Lock() - if index, ok := seen[key]; ok { - mergeDiscoveredService(&discovered[index], svc) - } else if len(discovered) < maxDiscoveredServices { - seen[key] = len(discovered) - discovered = append(discovered, svc) - } - mu.Unlock() - } - } - - if entriesSeen == maxBrowseEntries { - cancelBrowse() - } - } - }() - - errBrowse := zeroconf.Browse(browseCtx, serviceType, domain, entries, b.options...) - if errBrowse != nil { - // zeroconf may already have closed entries after a runtime error. - closeBrowseEntries(entries) - <-doneCh - return nil, fmt.Errorf("discovery: browse query failed: %w", errBrowse) - } - - <-doneCh - - return discovered, nil -} - -func closeBrowseEntries(entries chan *zeroconf.ServiceEntry) { - defer func() { - _ = recover() - }() - close(entries) -} - -func discoveredServiceKey(svc DiscoveredService) string { - return strings.Join([]string{svc.InstanceName, svc.ServiceType, svc.Domain}, "\x00") -} - -func mergeDiscoveredService(dst *DiscoveredService, src DiscoveredService) { - if dst == nil { - return - } - if src.Host != "" { - dst.Host = src.Host - } - if src.Port != 0 { - dst.Port = src.Port - } - dst.IPv4 = appendUniqueIPs(dst.IPv4, src.IPv4) - dst.IPv6 = appendUniqueIPs(dst.IPv6, src.IPv6) - if src.Product != "" { - dst.Product = src.Product - } - if src.Version != "" { - dst.Version = src.Version - } - if src.NodeRole != "" { - dst.NodeRole = src.NodeRole - } - if dst.RawTXT == nil { - dst.RawTXT = make(map[string]string) - } - mergeRawTXTRecords(dst.RawTXT, src.RawTXT) - if value, ok := dst.RawTXT["auth_required"]; ok { - dst.AuthRequired = strings.EqualFold(strings.TrimSpace(value), "true") - } - dst.AuthMethods = appendUniqueStrings(dst.AuthMethods, src.AuthMethods) - dst.Protocols = appendUniqueStrings(dst.Protocols, src.Protocols) - dst.Features = appendUniqueStrings(dst.Features, src.Features) - if dst.Endpoints == nil { - dst.Endpoints = make(map[string]string) - } - for key, value := range src.Endpoints { - dst.Endpoints[key] = value - } -} - -func mergeRawTXTRecords(dst, src map[string]string) { - totalBytes := rawTXTMapBytes(dst) - for key, value := range src { - if key == "" { - continue - } - replacedBytes := 0 - _, exists := dst[key] - if exists { - replacedBytes = rawTXTRecordBytes(key, dst[key]) - } else if len(dst) >= maxBrowseTXTRecords { - continue - } - nextBytes := totalBytes - replacedBytes + rawTXTRecordBytes(key, value) - if nextBytes > maxBrowseTXTBytes { - continue - } - dst[key] = value - totalBytes = nextBytes - } -} - -func rawTXTMapBytes(records map[string]string) int { - totalBytes := 0 - for key, value := range records { - totalBytes += rawTXTRecordBytes(key, value) + var errQuery error + serviceType, domain, errQuery = normalizeBrowseQuery(serviceType, domain) + if errQuery != nil { + return nil, errQuery } - return totalBytes -} - -func rawTXTRecordBytes(key, value string) int { - return len(key) + len(value) + 1 -} - -func appendUniqueIPs(dst, src []net.IP) []net.IP { - for _, ip := range src { - if ip == nil || len(dst) >= maxDiscoveredAddresses { - continue - } - duplicate := false - for _, existing := range dst { - if existing.Equal(ip) { - duplicate = true - break - } - } - if !duplicate { - dst = append(dst, append(net.IP(nil), ip...)) + ifaces := b.ifaces + if len(ifaces) == 0 { + var errFilter error + ifaces, errFilter = FilterInterfaces(nil, nil) + if errFilter != nil { + return nil, fmt.Errorf("discovery: filter browse interfaces: %w", errFilter) } } - return dst + return browseLAN(ctx, ifaces, serviceType, domain) } func appendUniqueStrings(dst, src []string) []string { @@ -471,29 +2808,13 @@ func (b *ZeroconfBrowser) BrowseWithFallbackServiceType(ctx context.Context, ser return append(cpaGateways, otherGateways...), nil } -// sanitizeEndpointPath validates that an endpoint path is a safe relative API path -// starting with '/' and containing no scheme, domain, protocol-relative prefixes, or traversal elements. -func sanitizeEndpointPath(p string) string { - p = strings.TrimSpace(p) - if !strings.HasPrefix(p, "/") || strings.HasPrefix(p, "//") || strings.Contains(p, `\`) || strings.Contains(p, "://") || strings.Contains(p, "..") { - return "" - } - for _, r := range p { - if r < 32 || r == 127 { - return "" - } - } - return p -} - -func entryToDiscovered(e *zeroconf.ServiceEntry) DiscoveredService { - parsed := ParseTXTRecords(e.Text) - +// entryToDiscovered converts a zeroconf entry and reports whether its TXT +// record satisfies draft-su-ai-gateway-discovery-00 (or the legacy CPA schema). +func entryToDiscovered(e *zeroconf.ServiceEntry) (DiscoveredService, bool) { port := e.Port if port < 1 || port > 65535 { port = 0 } - svc := DiscoveredService{ InstanceName: e.Instance, ServiceType: e.Service, @@ -502,49 +2823,22 @@ func entryToDiscovered(e *zeroconf.ServiceEntry) DiscoveredService { Port: port, IPv4: filterUsableIPs(e.AddrIPv4), IPv6: filterUsableIPs(e.AddrIPv6), - Product: parsed["product"], - Version: parsed["version"], - NodeRole: parsed["node_role"], - RawTXT: parsed, - Endpoints: make(map[string]string), - } - - if parsed["auth_required"] == "true" { - svc.AuthRequired = true - } - if methods := parsed["auth_methods"]; methods != "" { - svc.AuthMethods = parseTXTList(methods) - } - if protos := parsed["protocols"]; protos != "" { - svc.Protocols = parseTXTList(protos) - } - if feats := parsed["features"]; feats != "" { - svc.Features = parseTXTList(feats) } - - if v, ok := parsed["api_openai"]; ok { - if clean := sanitizeEndpointPath(v); clean != "" { - svc.Endpoints["openai"] = clean - } - } - if v, ok := parsed["api_anthropic"]; ok { - if clean := sanitizeEndpointPath(v); clean != "" { - svc.Endpoints["anthropic"] = clean - } + if e.HostName == "." || !browseEntryWithinLimits(e) { + return svc, false } - if v, ok := parsed["api_gemini"]; ok { - if clean := sanitizeEndpointPath(v); clean != "" { - svc.Endpoints["gemini"] = clean - } + parsed := ParseTXTRecords(e.Text) + if _, draft := parsed[txtKeyVersion]; draft && !validateDraftTXTRecords(e.Text) { + return svc, false } - - return svc + usable := applyTXTMetadata(&svc, parsed) + return svc, usable } func filterUsableIPs(ips []net.IP) []net.IP { usable := make([]net.IP, 0, min(len(ips), maxDiscoveredAddresses)) for _, ip := range ips { - if ip == nil || ip.IsLoopback() || ip.IsUnspecified() { + if ip == nil || ip.To16() == nil || ip.IsLoopback() || ip.IsUnspecified() || ip.IsMulticast() || ip.Equal(net.IPv4bcast) { continue } usable = append(usable, append(net.IP(nil), ip...)) diff --git a/internal/githubauth/token.go b/internal/githubauth/token.go new file mode 100644 index 000000000..fe4cafb66 --- /dev/null +++ b/internal/githubauth/token.go @@ -0,0 +1,43 @@ +// Package githubauth provides shared GitHub API credentials. +package githubauth + +import ( + "net/url" + "os" + "strings" + "sync/atomic" +) + +var configuredToken atomic.Pointer[string] + +// SetToken replaces the global configured token, including clearing it on reload. +func SetToken(token string) { + token = strings.TrimSpace(token) + configuredToken.Store(&token) +} + +// ResolveToken prefers the server configuration, then the legacy environment settings. +func ResolveToken() string { + if token := configuredToken.Load(); token != nil && *token != "" { + return *token + } + for _, name := range []string{"GITHUB_TOKEN", "github_token"} { + if token := strings.TrimSpace(os.Getenv(name)); token != "" { + return token + } + } + gitURL := strings.ToLower(strings.TrimSpace(os.Getenv("GITSTORE_GIT_URL"))) + if !strings.Contains(gitURL, "github.com") { + return "" + } + return strings.TrimSpace(os.Getenv("GITSTORE_GIT_TOKEN")) +} + +// TokenForURL never exposes the shared credential outside the HTTPS GitHub API. +func TokenForURL(rawURL string) string { + u, err := url.Parse(rawURL) + if err != nil || !strings.EqualFold(u.Scheme, "https") || !strings.EqualFold(u.Hostname(), "api.github.com") || u.User != nil || (u.Port() != "" && u.Port() != "443") { + return "" + } + return ResolveToken() +} diff --git a/internal/httpwire/flush.go b/internal/httpwire/flush.go new file mode 100644 index 000000000..36735953f --- /dev/null +++ b/internal/httpwire/flush.go @@ -0,0 +1,38 @@ +package httpwire + +import ( + "net/http" + "reflect" +) + +// FlushResponse prefers error-aware flushing and preserves legacy middleware. +// Only Gin's raw void Flush is bypassed so transport errors remain observable. +func FlushResponse(w http.ResponseWriter) error { + for { + if flusher, ok := w.(interface{ FlushError() error }); ok { + return flusher.FlushError() + } + if flusher, ok := w.(http.Flusher); ok { + if !isGinResponseWriter(w) { + flusher.Flush() + return nil + } + // Match Gin's Flush header accounting before reaching its transport. + w.(interface{ WriteHeaderNow() }).WriteHeaderNow() + } + unwrapper, ok := w.(interface{ Unwrap() http.ResponseWriter }) + if !ok { + return http.NewResponseController(w).Flush() + } + w = unwrapper.Unwrap() + } +} + +func isGinResponseWriter(w http.ResponseWriter) bool { + // Gin's concrete writer is unexported. Interface checks also match wrappers + // embedding gin.ResponseWriter, whose Flush may drain buffers or compressors. + // Restrict this exception to the exact package/type; never inspect its fields. + t := reflect.TypeOf(w) + return t != nil && t.Kind() == reflect.Pointer && + t.Elem().PkgPath() == "github.com/gin-gonic/gin" && t.Elem().Name() == "responseWriter" +} diff --git a/internal/logging/cpa_trace.go b/internal/logging/cpa_trace.go index bc1d243a8..fd2df4c55 100644 --- a/internal/logging/cpa_trace.go +++ b/internal/logging/cpa_trace.go @@ -6,6 +6,7 @@ import ( "time" "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v8/internal/httpwire" ) // CPATraceIDHeader is the downstream response header used to correlate requests with selected credentials. @@ -137,8 +138,13 @@ func (w *cpaTraceResponseWriter) WriteString(data string) (int, error) { } func (w *cpaTraceResponseWriter) Flush() { + _ = w.FlushError() +} + +// FlushError preserves trace headers while forwarding transport errors. +func (w *cpaTraceResponseWriter) FlushError() error { w.applyTraceHeader() - w.ResponseWriter.Flush() + return httpwire.FlushResponse(w.ResponseWriter) } func (w *cpaTraceResponseWriter) applyTraceHeader() { diff --git a/internal/logging/gin_logger.go b/internal/logging/gin_logger.go index a63d41894..804d74870 100644 --- a/internal/logging/gin_logger.go +++ b/internal/logging/gin_logger.go @@ -43,6 +43,11 @@ func GinLogrusLogger() gin.HandlerFunc { start := time.Now() path := c.Request.URL.Path raw := util.MaskSensitiveQuery(c.Request.URL.RawQuery) + // OAuth state is a callback capability; authorization codes and issued + // client IDs must not be included in management access logs either. + if strings.HasPrefix(path, "/v8/management/oauth/") { + raw = "" + } // Only generate request ID for AI API paths var requestID string @@ -89,7 +94,7 @@ func GinLogrusLogger() gin.HandlerFunc { if requestID == "" { requestID = "--------" } - logLine := fmt.Sprintf("%3d | %13v | %15s | %-7s \"%s\"", statusCode, latency, clientIP, method, path) + logLine := fmt.Sprintf("%3d | %13v | %15s | %-7s %q", statusCode, latency, clientIP, method, path) if creditsUsed(c) { logLine += " [credits]" } diff --git a/internal/logging/request_logger_writer.go b/internal/logging/request_logger_writer.go index cb15d5d55..b2bb16ddf 100644 --- a/internal/logging/request_logger_writer.go +++ b/internal/logging/request_logger_writer.go @@ -306,12 +306,18 @@ func (l *FileRequestLogger) generateFilename(url string, requestID ...string) st idPart = fmt.Sprintf("%d", id) } - return fmt.Sprintf("%s-%s-%s.log", sanitized, timestamp, idPart) + // Defense in depth: the sanitizer already removes both separators, but the + // final component must never be able to climb out of the log directory. + name := fmt.Sprintf("%s-%s-%s.log", sanitized, timestamp, idPart) + return filepath.Base(name) } // createUniqueLogFile atomically opens a unique log file within dir. If a file with filename already exists, // it avoids overwriting by injecting an incrementing sequence number before the trailing request ID component. func createUniqueLogFile(dir, filename string) (*os.File, string, error) { + if filename == "" || filename == "." || filename == ".." || strings.ContainsAny(filename, `/\`) || strings.Contains(filename, "..") || filepath.Base(filename) != filename { + return nil, "", fmt.Errorf("invalid log filename: %q", filename) + } ext := filepath.Ext(filename) base := strings.TrimSuffix(filename, ext) idx := strings.LastIndex(base, "-") @@ -353,8 +359,11 @@ func createUniqueLogFile(dir, filename string) (*os.File, string, error) { // Returns: // - string: A sanitized filename func (l *FileRequestLogger) sanitizeForFilename(path string) string { - // Replace slashes with hyphens + // Replace slashes with hyphens. Both separators are replaced because Windows + // treats "\\" as a path separator too: leaving it in place let a crafted + // request path escape the log directory via filepath.Join. sanitized := strings.ReplaceAll(path, "/", "-") + sanitized = strings.ReplaceAll(sanitized, "\\", "-") // Replace colons with hyphens sanitized = strings.ReplaceAll(sanitized, ":", "-") @@ -363,12 +372,17 @@ func (l *FileRequestLogger) sanitizeForFilename(path string) string { reg := regexp.MustCompile(`[<>:"|?*\s]`) sanitized = reg.ReplaceAllString(sanitized, "-") + // Never let a parent-directory sequence survive. + for strings.Contains(sanitized, "..") { + sanitized = strings.ReplaceAll(sanitized, "..", "-") + } + // Remove multiple consecutive hyphens reg = regexp.MustCompile(`-+`) sanitized = reg.ReplaceAllString(sanitized, "-") - // Remove leading/trailing hyphens - sanitized = strings.Trim(sanitized, "-") + // Remove leading/trailing hyphens and dots + sanitized = strings.Trim(sanitized, "-.") // Handle empty result if sanitized == "" { diff --git a/internal/managementasset/updater.go b/internal/managementasset/updater.go index c46670300..57aff6b87 100644 --- a/internal/managementasset/updater.go +++ b/internal/managementasset/updater.go @@ -18,6 +18,7 @@ import ( "time" "github.com/router-for-me/CLIProxyAPI/v8/internal/config" + "github.com/router-for-me/CLIProxyAPI/v8/internal/githubauth" "github.com/router-for-me/CLIProxyAPI/v8/internal/httpfetch" "github.com/router-for-me/CLIProxyAPI/v8/internal/util" sdkconfig "github.com/router-for-me/CLIProxyAPI/v8/sdk/config" @@ -350,7 +351,7 @@ func fetchLatestAsset(ctx context.Context, client *http.Client, releaseURL strin "Accept": "application/vnd.github+json", "User-Agent": httpUserAgent, } - if token := util.ResolveGitHubToken(); token != "" { + if token := githubauth.TokenForURL(releaseURL); token != "" { headers["Authorization"] = "Bearer " + token } diff --git a/internal/misc/antigravity_version.go b/internal/misc/antigravity_version.go index 679b815ac..1e77c2d57 100644 --- a/internal/misc/antigravity_version.go +++ b/internal/misc/antigravity_version.go @@ -18,10 +18,10 @@ import ( const ( // antigravityFallbackVersion is the client version reported when the hub // manifest has not been fetched yet or cannot be reached. Cloud Code rejects - // newer models for clients below 2.9.0, so this floor must stay at or above + // newer models for clients below 2.23.0, so this floor must stay at or above // that version. - antigravityFallbackVersion = "2.9.1" - antigravityHubPlatform = "darwin/arm64" + antigravityFallbackVersion = "2.24.0" + antigravityHubPlatform = "os_type=darwin; arch=arm64" antigravityVersionCacheTTL = 6 * time.Hour antigravityFetchTimeout = 10 * time.Second AntigravityNodeAPIClientUA = "google-api-nodejs-client/10.3.0" @@ -32,6 +32,15 @@ var ( antigravityHubLatestManifestURL = "https://antigravity-hub-auto-updater-974169037036.us-central1.run.app/manifest/latest-arm64-mac.yml" ) +// Changelists belong to specific Hub releases, not to the latest manifest in +// general. Update the fallback release and its entry together. The 2.24.0 value +// was verified in the official darwin-arm archive's webm_encoder build metadata. +// Unknown releases keep their online version without claiming another build's CL. +var antigravityHubChangelists = map[string]string{ + "2.23.0": "998057262", + "2.24.0": "1000335682", +} + type antigravityHubUpdaterManifest struct { Version string `yaml:"version"` } @@ -114,7 +123,12 @@ func AntigravityLatestVersion() string { // AntigravityUserAgent returns the User-Agent string used by the Antigravity Hub family. func AntigravityUserAgent() string { - return fmt.Sprintf("antigravity/hub/%s %s", AntigravityLatestVersion(), antigravityHubPlatform) + version := AntigravityLatestVersion() + metadata := "aidev_client; " + antigravityHubPlatform + if changelist := antigravityHubChangelists[version]; changelist != "" { + metadata += "; cl=" + changelist + } + return fmt.Sprintf("antigravity/hub/%s (%s)", version, metadata) } func isAntigravityFamilyUserAgent(lower string) bool { @@ -138,13 +152,13 @@ func antigravityBaseUserAgent(userAgent string) string { return userAgent } -// AntigravityRequestUserAgent returns the short Antigravity runtime UA used by +// AntigravityRequestUserAgent returns the Antigravity runtime UA used by // generate/stream/model-list requests. func AntigravityRequestUserAgent(userAgent string) string { return antigravityBaseUserAgent(userAgent) } -// AntigravityLoadCodeAssistUserAgent returns the short Antigravity UA used by +// AntigravityLoadCodeAssistUserAgent returns the Antigravity runtime UA used by // loadCodeAssist requests. func AntigravityLoadCodeAssistUserAgent(userAgent string) string { return AntigravityRequestUserAgent(userAgent) diff --git a/internal/misc/header_utils.go b/internal/misc/header_utils.go index 0c3abbf4b..bf50b6d97 100644 --- a/internal/misc/header_utils.go +++ b/internal/misc/header_utils.go @@ -74,6 +74,14 @@ func EnsureHeader(target http.Header, source http.Header, key, defaultValue stri target.Set(key, val) return } + for k, v := range source { + if strings.EqualFold(k, key) && len(v) > 0 { + if val := strings.TrimSpace(v[0]); val != "" { + target.Set(key, val) + return + } + } + } } if strings.TrimSpace(target.Get(key)) != "" { return diff --git a/internal/misc/oauth.go b/internal/misc/oauth.go index 88be2eefe..11a453649 100644 --- a/internal/misc/oauth.go +++ b/internal/misc/oauth.go @@ -26,6 +26,7 @@ func GenerateRandomState() (string, error) { type OAuthCallback struct { Code string State string + ClientID string Error string ErrorDescription string } @@ -76,6 +77,7 @@ func ParseOAuthCallback(input string) (*OAuthCallback, error) { query := parsedURL.Query() code := strings.TrimSpace(query.Get("code")) state := strings.TrimSpace(query.Get("state")) + clientID := strings.TrimSpace(query.Get("client_id")) errCode := strings.TrimSpace(query.Get("error")) errDesc := strings.TrimSpace(query.Get("error_description")) @@ -87,6 +89,9 @@ func ParseOAuthCallback(input string) (*OAuthCallback, error) { if state == "" { state = strings.TrimSpace(fragQuery.Get("state")) } + if clientID == "" { + clientID = strings.TrimSpace(fragQuery.Get("client_id")) + } if errCode == "" { errCode = strings.TrimSpace(fragQuery.Get("error")) } @@ -114,6 +119,7 @@ func ParseOAuthCallback(input string) (*OAuthCallback, error) { return &OAuthCallback{ Code: code, State: state, + ClientID: clientID, Error: errCode, ErrorDescription: errDesc, }, nil diff --git a/internal/redisqueue/plugin.go b/internal/redisqueue/plugin.go index 10ecf27d4..c9d0a944f 100644 --- a/internal/redisqueue/plugin.go +++ b/internal/redisqueue/plugin.go @@ -53,6 +53,18 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec authType = "unknown" } apiKey := strings.TrimSpace(record.APIKey) + accessProvider := strings.TrimSpace(record.AccessProvider) + if accessProvider == "" { + accessProvider = strings.TrimSpace(coreusage.AccessProviderFromContext(ctx)) + } + isNativeKey := record.IsNativeKey + if !isNativeKey { + if nativeExplicit, ok := coreusage.IsNativeKeyFromContext(ctx); ok { + isNativeKey = nativeExplicit + } else if accessProvider != "" { + isNativeKey = coreusage.IsNativeAccessProvider(accessProvider) + } + } requestID := strings.TrimSpace(internallogging.GetRequestID(ctx)) traceID := strings.TrimSpace(record.TraceID) if traceID == "" { @@ -73,6 +85,7 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec if serviceTier == "" { serviceTier = coreusage.ServiceTierFromContext(ctx) } + effectiveServiceTier := strings.TrimSpace(record.EffectiveServiceTier) responseServiceTier := strings.TrimSpace(record.ResponseServiceTier) responseModel := strings.TrimSpace(record.ResponseModel) clientRequestMetadata := internallogging.GetClientRequestMetadata(ctx) @@ -133,28 +146,31 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec } payload, err := json.Marshal(queuedUsageDetail{ - requestDetail: detail, - AccountingVersion: coreusage.TokenAccountingSchemaVersion, - TokenBreakdown: usageDetail.TokenBreakdown, - Provider: provider, - ExecutorType: executorType, - Model: modelName, - Alias: aliasName, - Endpoint: resolveEndpoint(ctx), - AuthType: authType, - APIKey: apiKey, - RequestID: requestID, - ExecutionID: executionID, - TraceID: traceID, - SessionID: sessionID, - ParentSessionID: parentSessionID, - NodeKind: strings.TrimSpace(clientRequestMetadata.NodeKind), - IsFork: clientRequestMetadata.IsFork, - IsCompaction: clientRequestMetadata.IsCompaction, - ReasoningEffort: reasoningEffort, - ServiceTier: serviceTier, - ResponseServiceTier: responseServiceTier, - ResponseModel: responseModel, + requestDetail: detail, + AccountingVersion: coreusage.TokenAccountingSchemaVersion, + TokenBreakdown: usageDetail.TokenBreakdown, + Provider: provider, + ExecutorType: executorType, + Model: modelName, + Alias: aliasName, + Endpoint: resolveEndpoint(ctx), + AuthType: authType, + APIKey: apiKey, + IsNativeKey: isNativeKey, + AccessProvider: accessProvider, + RequestID: requestID, + ExecutionID: executionID, + TraceID: traceID, + SessionID: sessionID, + ParentSessionID: parentSessionID, + NodeKind: strings.TrimSpace(clientRequestMetadata.NodeKind), + IsFork: clientRequestMetadata.IsFork, + IsCompaction: clientRequestMetadata.IsCompaction, + ReasoningEffort: reasoningEffort, + ServiceTier: serviceTier, + EffectiveServiceTier: effectiveServiceTier, + ResponseServiceTier: responseServiceTier, + ResponseModel: responseModel, }) if err != nil { return @@ -164,27 +180,30 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec type queuedUsageDetail struct { requestDetail - AccountingVersion int `json:"accounting_version"` - TokenBreakdown coreusage.TokenBreakdown `json:"token_breakdown"` - Provider string `json:"provider"` - ExecutorType string `json:"executor_type"` - Model string `json:"model"` - Alias string `json:"alias"` - Endpoint string `json:"endpoint"` - AuthType string `json:"auth_type"` - APIKey string `json:"api_key"` - RequestID string `json:"request_id"` - ExecutionID string `json:"execution_id,omitempty"` - TraceID string `json:"trace_id,omitempty"` - SessionID string `json:"session_id,omitempty"` - ParentSessionID string `json:"parent_session_id,omitempty"` - NodeKind string `json:"node_kind,omitempty"` - IsFork bool `json:"is_fork,omitempty"` - IsCompaction bool `json:"is_compaction,omitempty"` - ReasoningEffort string `json:"reasoning_effort"` - ServiceTier string `json:"service_tier"` - ResponseServiceTier string `json:"response_service_tier,omitempty"` - ResponseModel string `json:"response_model,omitempty"` + AccountingVersion int `json:"accounting_version"` + TokenBreakdown coreusage.TokenBreakdown `json:"token_breakdown"` + Provider string `json:"provider"` + ExecutorType string `json:"executor_type"` + Model string `json:"model"` + Alias string `json:"alias"` + Endpoint string `json:"endpoint"` + AuthType string `json:"auth_type"` + APIKey string `json:"api_key"` + IsNativeKey bool `json:"is_native_key"` + AccessProvider string `json:"access_provider,omitempty"` + RequestID string `json:"request_id"` + ExecutionID string `json:"execution_id,omitempty"` + TraceID string `json:"trace_id,omitempty"` + SessionID string `json:"session_id,omitempty"` + ParentSessionID string `json:"parent_session_id,omitempty"` + NodeKind string `json:"node_kind,omitempty"` + IsFork bool `json:"is_fork,omitempty"` + IsCompaction bool `json:"is_compaction,omitempty"` + ReasoningEffort string `json:"reasoning_effort"` + ServiceTier string `json:"service_tier"` + EffectiveServiceTier string `json:"effective_service_tier,omitempty"` + ResponseServiceTier string `json:"response_service_tier,omitempty"` + ResponseModel string `json:"response_model,omitempty"` } type requestDetail struct { @@ -220,12 +239,14 @@ type tokenStats struct { type failDetail struct { StatusCode int `json:"status_code"` Body string `json:"body"` + Code string `json:"code,omitempty"` } func resolveFail(ctx context.Context, record coreusage.Record, failed bool) failDetail { fail := failDetail{ StatusCode: record.Fail.StatusCode, Body: strings.TrimSpace(record.Fail.Body), + Code: strings.TrimSpace(record.Fail.Code), } if !failed { return failDetail{StatusCode: 200} diff --git a/internal/registry/catalog_config.go b/internal/registry/catalog_config.go new file mode 100644 index 000000000..65d8c9e1e --- /dev/null +++ b/internal/registry/catalog_config.go @@ -0,0 +1,29 @@ +package registry + +import ( + "fmt" + "net/url" + "path/filepath" +) + +// CatalogSources overrides the source of each model catalog independently. +// Empty values retain the process default (embedded or official remote catalogs). +type CatalogSources struct { + Catalog string `yaml:"catalog" json:"catalog,omitempty"` + CodexCatalog string `yaml:"codex-catalog" json:"codex-catalog,omitempty"` + DevinCatalog string `yaml:"devin-catalog" json:"devin-catalog,omitempty"` +} + +// Validate rejects relative paths and unsupported URL schemes. +func (m CatalogSources) Validate() error { + for name, source := range map[string]string{"catalog": m.Catalog, "codex-catalog": m.CodexCatalog, "devin-catalog": m.DevinCatalog} { + if source == "" || filepath.IsAbs(source) { + continue + } + parsed, errParse := url.Parse(source) + if errParse != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") || parsed.Hostname() == "" { + return fmt.Errorf("models.%s must be an http(s) URL or an absolute local path", name) + } + } + return nil +} diff --git a/internal/registry/catalog_sources.go b/internal/registry/catalog_sources.go new file mode 100644 index 000000000..f89a153f2 --- /dev/null +++ b/internal/registry/catalog_sources.go @@ -0,0 +1,265 @@ +package registry + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "sync" + "time" + + log "github.com/sirupsen/logrus" +) + +const embeddedCatalogSource = "embed" + +// catalogUpdater serializes source changes with publication, not network I/O. +// Cancellation saves work; the generation check also protects against readers +// that finish after cancellation (including local file reads). +type catalogUpdater struct { + mu sync.Mutex + source string + generation uint64 + cancel context.CancelFunc + ctx context.Context + parentDone <-chan struct{} + fetch func(context.Context, string) ([]byte, error) + publish func([]byte) ([]string, error) +} + +func (u *catalogUpdater) configure(ctx context.Context, source string) { + u.mu.Lock() + defer u.mu.Unlock() + if u.cancel != nil && u.source == source && u.parentDone == ctx.Done() && u.ctx.Err() == nil { + return + } + if u.cancel != nil { + u.cancel() + } + u.generation++ + u.source = source + refreshCtx, cancel := context.WithCancel(ctx) + u.cancel = cancel + u.ctx = refreshCtx + u.parentDone = ctx.Done() + if source == "disabled" { + return + } + generation := u.generation + go func() { + u.refresh(refreshCtx, source, generation) + if source == embeddedCatalogSource { + return + } + ticker := time.NewTicker(ModelsRefreshInterval) + defer ticker.Stop() + for { + select { + case <-refreshCtx.Done(): + return + case <-ticker.C: + u.refresh(refreshCtx, source, generation) + } + } + }() +} + +func (u *catalogUpdater) refresh(ctx context.Context, source string, generation uint64) { + data, errFetch := u.fetch(ctx, source) + if errFetch != nil { + if ctx.Err() == nil { + log.Warn("model catalog refresh failed; keeping last valid catalog") + } + return + } + u.mu.Lock() + if generation != u.generation || ctx.Err() != nil { + u.mu.Unlock() + return + } + changed, errPublish := u.publish(data) + u.mu.Unlock() + if errPublish != nil { + log.Warn("model catalog rejected; keeping last valid catalog") + return + } + notifyModelRefresh(changed) +} + +func readCatalogSource(ctx context.Context, source string) ([]byte, error) { + var reader io.ReadCloser + if filepath.IsAbs(source) { + file, errOpen := os.Open(source) + if errOpen != nil { + return nil, errOpen + } + reader = file + } else { + req, errRequest := http.NewRequestWithContext(ctx, http.MethodGet, source, nil) + if errRequest != nil { + return nil, errRequest + } + resp, errDo := http.DefaultClient.Do(req) + if errDo != nil { + return nil, errDo + } + reader = resp.Body + if resp.StatusCode != http.StatusOK { + if errClose := reader.Close(); errClose != nil { + log.Debug("catalog response close failed") + } + return nil, fmt.Errorf("catalog HTTP status %d", resp.StatusCode) + } + } + defer func() { + if errClose := reader.Close(); errClose != nil { + log.Debug("catalog reader close failed") + } + }() + data, errRead := io.ReadAll(io.LimitReader(reader, maxCodexClientModelsSize+1)) + if errRead != nil { + return nil, errRead + } + if len(data) > maxCodexClientModelsSize { + return nil, fmt.Errorf("catalog exceeds size limit") + } + return data, nil +} + +func catalogFetcher(embedded []byte, urls []string, validate func([]byte) error) func(context.Context, string) ([]byte, error) { + return func(ctx context.Context, source string) ([]byte, error) { + if source == embeddedCatalogSource { + return embedded, nil + } + sources := urls + if source != "" { + sources = []string{source} + } + for _, candidate := range sources { + data, errRead := readCatalogSource(ctx, candidate) + if errRead == nil && validate(data) == nil { + return data, nil + } + if ctx.Err() != nil { + break + } + } + return nil, fmt.Errorf("no valid model catalog source") + } +} + +func validateCatalogBytes(data []byte) error { + var parsed staticModelsJSON + if errUnmarshal := json.Unmarshal(data, &parsed); errUnmarshal != nil { + return errUnmarshal + } + return validateModelsCatalog(&parsed) +} + +func publishCatalogBytes(data []byte) ([]string, error) { + var parsed staticModelsJSON + if errUnmarshal := json.Unmarshal(data, &parsed); errUnmarshal != nil { + return nil, errUnmarshal + } + if errValidate := validateModelsCatalog(&parsed); errValidate != nil { + return nil, errValidate + } + old := getModels() + if len(parsed.Meta) == 0 && old != nil { + parsed.Meta = old.Meta + } + changed := detectChangedProviders(old, &parsed) + modelsCatalogStore.mu.Lock() + modelsCatalogStore.data = &parsed + modelsCatalogStore.mu.Unlock() + return changed, nil +} + +var generalCatalogUpdater = &catalogUpdater{ + fetch: catalogFetcher(embeddedModelsJSON, modelsURLs, validateCatalogBytes), + publish: publishCatalogBytes, +} +var codexCatalogUpdater = &catalogUpdater{ + fetch: catalogFetcher(embeddedCodexClientModelsJSON, codexClientModelsURLs, ValidateCodexClientModelsJSON), + publish: func(data []byte) ([]string, error) { + _, errLoad := loadCodexClientModelsFromBytes(data, "catalog") + return nil, errLoad + }, +} +var devinCatalogUpdater = &catalogUpdater{ + fetch: catalogFetcher(embeddedDevinModelsJSON, devinModelsURLs, validateDevinCatalogBytes), + publish: func(data []byte) ([]string, error) { + changed, errLoad := loadDevinModelsFromBytes(data, "catalog") + if errLoad == nil && changed { + return []string{"devin"}, nil + } + return nil, errLoad + }, +} + +var catalogRuntime struct { + sync.Mutex + ctx context.Context + local bool +} + +// SetLocalModelCatalogs sets the process default before the service starts. +// Explicit catalog sources still take precedence over this default. +func SetLocalModelCatalogs(local bool) { + catalogRuntime.Lock() + defer catalogRuntime.Unlock() + catalogRuntime.local = local +} + +// StartModelCatalogUpdaters binds process-wide readers to the service lifetime. +func StartModelCatalogUpdaters(ctx context.Context, sources CatalogSources, home bool) { + catalogRuntime.Lock() + defer catalogRuntime.Unlock() + catalogRuntime.ctx = ctx + configureCatalogs(ctx, sources, catalogRuntime.local, home) +} + +// UpdateModelCatalogSources updates an already-started process policy on reload. +// Updates before service startup or after shutdown do not start readers. +func UpdateModelCatalogSources(sources CatalogSources, home bool) { + catalogRuntime.Lock() + defer catalogRuntime.Unlock() + if catalogRuntime.ctx != nil && catalogRuntime.ctx.Err() == nil { + configureCatalogs(catalogRuntime.ctx, sources, catalogRuntime.local, home) + } +} + +func configureCatalogs(ctx context.Context, sources CatalogSources, local, home bool) { + if errValidate := sources.Validate(); errValidate != nil { + log.WithError(errValidate).Warn("invalid catalog sources") + return + } + sources = effectiveCatalogSources(sources, local, home) + generalCatalogUpdater.configure(ctx, sources.Catalog) + codexCatalogUpdater.configure(ctx, sources.CodexCatalog) + devinCatalogUpdater.configure(ctx, sources.DevinCatalog) +} + +func effectiveCatalogSources(sources CatalogSources, local, home bool) CatalogSources { + selectSource := func(source string) string { + if source == "" && local { + return embeddedCatalogSource + } + return source + } + general := selectSource(sources.Catalog) + devin := selectSource(sources.DevinCatalog) + if home { + general = "disabled" + devin = "disabled" + } + return CatalogSources{Catalog: general, CodexCatalog: selectSource(sources.CodexCatalog), DevinCatalog: devin} +} + +func validateDevinCatalogBytes(data []byte) error { + _, errValidate := ValidateDevinModelsJSON(data) + return errValidate +} diff --git a/internal/registry/codex_client_models.go b/internal/registry/codex_client_models.go index 370abf296..46445a8cc 100644 --- a/internal/registry/codex_client_models.go +++ b/internal/registry/codex_client_models.go @@ -96,9 +96,6 @@ func ValidateCodexClientModelsJSON(data []byte) error { return fmt.Errorf("Codex client model catalog model %q: %w", slug, err) } } - if _, ok := seen["gpt-5.5"]; !ok { - return fmt.Errorf("Codex client model catalog is missing default template %q", "gpt-5.5") - } return nil } diff --git a/internal/registry/codex_client_models_updater.go b/internal/registry/codex_client_models_updater.go index c556daad0..f7b7b0ff7 100644 --- a/internal/registry/codex_client_models_updater.go +++ b/internal/registry/codex_client_models_updater.go @@ -2,10 +2,6 @@ package registry import ( "context" - "io" - "net/http" - "sync" - "time" log "github.com/sirupsen/logrus" ) @@ -17,31 +13,11 @@ var codexClientModelsURLs = []string{ "https://models.router-for.me/codex_client_models.json", } -var codexClientModelsUpdaterOnce sync.Once - // StartCodexClientModelsUpdater starts a background updater that fetches the // Codex client model catalog immediately and then refreshes it every 3 hours. // Safe to call multiple times; only one updater will run. func StartCodexClientModelsUpdater(ctx context.Context) { - codexClientModelsUpdaterOnce.Do(func() { - go runCodexClientModelsUpdater(ctx) - }) -} - -func runCodexClientModelsUpdater(ctx context.Context) { - tryRefreshCodexClientModels(ctx, "startup Codex client model refresh") - - ticker := time.NewTicker(modelsRefreshInterval) - defer ticker.Stop() - log.Infof("periodic Codex client model refresh started (interval=%s)", modelsRefreshInterval) - for { - select { - case <-ctx.Done(): - return - case <-ticker.C: - tryRefreshCodexClientModels(ctx, "periodic Codex client model refresh") - } - } + codexCatalogUpdater.configure(ctx, "") } func tryRefreshCodexClientModels(ctx context.Context, label string) { @@ -64,51 +40,14 @@ func tryRefreshCodexClientModels(ctx context.Context, label string) { } func fetchCodexClientModelsFromRemote(ctx context.Context) ([]byte, string) { - client := &http.Client{Timeout: modelsFetchTimeout} - for _, sourceURL := range codexClientModelsURLs { - reqCtx, cancel := context.WithTimeout(ctx, modelsFetchTimeout) - req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, sourceURL, nil) - if err != nil { - cancel() - log.Debugf("Codex client models fetch request creation failed for %s: %v", sourceURL, err) - continue - } - - resp, err := client.Do(req) - if err != nil { - cancel() - log.Debugf("Codex client models fetch failed from %s: %v", sourceURL, err) - continue - } - if resp.StatusCode != http.StatusOK { - if errClose := resp.Body.Close(); errClose != nil { - log.Debugf("Codex client models response close failed for %s: %v", sourceURL, errClose) - } - cancel() - log.Debugf("Codex client models fetch returned %d from %s", resp.StatusCode, sourceURL) - continue - } - - data, errRead := io.ReadAll(io.LimitReader(resp.Body, maxCodexClientModelsSize+1)) - errClose := resp.Body.Close() - cancel() - if errRead != nil { - log.Debugf("Codex client models fetch read error from %s: %v", sourceURL, errRead) - continue - } - if errClose != nil { - log.Debugf("Codex client models response close failed for %s: %v", sourceURL, errClose) - continue - } - if len(data) > maxCodexClientModelsSize { - log.Warnf("Codex client models fetch from %s exceeded %d bytes", sourceURL, maxCodexClientModelsSize) - continue + for _, source := range codexClientModelsURLs { + data, errRead := readCatalogSource(ctx, source) + if errRead == nil && ValidateCodexClientModelsJSON(data) == nil { + return data, source } - if err := ValidateCodexClientModelsJSON(data); err != nil { - log.Warnf("Codex client models validate failed from %s: %v", sourceURL, err) - continue + if ctx.Err() != nil { + break } - return data, sourceURL } return nil, "" } diff --git a/internal/registry/devin_models.go b/internal/registry/devin_models.go index 59af4aefc..5e1936555 100644 --- a/internal/registry/devin_models.go +++ b/internal/registry/devin_models.go @@ -215,6 +215,15 @@ var devinCompoundSuffixes = []struct { {suffix: "-thinking", effort: ""}, {suffix: "-max-1m", effort: "max", readd: "-1m"}, {suffix: "-none-1m", effort: "none", readd: "-1m"}, + // Lightning tier keeps the effort before the -lightning suffix upstream + // (e.g. swe-2-high-lightning), so it aggregates as -lightning. + {suffix: "-minimal-lightning", effort: "minimal", readd: "-lightning"}, + {suffix: "-low-lightning", effort: "low", readd: "-lightning"}, + {suffix: "-medium-lightning", effort: "medium", readd: "-lightning"}, + {suffix: "-high-lightning", effort: "high", readd: "-lightning"}, + {suffix: "-xhigh-lightning", effort: "xhigh", readd: "-lightning"}, + {suffix: "-max-lightning", effort: "max", readd: "-lightning"}, + {suffix: "-none-lightning", effort: "none", readd: "-lightning"}, } var devinSimpleEffortSuffixes = []struct { @@ -237,7 +246,7 @@ var devinDisplayNameSuffixes = []string{ " Low Thinking", " Medium Thinking", " High Thinking", " XHigh Thinking", " Max Thinking", " No Thinking", " Low", " Medium", " High", " XHigh", " Max", " None", " Minimal", - " Thinking", " Fast", + " Thinking", " Fast", " Lightning", } var devinLevelOrder = map[string]int{ @@ -408,6 +417,13 @@ func aggregateDevinModels(models []*ModelInfo) []*ModelInfo { entry := aggregated[id] m := entry.model + // Lightning display names strip both the effort word and "Lightning"; + // re-append it so devin/swe-2-lightning shows as "SWE-2 Lightning". + if strings.HasSuffix(strings.ToLower(strings.TrimPrefix(id, "devin/")), "-lightning") && + !strings.HasSuffix(strings.ToLower(m.DisplayName), "lightning") { + m.DisplayName = strings.TrimSpace(m.DisplayName + " Lightning") + } + if len(entry.levels) > 0 { lvls := make([]string, 0, len(entry.levels)) for l := range entry.levels { diff --git a/internal/registry/devin_models_updater.go b/internal/registry/devin_models_updater.go index fd4d1cd78..5ed460d68 100644 --- a/internal/registry/devin_models_updater.go +++ b/internal/registry/devin_models_updater.go @@ -2,10 +2,6 @@ package registry import ( "context" - "io" - "net/http" - "sync" - "time" log "github.com/sirupsen/logrus" ) @@ -17,31 +13,11 @@ var devinModelsURLs = []string{ "https://models.router-for.me/devin_models.json", } -var devinModelsUpdaterOnce sync.Once - // StartDevinModelsUpdater starts a background updater that fetches the // Devin model catalog immediately and refreshes it every 3 hours. // Safe to call multiple times; only one updater runs. func StartDevinModelsUpdater(ctx context.Context) { - devinModelsUpdaterOnce.Do(func() { - go runDevinModelsUpdater(ctx) - }) -} - -func runDevinModelsUpdater(ctx context.Context) { - tryRefreshDevinModels(ctx, "startup Devin model refresh") - - ticker := time.NewTicker(modelsRefreshInterval) - defer ticker.Stop() - log.Infof("periodic Devin model refresh started (interval=%s)", modelsRefreshInterval) - for { - select { - case <-ctx.Done(): - return - case <-ticker.C: - tryRefreshDevinModels(ctx, "periodic Devin model refresh") - } - } + devinCatalogUpdater.configure(ctx, "") } func tryRefreshDevinModels(ctx context.Context, label string) { @@ -64,45 +40,14 @@ func tryRefreshDevinModels(ctx context.Context, label string) { } func fetchDevinModelsFromRemote(ctx context.Context) ([]byte, string) { - client := &http.Client{Timeout: modelsFetchTimeout} - for _, sourceURL := range devinModelsURLs { - reqCtx, cancel := context.WithTimeout(ctx, modelsFetchTimeout) - req, errReq := http.NewRequestWithContext(reqCtx, http.MethodGet, sourceURL, nil) - if errReq != nil { - cancel() - log.Warnf("devin models updater: invalid request for %s: %v", sourceURL, errReq) - continue + for _, source := range devinModelsURLs { + data, errRead := readCatalogSource(ctx, source) + if errRead == nil && validateDevinCatalogBytes(data) == nil { + return data, source } - - resp, errDo := client.Do(req) - if errDo != nil { - cancel() - log.Warnf("devin models updater: fetch failed from %s: %v", sourceURL, errDo) - continue + if ctx.Err() != nil { + break } - - if resp.StatusCode != http.StatusOK { - if errClose := resp.Body.Close(); errClose != nil { - log.Warnf("devin models updater: response close failed for %s: %v", sourceURL, errClose) - } - cancel() - log.Warnf("devin models updater: unexpected status %d from %s", resp.StatusCode, sourceURL) - continue - } - - body, errRead := io.ReadAll(io.LimitReader(resp.Body, maxDevinModelsSize)) - errClose := resp.Body.Close() - cancel() - if errRead != nil { - log.Warnf("devin models updater: read failed from %s: %v", sourceURL, errRead) - continue - } - if errClose != nil { - log.Warnf("devin models updater: response close failed for %s: %v", sourceURL, errClose) - continue - } - - return body, sourceURL } return nil, "" } diff --git a/internal/registry/model_definitions.go b/internal/registry/model_definitions.go index 126ee4795..b1681e779 100644 --- a/internal/registry/model_definitions.go +++ b/internal/registry/model_definitions.go @@ -18,6 +18,8 @@ const ( xaiBuiltinVideoModelID = "grok-imagine-video" xaiBuiltinVideo15ModelID = "grok-imagine-video-1.5" xaiBuiltinVideo15PreviewID = "grok-imagine-video-1.5-preview" + xaiBuiltinSpeechModelID = "grok-tts" + xaiBuiltinSpeechVoiceModelID = "grok-voice-tts-1.0" ) // staticModelsJSON mirrors the top-level structure of models.json. @@ -273,10 +275,10 @@ func WithCodexBuiltins(models []*ModelInfo) []*ModelInfo { }) } -// WithXAIBuiltins injects hard-coded xAI image/video model definitions that should -// not depend on remote models.json updates. +// WithXAIBuiltins injects hard-coded xAI image, video, and speech model definitions +// that should not depend on remote models.json updates. func WithXAIBuiltins(models []*ModelInfo) []*ModelInfo { - return upsertModelInfos(models, xaiBuiltinImageModelInfo(), xaiBuiltinImageQualityModelInfo(), xaiBuiltinImage20ModelInfo(), xaiBuiltinVideoModelInfo(), xaiBuiltinVideo15ModelInfo(), xaiBuiltinVideo15PreviewModelInfo()) + return upsertModelInfos(models, xaiBuiltinImageModelInfo(), xaiBuiltinImageQualityModelInfo(), xaiBuiltinImage20ModelInfo(), xaiBuiltinVideoModelInfo(), xaiBuiltinVideo15ModelInfo(), xaiBuiltinVideo15PreviewModelInfo(), xaiBuiltinSpeechModelInfo(), xaiBuiltinSpeechVoiceModelInfo()) } func normalizeAntigravityCapabilityModelID(modelID string) string { @@ -425,6 +427,32 @@ func xaiBuiltinVideo15PreviewModelInfo() *ModelInfo { } } +func xaiBuiltinSpeechModelInfo() *ModelInfo { + return &ModelInfo{ + ID: xaiBuiltinSpeechModelID, + Object: "model", + Created: 1773619200, // 2026-03-16 + OwnedBy: "xai", + Type: "xai", + DisplayName: "Grok TTS", + Name: xaiBuiltinSpeechModelID, + Description: "xAI Grok unary text-to-speech model.", + } +} + +func xaiBuiltinSpeechVoiceModelInfo() *ModelInfo { + return &ModelInfo{ + ID: xaiBuiltinSpeechVoiceModelID, + Object: "model", + Created: 1773619200, // 2026-03-16 + OwnedBy: "xai", + Type: "xai", + DisplayName: "Grok Voice TTS 1.0", + Name: xaiBuiltinSpeechVoiceModelID, + Description: "xAI Grok unary text-to-speech model.", + } +} + func upsertModelInfos(models []*ModelInfo, extras ...*ModelInfo) []*ModelInfo { if len(extras) == 0 { return models @@ -535,11 +563,21 @@ func LookupStaticModelInfoByChannel(modelID, channel string) *ModelInfo { if modelID == "" { return nil } - for _, model := range GetStaticModelDefinitionsByChannel(channel) { + models := GetStaticModelDefinitionsByChannel(channel) + for _, model := range models { if model != nil && model.ID == modelID { return cloneModelInfo(model) } } + if strings.EqualFold(strings.TrimSpace(channel), "claude") { + for _, familyID := range claudeFamilyModelIDs(modelID) { + for _, model := range models { + if model != nil && model.ID == familyID { + return cloneModelInfo(model) + } + } + } + } return nil } @@ -565,9 +603,8 @@ func LookupStaticModelInfo(modelID string) *ModelInfo { data.Kimi, data.Antigravity, data.XAI, - data.Devin, - staticDevinModels, - data.Meta, + GetDevinModels(), + GetMetaModels(), } for _, models := range allModels { for _, m := range models { @@ -577,5 +614,45 @@ func LookupStaticModelInfo(modelID string) *ModelInfo { } } + // Dated and bracketed Claude variants (claude-sonnet-5-5-20261001, + // claude-opus-5-5[1m]) share the capabilities of their catalog entry. + for _, familyID := range claudeFamilyModelIDs(modelID) { + for _, m := range data.Claude { + if m != nil && m.ID == familyID { + return cloneModelInfo(m) + } + } + } + return nil } + +// claudeFamilyModelIDs returns the catalog IDs a Claude model variant may belong +// to: without a "[...]" context suffix, then also without a "-YYYYMMDD" date. +func claudeFamilyModelIDs(modelID string) []string { + if !strings.HasPrefix(modelID, "claude-") { + return nil + } + var ids []string + base := modelID + if i := strings.IndexByte(base, '['); i > 0 && strings.HasSuffix(base, "]") { + base = base[:i] + ids = append(ids, base) + } + if n := len(base); n > 9 && base[n-9] == '-' && isEightDigits(base[n-8:]) { + ids = append(ids, base[:n-9]) + } + return ids +} + +func isEightDigits(s string) bool { + if len(s) != 8 { + return false + } + for _, r := range s { + if r < '0' || r > '9' { + return false + } + } + return true +} diff --git a/internal/registry/model_registry.go b/internal/registry/model_registry.go index 8a4bf7c15..0c0a84319 100644 --- a/internal/registry/model_registry.go +++ b/internal/registry/model_registry.go @@ -30,6 +30,10 @@ type NativeCapabilities struct { // WebSearch reports explicit per-model support for native web search. // nil means the catalog has not established support either way. WebSearch *bool `json:"web_search,omitempty"` + // ToolSearch reports whether every credential behind a Home route can take + // Codex tool_search as Anthropic native tool search. Home computes it per + // route; nil means unknown and is treated as unsupported. + ToolSearch *bool `json:"tool_search,omitempty"` } // ModelInfo represents information about an available model @@ -146,6 +150,9 @@ type ThinkingSupport struct { ZeroAllowed bool `json:"zero_allowed,omitempty" yaml:"zero-allowed,omitempty"` // DynamicAllowed indicates whether -1 is a valid value (dynamic thinking budget). DynamicAllowed bool `json:"dynamic_allowed,omitempty" yaml:"dynamic-allowed,omitempty"` + // DisableType is the Claude thinking.type sent when thinking is turned off. + // Empty means "disabled"; claude-sonnet-5-5 needs "between_tools". + DisableType string `json:"disable_type,omitempty" yaml:"disable-type,omitempty"` // Levels defines discrete reasoning effort levels (e.g., "low", "medium", "high"). // When set, the model uses level-based reasoning instead of token budgets. Levels []string `json:"levels,omitempty" yaml:"levels,omitempty"` @@ -434,6 +441,58 @@ func (r *ModelRegistry) triggerModelsUnregistered(provider, clientID string) { func (r *ModelRegistry) RegisterClient(clientID, clientProvider string, models []*ModelInfo) { r.mutex.Lock() defer r.mutex.Unlock() + r.registerClientLocked(clientID, clientProvider, models, false) +} + +// RegisterClientIfEpochCurrent registers a client's models only if its +// registration epoch still equals expectedEpoch and publish is non-nil. The +// fence evaluation and the registration runs atomically under the registry +// mutex. publish must be fast and must not call back into the registry or +// block. The returned value reports whether the registration was applied. +func (r *ModelRegistry) RegisterClientIfEpochCurrent(clientID, clientProvider string, expectedEpoch uint64, models []*ModelInfo, publish func() bool) bool { + r.mutex.Lock() + defer r.mutex.Unlock() + if r.clientEpochs[clientID] != expectedEpoch || publish == nil || !publish() { + return false + } + r.registerClientLocked(clientID, clientProvider, models, false) + return true +} + +// UnregisterClientIfEpochCurrent removes a client only if its registration +// epoch still equals expectedEpoch and publish is non-nil. The fence +// evaluation and the removal runs atomically under the registry mutex. +// publish must be fast and must not call back into the registry or block. +// The returned value reports whether the removal was applied. +func (r *ModelRegistry) UnregisterClientIfEpochCurrent(clientID string, expectedEpoch uint64, publish func() bool) bool { + r.mutex.Lock() + defer r.mutex.Unlock() + if r.clientEpochs[clientID] != expectedEpoch || publish == nil || !publish() { + return false + } + r.unregisterClientInternal(clientID) + r.invalidateAvailableModelsCacheLocked() + return true +} + +// ReplaceClientModels replaces a client's models only if its registration epoch matches. +// Epoch zero and epochs for empty catalogs are valid. A successful replacement advances +// the epoch and preserves quota, suspension state and the projection generation +// watermark for an existing registration with the same provider. +// Empty models unregister the client; a mismatched epoch leaves the registry unchanged. +// The successful epoch is returned under the same lock as the replacement. +func (r *ModelRegistry) ReplaceClientModels(clientID, clientProvider string, expectedEpoch uint64, models []*ModelInfo) (uint64, bool) { + r.mutex.Lock() + defer r.mutex.Unlock() + if r.clientEpochs[clientID] != expectedEpoch { + return 0, false + } + r.registerClientLocked(clientID, clientProvider, models, true) + return r.clientEpochs[clientID], true +} + +// registerClientLocked reconciles a client's models while the registry mutex is held. +func (r *ModelRegistry) registerClientLocked(clientID, clientProvider string, models []*ModelInfo, preserveState bool) { r.ensureAvailableModelsCacheLocked() if r.clientGenerations == nil { @@ -472,16 +531,19 @@ func (r *ModelRegistry) RegisterClient(clientID, clientProvider string, models [ return } - // Monotonically increment client registration epoch and reset generation to 0. + oldModels, hadExisting := r.clientModels[clientID] + oldProvider := r.clientProviders[clientID] + providerChanged := oldProvider != provider + + // Retaining scheduling state also retains its projection watermark. An older + // request result may read the new epoch after capturing its auth snapshot. r.clientEpochs[clientID]++ - r.clientGenerations[clientID] = uint64(0) + if !preserveState || !hadExisting || providerChanged { + r.clientGenerations[clientID] = uint64(0) + } r.registrationEpoch.Add(1) now := time.Now() - - oldModels, hadExisting := r.clientModels[clientID] - oldProvider := r.clientProviders[clientID] - providerChanged := oldProvider != provider if !hadExisting { // Pure addition path. for _, modelID := range rawModelIDs { @@ -562,7 +624,7 @@ func (r *ModelRegistry) RegisterClient(clientID, clientProvider string, models [ for _, id := range removed { oldCount := oldCounts[id] for i := 0; i < oldCount; i++ { - r.removeModelRegistration(clientID, id, oldProvider, now) + r.removeModelRegistration(clientID, id, oldProvider, now, false) } } @@ -573,7 +635,7 @@ func (r *ModelRegistry) RegisterClient(clientID, clientProvider string, models [ } overage := oldCount - newCount for i := 0; i < overage; i++ { - r.removeModelRegistration(clientID, id, oldProvider, now) + r.removeModelRegistration(clientID, id, oldProvider, now, preserveState && !providerChanged) } } @@ -613,13 +675,10 @@ func (r *ModelRegistry) RegisterClient(clientID, clientProvider string, models [ reg.InfoByProvider[oldProvider].SupportsWebSearch = r.hasClientSupportingWebSearchLocked(id, oldProvider, clientID) } reg.LastUpdated = now - // Re-registering an existing client/model binding starts a fresh registry - // snapshot for that binding. Cooldown and suspension are transient - // scheduling state and must not survive this reconciliation step. - if reg.QuotaExceededClients != nil { + // Only conditional replacements preserve scheduling state for retained + // bindings with the same provider. Ordinary registration starts fresh. + if !preserveState || providerChanged || oldCounts[id] == 0 { delete(reg.QuotaExceededClients, clientID) - } - if reg.SuspendedClients != nil { delete(reg.SuspendedClients, clientID) } if providerChanged && provider != "" { @@ -713,17 +772,16 @@ func (r *ModelRegistry) addModelRegistration(modelID, provider string, model *Mo log.Debugf("Registered new model %s from provider %s", modelID, provider) } -func (r *ModelRegistry) removeModelRegistration(clientID, modelID, provider string, now time.Time) { +// removeModelRegistration may preserve state when only a duplicate binding is removed. +func (r *ModelRegistry) removeModelRegistration(clientID, modelID, provider string, now time.Time, preserveState bool) { registration, exists := r.models[modelID] if !exists { return } registration.Count-- registration.LastUpdated = now - if registration.QuotaExceededClients != nil { + if !preserveState { delete(registration.QuotaExceededClients, clientID) - } - if registration.SuspendedClients != nil { delete(registration.SuspendedClients, clientID) } if registration.Count < 0 { @@ -1281,6 +1339,16 @@ func (r *ModelRegistry) GetAvailableModels(handlerType string) []map[string]any return models } +// isQuotaSuspension reports whether a suspension is a quota cooldown. A model in a +// quota cooldown stays listed: the limit is temporary, and a request meanwhile gets a +// 429 with its retry time instead of an unknown model. A credential-wide quota +// ("credential_quota") suspends every model of that credential, so it is a quota +// cooldown too; counting it as another failure hides all of them when no other +// credential serves them. +func isQuotaSuspension(reason string) bool { + return strings.EqualFold(reason, "quota") || strings.EqualFold(reason, "credential_quota") +} + func modelRegistrationAvailability(registration *ModelRegistration, now time.Time) (bool, time.Time) { if registration == nil { return false, time.Time{} @@ -1304,23 +1372,23 @@ func modelRegistrationAvailability(registration *ModelRegistration, now time.Tim cooldownSuspended := 0 otherSuspended := 0 - quotaAndOtherSuspended := 0 + quotaAndSuspended := 0 if registration.SuspendedClients != nil { for clientID, reason := range registration.SuspendedClients { - if strings.EqualFold(reason, "quota") { + if quotaTime := registration.QuotaExceededClients[clientID]; quotaTime != nil && now.Before(quotaTime.Add(modelQuotaExceededWindow)) { + quotaAndSuspended++ + } + if isQuotaSuspension(reason) { cooldownSuspended++ continue } otherSuspended++ - if quotaTime := registration.QuotaExceededClients[clientID]; quotaTime != nil && now.Before(quotaTime.Add(modelQuotaExceededWindow)) { - quotaAndOtherSuspended++ - } } } - // A credential-wide quota can mark the same client both quota-exceeded and + // A quota cooldown can mark the same client both quota-exceeded and // suspended. Count that unavailable client only once. - effectiveClients := availableClients - expiredClients - otherSuspended + quotaAndOtherSuspended + effectiveClients := availableClients - expiredClients - cooldownSuspended - otherSuspended + quotaAndSuspended if effectiveClients < 0 { effectiveClients = 0 } @@ -1481,7 +1549,7 @@ func (r *ModelRegistry) GetAvailableModelsByProvider(provider string) []*ModelIn expiredClients := 0 cooldownSuspended := 0 otherSuspended := 0 - quotaAndOtherSuspended := 0 + quotaAndSuspended := 0 if ok && registration != nil { if registration.QuotaExceededClients != nil { for clientID, quotaTime := range registration.QuotaExceededClients { @@ -1504,20 +1572,20 @@ func (r *ModelRegistry) GetAvailableModelsByProvider(provider string) []*ModelIn if p, okProvider := r.clientProviders[clientID]; !okProvider || p != provider { continue } - if strings.EqualFold(reason, "quota") { + if quotaTime := registration.QuotaExceededClients[clientID]; quotaTime != nil && now.Before(quotaTime.Add(modelQuotaExceededWindow)) { + quotaAndSuspended++ + } + if isQuotaSuspension(reason) { cooldownSuspended++ continue } otherSuspended++ - if quotaTime := registration.QuotaExceededClients[clientID]; quotaTime != nil && now.Before(quotaTime.Add(modelQuotaExceededWindow)) { - quotaAndOtherSuspended++ - } } } } availableClients := entry.count - effectiveClients := availableClients - expiredClients - otherSuspended + quotaAndOtherSuspended + effectiveClients := availableClients - expiredClients - cooldownSuspended - otherSuspended + quotaAndSuspended if effectiveClients < 0 { effectiveClients = 0 } @@ -1735,11 +1803,19 @@ func (r *ModelRegistry) convertModelToMap(model *ModelInfo, handlerType string) if model.Description != "" { result["description"] = model.Description } - if model.InputTokenLimit > 0 { - result["inputTokenLimit"] = model.InputTokenLimit + inputTokenLimit := model.InputTokenLimit + if inputTokenLimit <= 0 { + inputTokenLimit = model.ContextLength + } + if inputTokenLimit > 0 { + result["inputTokenLimit"] = inputTokenLimit + } + outputTokenLimit := model.OutputTokenLimit + if outputTokenLimit <= 0 { + outputTokenLimit = model.MaxCompletionTokens } - if model.OutputTokenLimit > 0 { - result["outputTokenLimit"] = model.OutputTokenLimit + if outputTokenLimit > 0 { + result["outputTokenLimit"] = outputTokenLimit } if len(model.SupportedGenerationMethods) > 0 { result["supportedGenerationMethods"] = append([]string(nil), model.SupportedGenerationMethods...) diff --git a/internal/registry/model_updater.go b/internal/registry/model_updater.go index 5a21f4882..59a0329b5 100644 --- a/internal/registry/model_updater.go +++ b/internal/registry/model_updater.go @@ -5,8 +5,6 @@ import ( _ "embed" "encoding/json" "fmt" - "io" - "net/http" "reflect" "strings" "sync" @@ -15,10 +13,8 @@ import ( log "github.com/sirupsen/logrus" ) -const ( - modelsFetchTimeout = 30 * time.Second - modelsRefreshInterval = 3 * time.Hour -) +// ModelsRefreshInterval is the shared cadence for model catalogs and account entitlements. +const ModelsRefreshInterval = 3 * time.Hour var modelsURLs = []string{ "https://raw.githubusercontent.com/router-for-me/models/refs/heads/main/models.json", @@ -35,8 +31,6 @@ type modelStore struct { var modelsCatalogStore = &modelStore{} -var updaterOnce sync.Once - // ModelRefreshCallback is invoked when startup or periodic model refresh detects changes. // changedProviders contains the provider names whose model definitions changed. type ModelRefreshCallback func(changedProviders []string) @@ -76,122 +70,7 @@ func init() { // immediately on startup and then refreshes the model catalog every 3 hours. // Safe to call multiple times; only one updater will run. func StartModelsUpdater(ctx context.Context) { - updaterOnce.Do(func() { - go runModelsUpdater(ctx) - }) -} - -func runModelsUpdater(ctx context.Context) { - tryStartupRefresh(ctx) - periodicRefresh(ctx) -} - -func periodicRefresh(ctx context.Context) { - ticker := time.NewTicker(modelsRefreshInterval) - defer ticker.Stop() - log.Infof("periodic model refresh started (interval=%s)", modelsRefreshInterval) - for { - select { - case <-ctx.Done(): - return - case <-ticker.C: - tryPeriodicRefresh(ctx) - } - } -} - -// tryPeriodicRefresh fetches models from remote, compares with the current -// catalog, and notifies the registered callback if any provider changed. -func tryPeriodicRefresh(ctx context.Context) { - tryRefreshModels(ctx, "periodic model refresh") -} - -// tryStartupRefresh fetches models from remote in the background during -// process startup. It uses the same change detection as periodic refresh so -// existing auth registrations can be updated after the callback is registered. -func tryStartupRefresh(ctx context.Context) { - tryRefreshModels(ctx, "startup model refresh") -} - -func tryRefreshModels(ctx context.Context, label string) { - oldData := getModels() - - parsed, url := fetchModelsFromRemote(ctx) - if parsed == nil { - log.Warnf("%s: fetch failed from all URLs, keeping current data", label) - return - } - - if len(parsed.Meta) == 0 && oldData != nil && len(oldData.Meta) > 0 { - parsed.Meta = oldData.Meta - } - - // Detect changes before updating store. - changed := detectChangedProviders(oldData, parsed) - - // Update store with new data regardless. - modelsCatalogStore.mu.Lock() - modelsCatalogStore.data = parsed - modelsCatalogStore.mu.Unlock() - - if len(changed) == 0 { - log.Infof("%s completed from %s, no changes detected", label, url) - return - } - - log.Infof("%s completed from %s, changes detected for providers: %v", label, url, changed) - notifyModelRefresh(changed) -} - -// fetchModelsFromRemote tries all remote URLs and returns the parsed model catalog -// along with the URL it was fetched from. Returns (nil, "") if all fetches fail. -func fetchModelsFromRemote(ctx context.Context) (*staticModelsJSON, string) { - client := &http.Client{Timeout: modelsFetchTimeout} - for _, url := range modelsURLs { - reqCtx, cancel := context.WithTimeout(ctx, modelsFetchTimeout) - req, err := http.NewRequestWithContext(reqCtx, "GET", url, nil) - if err != nil { - cancel() - log.Debugf("models fetch request creation failed for %s: %v", url, err) - continue - } - - resp, err := client.Do(req) - if err != nil { - cancel() - log.Debugf("models fetch failed from %s: %v", url, err) - continue - } - - if resp.StatusCode != 200 { - resp.Body.Close() - cancel() - log.Debugf("models fetch returned %d from %s", resp.StatusCode, url) - continue - } - - data, err := io.ReadAll(resp.Body) - resp.Body.Close() - cancel() - - if err != nil { - log.Debugf("models fetch read error from %s: %v", url, err) - continue - } - - var parsed staticModelsJSON - if err := json.Unmarshal(data, &parsed); err != nil { - log.Warnf("models parse failed from %s: %v", url, err) - continue - } - if err := validateModelsCatalog(&parsed); err != nil { - log.Warnf("models validate failed from %s: %v", url, err) - continue - } - - return &parsed, url - } - return nil, "" + generalCatalogUpdater.configure(ctx, "") } // detectChangedProviders compares two model catalogs and returns provider names diff --git a/internal/store/gitstore.go b/internal/store/gitstore.go index 7bf5e2055..bf645fbb5 100644 --- a/internal/store/gitstore.go +++ b/internal/store/gitstore.go @@ -45,6 +45,8 @@ type GitTokenStore struct { username string password string lastGC time.Time + pullFn func(ctx context.Context, worktree *git.Worktree, pullOpts *git.PullOptions) error + pushFn func(ctx context.Context, repo *git.Repository, pushOpts *git.PushOptions) error } type resolvedRemoteBranch struct { @@ -115,7 +117,17 @@ func (s *GitTokenStore) EnsureRepository() error { return s.ensureRepositoryLocked() } -func (s *GitTokenStore) ensureRepositoryLocked() (errResult error) { +func (s *GitTokenStore) ensureRepositoryLocked() error { + return s.ensureRepositoryContextLocked(context.Background()) +} + +func (s *GitTokenStore) ensureRepositoryContextLocked(ctx context.Context) (errResult error) { + if ctx == nil { + ctx = context.Background() + } + if errCtx := ctx.Err(); errCtx != nil { + return errCtx + } s.dirLock.Lock() if s.remote == "" { s.dirLock.Unlock() @@ -150,7 +162,7 @@ func (s *GitTokenStore) ensureRepositoryLocked() (errResult error) { if s.branch != "" { cloneOpts.ReferenceName = plumbing.NewBranchReferenceName(s.branch) } - if cloned, errClone := git.PlainClone(repoDir, cloneOpts); errClone != nil { + if cloned, errClone := git.PlainCloneContext(ctx, repoDir, cloneOpts); errClone != nil { if errors.Is(errClone, transport.ErrEmptyRemoteRepository) { _ = os.RemoveAll(gitDir) repo, errInit := git.PlainInit(repoDir, false) @@ -245,7 +257,7 @@ func (s *GitTokenStore) ensureRepositoryLocked() (errResult error) { s.dirLock.Unlock() return fmt.Errorf("git token store: verify repository before pull: %w", errVerify) } - errRecover := s.recoverRepositoryLocked(repoDir, authMethod, repo, nil, nil, (*git.Repository).Close, os.Rename) + errRecover := s.recoverRepositoryContextLocked(ctx, repoDir, authMethod, repo, nil, nil, (*git.Repository).Close, os.Rename) repo = nil if errRecover != nil { s.dirLock.Unlock() @@ -263,13 +275,13 @@ func (s *GitTokenStore) ensureRepositoryLocked() (errResult error) { } } if s.branch != "" { - if errCheckout := s.checkoutConfiguredBranch(repo, worktree, authMethod); errCheckout != nil { + if errCheckout := s.checkoutConfiguredBranch(ctx, repo, worktree, authMethod); errCheckout != nil { s.dirLock.Unlock() return errCheckout } } else { // When branch is unset, ensure the working tree follows the remote default branch - if err := checkoutRemoteDefaultBranch(repo, worktree, authMethod); err != nil { + if err := checkoutRemoteDefaultBranch(ctx, repo, worktree, authMethod); err != nil { if !shouldFallbackToCurrentBranch(repo, err) { s.dirLock.Unlock() return fmt.Errorf("git token store: checkout remote default: %w", err) @@ -304,7 +316,13 @@ func (s *GitTokenStore) ensureRepositoryLocked() (errResult error) { return fmt.Errorf("git token store: inspect worktree before pull: %w", errDirtyPaths) } repositoryRecovered := false - if errPull := worktree.Pull(pullOpts); errPull != nil { + var errPull error + if s.pullFn != nil { + errPull = s.pullFn(ctx, worktree, pullOpts) + } else { + errPull = worktree.PullContext(ctx, pullOpts) + } + if errPull != nil { switch { case errors.Is(errPull, git.NoErrAlreadyUpToDate): if errReset := resetIndexToHead(repo, worktree); errReset != nil { @@ -312,7 +330,7 @@ func (s *GitTokenStore) ensureRepositoryLocked() (errResult error) { s.dirLock.Unlock() return fmt.Errorf("git token store: repair index after up-to-date pull: %w", errReset) } - errRecover := s.recoverRepositoryLocked(repoDir, authMethod, repo, prePullTree, dirtyPaths, (*git.Repository).Close, os.Rename) + errRecover := s.recoverRepositoryContextLocked(ctx, repoDir, authMethod, repo, prePullTree, dirtyPaths, (*git.Repository).Close, os.Rename) repo = nil if errRecover != nil { s.dirLock.Unlock() @@ -330,7 +348,7 @@ func (s *GitTokenStore) ensureRepositoryLocked() (errResult error) { s.dirLock.Unlock() return fmt.Errorf("git token store: reconcile remote changes: %w", errReconcile) } - errRecover := s.recoverRepositoryLocked(repoDir, authMethod, repo, prePullTree, dirtyPaths, (*git.Repository).Close, os.Rename) + errRecover := s.recoverRepositoryContextLocked(ctx, repoDir, authMethod, repo, prePullTree, dirtyPaths, (*git.Repository).Close, os.Rename) repo = nil if errRecover != nil { s.dirLock.Unlock() @@ -348,7 +366,7 @@ func (s *GitTokenStore) ensureRepositoryLocked() (errResult error) { } // Ignore missing references only when following the remote default branch. case isRepositoryCorruptionError(errPull): - errRecover := s.recoverRepositoryLocked(repoDir, authMethod, repo, prePullTree, dirtyPaths, (*git.Repository).Close, os.Rename) + errRecover := s.recoverRepositoryContextLocked(ctx, repoDir, authMethod, repo, prePullTree, dirtyPaths, (*git.Repository).Close, os.Rename) repo = nil if errRecover != nil { s.dirLock.Unlock() @@ -366,7 +384,7 @@ func (s *GitTokenStore) ensureRepositoryLocked() (errResult error) { s.dirLock.Unlock() return fmt.Errorf("git token store: verify repository after pull: %w", errVerify) } - errRecover := s.recoverRepositoryLocked(repoDir, authMethod, repo, prePullTree, dirtyPaths, (*git.Repository).Close, os.Rename) + errRecover := s.recoverRepositoryContextLocked(ctx, repoDir, authMethod, repo, prePullTree, dirtyPaths, (*git.Repository).Close, os.Rename) repo = nil if errRecover != nil { s.dirLock.Unlock() @@ -396,7 +414,7 @@ func (s *GitTokenStore) ensureRepositoryLocked() (errResult error) { } s.dirLock.Unlock() if len(initPaths) > 0 { - if errCommit := s.commitAndPushInitialLocked("Initialize git token store", initPaths...); errCommit != nil { + if errCommit := s.commitAndPushWithOptionsLocked(ctx, "Initialize git token store", true, initPaths...); errCommit != nil { return errCommit } } @@ -433,7 +451,7 @@ func (s *GitTokenStore) Save(ctx context.Context, auth *cliproxyauth.Auth) (stri } } - if err = s.ensureRepositoryLocked(); err != nil { + if err = s.ensureRepositoryContextLocked(ctx); err != nil { return "", err } relPath, errRel := s.relativeToRepo(path) @@ -444,41 +462,47 @@ func (s *GitTokenStore) Save(ctx context.Context, auth *cliproxyauth.Auth) (stri return "", fmt.Errorf("auth filestore: create dir failed: %w", err) } - switch { - case auth.Storage != nil: - if auth.Metadata == nil { - auth.Metadata = make(map[string]any) - } - auth.Metadata["disabled"] = auth.Disabled - if setter, ok := auth.Storage.(interface{ SetMetadata(map[string]any) }); ok { - setter.SetMetadata(auth.Metadata) - } - if err = auth.Storage.SaveTokenToFile(path); err != nil { - return "", err - } - case auth.Metadata != nil: - auth.Metadata["disabled"] = auth.Disabled - raw, errMarshal := json.Marshal(auth.Metadata) - if errMarshal != nil { - return "", fmt.Errorf("auth filestore: marshal metadata failed: %w", errMarshal) - } - contentsMatch := false - if existing, errRead := os.ReadFile(path); errRead == nil { - contentsMatch = jsonEqual(existing, raw) - } else if !os.IsNotExist(errRead) { - return "", fmt.Errorf("auth filestore: read existing failed: %w", errRead) - } - if !contentsMatch { - tmp := path + ".tmp" - if errWrite := os.WriteFile(tmp, raw, 0o600); errWrite != nil { - return "", fmt.Errorf("auth filestore: write temp failed: %w", errWrite) + if handled, errCodex := cliproxyauth.PersistCodexAuthFile(ctx, path, auth); handled { + if errCodex != nil { + return "", errCodex + } + } else { + switch { + case auth.Storage != nil: + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) + } + auth.Metadata["disabled"] = auth.Disabled + if setter, ok := auth.Storage.(interface{ SetMetadata(map[string]any) }); ok { + setter.SetMetadata(auth.Metadata) + } + if err = auth.Storage.SaveTokenToFile(path); err != nil { + return "", err + } + case auth.Metadata != nil: + auth.Metadata["disabled"] = auth.Disabled + raw, errMarshal := json.Marshal(auth.Metadata) + if errMarshal != nil { + return "", fmt.Errorf("auth filestore: marshal metadata failed: %w", errMarshal) + } + contentsMatch := false + if existing, errRead := os.ReadFile(path); errRead == nil { + contentsMatch = jsonEqual(existing, raw) + } else if !os.IsNotExist(errRead) { + return "", fmt.Errorf("auth filestore: read existing failed: %w", errRead) } - if errRename := os.Rename(tmp, path); errRename != nil { - return "", fmt.Errorf("auth filestore: rename failed: %w", errRename) + if !contentsMatch { + tmp := path + ".tmp" + if errWrite := os.WriteFile(tmp, raw, 0o600); errWrite != nil { + return "", fmt.Errorf("auth filestore: write temp failed: %w", errWrite) + } + if errRename := os.Rename(tmp, path); errRename != nil { + return "", fmt.Errorf("auth filestore: rename failed: %w", errRename) + } } + default: + return "", fmt.Errorf("auth filestore: nothing to persist for %s", auth.ID) } - default: - return "", fmt.Errorf("auth filestore: nothing to persist for %s", auth.ID) } if auth.Attributes == nil { @@ -495,7 +519,7 @@ func (s *GitTokenStore) Save(ctx context.Context, auth *cliproxyauth.Auth) (stri if strings.TrimSpace(messageID) == "" { messageID = filepath.Base(path) } - if errCommit := s.commitAndPushLocked(fmt.Sprintf("Update auth %s", strings.TrimSpace(messageID)), relPath); errCommit != nil { + if errCommit := s.commitAndPushContextLocked(ctx, fmt.Sprintf("Update auth %s", strings.TrimSpace(messageID)), relPath); errCommit != nil { return "", errCommit } @@ -874,7 +898,7 @@ func (s *GitTokenStore) relativeToRepo(path string) (string, error) { return rel, nil } -func (s *GitTokenStore) checkoutConfiguredBranch(repo *git.Repository, worktree *git.Worktree, authMethod []client.Option) error { +func (s *GitTokenStore) checkoutConfiguredBranch(ctx context.Context, repo *git.Repository, worktree *git.Worktree, authMethod []client.Option) error { branchRefName := plumbing.NewBranchReferenceName(s.branch) headRef, errHead := repo.Head() switch { @@ -890,18 +914,18 @@ func (s *GitTokenStore) checkoutConfiguredBranch(repo *git.Repository, worktree return fmt.Errorf("git token store: checkout branch %s: %w", s.branch, err) } else if !errors.Is(errRef, plumbing.ErrReferenceNotFound) { return fmt.Errorf("git token store: inspect branch %s: %w", s.branch, errRef) - } else if err := s.checkoutConfiguredRemoteTrackingBranch(repo, worktree, branchRefName, authMethod); err != nil { + } else if err := s.checkoutConfiguredRemoteTrackingBranch(ctx, repo, worktree, branchRefName, authMethod); err != nil { return fmt.Errorf("git token store: checkout branch %s: %w", s.branch, err) } return nil } -func (s *GitTokenStore) checkoutConfiguredRemoteTrackingBranch(repo *git.Repository, worktree *git.Worktree, branchRefName plumbing.ReferenceName, authMethod []client.Option) error { +func (s *GitTokenStore) checkoutConfiguredRemoteTrackingBranch(ctx context.Context, repo *git.Repository, worktree *git.Worktree, branchRefName plumbing.ReferenceName, authMethod []client.Option) error { remoteRefName := plumbing.ReferenceName("refs/remotes/origin/" + s.branch) remoteRef, err := repo.Reference(remoteRefName, true) if errors.Is(err, plumbing.ErrReferenceNotFound) { - if errSync := syncRemoteReferences(repo, authMethod); errSync != nil { + if errSync := syncRemoteReferences(ctx, repo, authMethod); errSync != nil { return fmt.Errorf("sync remote refs: %w", errSync) } remoteRef, err = repo.Reference(remoteRefName, true) @@ -928,8 +952,8 @@ func (s *GitTokenStore) checkoutConfiguredRemoteTrackingBranch(repo *git.Reposit return nil } -func syncRemoteReferences(repo *git.Repository, authMethod []client.Option) error { - if err := repo.Fetch(&git.FetchOptions{ClientOptions: authMethod, RemoteName: "origin"}); err != nil && !errors.Is(err, git.NoErrAlreadyUpToDate) { +func syncRemoteReferences(ctx context.Context, repo *git.Repository, authMethod []client.Option) error { + if err := repo.FetchContext(ctx, &git.FetchOptions{ClientOptions: authMethod, RemoteName: "origin"}); err != nil && !errors.Is(err, git.NoErrAlreadyUpToDate) { return err } return nil @@ -937,15 +961,15 @@ func syncRemoteReferences(repo *git.Repository, authMethod []client.Option) erro // resolveRemoteDefaultBranch queries the origin remote to determine the remote's default branch // (the target of HEAD) and returns the corresponding local branch reference name (e.g. refs/heads/master). -func resolveRemoteDefaultBranch(repo *git.Repository, authMethod []client.Option) (resolvedRemoteBranch, error) { - if err := syncRemoteReferences(repo, authMethod); err != nil { +func resolveRemoteDefaultBranch(ctx context.Context, repo *git.Repository, authMethod []client.Option) (resolvedRemoteBranch, error) { + if err := syncRemoteReferences(ctx, repo, authMethod); err != nil { return resolvedRemoteBranch{}, fmt.Errorf("resolve remote default: sync remote refs: %w", err) } remote, err := repo.Remote("origin") if err != nil { return resolvedRemoteBranch{}, fmt.Errorf("resolve remote default: get remote: %w", err) } - refs, err := remote.List(&git.ListOptions{ClientOptions: authMethod}) + refs, err := remote.ListContext(ctx, &git.ListOptions{ClientOptions: authMethod}) if err != nil { if resolved, ok := resolveRemoteDefaultBranchFromLocal(repo); ok { return resolved, nil @@ -1179,7 +1203,11 @@ func closeRecoveryRepositories(closeRepository func(*git.Repository) error, base return errClose } -func (s *GitTokenStore) recoverRepositoryLocked(repoDir string, authMethod []client.Option, callerRepo *git.Repository, baselineTree *object.Tree, dirtyPaths map[string]struct{}, closeRepository func(*git.Repository) error, renameWorktreeEntry func(string, string) error) (errRecovery error) { +func (s *GitTokenStore) recoverRepositoryLocked(repoDir string, authMethod []client.Option, callerRepo *git.Repository, baselineTree *object.Tree, dirtyPaths map[string]struct{}, closeRepository func(*git.Repository) error, renameWorktreeEntry func(string, string) error) error { + return s.recoverRepositoryContextLocked(context.Background(), repoDir, authMethod, callerRepo, baselineTree, dirtyPaths, closeRepository, renameWorktreeEntry) +} + +func (s *GitTokenStore) recoverRepositoryContextLocked(ctx context.Context, repoDir string, authMethod []client.Option, callerRepo *git.Repository, baselineTree *object.Tree, dirtyPaths map[string]struct{}, closeRepository func(*git.Repository) error, renameWorktreeEntry func(string, string) error) (errRecovery error) { parentDir := filepath.Dir(repoDir) recoveryRoot, errTemp := os.MkdirTemp(parentDir, ".gitstore-recovery-") if errTemp != nil { @@ -1231,7 +1259,7 @@ func (s *GitTokenStore) recoverRepositoryLocked(repoDir string, authMethod []cli if s.branch != "" { cloneOpts.ReferenceName = plumbing.NewBranchReferenceName(s.branch) } - clonedRepo, errClone := git.PlainClone(cloneDir, cloneOpts) + clonedRepo, errClone := git.PlainCloneContext(ctx, cloneDir, cloneOpts) if errClone != nil { if baselineRepo != nil { if errClose := closeRepository(baselineRepo); errClose != nil { @@ -1524,7 +1552,12 @@ func rollbackRecoveredGitDirectory(gitDir, backupGitDir string) error { } func isRepositoryCorruptionError(err error) bool { - return errors.Is(err, dotgit.ErrPackfileNotFound) || errors.Is(err, plumbing.ErrObjectNotFound) + if errors.Is(err, dotgit.ErrPackfileNotFound) || + errors.Is(err, dotgit.ErrIdxNotFound) || + errors.Is(err, plumbing.ErrObjectNotFound) { + return true + } + return errors.Is(err, fs.ErrNotExist) && err != nil && strings.Contains(err.Error(), "packhandle:") } func verifyRepositoryHead(repo *git.Repository) (result error) { @@ -1607,8 +1640,8 @@ func shouldFallbackToCurrentBranch(repo *git.Repository, err error) bool { // checkoutRemoteDefaultBranch ensures the working tree is checked out to the remote's default branch // (the branch target of origin/HEAD). If the local branch does not exist it will be created to track // the remote branch. -func checkoutRemoteDefaultBranch(repo *git.Repository, worktree *git.Worktree, authMethod []client.Option) error { - resolved, err := resolveRemoteDefaultBranch(repo, authMethod) +func checkoutRemoteDefaultBranch(ctx context.Context, repo *git.Repository, worktree *git.Worktree, authMethod []client.Option) error { + resolved, err := resolveRemoteDefaultBranch(ctx, repo, authMethod) if err != nil { return err } @@ -1656,14 +1689,24 @@ func checkoutRemoteDefaultBranch(repo *git.Repository, worktree *git.Worktree, a } func (s *GitTokenStore) commitAndPushLocked(message string, relPaths ...string) error { - return s.commitAndPushWithOptionsLocked(message, false, relPaths...) + return s.commitAndPushWithOptionsLocked(context.Background(), message, false, relPaths...) +} + +func (s *GitTokenStore) commitAndPushContextLocked(ctx context.Context, message string, relPaths ...string) error { + return s.commitAndPushWithOptionsLocked(ctx, message, false, relPaths...) } func (s *GitTokenStore) commitAndPushInitialLocked(message string, relPaths ...string) error { - return s.commitAndPushWithOptionsLocked(message, true, relPaths...) + return s.commitAndPushWithOptionsLocked(context.Background(), message, true, relPaths...) } -func (s *GitTokenStore) commitAndPushWithOptionsLocked(message string, allowMissingRemote bool, relPaths ...string) (errResult error) { +func (s *GitTokenStore) commitAndPushWithOptionsLocked(ctx context.Context, message string, allowMissingRemote bool, relPaths ...string) (errResult error) { + if ctx == nil { + ctx = context.Background() + } + if errCtx := ctx.Err(); errCtx != nil { + return errCtx + } repoDir := s.repoDirSnapshot() if repoDir == "" { return fmt.Errorf("git token store: repository path not configured") @@ -1769,7 +1812,7 @@ func (s *GitTokenStore) commitAndPushWithOptionsLocked(message string, allowMiss if errRewrite := s.rewriteHeadAsSingleCommit(repo, headRef.Name(), commitHash, message, signature); errRewrite != nil { return errRewrite } - if errPush := s.pushRepositoryLocked(repo, repoDir, allowMissingRemote); errPush != nil { + if errPush := s.pushRepositoryLocked(ctx, repo, repoDir, allowMissingRemote); errPush != nil { if baseRef == nil { return errPush } @@ -1854,7 +1897,10 @@ func restoreHeadAndIndex(repo *git.Repository, worktree *git.Worktree, head *plu return worktree.Reset(&git.ResetOptions{Mode: git.MixedReset, Commit: head.Hash()}) } -func (s *GitTokenStore) pushRepositoryLocked(repo *git.Repository, repoDir string, allowMissingRemote bool) error { +func (s *GitTokenStore) pushRepositoryLocked(ctx context.Context, repo *git.Repository, repoDir string, allowMissingRemote bool) error { + if ctx == nil { + ctx = context.Background() + } if repo == nil { return fmt.Errorf("git token store: repository is nil") } @@ -1885,7 +1931,13 @@ func (s *GitTokenStore) pushRepositoryLocked(repo *git.Repository, repoDir strin default: return fmt.Errorf("git token store: inspect remote tracking branch %s: %w", remoteName, errRemote) } - if errPush := repo.Push(pushOpts); errPush != nil { + var errPush error + if s.pushFn != nil { + errPush = s.pushFn(ctx, repo, pushOpts) + } else { + errPush = repo.PushContext(ctx, pushOpts) + } + if errPush != nil { if !errors.Is(errPush, git.NoErrAlreadyUpToDate) { return fmt.Errorf("git token store: push: %w", errPush) } diff --git a/internal/store/objectstore.go b/internal/store/objectstore.go index afa9e8938..39b8b0898 100644 --- a/internal/store/objectstore.go +++ b/internal/store/objectstore.go @@ -17,6 +17,7 @@ import ( "github.com/minio/minio-go/v7" "github.com/minio/minio-go/v7/pkg/credentials" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" "github.com/router-for-me/CLIProxyAPI/v8/internal/misc" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" log "github.com/sirupsen/logrus" @@ -43,12 +44,13 @@ type ObjectStoreConfig struct { // ObjectTokenStore persists configuration and authentication metadata using an S3-compatible object storage backend. // Files are mirrored to a local workspace so existing file-based flows continue to operate. type ObjectTokenStore struct { - client *minio.Client - cfg ObjectStoreConfig - spoolRoot string - configPath string - authDir string - mu sync.Mutex + client *minio.Client + cfg ObjectStoreConfig + spoolRoot string + configPath string + authDir string + mu sync.Mutex + putObjectFn func(ctx context.Context, key string, data []byte, contentType string) error } // NewObjectTokenStore initializes an object storage backed token store. @@ -189,40 +191,47 @@ func (s *ObjectTokenStore) Save(ctx context.Context, auth *cliproxyauth.Auth) (s return "", fmt.Errorf("object store: create auth directory: %w", err) } - switch { - case auth.Storage != nil: - if auth.Metadata == nil { - auth.Metadata = make(map[string]any) - } - auth.Metadata["disabled"] = auth.Disabled - if setter, ok := auth.Storage.(interface{ SetMetadata(map[string]any) }); ok { - setter.SetMetadata(auth.Metadata) - } - if err = auth.Storage.SaveTokenToFile(path); err != nil { - return "", err + if handled, errCodex := cliproxyauth.PersistCodexAuthFile(ctx, path, auth); handled { + if errCodex != nil { + return "", errCodex } - case auth.Metadata != nil: - auth.Metadata["disabled"] = auth.Disabled - raw, errMarshal := json.Marshal(auth.Metadata) - if errMarshal != nil { - return "", fmt.Errorf("object store: marshal metadata: %w", errMarshal) - } - if existing, errRead := os.ReadFile(path); errRead == nil { - if jsonEqual(existing, raw) { - return path, nil + } else { + switch { + case auth.Storage != nil: + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) } - } else if errRead != nil && !errors.Is(errRead, fs.ErrNotExist) { - return "", fmt.Errorf("object store: read existing metadata: %w", errRead) - } - tmp := path + ".tmp" - if errWrite := os.WriteFile(tmp, raw, 0o600); errWrite != nil { - return "", fmt.Errorf("object store: write temp auth file: %w", errWrite) - } - if errRename := os.Rename(tmp, path); errRename != nil { - return "", fmt.Errorf("object store: rename auth file: %w", errRename) + auth.Metadata["disabled"] = auth.Disabled + if setter, ok := auth.Storage.(interface{ SetMetadata(map[string]any) }); ok { + setter.SetMetadata(auth.Metadata) + } + if err = auth.Storage.SaveTokenToFile(path); err != nil { + return "", err + } + case auth.Metadata != nil: + auth.Metadata["disabled"] = auth.Disabled + raw, errMarshal := json.Marshal(auth.Metadata) + if errMarshal != nil { + return "", fmt.Errorf("object store: marshal metadata: %w", errMarshal) + } + contentsMatch := false + if existing, errRead := os.ReadFile(path); errRead == nil { + contentsMatch = jsonEqual(existing, raw) + } else if errRead != nil && !errors.Is(errRead, fs.ErrNotExist) { + return "", fmt.Errorf("object store: read existing metadata: %w", errRead) + } + if !contentsMatch { + tmp := path + ".tmp" + if errWrite := os.WriteFile(tmp, raw, 0o600); errWrite != nil { + return "", fmt.Errorf("object store: write temp auth file: %w", errWrite) + } + if errRename := os.Rename(tmp, path); errRename != nil { + return "", fmt.Errorf("object store: rename auth file: %w", errRename) + } + } + default: + return "", fmt.Errorf("object store: nothing to persist for %s", auth.ID) } - default: - return "", fmt.Errorf("object store: nothing to persist for %s", auth.ID) } if auth.Attributes == nil { @@ -386,6 +395,12 @@ func (s *ObjectTokenStore) syncConfigFromBucket(ctx context.Context, example str } } } + // Seed the bucket with the file the startup load leaves behind, so a plaintext + // management secret never reaches it. Fail instead of skipping: the startup load + // could turn a refused file into one that a later start would upload. + if errPrepare := config.PrepareConfigFileForUpload(s.configPath); errPrepare != nil { + return fmt.Errorf("object store: prepare local config for seeding: %w", errPrepare) + } data, errRead := os.ReadFile(s.configPath) if errRead != nil { return fmt.Errorf("object store: read local config: %w", errRead) @@ -488,9 +503,15 @@ func (s *ObjectTokenStore) deleteAuthObject(ctx context.Context, path string) er } func (s *ObjectTokenStore) putObject(ctx context.Context, key string, data []byte, contentType string) error { + if s.putObjectFn != nil { + return s.putObjectFn(ctx, key, data, contentType) + } if len(data) == 0 { return s.deleteObject(ctx, key) } + if s.client == nil { + return fmt.Errorf("object store: client not configured") + } fullKey := s.prefixedKey(key) reader := bytes.NewReader(data) _, err := s.client.PutObject(ctx, s.cfg.Bucket, fullKey, reader, int64(len(data)), minio.PutObjectOptions{ diff --git a/internal/store/postgresstore.go b/internal/store/postgresstore.go index eb8c0a173..69bc318bd 100644 --- a/internal/store/postgresstore.go +++ b/internal/store/postgresstore.go @@ -14,6 +14,7 @@ import ( "time" _ "github.com/jackc/pgx/v5/stdlib" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" "github.com/router-for-me/CLIProxyAPI/v8/internal/misc" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/auth" log "github.com/sirupsen/logrus" @@ -46,6 +47,7 @@ type PostgresStore struct { authDir string cooldownStore *postgresCooldownStateStore mu sync.Mutex + persistAuthFn func(ctx context.Context, relID string, data []byte) error } // NewPostgresStore establishes a connection to PostgreSQL and prepares the local workspace. @@ -240,40 +242,47 @@ func (s *PostgresStore) Save(ctx context.Context, auth *cliproxyauth.Auth) (stri return "", fmt.Errorf("postgres store: create auth directory: %w", err) } - switch { - case auth.Storage != nil: - if auth.Metadata == nil { - auth.Metadata = make(map[string]any) - } - auth.Metadata["disabled"] = auth.Disabled - if setter, ok := auth.Storage.(interface{ SetMetadata(map[string]any) }); ok { - setter.SetMetadata(auth.Metadata) - } - if err = auth.Storage.SaveTokenToFile(path); err != nil { - return "", err + if handled, errCodex := cliproxyauth.PersistCodexAuthFile(ctx, path, auth); handled { + if errCodex != nil { + return "", errCodex } - case auth.Metadata != nil: - auth.Metadata["disabled"] = auth.Disabled - raw, errMarshal := json.Marshal(auth.Metadata) - if errMarshal != nil { - return "", fmt.Errorf("postgres store: marshal metadata: %w", errMarshal) - } - if existing, errRead := os.ReadFile(path); errRead == nil { - if jsonEqual(existing, raw) { - return path, nil + } else { + switch { + case auth.Storage != nil: + if auth.Metadata == nil { + auth.Metadata = make(map[string]any) } - } else if errRead != nil && !errors.Is(errRead, fs.ErrNotExist) { - return "", fmt.Errorf("postgres store: read existing metadata: %w", errRead) - } - tmp := path + ".tmp" - if errWrite := os.WriteFile(tmp, raw, 0o600); errWrite != nil { - return "", fmt.Errorf("postgres store: write temp auth file: %w", errWrite) - } - if errRename := os.Rename(tmp, path); errRename != nil { - return "", fmt.Errorf("postgres store: rename auth file: %w", errRename) + auth.Metadata["disabled"] = auth.Disabled + if setter, ok := auth.Storage.(interface{ SetMetadata(map[string]any) }); ok { + setter.SetMetadata(auth.Metadata) + } + if err = auth.Storage.SaveTokenToFile(path); err != nil { + return "", err + } + case auth.Metadata != nil: + auth.Metadata["disabled"] = auth.Disabled + raw, errMarshal := json.Marshal(auth.Metadata) + if errMarshal != nil { + return "", fmt.Errorf("postgres store: marshal metadata: %w", errMarshal) + } + contentsMatch := false + if existing, errRead := os.ReadFile(path); errRead == nil { + contentsMatch = jsonEqual(existing, raw) + } else if errRead != nil && !errors.Is(errRead, fs.ErrNotExist) { + return "", fmt.Errorf("postgres store: read existing metadata: %w", errRead) + } + if !contentsMatch { + tmp := path + ".tmp" + if errWrite := os.WriteFile(tmp, raw, 0o600); errWrite != nil { + return "", fmt.Errorf("postgres store: write temp auth file: %w", errWrite) + } + if errRename := os.Rename(tmp, path); errRename != nil { + return "", fmt.Errorf("postgres store: rename auth file: %w", errRename) + } + } + default: + return "", fmt.Errorf("postgres store: nothing to persist for %s", auth.ID) } - default: - return "", fmt.Errorf("postgres store: nothing to persist for %s", auth.ID) } if auth.Attributes == nil { @@ -462,6 +471,12 @@ func (s *PostgresStore) syncConfigFromDatabase(ctx context.Context, exampleConfi } } } + // Seed the database with the file the startup load leaves behind, so a plaintext + // management secret never reaches it. Fail instead of skipping: the startup load + // could turn a refused file into one that a later start would upload. + if errPrepare := config.PrepareConfigFileForUpload(s.configPath); errPrepare != nil { + return fmt.Errorf("postgres store: prepare local config for seeding: %w", errPrepare) + } data, errRead := os.ReadFile(s.configPath) if errRead != nil { return fmt.Errorf("postgres store: read local config: %w", errRead) @@ -551,6 +566,12 @@ func (s *PostgresStore) upsertAuthRecord(ctx context.Context, relID, path string } func (s *PostgresStore) persistAuth(ctx context.Context, relID string, data []byte) error { + if s.persistAuthFn != nil { + return s.persistAuthFn(ctx, relID, data) + } + if s.db == nil { + return fmt.Errorf("postgres store: database not configured") + } jsonPayload := json.RawMessage(data) query := fmt.Sprintf(` INSERT INTO %s (id, content, created_at, updated_at)