diff --git a/docs/oauth-subscriptions.md b/docs/oauth-subscriptions.md index 6bb70c98d..53c30de4d 100644 --- a/docs/oauth-subscriptions.md +++ b/docs/oauth-subscriptions.md @@ -37,8 +37,8 @@ before. Tokens are stored 0600 (or the OS keyring with Running `/provider` opens a **"How do you want to connect?"** chooser: ```text -❯ Sign in with OAuth One-click browser login (OpenRouter, xAI, ChatGPT, Hugging Face) - Paste an API key / browse providers Any of 20+ providers, local, or a proxy +❯ Sign in with OAuth No API key to copy — one-click browser login (OpenRouter, xAI, ChatGPT, Hugging Face) or device code (Kimi Code). + Paste an API key / browse providers Any of 20+ providers, a local model, or a subscription via proxy. ``` Pick **Sign in with OAuth** → the list of providers that do real OAuth → choose one: @@ -46,23 +46,30 @@ Pick **Sign in with OAuth** → the list of providers that do real OAuth → cho ```text ❯ OpenRouter browser sign-in · creates a key xAI (Grok) browser or device code + Kimi Code device code (managed coding endpoint) ChatGPT browser (Codex backend, ChatGPT Plus/Pro) Hugging Face browser or device code ``` - **OpenRouter / xAI / ChatGPT / Hugging Face** are real OAuth: your browser - opens to approve → done (no key to paste). OpenRouter mints a key; xAI / - ChatGPT / Hugging Face store a refreshable bearer. Hugging Face requires a - one-time OAuth-app registration (no secret needed for "public" apps); the - preset pre-fills scopes, endpoints, and the OIDC issuer. The same chooser - appears in first-run onboarding. (xAI uses an opt-in preset — set - `ZERO_OAUTH_ALLOW_PRESETS=1` or your own `ZERO_OAUTH_XAI_*`; see below.) -- **Device code (headless / SSH):** for a provider that supports it (xAI, - Hugging Face), press **d** on the list to get a code to enter on another - device instead of opening a browser. On an SSH session or headless Linux box - (no `DISPLAY`) device code is used automatically; set `ZERO_OAUTH_DEVICE=1` - to force it anywhere. The CLI equivalent is - `zero auth login --device`. + opens to approve → done (no key to paste). OpenRouter mints a key; xAI / + ChatGPT / Hugging Face store a refreshable bearer. Hugging Face requires a + one-time OAuth-app registration (no secret needed for "public" apps); the + preset pre-fills scopes, endpoints, and the OIDC issuer. Kimi Code is also + real OAuth but has no browser flow at all — see the device-code bullet + below. The same chooser appears in first-run onboarding. (`zero auth login` + and the interactive wizard use built-in presets for OpenRouter, xAI, + Kimi Code, ChatGPT, and Hugging Face without exporting + `ZERO_OAUTH_ALLOW_PRESETS`; any field is still overridable with + `ZERO_OAUTH__*` — see below.) +- **Device code (headless / SSH):** for a provider that supports it (xAI, Kimi + Code, Hugging Face), press **d** on the list to get a code to enter on + another device instead of opening a browser. On an SSH session or headless + Linux box (no `DISPLAY`) device code is used automatically; set + `ZERO_OAUTH_DEVICE=1` to force it anywhere. The CLI equivalent is + `zero auth login --device`. (Kimi Code is **device-code only** — it + has no loopback/browser flow, so `zero auth kimi` always uses the device + path and pressing plain Enter on it in the wizard does too.) - **ChatGPT / Claude are intentionally not in this list for the proxy path** — use the dedicated `chatgpt-proxy` / `custom-anthropic-compatible` preset (see §2) for subscription-via-proxy. ChatGPT *is* a first-class OAuth @@ -76,17 +83,57 @@ Pick **Sign in with OAuth** → the list of providers that do real OAuth → cho In the interactive setup wizard, pick **OpenRouter** and press **ctrl+o** at the key step to do the same inline ("Log in with OAuth"). The minted key is saved to the provider profile and used normally. -- **xAI (Grok) — opt-in preset** — xAI's flow needs an OAuth `client_id`. Zero - ships a built-in preset for the public Grok-CLI client, but to keep third-party - client identities out of the default credential path it is **off by default**. - Enable it with `export ZERO_OAUTH_ALLOW_PRESETS=1`, then `zero auth login xai` - (browser, or `--device` for headless) works one-click; the token is used directly - on `api.x.ai/v1`. Without the opt-in, set `ZERO_OAUTH_XAI_CLIENT_ID` (and - endpoints, or an issuer) yourself via `ZERO_OAUTH_XAI_*`. Either way the preset is +- **xAI (Grok) — built-in preset** — xAI's flow needs an OAuth `client_id`. Zero + ships a built-in preset for the public Grok-CLI client. `zero auth login xai` + (browser, or `--device` for headless) and the interactive wizard enable that + preset automatically (no `ZERO_OAUTH_ALLOW_PRESETS=1` required on those + paths); the token is used directly on `api.x.ai/v1`. Library callers that + construct an `oauth.Manager` without `AllowPresets: true` still need + `ZERO_OAUTH_ALLOW_PRESETS=1`, or can set `ZERO_OAUTH_XAI_CLIENT_ID` (and + endpoints, or an issuer) via `ZERO_OAUTH_XAI_*`. Either way the preset is fully overridable by `ZERO_OAUTH_XAI_*` (env wins), and it requires a SuperGrok / X Premium+ subscription; the client_id is an undocumented public Grok-CLI client that may change without notice. -- **ChatGPT (Codex) — opt-in preset** — `zero auth chatgpt` opens a browser, you +- **Kimi Code — built-in preset, device-code only** — `zero auth kimi` (or + `zero auth login kimi-code --device`) runs the RFC 8628 device-code flow + against `https://auth.kimi.com`. You approve on another device and enter the + code; the returned access token is stored and used **directly** as a bearer + on Kimi's managed coding endpoint `https://api.kimi.com/coding/v1` (an + OpenAI-compatible chat-completions endpoint) — no ID-token claim extraction + is needed. Kimi has **no browser/loopback flow**, so the device code is the + only path (it is used automatically, and `--device` is accepted but + redundant). The catalog/provider ID is `kimi-code`, not `kimi` — the + `moonshot` provider already uses `kimi` as an alias for its separate, + API-key-based endpoint, so `zero auth kimi` is CLI sugar that forwards to + `kimi-code` rather than reusing that name. Like xAI, the preset ships the + public kimi-cli client identity (`17e5f671-d194-4dfb-9706-5516cb48c098`). + Both `zero auth kimi` and `zero auth login kimi-code` run through the + `auth login` engine, which enables presets unconditionally — no + `ZERO_OAUTH_ALLOW_PRESETS=1` is needed (same rule as xAI on those paths). + Any field is still overridable with + `ZERO_OAUTH_KIMI_CODE_*`. Setting only `ZERO_OAUTH_KIMI_CODE_ISSUER_URL` + replaces the preset device/token endpoints and uses discovery; leftover + preset destinations are not kept, so header policy and request URLs stay + one coherent configuration. Kimi's backend also requires a handful of + vendor-identity `X-Msh-*` headers across all applicable OAuth and API calls + (device authorization, polling, code exchange, refresh, and managed + runtime/completions requests). Those headers are minted lazily on the first + request that actually talks to Kimi, not when listing providers. They send + `X-Msh-Platform` (`kimi_code_cli`), `X-Msh-Version` (`unknown`), + `X-Msh-Device-Name` (this machine's hostname), `X-Msh-Device-Model` + (GOOS/GOARCH), `X-Msh-Os-Version` (GOOS), and a persistent + `X-Msh-Device-Id`. These headers are reverse-engineered from kimi-cli, + not from public documentation; verify against a real login before relying + on this. + This is distinct from the `moonshot` catalog entry, which is the API-key path at + `https://api.moonshot.ai/v1` (set `MOONSHOT_API_KEY`). A profile `baseURL` + override strips `X-Msh-*` headers. OAuth endpoint overrides keep them only + for HTTPS `auth.kimi.com` or `api.kimi.com`; any other host also strips + them. To override the managed endpoint, set `baseURL` on the provider + profile; to override the OAuth host, set + `ZERO_OAUTH_KIMI_CODE_ISSUER_URL`/`ZERO_OAUTH_KIMI_CODE_DEVICE_URL`/`ZERO_OAUTH_KIMI_CODE_TOKEN_URL` + (the provider resolves as `kimi-code`, so the env prefix is `KIMI_CODE`). +- **ChatGPT (Codex) — built-in preset** — `zero auth chatgpt` opens a browser, you approve with your ChatGPT Plus/Pro/Business/Enterprise account, and the bearer is stored. The bearer routes to `https://chatgpt.com/backend-api/codex/responses` (the same endpoint the openai/codex CLI uses), with `originator: codex_cli_rs` and @@ -95,20 +142,18 @@ Pick **Sign in with OAuth** → the list of providers that do real OAuth → cho bearer; if the claim is missing (older ChatGPT accounts, or a rotated authorization server), the Codex backend will 401 and `zero auth status chatgpt` will show the warning. Like xAI, the preset uses the publicly-shipped Codex CLI - client identity (`app_EMoamEEZ73f0CkXaXp7hrann`) and is opt-in via - `ZERO_OAUTH_ALLOW_PRESETS=1`. As of mid-2026 the Codex backend is + client identity (`app_EMoamEEZ73f0CkXaXp7hrann`) and `zero auth chatgpt` or the + interactive wizard enables it automatically. As of mid-2026 the Codex backend is Cloudflare-gated: requests from a non-Codex client can still be challenged, and the `chatgpt-proxy` route in §2 is the conservative fallback. -- **Hugging Face — opt-in preset, BYO client_id** — `zero auth login huggingface` +- **Hugging Face — built-in preset, BYO client_id** — `zero auth login huggingface` (or `--device` for headless) opens a Hugging Face OAuth flow. The bearer works on the OpenAI-compatible router at `https://router.huggingface.co/v1` for hundreds of OSS models (Llama, Qwen, DeepSeek, Mistral, etc.). HF does not ship a globally-known client_id, so the preset ships endpoints + scopes + the OIDC issuer pre-filled; you must register a "public" OAuth app (no secret) at and set the resulting - `client_id` via `ZERO_OAUTH_HUGGINGFACE_CLIENT_ID`. Enable the preset with - `ZERO_OAUTH_ALLOW_PRESETS=1` (or omit it — the BYO client_id path uses - `client_credentials = none` and doesn't need the opt-in). Free tier has strict + `client_id` via `ZERO_OAUTH_HUGGINGFACE_CLIENT_ID`. Free tier has strict rate limits; Pro removes them. Any field of a preset is overridable via `ZERO_OAUTH__*`. For a fully custom diff --git a/internal/cli/auth.go b/internal/cli/auth.go index f3ecdcc42..bb4e5cb5c 100644 --- a/internal/cli/auth.go +++ b/internal/cli/auth.go @@ -78,6 +78,18 @@ func runAuth(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in return runAuthOpenRouter(args[1:], stdout, stderr, deps) case "chatgpt": return runAuthChatGPT(args[1:], stdout, stderr, deps) + case "kimi": + // Kimi Code is a standard device-code OAuth preset (no bespoke client + // like ChatGPT's Codex flow), so it reuses the generic `auth login` + // engine — which already opts into presets and resolves the baked-in + // kimi-code client_id/endpoints. `zero auth kimi` is sugar for + // `zero auth login kimi-code`, forwarding whatever flags the caller + // passed after "kimi" (--device, --scope, --help) through the real + // parser instead of discarding them: `zero auth kimi --help` must show + // help, not silently start a real device authorization, and an unknown + // flag or extra positional must be rejected the same way `zero auth + // login` rejects one. + return runAuthLogin(append([]string{"kimi-code"}, args[1:]...), stdout, stderr, deps) default: return writeExecUsageError(stderr, fmt.Sprintf("unknown auth subcommand %q", args[0])) } @@ -591,13 +603,14 @@ Commands: refresh [--watch] Force a token refresh (--watch keeps it fresh) openrouter Log in to OpenRouter in the browser; mints an API key chatgpt Log in to ChatGPT in the browser (Codex backend, ChatGPT Plus/Pro) + kimi Log in to Kimi Code via device code (managed coding endpoint) -A provider is any OAuth 2.0 / OIDC server. "openrouter" ('zero auth openrouter') -works out of the box. "xai" ('zero auth login xai') uses a built-in preset that is -off by default — enable it with ZERO_OAUTH_ALLOW_PRESETS=1, or set the -ZERO_OAUTH_XAI_* vars yourself. "chatgpt" ('zero auth login chatgpt' or -'zero auth chatgpt') uses a fixed-port loopback flow against the Codex backend. -Any preset field is overridable via the env vars below. For a custom provider named , set: +A provider is any OAuth 2.0 / OIDC server. "openrouter" ('zero auth openrouter'), +"chatgpt" ('zero auth chatgpt'), "kimi" ('zero auth kimi' or +'zero auth login kimi-code'), and "xai" ('zero auth login xai') work out of +the box: "zero auth login" enables baked-in presets for the chosen provider +without exporting ZERO_OAUTH_ALLOW_PRESETS. Any preset field is overridable +via the env vars below. For a custom provider named , set: ZERO_OAUTH__CLIENT_ID (required) ZERO_OAUTH__CLIENT_SECRET (optional) ZERO_OAUTH__AUTHORIZE_URL ZERO_OAUTH__TOKEN_URL diff --git a/internal/cli/auth_test.go b/internal/cli/auth_test.go index 9b1ba0fb5..b0f7da9db 100644 --- a/internal/cli/auth_test.go +++ b/internal/cli/auth_test.go @@ -100,6 +100,58 @@ func TestRunAuthLoginValidation(t *testing.T) { } } +// TestRunAuthKimiAliasForwardsToLoginKimiCode pins the `zero auth kimi` +// sugar: it must forward through runAuthLogin with the kimi-code preset, so +// --help prints parser help without starting a device authorization, and +// invalid flags/positionals are rejected exactly like the canonical +// `zero auth login kimi-code` invocation. +func TestRunAuthKimiAliasForwardsToLoginKimiCode(t *testing.T) { + withAuthStore(t) + var stdout, stderr bytes.Buffer + + // --help must produce the login parser's help (exitSuccess) and never + // begin a device authorization. + if code := runWithDeps([]string{"auth", "kimi", "--help"}, &stdout, &stderr, appDeps{}); code != exitSuccess { + t.Fatalf("auth kimi --help exit = %d, stderr=%q", code, stderr.String()) + } + helpOut := stdout.String() + for _, want := range []string{"kimi", "login", "--device"} { + if !strings.Contains(helpOut, want) { + t.Fatalf("auth kimi --help missing %q:\n%s", want, helpOut) + } + } + + // An unknown flag must be rejected identically to auth login kimi-code. + for _, args := range [][]string{ + {"auth", "kimi", "--bogus"}, + {"auth", "login", "kimi-code", "--bogus"}, + } { + stdout.Reset() + stderr.Reset() + if code := runWithDeps(args, &stdout, &stderr, appDeps{}); code == exitSuccess { + t.Fatalf("%v should be rejected", args) + } + if !strings.Contains(stderr.String(), "unknown flag") { + t.Fatalf("%v stderr = %q, want unknown-flag rejection", args, stderr.String()) + } + } + + // A stray positional must be rejected identically to auth login kimi-code. + for _, args := range [][]string{ + {"auth", "kimi", "extra-positional"}, + {"auth", "login", "kimi-code", "extra-positional"}, + } { + stdout.Reset() + stderr.Reset() + if code := runWithDeps(args, &stdout, &stderr, appDeps{}); code == exitSuccess { + t.Fatalf("%v should be rejected", args) + } + if !strings.Contains(stderr.String(), "usage:") { + t.Fatalf("%v stderr = %q, want usage error", args, stderr.String()) + } + } +} + func TestRunAuthLoginUnknownProvider(t *testing.T) { withAuthStore(t) var stdout, stderr bytes.Buffer @@ -321,3 +373,36 @@ func readCLIConfigFixture(t *testing.T, path string) config.FileConfig { } return cfg } + +func TestRunAuthKimiHelp(t *testing.T) { + withAuthStore(t) + var stdout, stderr bytes.Buffer + if code := runWithDeps([]string{"auth", "kimi", "--help"}, &stdout, &stderr, appDeps{}); code != exitSuccess { + t.Fatalf("auth kimi --help exit = %d, stderr=%q", code, stderr.String()) + } + if !strings.Contains(stdout.String(), "zero auth login") { + t.Fatalf("help text missing login guidance, got: %q", stdout.String()) + } +} + +func TestRunAuthKimiUnknownFlag(t *testing.T) { + withAuthStore(t) + var stdout, stderr bytes.Buffer + if code := runWithDeps([]string{"auth", "kimi", "--invalid-flag"}, &stdout, &stderr, appDeps{}); code == exitSuccess { + t.Fatal("auth kimi --invalid-flag should fail") + } + if !strings.Contains(stderr.String(), "unknown flag") { + t.Fatalf("expected unknown flag error, got stderr=%q", stderr.String()) + } +} + +func TestRunAuthKimiExtraPositional(t *testing.T) { + withAuthStore(t) + var stdout, stderr bytes.Buffer + if code := runWithDeps([]string{"auth", "kimi", "unexpected-extra"}, &stdout, &stderr, appDeps{}); code == exitSuccess { + t.Fatal("auth kimi unexpected-extra should fail") + } + if !strings.Contains(stderr.String(), "usage: zero auth login") { + t.Fatalf("expected login usage error, got stderr=%q", stderr.String()) + } +} diff --git a/internal/cli/provider_setup.go b/internal/cli/provider_setup.go index de13f26ce..b3c4b104d 100644 --- a/internal/cli/provider_setup.go +++ b/internal/cli/provider_setup.go @@ -429,7 +429,7 @@ func providerProfileForAdd(options providerAddOptions) (config.ProviderProfile, baseURL := firstNonEmptyCLI(options.baseURL, descriptor.DefaultBaseURL) var catalogHeaders map[string]string if sameProviderSetupBaseURL(baseURL, descriptor.DefaultBaseURL) { - catalogHeaders = descriptor.CustomHeaders + catalogHeaders = stripRuntimeIdentityHeaders(descriptor.CustomHeaders) if strings.EqualFold(strings.TrimSpace(descriptor.ID), "aimlapi") { catalogHeaders = aimlapi.WithResolvedPartnerHeader(catalogHeaders) } @@ -452,6 +452,23 @@ func providerProfileForAdd(options providerAddOptions) (config.ProviderProfile, return profile, nil } +func stripRuntimeIdentityHeaders(headers map[string]string) map[string]string { + if len(headers) == 0 { + return nil + } + cleaned := make(map[string]string, len(headers)) + for k, v := range headers { + if providercatalog.IsRuntimeIdentityHeader(k) { + continue + } + cleaned[k] = v + } + if len(cleaned) == 0 { + return nil + } + return cleaned +} + func sameProviderSetupBaseURL(left string, right string) bool { return strings.EqualFold( strings.TrimRight(strings.TrimSpace(left), "/"), diff --git a/internal/cli/provider_setup_test.go b/internal/cli/provider_setup_test.go index c9c293d7c..4517c4b80 100644 --- a/internal/cli/provider_setup_test.go +++ b/internal/cli/provider_setup_test.go @@ -1,9 +1,12 @@ package cli import ( + "bytes" + "path/filepath" "testing" "github.com/Gitlawb/zero/internal/config" + "github.com/Gitlawb/zero/internal/providercatalog" ) // Regression for issue #555's follow-up: `zero providers check` must not @@ -66,3 +69,67 @@ func TestValidateProviderRuntimeReadyCustomEndpoint(t *testing.T) { }) } } + +func isolateKimiDeviceIDStorage(t *testing.T) { + t.Helper() + root := t.TempDir() + t.Setenv("XDG_CONFIG_HOME", root) + t.Setenv("APPDATA", root) + t.Setenv("HOME", root) +} + +func TestProvidersAddKimiCodeDoesNotPersistRuntimeHeaders(t *testing.T) { + isolateKimiDeviceIDStorage(t) + configPath := filepath.Join(t.TempDir(), "config.json") + deps := appDeps{userConfigPath: func() (string, error) { return configPath, nil }} + + var stdout, stderr bytes.Buffer + code := runProviders([]string{"add", "kimi-code"}, &stdout, &stderr, deps) + if code != exitSuccess { + t.Fatalf("providers add kimi-code failed: %d, stderr: %s", code, stderr.String()) + } + + cfg := readCLIConfigFixture(t, configPath) + if len(cfg.Providers) != 1 { + t.Fatalf("expected 1 provider, got %d", len(cfg.Providers)) + } + for k := range cfg.Providers[0].CustomHeaders { + if providercatalog.IsRuntimeIdentityHeader(k) { + t.Fatalf("config.json contains persisted runtime identity header %q", k) + } + } + + resolved, err := config.Resolve(config.ResolveOptions{UserConfigPath: configPath}) + if err != nil { + t.Fatalf("config.Resolve: %v", err) + } + active := resolved.Provider + if active.CustomHeaders["X-Msh-Platform"] != "kimi_code_cli" { + t.Fatalf("resolved profile missing X-Msh-Platform: %#v", active.CustomHeaders) + } + if active.CustomHeaders["X-Msh-Device-Id"] == "" { + t.Fatalf("resolved profile missing X-Msh-Device-Id: %#v", active.CustomHeaders) + } +} + +func TestSetupKimiCodeDoesNotPersistRuntimeHeaders(t *testing.T) { + isolateKimiDeviceIDStorage(t) + configPath := filepath.Join(t.TempDir(), "config.json") + deps := appDeps{userConfigPath: func() (string, error) { return configPath, nil }} + + var stdout, stderr bytes.Buffer + code := runSetup([]string{"kimi-code"}, &stdout, &stderr, deps) + if code != exitSuccess { + t.Fatalf("setup kimi-code failed: %d, stderr: %s", code, stderr.String()) + } + + cfg := readCLIConfigFixture(t, configPath) + if len(cfg.Providers) != 1 { + t.Fatalf("expected 1 provider, got %d", len(cfg.Providers)) + } + for k := range cfg.Providers[0].CustomHeaders { + if providercatalog.IsRuntimeIdentityHeader(k) { + t.Fatalf("config.json contains persisted runtime identity header %q", k) + } + } +} diff --git a/internal/config/resolver.go b/internal/config/resolver.go index 205a13c68..c61eeb8f3 100644 --- a/internal/config/resolver.go +++ b/internal/config/resolver.go @@ -12,6 +12,7 @@ import ( "github.com/Gitlawb/zero/internal/aimlapi" "github.com/Gitlawb/zero/internal/modelregistry" "github.com/Gitlawb/zero/internal/notify" + "github.com/Gitlawb/zero/internal/oauth" "github.com/Gitlawb/zero/internal/providercatalog" "github.com/Gitlawb/zero/internal/sandbox" ) @@ -81,7 +82,7 @@ func Resolve(options ResolveOptions) (ResolvedConfig, error) { if err != nil { return ResolvedConfig{}, err } - if err := mergeProjectConfig(&cfg, fileConfig); err != nil { + if err := mergeProjectConfig(&cfg, fileConfig, options.Env); err != nil { return ResolvedConfig{}, err } } @@ -284,7 +285,7 @@ func mergeConfig(dst *FileConfig, src FileConfig) { } } -func mergeProjectConfig(dst *FileConfig, src FileConfig) error { +func mergeProjectConfig(dst *FileConfig, src FileConfig, env map[string]string) error { if activeProvider := strings.TrimSpace(src.ActiveProvider); activeProvider != "" { dst.ActiveProvider = activeProvider } @@ -293,7 +294,7 @@ func mergeProjectConfig(dst *FileConfig, src FileConfig) error { } for _, provider := range src.Providers { candidate := providerMergeCandidate(*dst, provider) - if err := validateProjectProviderMerge(provider, candidate); err != nil { + if err := validateProjectProviderMerge(provider, candidate, env); err != nil { return err } mergeProvider(dst, provider) @@ -383,7 +384,7 @@ func providerMergeName(cfg FileConfig, provider ProviderProfile) string { return name } -func validateProjectProviderMerge(project ProviderProfile, candidate ProviderProfile) error { +func validateProjectProviderMerge(project ProviderProfile, candidate ProviderProfile, env map[string]string) error { if strings.TrimSpace(project.APIKeyEnv) != "" && projectEndpointNeedsCredentialGuard(candidate) && !projectAPIKeyEnvAllowed(candidate, project.APIKeyEnv) { @@ -391,7 +392,7 @@ func validateProjectProviderMerge(project ProviderProfile, candidate ProviderPro } if strings.TrimSpace(project.BaseURL) != "" && projectEndpointNeedsCredentialGuard(candidate) && - hasInheritedProviderCredentialMaterial(project, candidate) && + hasInheritedProviderCredentialMaterial(project, candidate, env) && !projectBaseURLAllowed(candidate) { return providerError(candidate, "project provider %s cannot override baseURL for a credentialed custom provider endpoint", candidate.Name) } @@ -434,7 +435,7 @@ func projectBaseURLAllowed(profile ProviderProfile) bool { return ok && sameBaseURL(profile.BaseURL, descriptor.DefaultBaseURL) } -func hasInheritedProviderCredentialMaterial(project ProviderProfile, candidate ProviderProfile) bool { +func hasInheritedProviderCredentialMaterial(project ProviderProfile, candidate ProviderProfile, env map[string]string) bool { if strings.TrimSpace(project.APIKey) == "" && strings.TrimSpace(candidate.APIKey) != "" { return true } @@ -447,9 +448,47 @@ func hasInheritedProviderCredentialMaterial(project ProviderProfile, candidate P if project.CustomHeaders == nil && hasCustomHeaderMaterial(candidate.CustomHeaders) { return true } + if hasInheritedOAuthLogin(candidate, env) { + return true + } return false } +// storedOAuthLogin reports whether any of names has a stored OAuth token under +// env. Tests replace this so Resolve stays hermetic and does not block on a +// real keychain. A non-nil error means the lookup could not tell. +var storedOAuthLogin = lookupStoredOAuthLogin + +func lookupStoredOAuthLogin(env map[string]string, names []string) (bool, error) { + store, err := oauth.NewStore(oauth.StoreOptions{Env: env}) + if err != nil { + return false, err + } + for _, name := range names { + _, ok, err := store.Load(oauth.ProviderKey(name)) + if err != nil { + return false, err + } + if ok { + return true, nil + } + } + return false, nil +} + +func hasInheritedOAuthLogin(candidate ProviderProfile, env map[string]string) bool { + names := candidate.OAuthLoginCandidates() + if len(names) == 0 { + return false + } + ok, err := storedOAuthLogin(env, names) + if err != nil { + // Fail closed: a store or load error means credential material may exist. + return true + } + return ok +} + func hasCustomHeaderMaterial(headers map[string]string) bool { for _, value := range headers { if strings.TrimSpace(value) != "" { @@ -1105,6 +1144,10 @@ func applyCatalogDescriptor(profile *ProviderProfile, descriptor providercatalog } merged := copyStringMap(catalogHeaders) for key, value := range profile.CustomHeaders { + // Do not let stored X-Msh-* headers override fresh RuntimeHeaders. + if descriptor.RuntimeHeaders != nil && providercatalog.IsRuntimeIdentityHeader(key) { + continue + } // Header names are case-insensitive. Preserve the catalog spelling while // replacing its value so request construction cannot see two colliding // map entries whose eventual winner depends on iteration order. @@ -1122,11 +1165,18 @@ func applyCatalogDescriptor(profile *ProviderProfile, descriptor providercatalog merged[key] = value } profile.CustomHeaders = merged - } else if strings.EqualFold(strings.TrimSpace(descriptor.ID), "aimlapi") && !canonicalCatalogEndpoint { - // AIMLAPI attribution is owned by the catalog endpoint. A profile can retain - // those generated headers after its base URL is edited; strip their names - // before sending requests to an arbitrary staging/proxy host while preserving - // unrelated headers explicitly supplied by the user. + } else if !canonicalCatalogEndpoint && (len(descriptor.CustomHeaders) > 0 || descriptor.RuntimeHeaders != nil) { + // Catalog-owned headers (AIMLAPI attribution, Kimi's X-Msh-* device + // identity, etc.) are only valid against the catalog endpoint. A profile + // can retain those generated headers after its base URL is edited; strip + // their names before sending requests to an arbitrary staging/proxy host + // while preserving unrelated headers explicitly supplied by the user. + // + // Gate on RuntimeHeaders as well as CustomHeaders: kimi-code's identity + // headers are produced lazily by RuntimeHeaders at Get/Require time. If + // that map is empty (or a caller holds a listing descriptor that never + // ran RuntimeHeaders), we still must strip any persisted X-Msh-* keys + // rather than forwarding them to a retargeted host. for profileKey := range profile.CustomHeaders { for catalogKey := range descriptor.CustomHeaders { if strings.EqualFold(profileKey, catalogKey) { @@ -1134,6 +1184,15 @@ func applyCatalogDescriptor(profile *ProviderProfile, descriptor providercatalog break } } + // Only a descriptor that mints runtime headers owns the X-Msh-* + // namespace, so only it may drop keys by prefix. This branch is also + // reached by descriptors that merely carry static CustomHeaders + // (AIMLAPI attribution, say), and those must not delete a user's own + // X-Msh-* header they had nothing to do with. Matches the same + // RuntimeHeaders test on the canonical-endpoint path above. + if descriptor.RuntimeHeaders != nil && providercatalog.IsRuntimeIdentityHeader(profileKey) { + delete(profile.CustomHeaders, profileKey) + } } } } diff --git a/internal/config/resolver_test.go b/internal/config/resolver_test.go index 13038664e..1ab9c2515 100644 --- a/internal/config/resolver_test.go +++ b/internal/config/resolver_test.go @@ -9,6 +9,7 @@ import ( "testing" "github.com/Gitlawb/zero/internal/modelregistry" + "github.com/Gitlawb/zero/internal/oauth" "github.com/Gitlawb/zero/internal/providercatalog" ) @@ -492,6 +493,120 @@ func TestResolveRejectsProjectBaseURLOverrideWithInheritedCompatibleCredentials( } } +func TestResolveRejectsProjectBaseURLOverrideWithStoredKimiOAuth(t *testing.T) { + isolateKimiDeviceIDStorage(t) + tokenPath := filepath.Join(t.TempDir(), "oauth-tokens.json") + store, err := oauth.NewStore(oauth.StoreOptions{FilePath: tokenPath}) + if err != nil { + t.Fatalf("NewStore: %v", err) + } + if err := store.Save(oauth.ProviderKey("kimi-code"), oauth.Token{AccessToken: "kimi-secret-token"}); err != nil { + t.Fatalf("seed Kimi login: %v", err) + } + + // Poison process env so a store built from ambient env would miss the token + // and fail open. The lookup must use ResolveOptions.Env instead. + t.Setenv("ZERO_OAUTH_TOKENS_PATH", filepath.Join(t.TempDir(), "ambient-empty.json")) + t.Setenv("ZERO_OAUTH_STORAGE", "file") + + userPath := writeConfig(t, `{ + "activeProvider": "kimi-code", + "providers": [{ + "name": "kimi-code", + "catalogID": "kimi-code" + }] + }`) + projectPath := writeConfig(t, `{ + "providers": [{ + "name": "kimi-code", + "baseURL": "https://attacker.example/v1" + }] + }`) + + _, err = Resolve(ResolveOptions{ + UserConfigPath: userPath, + ProjectConfigPath: projectPath, + Env: map[string]string{ + "ZERO_OAUTH_TOKENS_PATH": tokenPath, + "ZERO_OAUTH_STORAGE": "file", + }, + }) + if err == nil { + t.Fatal("Resolve() error = nil, want project baseURL override rejection for stored OAuth") + } + if strings.Contains(err.Error(), "kimi-secret-token") { + t.Fatalf("error leaked OAuth token: %q", err.Error()) + } + if !strings.Contains(err.Error(), "cannot override baseURL") { + t.Fatalf("error = %q, want project baseURL override rejection", err.Error()) + } +} + +func TestResolveRejectsProjectBaseURLWhenOAuthLookupErrors(t *testing.T) { + isolateKimiDeviceIDStorage(t) + userPath, projectPath := writeKimiProjectBaseURLOverride(t) + + t.Run("injected lookup error", func(t *testing.T) { + prev := storedOAuthLogin + var sawEnv map[string]string + storedOAuthLogin = func(env map[string]string, names []string) (bool, error) { + sawEnv = env + return false, errors.New("keychain unavailable") + } + t.Cleanup(func() { storedOAuthLogin = prev }) + + env := map[string]string{"ZERO_OAUTH_TOKENS_PATH": filepath.Join(t.TempDir(), "unused.json")} + _, err := Resolve(ResolveOptions{ + UserConfigPath: userPath, + ProjectConfigPath: projectPath, + Env: env, + }) + if err == nil { + t.Fatal("Resolve() error = nil, want fail-closed rejection when OAuth lookup errors") + } + if !strings.Contains(err.Error(), "cannot override baseURL") { + t.Fatalf("error = %q, want project baseURL override rejection", err.Error()) + } + if sawEnv["ZERO_OAUTH_TOKENS_PATH"] != env["ZERO_OAUTH_TOKENS_PATH"] { + t.Fatalf("storedOAuthLogin env = %#v, want ResolveOptions.Env", sawEnv) + } + }) + + t.Run("store construction error", func(t *testing.T) { + _, err := Resolve(ResolveOptions{ + UserConfigPath: userPath, + ProjectConfigPath: projectPath, + Env: map[string]string{ + "ZERO_OAUTH_STORAGE": "not-a-backend", + }, + }) + if err == nil { + t.Fatal("Resolve() error = nil, want fail-closed rejection when OAuth store construction fails") + } + if !strings.Contains(err.Error(), "cannot override baseURL") { + t.Fatalf("error = %q, want project baseURL override rejection", err.Error()) + } + }) +} + +func writeKimiProjectBaseURLOverride(t *testing.T) (userPath string, projectPath string) { + t.Helper() + userPath = writeConfig(t, `{ + "activeProvider": "kimi-code", + "providers": [{ + "name": "kimi-code", + "catalogID": "kimi-code" + }] + }`) + projectPath = writeConfig(t, `{ + "providers": [{ + "name": "kimi-code", + "baseURL": "https://attacker.example/v1" + }] + }`) + return userPath, projectPath +} + func TestResolveAllowsProjectCatalogAPIKeyEnvOnCatalogEndpoint(t *testing.T) { path := writeConfig(t, `{ "activeProvider": "openrouter", @@ -1577,6 +1692,156 @@ func TestApplyCatalogDescriptorStripsAimlapiAttributionFromRetargetedProfile(t * } } +// TestApplyCatalogDescriptorKeepsUserXMshOnNonKimiRetarget pins that the +// X-Msh-* namespace is only strippable by a descriptor that actually mints +// those headers. AIMLAPI reaches the same retarget branch because it carries +// static CustomHeaders, but it has no RuntimeHeaders and no claim on X-Msh-*, +// so a header the user set themselves must survive. +func TestApplyCatalogDescriptorKeepsUserXMshOnNonKimiRetarget(t *testing.T) { + descriptor, err := providercatalog.Require("aimlapi") + if err != nil { + t.Fatal(err) + } + if descriptor.RuntimeHeaders != nil { + t.Fatal("setup: aimlapi is expected to have no RuntimeHeaders; pick another non-runtime descriptor") + } + profile := ProviderProfile{ + BaseURL: "https://proxy.example.test/v1", + CustomHeaders: map[string]string{ + "x-aimlapi-partner-id": "persisted-partner", + "X-Msh-Trace": "user-supplied", + }, + } + + applyCatalogDescriptor(&profile, descriptor, true) + + if profile.CustomHeaders["X-Msh-Trace"] != "user-supplied" { + t.Fatalf("a descriptor that does not own X-Msh-* deleted the user's own header: %#v", profile.CustomHeaders) + } + // The descriptor's own catalog headers must still be stripped. + for key := range profile.CustomHeaders { + if strings.HasPrefix(strings.ToLower(key), "x-aimlapi-") { + t.Fatalf("catalog attribution survived retargeting: %#v", profile.CustomHeaders) + } + } +} + +// isolateKimiDeviceIDStorage redirects os.UserConfigDir() to a throwaway temp +// dir before the kimi-code descriptor's RuntimeHeaders (kimiidentity.Headers) +// run, so tests never create or touch the real kimi-device-id file. Sets all +// three env vars os.UserConfigDir may consult so isolation is portable. +// DeviceID is path-keyed, so no separate cache reset is required. +func isolateKimiDeviceIDStorage(t *testing.T) { + t.Helper() + root := t.TempDir() + t.Setenv("XDG_CONFIG_HOME", root) + t.Setenv("APPDATA", root) + t.Setenv("HOME", root) +} + +func TestApplyCatalogDescriptorStripsKimiIdentityFromRetargetedProfile(t *testing.T) { + isolateKimiDeviceIDStorage(t) + descriptor, err := providercatalog.Require("kimi-code") + if err != nil { + t.Fatal(err) + } + profile := ProviderProfile{ + BaseURL: "https://proxy.example.test/v1", + CustomHeaders: map[string]string{ + "x-msh-platform": "kimi_code_cli", + "X-Msh-Device-Id": "persisted-device-id", + "X-Environment": "staging", + }, + } + + applyCatalogDescriptor(&profile, descriptor, true) + + for key := range profile.CustomHeaders { + if strings.HasPrefix(strings.ToLower(key), "x-msh-") { + t.Fatalf("kimi identity headers survived retargeting: %#v", profile.CustomHeaders) + } + } + if profile.CustomHeaders["X-Environment"] != "staging" { + t.Fatalf("user header was removed: %#v", profile.CustomHeaders) + } +} + +// TestApplyCatalogDescriptorStripsXMshWhenRuntimeHeadersOnly covers the case +// the x-msh- prefix branch exists for: a listing-style descriptor that has +// RuntimeHeaders set but empty CustomHeaders. The EqualFold catalog-key loop +// alone cannot strip persisted identity headers when CustomHeaders is empty. +func TestApplyCatalogDescriptorStripsXMshWhenRuntimeHeadersOnly(t *testing.T) { + isolateKimiDeviceIDStorage(t) + descriptor := providercatalog.Descriptor{ + ID: "kimi-code", + DefaultBaseURL: "https://api.kimi.com/coding/v1", + RuntimeHeaders: func() map[string]string { + return map[string]string{"X-Msh-Device-Id": "fresh-from-runtime"} + }, + // CustomHeaders deliberately empty — listing clones never run RuntimeHeaders. + } + profile := ProviderProfile{ + BaseURL: "https://proxy.example.test/v1", + CustomHeaders: map[string]string{ + "X-Msh-Device-Id": "stale-persisted", + "X-Msh-Device-Name": "hostname-from-disk", + "X-User-Agent": "keep-me", + }, + } + + applyCatalogDescriptor(&profile, descriptor, true) + + for key := range profile.CustomHeaders { + if strings.HasPrefix(strings.ToLower(key), "x-msh-") { + t.Fatalf("x-msh- header survived retargeting with empty catalog CustomHeaders: %#v", profile.CustomHeaders) + } + } + if profile.CustomHeaders["X-User-Agent"] != "keep-me" { + t.Fatalf("user header was removed: %#v", profile.CustomHeaders) + } +} + +// TestApplyCatalogDescriptorFreshRuntimeHeadersWinOverStalePersisted pins the +// stored-header override guard: a stale X-Msh-Device-Id in config.json must +// lose to the freshly minted RuntimeHeaders value on the canonical endpoint. +func TestApplyCatalogDescriptorFreshRuntimeHeadersWinOverStalePersisted(t *testing.T) { + isolateKimiDeviceIDStorage(t) + const freshID = "fresh-runtime-device-id" + descriptor := providercatalog.Descriptor{ + ID: "kimi-code", + DefaultBaseURL: "https://api.kimi.com/coding/v1", + CustomHeaders: map[string]string{ + "X-Msh-Platform": "kimi_code_cli", + "X-Msh-Device-Id": freshID, + }, + RuntimeHeaders: func() map[string]string { + return map[string]string{ + "X-Msh-Platform": "kimi_code_cli", + "X-Msh-Device-Id": freshID, + } + }, + } + profile := ProviderProfile{ + BaseURL: "https://api.kimi.com/coding/v1", + CustomHeaders: map[string]string{ + "X-Msh-Device-Id": "stale-persisted-device-id", + "X-User-Agent": "keep-me", + }, + } + + applyCatalogDescriptor(&profile, descriptor, false) + + if got := profile.CustomHeaders["X-Msh-Device-Id"]; got != freshID { + t.Fatalf("X-Msh-Device-Id = %q, want fresh runtime value %q (stale persist must not win)", got, freshID) + } + if profile.CustomHeaders["X-User-Agent"] != "keep-me" { + t.Fatalf("user header was removed: %#v", profile.CustomHeaders) + } + if profile.CustomHeaders["X-Msh-Platform"] != "kimi_code_cli" { + t.Fatalf("catalog platform header missing: %#v", profile.CustomHeaders) + } +} + func TestResolveProviderProfileParseThinkTagsFalseAlias(t *testing.T) { path := writeConfig(t, `{ "activeProvider": "custom", @@ -2087,7 +2352,7 @@ func TestMergeProjectConfigIgnoresAdditionalWriteRoots(t *testing.T) { dst.Sandbox.AdditionalWriteRoots = []string{"/global/one"} src := FileConfig{} src.Sandbox.AdditionalWriteRoots = []string{"/repo/sneaky"} - if err := mergeProjectConfig(&dst, src); err != nil { + if err := mergeProjectConfig(&dst, src, nil); err != nil { t.Fatalf("mergeProjectConfig: %v", err) } if !reflect.DeepEqual(dst.Sandbox.AdditionalWriteRoots, []string{"/global/one"}) { @@ -2231,7 +2496,7 @@ func TestCrossSessionInboundProjectConfigCanOnlyTighten(t *testing.T) { for _, test := range tests { t.Run(test.name, func(t *testing.T) { cfg := FileConfig{CrossSessionInbound: test.user} - if err := mergeProjectConfig(&cfg, FileConfig{CrossSessionInbound: test.project}); err != nil { + if err := mergeProjectConfig(&cfg, FileConfig{CrossSessionInbound: test.project}, nil); err != nil { t.Fatal(err) } if cfg.CrossSessionInbound != test.want { diff --git a/internal/kimiidentity/export_test.go b/internal/kimiidentity/export_test.go new file mode 100644 index 000000000..a7686397c --- /dev/null +++ b/internal/kimiidentity/export_test.go @@ -0,0 +1,59 @@ +package kimiidentity + +import ( + "context" + "os" + "testing" + "time" +) + +// IsolateDeviceIDStorage redirects the env vars os.UserConfigDir consults so +// subsequent DeviceID/Headers calls store under root for the duration of t +// (and any nested tests that inherit the env). Sets XDG_CONFIG_HOME, APPDATA, +// and HOME together so the redirect is portable across Windows, macOS, and +// Linux without GOOS branching. +// +// DeviceID's cache is path-keyed, so no separate cache reset is required: +// once the config root changes, the next DeviceID/Headers call reloads. +// +// Cross-package tests cannot call this (export_test.go is package-local to +// go test of kimiidentity); they should set the same three env keys via +// t.Setenv on a t.TempDir() before invoking Headers/DeviceID or +// providercatalog.Get("kimi-code"). +func IsolateDeviceIDStorage(t *testing.T) { + t.Helper() + root := t.TempDir() + t.Setenv("XDG_CONFIG_HOME", root) + t.Setenv("APPDATA", root) + t.Setenv("HOME", root) +} + +// SetBeforeRenameHook sets a hook invoked immediately before the staged +// device-id file is renamed into place. Returns a cleanup function. +func SetBeforeRenameHook(hook func()) func() { + prev := beforeRenameHook + beforeRenameHook = hook + return func() { + beforeRenameHook = prev + } +} + +func SetDeviceIDMaxWait(d time.Duration) func() { + prev := deviceIDMaxWait + deviceIDMaxWait = d + return func() { deviceIDMaxWait = prev } +} + +// SetReadDeviceLock replaces the lock-file reader. Tests use this to inject +// a transient Windows-style read failure while a holder is publishing. +func SetReadDeviceLock(fn func(root *os.Root, name string) ([]byte, error)) func() { + prev := readDeviceLock + if fn != nil { + readDeviceLock = fn + } + return func() { readDeviceLock = prev } +} + +func LoadOrCreateDeviceIDAtContext(ctx context.Context, path string) string { + return loadOrCreateDeviceIDAtContext(ctx, path) +} diff --git a/internal/kimiidentity/kimiidentity.go b/internal/kimiidentity/kimiidentity.go new file mode 100644 index 000000000..edf900ec4 --- /dev/null +++ b/internal/kimiidentity/kimiidentity.go @@ -0,0 +1,469 @@ +// Package kimiidentity builds the X-Msh-* vendor-identity headers Kimi +// Code's backend requires on every request — OAuth device authorization, +// token polling, refresh, AND managed-endpoint model calls. It is shared by +// both internal/oauth (login/refresh) and internal/providercatalog (the +// kimi-code descriptor's CustomHeaders, applied to runtime completions) so +// they send the SAME identity: a login accepted under one device identity +// and completions sent under another (or under none) is rejected by the +// backend. +// +// Header names and general shape are reverse-engineered from the +// open-source kimi-cli client (src/kimi_cli/auth/oauth.py, _common_headers); +// Kimi has no published public API documentation for this, so these values +// are a best-effort match, not a verified spec. +package kimiidentity + +import ( + "context" + "crypto/rand" + "errors" + "fmt" + "os" + "path/filepath" + "runtime" + "strconv" + "strings" + "sync" + "time" + + "github.com/Gitlawb/zero/internal/lockutil" +) + +// Headers returns the X-Msh-* vendor-identity headers, including the stable +// per-device identifier. +// +// X-Msh-Platform is "kimi_code_cli". That is the value Moonshot's own Kimi +// Code CLI sends (packages/oauth/src/identity.ts, KIMI_CODE_PLATFORM) as of +// its oauth package changelog entry correcting the header from an earlier +// "kimi-code-cli" typo (PR MoonshotAI/kimi-code#52, commit 064343a); the +// older, separate open-source kimi-cli client instead hardcodes "kimi_cli". +// Kimi's coding/v1 endpoint documents a client whitelist ("Kimi CLI, Claude +// Code, Roo Code, ..."); sending the wrong platform value risks the managed +// endpoint rejecting completions even after a successful login. +func Headers() map[string]string { + hostname, err := os.Hostname() + if err != nil || strings.TrimSpace(hostname) == "" { + hostname = "unknown-host" + } + return map[string]string{ + "X-Msh-Platform": "kimi_code_cli", + "X-Msh-Version": "unknown", + "X-Msh-Device-Name": asciiHeaderValue(hostname), + "X-Msh-Device-Model": asciiHeaderValue(runtime.GOOS + " " + runtime.GOARCH), + "X-Msh-Os-Version": runtime.GOOS, + "X-Msh-Device-Id": DeviceID(), + } +} + +var ( + deviceIDMu sync.Mutex + cachedDevicePath string + cachedDeviceID string +) + +// DeviceID returns the persistent device identifier sent as X-Msh-Device-Id. +// Kimi Code's own CLI persists this to ~/.kimi/device_id so the same value +// follows a device across logins, refreshes, and model calls; mirroring +// that, the ID is stored under the user config dir (zero/kimi-device-id) and +// minted once on first use. Acquisition is bounded: an existing valid ID is +// returned, a live publisher is waited on only up to deviceIDMaxWait, a +// proven-dead or expired lease is reclaimed, and unreadable/unwritable +// storage, cancellation, or a live holder that outlasts the wait all return +// a process-local ID without overwriting a file this process does not own. +// +// The cache is keyed by the resolved storage path so tests that redirect +// os.UserConfigDir (via XDG_CONFIG_HOME / APPDATA / HOME) pick up a fresh +// identity without a separate test-only reset hook. +func DeviceID() string { + deviceIDMu.Lock() + defer deviceIDMu.Unlock() + path := deviceIDPath() + if cachedDeviceID != "" && cachedDevicePath == path { + return cachedDeviceID + } + id := loadOrCreateDeviceIDAt(path) + cachedDevicePath = path + cachedDeviceID = id + return id +} + +// loadOrCreateDeviceIDAt is the real load-or-create logic behind DeviceID, +// parameterized by the storage path so tests can exercise production code +// directly (env var indirection through os.UserConfigDir is not portable to +// redirect in tests). It reads an existing UUID if present, otherwise mints +// one and persists it exclusively (see the concurrency note below). +// +// path must be of the form /zero/kimi-device-id. All file +// operations bind to an opened handle and then the zero/ +// subdirectory so a symlink at zero cannot redirect device-id, lock, or +// temporary-file traffic outside the configuration root. +func loadOrCreateDeviceIDAt(path string) string { + return loadOrCreateDeviceIDAtContext(context.Background(), path) +} + +// loadOrCreateDeviceIDAtContext is loadOrCreateDeviceIDAt with a caller +// context. Cancellation returns a process-local id without touching a live +// holder's persisted file. +func loadOrCreateDeviceIDAtContext(ctx context.Context, path string) string { + if path == "" { + return generateDeviceID() + } + root, name, err := openDeviceIDDir(path) + if err != nil { + return generateDeviceID() + } + defer root.Close() + if id := readValidDeviceID(root, name); id != "" { + return id + } + return publishOrAdoptDeviceID(ctx, root, name, generateDeviceID()) +} + +// openDeviceIDDir opens the zero/ directory under the configuration root for +// path (/zero/) using rooted, traversal-resistant handles. +// A zero component that is a symlink escaping the config root is rejected by +// Root.OpenRoot rather than followed into attacker-controlled storage. +func openDeviceIDDir(path string) (*os.Root, string, error) { + name := filepath.Base(path) + zeroDir := filepath.Dir(path) + configDir := filepath.Dir(zeroDir) + zeroName := filepath.Base(zeroDir) + if name == "" || name == "." || zeroName == "" || zeroName == "." || configDir == "" || configDir == "." { + return nil, "", fmt.Errorf("kimiidentity: invalid device-id path %q", path) + } + if err := os.MkdirAll(configDir, 0o700); err != nil { + return nil, "", err + } + cfgRoot, err := os.OpenRoot(configDir) + if err != nil { + return nil, "", err + } + // Best-effort create; Exist is fine. Root refuses a zero symlink that + // points outside configDir on the subsequent OpenRoot. + _ = cfgRoot.Mkdir(zeroName, 0o700) + zeroRoot, err := cfgRoot.OpenRoot(zeroName) + _ = cfgRoot.Close() + if err != nil { + return nil, "", err + } + return zeroRoot, name, nil +} + +var ( + beforeRenameHook func() + readDeviceLock = func(root *os.Root, name string) ([]byte, error) { + return root.ReadFile(name) + } + deviceIDNow = time.Now + deviceIDMaxWait = 5 * time.Second + deviceIDLeaseTTL = 10 * time.Second +) + +type publishOutcome int + +const ( + publishOK publishOutcome = iota + publishContended + publishPersistFailed +) + +func publishDeviceIDAsHolder(root *os.Root, name, lockName, id, ownerToken string) (string, publishOutcome) { + tmpLockName := fmt.Sprintf("%s.tmp.%s", lockName, ownerToken) + if err := root.WriteFile(tmpLockName, []byte(ownerToken+"\n"), 0o600); err != nil { + return "", publishPersistFailed + } + defer func() { _ = root.Remove(tmpLockName) }() + + if err := root.Link(tmpLockName, lockName); err != nil { + if errors.Is(err, os.ErrExist) { + return "", publishContended + } + // Fallback for filesystems where Link is unsupported: + lock, oerr := root.OpenFile(lockName, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if oerr != nil { + if errors.Is(oerr, os.ErrExist) { + return "", publishContended + } + return "", publishPersistFailed + } + if _, werr := lock.WriteString(ownerToken + "\n"); werr != nil { + _ = lock.Close() + _ = root.Remove(lockName) + return "", publishPersistFailed + } + if serr := lock.Sync(); serr != nil { + _ = lock.Close() + _ = root.Remove(lockName) + return "", publishPersistFailed + } + _ = lock.Close() + } + + defer func() { + if curRaw, rerr := root.ReadFile(lockName); rerr == nil && strings.TrimSpace(string(curRaw)) == ownerToken { + _ = root.Remove(lockName) + } + }() + + if existingID := readValidDeviceID(root, name); existingID != "" { + return existingID, publishOK + } + + tmpName := tmpDeviceIDName(name) + if err := writeDeviceIDFile(root, tmpName, id); err != nil { + if existingID := readValidDeviceIDWithRetry(root, name); existingID != "" { + return existingID, publishOK + } + return "", publishPersistFailed + } + defer func() { _ = root.Remove(tmpName) }() + + if beforeRenameHook != nil { + beforeRenameHook() + } + + if err := root.Rename(tmpName, name); err != nil { + if existingID := readValidDeviceIDWithRetry(root, name); existingID != "" { + return existingID, publishOK + } + return "", publishPersistFailed + } + if existingID := readValidDeviceID(root, name); existingID != "" { + return existingID, publishOK + } + return id, publishOK +} + +func publishOrAdoptDeviceID(ctx context.Context, root *os.Root, name, id string) string { + if existingID := readValidDeviceID(root, name); existingID != "" { + return existingID + } + if ctx == nil { + ctx = context.Background() + } + lockName := name + ".lock" + ownerToken := fmt.Sprintf("%d.%d", os.Getpid(), deviceIDNow().UnixNano()) + deadline := deviceIDNow().Add(deviceIDMaxWait) + + const pollInterval = 10 * time.Millisecond + + for { + if existingID := readValidDeviceID(root, name); existingID != "" { + return existingID + } + if ctx.Err() != nil { + return id + } + + publishedID, outcome := publishDeviceIDAsHolder(root, name, lockName, id, ownerToken) + switch outcome { + case publishOK: + return publishedID + case publishPersistFailed: + if existingID := readValidDeviceID(root, name); existingID != "" { + return existingID + } + return id + } + + if existingID := readValidDeviceID(root, name); existingID != "" { + return existingID + } + + raw, rerr := readDeviceLock(root, lockName) + switch { + case rerr != nil && !errors.Is(rerr, os.ErrNotExist): + // Windows often returns sharing-violation or access-denied + // while the holder is deleting the lock after publish. That is + // contention, not a dead store: returning a process-local id + // here makes concurrent first-use callers diverge. + if existingID := readValidDeviceID(root, name); existingID != "" { + return existingID + } + if !deviceIDNow().Before(deadline) || ctx.Err() != nil { + return id + } + time.Sleep(pollInterval) + continue + case rerr == nil && lockHolderAlive(raw): + if !deviceIDNow().Before(deadline) || ctx.Err() != nil { + if existingID := readValidDeviceID(root, name); existingID != "" { + return existingID + } + return id + } + time.Sleep(pollInterval) + continue + } + + reclaimed, rerr := reclaimDeadRepairLock(root, lockName) + if rerr != nil { + return id + } + if reclaimed { + continue + } + if !deviceIDNow().Before(deadline) || ctx.Err() != nil { + return id + } + time.Sleep(pollInterval) + } +} + +// reclaimDeadRepairLock renames the repair lock aside and keeps it only when +// the holder is proven dead (or the lock is empty/corrupt and therefore not a +// live lease). Uses lockutil's rooted reclaim so only one racer wins the +// rename-aside, a live holder's lock is restored rather than stolen, and every +// rename/read/remove stays inside the opened root handle instead of re-walking +// root.Name()+lockName as plain paths (a symlink or reparse point swapped in +// under the lock name after the root was opened cannot redirect them). +func reclaimDeadRepairLock(root *os.Root, lockName string) (bool, error) { + suffix := fmt.Sprintf("%d.%d", os.Getpid(), time.Now().UnixNano()) + return lockutil.ReclaimStaleLockRooted(root, lockName, suffix, lockHolderAlive) +} + +// lockHolderAlive reports whether the repair-lock contents still represent a +// live holder. Token format is ".". Empty or unparseable contents +// are treated as dead (abandoned claim) so a crashed mid-write holder can be +// recovered. A parseable live PID is not enough: the lease also expires after +// deviceIDLeaseTTL so a reused PID cannot pin the lock forever. +func lockHolderAlive(raw []byte) bool { + pid, issued, ok := parseLockToken(strings.TrimSpace(string(raw))) + if !ok || pid <= 0 { + return false + } + if deviceIDNow().Sub(issued) > deviceIDLeaseTTL { + return false + } + return processAlive(pid) +} + +func parseLockToken(token string) (pid int, issued time.Time, ok bool) { + if token == "" { + return 0, time.Time{}, false + } + dot := strings.IndexByte(token, '.') + if dot <= 0 || dot == len(token)-1 { + return 0, time.Time{}, false + } + pid, err := strconv.Atoi(token[:dot]) + if err != nil { + return 0, time.Time{}, false + } + nano, err := strconv.ParseInt(token[dot+1:], 10, 64) + if err != nil || nano <= 0 { + return 0, time.Time{}, false + } + return pid, time.Unix(0, nano), true +} + +// writeDeviceIDFile writes a complete id+"\n" to root/name, checking write, +// sync, and close errors. On any failure the partial file is removed. +func writeDeviceIDFile(root *os.Root, name, id string) error { + f, err := root.OpenFile(name, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600) + if err != nil { + return err + } + if _, err := f.WriteString(id + "\n"); err != nil { + _ = f.Close() + _ = root.Remove(name) + return err + } + if err := f.Sync(); err != nil { + _ = f.Close() + _ = root.Remove(name) + return err + } + if err := f.Close(); err != nil { + _ = root.Remove(name) + return err + } + return nil +} + +func tmpDeviceIDName(name string) string { + return fmt.Sprintf("%s.tmp.%d.%d", name, os.Getpid(), time.Now().UnixNano()) +} + +// readValidDeviceID returns a UUID from root/name, or "" if missing/invalid. +func readValidDeviceID(root *os.Root, name string) string { + raw, err := root.ReadFile(name) + if err != nil { + return "" + } + if id := strings.TrimSpace(string(raw)); isUUID(id) { + return id + } + return "" +} + +// readValidDeviceIDWithRetry re-reads briefly so a process that lost the +// exclusive create can adopt the winner even if it observed the file before +// the winner finished publishing the UUID. +func readValidDeviceIDWithRetry(root *os.Root, name string) string { + const attempts = 40 + const delay = 5 * time.Millisecond + for i := 0; i < attempts; i++ { + if id := readValidDeviceID(root, name); id != "" { + return id + } + time.Sleep(delay) + } + return "" +} + +func deviceIDPath() string { + configDir, err := os.UserConfigDir() + if err != nil || strings.TrimSpace(configDir) == "" { + return "" + } + return filepath.Join(configDir, "zero", "kimi-device-id") +} + +func generateDeviceID() string { + raw := make([]byte, 16) + if _, err := rand.Read(raw); err != nil { + return "00000000-0000-0000-0000-000000000000" + } + raw[6] = (raw[6] & 0x0f) | 0x40 // version 4 + raw[8] = (raw[8] & 0x3f) | 0x80 // variant 10 + return fmt.Sprintf("%x-%x-%x-%x-%x", raw[0:4], raw[4:6], raw[6:8], raw[8:10], raw[10:16]) +} + +func isUUID(s string) bool { + if len(s) != 36 { + return false + } + for i, r := range s { + switch i { + case 8, 13, 18, 23: + if r != '-' { + return false + } + default: + if (r < '0' || r > '9') && (r < 'a' || r > 'f') && (r < 'A' || r > 'F') { + return false + } + } + } + return true +} + +// asciiHeaderValue strips anything outside printable ASCII (0x20-0x7e). This +// mirrors a defensive fix kimi-cli itself needed: a raw platform-version +// string containing "#" broke an HTTP client's header validation on Linux +// (MoonshotAI/kimi-cli#1169) because HTTP header values must not contain +// control characters. +func asciiHeaderValue(s string) string { + var b strings.Builder + for _, r := range s { + if r >= 0x20 && r <= 0x7e { + b.WriteRune(r) + } + } + clean := strings.TrimSpace(b.String()) + if clean == "" { + return "unknown" + } + return clean +} diff --git a/internal/kimiidentity/kimiidentity_test.go b/internal/kimiidentity/kimiidentity_test.go new file mode 100644 index 000000000..e91e61b89 --- /dev/null +++ b/internal/kimiidentity/kimiidentity_test.go @@ -0,0 +1,760 @@ +package kimiidentity + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "runtime" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestHeadersIncludesDeviceIdentity(t *testing.T) { + IsolateDeviceIDStorage(t) + headers := Headers() + for _, key := range []string{ + "X-Msh-Platform", + "X-Msh-Version", + "X-Msh-Device-Name", + "X-Msh-Device-Model", + "X-Msh-Os-Version", + "X-Msh-Device-Id", + } { + if headers[key] == "" { + t.Fatalf("Headers()[%q] empty", key) + } + } + if headers["X-Msh-Platform"] != "kimi_code_cli" { + t.Fatalf("X-Msh-Platform = %q, want kimi_code_cli", headers["X-Msh-Platform"]) + } + if !isUUID(headers["X-Msh-Device-Id"]) { + t.Fatalf("X-Msh-Device-Id = %q, want UUID", headers["X-Msh-Device-Id"]) + } +} + +// TestDeviceIDReloadsWhenConfigRootChanges pins the path-keyed cache: after an +// identity is cached for one config root, redirecting os.UserConfigDir must +// mint (or load) a different root's id rather than returning the first. +func TestDeviceIDReloadsWhenConfigRootChanges(t *testing.T) { + root1 := t.TempDir() + t.Setenv("XDG_CONFIG_HOME", root1) + t.Setenv("APPDATA", root1) + t.Setenv("HOME", root1) + id1 := DeviceID() + if !isUUID(id1) { + t.Fatalf("DeviceID() under root1 = %q, want UUID", id1) + } + path1 := mustDeviceIDPath(t) + if raw, err := os.ReadFile(path1); err != nil { + t.Fatalf("read root1 device id: %v", err) + } else if got := strings.TrimSpace(string(raw)); got != id1 { + t.Fatalf("root1 file = %q, want %q", got, id1) + } + + root2 := t.TempDir() + t.Setenv("XDG_CONFIG_HOME", root2) + t.Setenv("APPDATA", root2) + t.Setenv("HOME", root2) + id2 := DeviceID() + if !isUUID(id2) { + t.Fatalf("DeviceID() under root2 = %q, want UUID", id2) + } + if id1 == id2 { + t.Fatalf("DeviceID reused first root's id %q after config root change", id1) + } + path2 := mustDeviceIDPath(t) + if path1 == path2 { + t.Fatalf("device id path did not change with config root: %q", path1) + } + if raw, err := os.ReadFile(path2); err != nil { + t.Fatalf("read root2 device id: %v", err) + } else if got := strings.TrimSpace(string(raw)); got != id2 { + t.Fatalf("root2 file = %q, want %q", got, id2) + } + // First root's file must still hold id1 (no clobber across roots). + if raw, err := os.ReadFile(path1); err != nil { + t.Fatalf("re-read root1 device id: %v", err) + } else if got := strings.TrimSpace(string(raw)); got != id1 { + t.Fatalf("root1 file changed after root2 mint: got %q, want %q", got, id1) + } +} + +func mustDeviceIDPath(t *testing.T) string { + t.Helper() + configDir, err := os.UserConfigDir() + if err != nil { + t.Fatalf("UserConfigDir: %v", err) + } + return filepath.Join(configDir, "zero", "kimi-device-id") +} + +func TestLoadOrCreateDeviceIDExclusiveCreate(t *testing.T) { + // Exercise the production loader directly via its path-parameterized + // helper. Concurrent first-use must converge on a single persisted ID: + // the O_EXCL loser reads back the winner's file instead of overwriting it. + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + + const workers = 8 + ids := make([]string, workers) + var wg sync.WaitGroup + wg.Add(workers) + for i := range workers { + go func(i int) { + defer wg.Done() + ids[i] = loadOrCreateDeviceIDAt(path) + }(i) + } + wg.Wait() + + winner := "" + for _, id := range ids { + if id == "" { + t.Fatal("worker returned empty id") + } + if winner == "" { + winner = id + continue + } + if id != winner { + t.Fatalf("workers diverged: got %q and %q", winner, id) + } + } + if !isUUID(winner) { + t.Fatalf("winner id %q is not a UUID", winner) + } + // The persisted file carries the winner exactly once. + if raw, err := os.ReadFile(path); err != nil { + t.Fatalf("read persisted id: %v", err) + } else if got := strings.TrimSpace(string(raw)); got != winner { + t.Fatalf("persisted id = %q, want %q", got, winner) + } +} + +// TestLoadOrCreateDeviceIDConvergesWhenLockReadIsTransient is the Windows +// exclusive-create regression: a holder deleting the lock after publish can +// make a concurrent ReadFile return a sharing-violation (not ErrNotExist). +// That must be treated as contention so every caller adopts the persisted id. +func TestLoadOrCreateDeviceIDConvergesWhenLockReadIsTransient(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + + pause := make(chan struct{}) + paused := make(chan struct{}) + var once sync.Once + cleanupHook := SetBeforeRenameHook(func() { + once.Do(func() { + close(paused) + <-pause + }) + }) + defer cleanupHook() + + var injected atomic.Bool + transientReadErr := errors.New("sharing violation") + cleanupRead := SetReadDeviceLock(func(root *os.Root, name string) ([]byte, error) { + if injected.CompareAndSwap(false, true) { + return nil, transientReadErr + } + return root.ReadFile(name) + }) + defer cleanupRead() + + var publisherID string + var pubWg sync.WaitGroup + pubWg.Add(1) + go func() { + defer pubWg.Done() + publisherID = loadOrCreateDeviceIDAt(path) + }() + <-paused + + const workers = 8 + ids := make([]string, workers) + var wg sync.WaitGroup + wg.Add(workers) + for i := range workers { + go func(i int) { + defer wg.Done() + ids[i] = loadOrCreateDeviceIDAt(path) + }(i) + } + // Let racers observe the lock (and the injected read failure) first. + time.Sleep(30 * time.Millisecond) + close(pause) + pubWg.Wait() + wg.Wait() + + if !isUUID(publisherID) { + t.Fatalf("publisher id %q is not a UUID", publisherID) + } + for i, id := range ids { + if id != publisherID { + t.Fatalf("worker %d returned %q, want publisher %q (all: %v)", i, id, publisherID, ids) + } + } + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read persisted id: %v", err) + } + if persisted := strings.TrimSpace(string(raw)); persisted != publisherID { + t.Fatalf("persisted %q, want publisher %q", persisted, publisherID) + } +} + +func TestLoadOrCreateDeviceIDReadsExisting(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + const existing = "aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee" + if err := os.WriteFile(path, []byte(existing+"\n"), 0o600); err != nil { + t.Fatal(err) + } + if got := loadOrCreateDeviceIDAt(path); got != existing { + t.Fatalf("loadOrCreateDeviceIDAt = %q, want existing %q", got, existing) + } +} + +// TestLoadOrCreateDeviceIDAdoptsWinnerAfterEmptyCreate covers the +// multi-process window where the lock winner has acquired the lock but not +// yet published the UUID. Concurrent callers must wait and adopt that UUID +// rather than each minting a divergent identity. +func TestLoadOrCreateDeviceIDAdoptsWinnerAfterEmptyCreate(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + lockPath := path + ".lock" + // Simulate the lock winner holding the lock before writing. + lockF, err := os.OpenFile(lockPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + t.Fatal(err) + } + ownerToken := fmt.Sprintf("%d.%d\n", os.Getpid(), time.Now().UnixNano()) + _, _ = lockF.WriteString(ownerToken) + _ = lockF.Sync() + + const winner = "11111111-2222-4333-8444-555555555555" + done := make(chan struct{}) + go func() { + defer close(done) + time.Sleep(30 * time.Millisecond) + _ = os.WriteFile(path, []byte(winner+"\n"), 0o600) + _ = lockF.Close() + _ = os.Remove(lockPath) + }() + + const workers = 4 + ids := make([]string, workers) + var wg sync.WaitGroup + wg.Add(workers) + for i := range workers { + go func(i int) { + defer wg.Done() + ids[i] = loadOrCreateDeviceIDAt(path) + }(i) + } + wg.Wait() + <-done + + for _, id := range ids { + if id != winner { + t.Fatalf("worker returned %q, want winner %q (all: %v)", id, winner, ids) + } + } +} + +// TestLoadOrCreateDeviceIDRepairsAbandonedEmptyFile covers the case where +// a previous process exclusive-created the path and died before writing a +// UUID. Callers must not permanently diverge: after the retry window the +// empty file is removed and a new exclusive create publishes a valid id. +func TestLoadOrCreateDeviceIDRepairsAbandonedEmptyFile(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + t.Fatal(err) + } + _ = f.Close() // abandoned: never written + + got := loadOrCreateDeviceIDAt(path) + if !isUUID(got) { + t.Fatalf("repaired id %q is not a UUID", got) + } + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read repaired file: %v", err) + } + if persisted := strings.TrimSpace(string(raw)); persisted != got { + t.Fatalf("persisted %q, want repaired %q", persisted, got) + } +} + +// TestLoadOrCreateDeviceIDConcurrentAbandonedFileRepairConverges covers +// multiple racing processes all finding the same abandoned/invalid file at +// once. Repair must be mutually exclusive: only one racer may remove and +// recreate the file, so every caller ends up with the same id and that id is +// exactly what is persisted (no caller returns an in-memory id that a later +// repair silently unlinked and replaced). +func TestLoadOrCreateDeviceIDConcurrentAbandonedFileRepairConverges(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + t.Fatal(err) + } + _ = f.Close() // abandoned: never written + + const workers = 16 + ids := make([]string, workers) + var wg sync.WaitGroup + wg.Add(workers) + for i := range workers { + go func(i int) { + defer wg.Done() + ids[i] = loadOrCreateDeviceIDAt(path) + }(i) + } + wg.Wait() + + winner := "" + for _, id := range ids { + if id == "" { + t.Fatal("worker returned empty id") + } + if winner == "" { + winner = id + continue + } + if id != winner { + t.Fatalf("workers diverged repairing abandoned file: got %q and %q", winner, id) + } + } + if !isUUID(winner) { + t.Fatalf("winner id %q is not a UUID", winner) + } + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read persisted id: %v", err) + } + if persisted := strings.TrimSpace(string(raw)); persisted != winner { + t.Fatalf("persisted %q, want winner %q", persisted, winner) + } +} + +func TestLoadOrCreateDeviceIDRepairsStaleRepairLock(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + t.Fatal(err) + } + _ = f.Close() // abandoned target file + + lockPath := path + ".lock" + // Empty lock contents are treated as abandoned (unparseable holder) and + // reclaimed. A live holder's "." token is not reclaimed. + lockF, err := os.OpenFile(lockPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + t.Fatal(err) + } + _ = lockF.Close() + + got := loadOrCreateDeviceIDAt(path) + if !isUUID(got) { + t.Fatalf("repaired id %q is not a UUID", got) + } + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read repaired file: %v", err) + } + if persisted := strings.TrimSpace(string(raw)); persisted != got { + t.Fatalf("persisted %q, want repaired %q", persisted, got) + } + if _, err := os.Stat(lockPath); !os.IsNotExist(err) { + t.Fatalf("repair lock should be cleaned up after reclaim+repair: err=%v", err) + } +} + +func TestLoadOrCreateDeviceIDRepairsDeadPIDRepairLock(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + t.Fatal(err) + } + _ = f.Close() + + lockPath := path + ".lock" + // Well-formed token with a non-positive PID is treated as dead (same as a + // crashed holder). Avoid inventing a high PID that might be live. + if err := os.WriteFile(lockPath, []byte("0.12345\n"), 0o600); err != nil { + t.Fatal(err) + } + + got := loadOrCreateDeviceIDAt(path) + if !isUUID(got) { + t.Fatalf("repaired id %q is not a UUID", got) + } + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read repaired file: %v", err) + } + if persisted := strings.TrimSpace(string(raw)); persisted != got { + t.Fatalf("persisted %q, want repaired %q", persisted, got) + } +} + +func TestLockHolderAlive(t *testing.T) { + if lockHolderAlive([]byte("")) { + t.Fatal("empty lock should not be treated as live") + } + if lockHolderAlive([]byte("not-a-token")) { + t.Fatal("unparseable lock should not be treated as live") + } + if lockHolderAlive([]byte("0.1")) { + t.Fatal("non-positive pid should not be treated as live") + } + // Our own PID must be treated as live so we never reclaim our own lock. + self := fmt.Sprintf("%d.%d", os.Getpid(), time.Now().UnixNano()) + if !lockHolderAlive([]byte(self)) { + t.Fatalf("self pid token %q should be live", self) + } + // A live PID with an expired timestamp is dead: PID reuse must not pin the lock. + stale := fmt.Sprintf("%d.%d", os.Getpid(), time.Now().Add(-time.Hour).UnixNano()) + if lockHolderAlive([]byte(stale)) { + t.Fatalf("expired lease %q should not be live", stale) + } +} + +func TestAsciiHeaderValueStripsNonPrintable(t *testing.T) { + if got := asciiHeaderValue("linux#6.1"); got != "linux#6.1" { + // printable ASCII including # is kept; the kimi-cli bug was a different + // control character path — ensure we still strip true controls. + t.Fatalf("got %q", got) + } + if got := asciiHeaderValue("a\nb\x00c"); got != "abc" { + t.Fatalf("got %q, want abc", got) + } + if got := asciiHeaderValue("\x01\x02"); got != "unknown" { + t.Fatalf("got %q, want unknown", got) + } +} + +// TestLoadOrCreateDeviceIDPausedPublisherConverges covers a publisher that is +// paused/descheduled before its rename: concurrent callers must wait for the +// lock holder to complete publication and adopt its published ID rather than +// treating it as abandoned and overwriting it. +func TestLoadOrCreateDeviceIDPausedPublisherConverges(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + + pause := make(chan struct{}) + paused := make(chan struct{}) + var once sync.Once + cleanup := SetBeforeRenameHook(func() { + once.Do(func() { + close(paused) + <-pause + }) + }) + defer cleanup() + + var publisherID string + var pubWg sync.WaitGroup + pubWg.Add(1) + go func() { + defer pubWg.Done() + publisherID = loadOrCreateDeviceIDAt(path) + }() + + // Wait until the publisher holds the lock and reaches the hook before rename. + <-paused + + const workers = 4 + ids := make([]string, workers) + var wg sync.WaitGroup + wg.Add(workers) + for i := range workers { + go func(i int) { + defer wg.Done() + ids[i] = loadOrCreateDeviceIDAt(path) + }(i) + } + + // Give racers a chance to observe the lock, then resume the publisher. + time.Sleep(30 * time.Millisecond) + close(pause) + + pubWg.Wait() + wg.Wait() + + if !isUUID(publisherID) { + t.Fatalf("publisher id %q is not a UUID", publisherID) + } + for i, id := range ids { + if id != publisherID { + t.Fatalf("worker %d returned %q, want publisher winner %q", i, id, publisherID) + } + } + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read persisted id: %v", err) + } + if persisted := strings.TrimSpace(string(raw)); persisted != publisherID { + t.Fatalf("persisted %q, want publisher winner %q", persisted, publisherID) + } +} + +// TestLoadOrCreateDeviceIDWaitsForSlowLiveRepairLockHolder ensures that when a live +// process owns the repair lock, a competing caller waits for the published ID +// instead of returning an unpersisted locally-generated ID. +func TestLoadOrCreateDeviceIDWaitsForSlowLiveRepairLockHolder(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + + lockPath := path + ".lock" + // Live owner token using our own PID. + ownerToken := fmt.Sprintf("%d.%d\n", os.Getpid(), time.Now().UnixNano()) + if err := os.WriteFile(lockPath, []byte(ownerToken), 0o600); err != nil { + t.Fatal(err) + } + + const publishedID = "33333333-4444-4555-8666-777777777777" + done := make(chan struct{}) + go func() { + defer close(done) + // Hold the lock longer than a quick retry window. + time.Sleep(100 * time.Millisecond) + if err := os.WriteFile(path, []byte(publishedID+"\n"), 0o600); err != nil { + t.Errorf("write published id: %v", err) + } + _ = os.Remove(lockPath) + }() + + got := loadOrCreateDeviceIDAt(path) + <-done + + if got != publishedID { + t.Fatalf("loadOrCreateDeviceIDAt = %q, want published %q", got, publishedID) + } + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read persisted device ID: %v", err) + } + if persisted := strings.TrimSpace(string(raw)); persisted != publishedID { + t.Fatalf("persisted device ID = %q, want %q", persisted, publishedID) + } +} + +// TestLiveLockHolderLeaseNeverOverwrittenByCompetitor pins that while a live lock +// holder owns the lease, competing callers must never publish or overwrite the +// device ID file without holding the lock. +func TestLiveLockHolderLeaseNeverOverwrittenByCompetitor(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + + lockPath := path + ".lock" + ownerToken := fmt.Sprintf("%d.%d\n", os.Getpid(), time.Now().UnixNano()) + if err := os.WriteFile(lockPath, []byte(ownerToken), 0o600); err != nil { + t.Fatal(err) + } + + const holderPublishedID = "11111111-2222-4333-8444-555555555555" + done := make(chan struct{}) + go func() { + defer close(done) + time.Sleep(150 * time.Millisecond) + if err := os.WriteFile(path, []byte(holderPublishedID+"\n"), 0o600); err != nil { + t.Errorf("write published ID: %v", err) + } + _ = os.Remove(lockPath) + }() + + got := loadOrCreateDeviceIDAt(path) + <-done + + if got != holderPublishedID { + t.Fatalf("competitor returned %q, want holder's %q", got, holderPublishedID) + } + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read persisted file: %v", err) + } + if persisted := strings.TrimSpace(string(raw)); persisted != holderPublishedID { + t.Fatalf("persisted file = %q, want holder's %q", persisted, holderPublishedID) + } +} + +func TestLoadOrCreateDeviceIDUnwritableDirReturnsLocalWithoutHang(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("directory write-bit is not enforced the same way on Windows") + } + if os.Geteuid() == 0 { + t.Skip("root can write a 0555 directory") + } + dir := t.TempDir() + zeroDir := filepath.Join(dir, "zero") + if err := os.MkdirAll(zeroDir, 0o700); err != nil { + t.Fatal(err) + } + if err := os.Chmod(zeroDir, 0o555); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = os.Chmod(zeroDir, 0o700) }) + path := filepath.Join(zeroDir, "kimi-device-id") + + done := make(chan string, 1) + go func() { done <- loadOrCreateDeviceIDAt(path) }() + var got string + select { + case got = <-done: + case <-time.After(time.Second): + t.Fatal("identity acquisition hung on unwritable storage") + } + if !isUUID(got) { + t.Fatalf("process-local id %q is not a UUID", got) + } + if _, err := os.Stat(path); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("unwritable storage must not persist a device id: stat err=%v", err) + } +} + +func TestLoadOrCreateDeviceIDReclaimsExpiredLivePIDLock(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + lockPath := path + ".lock" + // Our PID is live, but the timestamp is older than the lease TTL: this is + // the PID-reuse case, not a current holder. + stale := fmt.Sprintf("%d.%d\n", os.Getpid(), time.Now().Add(-time.Hour).UnixNano()) + if err := os.WriteFile(lockPath, []byte(stale), 0o600); err != nil { + t.Fatal(err) + } + + got := loadOrCreateDeviceIDAt(path) + if !isUUID(got) { + t.Fatalf("reclaimed id %q is not a UUID", got) + } + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read persisted id: %v", err) + } + if persisted := strings.TrimSpace(string(raw)); persisted != got { + t.Fatalf("persisted %q, want reclaimed %q", persisted, got) + } +} + +func TestLoadOrCreateDeviceIDLiveHolderPastDeadlineDoesNotOverwrite(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + const holderID = "aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee" + if err := os.WriteFile(path, []byte(holderID+"\n"), 0o600); err != nil { + t.Fatal(err) + } + lockPath := path + ".lock" + ownerToken := fmt.Sprintf("%d.%d\n", os.Getpid(), time.Now().UnixNano()) + if err := os.WriteFile(lockPath, []byte(ownerToken), 0o600); err != nil { + t.Fatal(err) + } + restore := SetDeviceIDMaxWait(0) + defer restore() + + got := loadOrCreateDeviceIDAt(path) + if got != holderID { + t.Fatalf("got %q, want existing holder id %q", got, holderID) + } + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read persisted id: %v", err) + } + if persisted := strings.TrimSpace(string(raw)); persisted != holderID { + t.Fatalf("persisted %q, want unchanged holder id %q", persisted, holderID) + } +} + +func TestLoadOrCreateDeviceIDLiveHolderPastDeadlineLeavesMissingFile(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + lockPath := path + ".lock" + ownerToken := fmt.Sprintf("%d.%d\n", os.Getpid(), time.Now().UnixNano()) + if err := os.WriteFile(lockPath, []byte(ownerToken), 0o600); err != nil { + t.Fatal(err) + } + restore := SetDeviceIDMaxWait(0) + defer restore() + + got := loadOrCreateDeviceIDAt(path) + if !isUUID(got) { + t.Fatalf("process-local id %q is not a UUID", got) + } + if _, err := os.Stat(path); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("live-holder deadline fallback must not publish: stat err=%v", err) + } +} + +func TestLoadOrCreateDeviceIDCancelReturnsLocalWithoutWrite(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "zero", "kimi-device-id") + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + lockPath := path + ".lock" + ownerToken := fmt.Sprintf("%d.%d\n", os.Getpid(), time.Now().UnixNano()) + if err := os.WriteFile(lockPath, []byte(ownerToken), 0o600); err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + got := LoadOrCreateDeviceIDAtContext(ctx, path) + if !isUUID(got) { + t.Fatalf("canceled id %q is not a UUID", got) + } + if _, err := os.Stat(path); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("cancellation must not persist a device id: stat err=%v", err) + } +} diff --git a/internal/kimiidentity/process_alive_posix.go b/internal/kimiidentity/process_alive_posix.go new file mode 100644 index 000000000..f833c5ac9 --- /dev/null +++ b/internal/kimiidentity/process_alive_posix.go @@ -0,0 +1,26 @@ +//go:build !windows + +package kimiidentity + +import ( + "errors" + "syscall" +) + +// processAlive reports whether pid is a live process on POSIX. Signal 0 does +// not deliver a signal; it only checks existence/permission. +// +// Fail closed: return false only when the process is established to be gone +// (ESRCH). nil and EPERM mean alive; any other error is treated as alive so +// an inconclusive liveness check never makes a repair lease reclaimable. +func processAlive(pid int) bool { + if pid <= 0 { + return false + } + err := syscall.Kill(pid, 0) + if err == nil || errors.Is(err, syscall.EPERM) { + return true + } + // Only ESRCH proves the process is dead. + return !errors.Is(err, syscall.ESRCH) +} diff --git a/internal/kimiidentity/process_alive_windows.go b/internal/kimiidentity/process_alive_windows.go new file mode 100644 index 000000000..5d927662b --- /dev/null +++ b/internal/kimiidentity/process_alive_windows.go @@ -0,0 +1,56 @@ +//go:build windows + +package kimiidentity + +import ( + "errors" + + "golang.org/x/sys/windows" +) + +// processAlive reports whether pid is a live process on Windows. It opens the +// process with SYNCHRONIZE and probes termination via WaitForSingleObject with +// a zero timeout. Exit-code checks are not used: a terminated process can +// legitimately return exit code 259 (STILL_ACTIVE), which would look alive. +// +// Fail closed on inconclusive results: only a proven-dead process makes a +// repair lease reclaimable. Access-denied means the process exists but we +// cannot open it (alive). ERROR_INVALID_PARAMETER is the usual free/invalid +// PID signal (dead). Any other OpenProcess failure, WaitForSingleObject +// failure, or unexpected wait result is treated as alive. +func processAlive(pid int) bool { + if pid <= 0 { + return false + } + handle, err := windows.OpenProcess(windows.SYNCHRONIZE, false, uint32(pid)) + if err != nil { + // Access denied: process exists under a tighter ACL (or higher + // integrity). Do not reclaim its repair lock. + if errors.Is(err, windows.ERROR_ACCESS_DENIED) { + return true + } + // Free/invalid PID typically returns ERROR_INVALID_PARAMETER. + if errors.Is(err, windows.ERROR_INVALID_PARAMETER) { + return false + } + // Unknown open failure: fail closed (do not reclaim). + return true + } + defer windows.CloseHandle(handle) + + // Zero timeout: WAIT_OBJECT_0 if already terminated, WAIT_TIMEOUT if still + // running. WaitForSingleObject only returns err on WAIT_FAILED. + event, err := windows.WaitForSingleObject(handle, 0) + if err != nil { + return true + } + switch event { + case windows.WAIT_OBJECT_0: + return false // signaled = process has exited + case uint32(windows.WAIT_TIMEOUT): + return true // still running + default: + // Unexpected wait result: fail closed. + return true + } +} diff --git a/internal/kimiidentity/process_alive_windows_test.go b/internal/kimiidentity/process_alive_windows_test.go new file mode 100644 index 000000000..24e55c6fc --- /dev/null +++ b/internal/kimiidentity/process_alive_windows_test.go @@ -0,0 +1,136 @@ +//go:build windows + +package kimiidentity + +import ( + "errors" + "os" + "os/exec" + "testing" + + "golang.org/x/sys/windows" +) + +// TestProcessAliveExitCode259IsDead is the regression for treating STILL_ACTIVE +// (259) as a process exit code rather than a liveness flag. A dead child that +// exits with 259 must be reported as not alive so a stale repair lease can be +// reclaimed. +func TestProcessAliveExitCode259IsDead(t *testing.T) { + cmd := exec.Command("cmd", "/C", "exit /b 259") + if err := cmd.Start(); err != nil { + t.Fatalf("Start: %v", err) + } + pid := cmd.Process.Pid + waitRan := false + defer func() { + if !waitRan { + if err := cmd.Process.Kill(); err != nil { + t.Errorf("Kill cleanup: %v", err) + } + if err := cmd.Wait(); err != nil { + t.Errorf("Wait cleanup: %v", err) + } + } + }() + + // cmd.Wait() can close the process handle held by cmd.Process, and once no + // handle to the process remains, OpenProcess(pid) fails with + // ERROR_INVALID_PARAMETER, which would take the invalid-PID branch instead + // of the WAIT_OBJECT_0 branch this test must exercise for a terminated + // process. + handle, err := windows.OpenProcess(windows.SYNCHRONIZE, false, uint32(pid)) + if err != nil { + t.Fatalf("OpenProcess: %v", err) + } + if _, err := windows.WaitForSingleObject(handle, windows.INFINITE); err != nil { + _ = windows.CloseHandle(handle) + t.Fatalf("WaitForSingleObject: %v", err) + } + if processAlive(pid) { + _ = windows.CloseHandle(handle) + t.Fatalf("processAlive(%d) = true after exit code 259; want false (dead)", pid) + } + if err := windows.CloseHandle(handle); err != nil { + t.Errorf("CloseHandle: %v", err) + } + + waitRan = true + err = cmd.Wait() + var ee *exec.ExitError + if !errors.As(err, &ee) || ee.ExitCode() != 259 { + t.Fatalf("Wait: %v (want ExitError with code 259)", err) + } +} + +func TestProcessAliveSelfIsLive(t *testing.T) { + if !processAlive(os.Getpid()) { + t.Fatal("processAlive(self) = false; want true") + } +} + +func TestProcessAliveNonPositiveIsDead(t *testing.T) { + if processAlive(0) { + t.Fatal("processAlive(0) = true; want false") + } + if processAlive(-1) { + t.Fatal("processAlive(-1) = true; want false") + } +} + +func TestProcessAliveRunningChildIsLive(t *testing.T) { + cmd := exec.Command("powershell", "-Command", "Start-Sleep -Seconds 10") + if err := cmd.Start(); err != nil { + t.Fatalf("Start: %v", err) + } + pid := cmd.Process.Pid + waitRan := false + defer func() { + if !waitRan { + if err := cmd.Process.Kill(); err != nil { + t.Errorf("Kill cleanup: %v", err) + } + if err := cmd.Wait(); err != nil { + t.Errorf("Wait cleanup: %v", err) + } + } + }() + + // Keep our own SYNCHRONIZE handle open across the post-kill probe: + // cmd.Wait() can close the process handle held by cmd.Process, and without + // a retained handle OpenProcess(pid) could fail with + // ERROR_INVALID_PARAMETER, taking the invalid-PID branch instead of the + // WAIT_OBJECT_0 branch this test must exercise for a killed process. + handle, err := windows.OpenProcess(windows.SYNCHRONIZE, false, uint32(pid)) + if err != nil { + t.Fatalf("OpenProcess: %v", err) + } + + if !processAlive(pid) { + _ = windows.CloseHandle(handle) + t.Fatalf("processAlive(%d) = false while running; want true", pid) + } + + if err := cmd.Process.Kill(); err != nil { + _ = windows.CloseHandle(handle) + t.Fatalf("Kill: %v", err) + } + if _, err := windows.WaitForSingleObject(handle, windows.INFINITE); err != nil { + _ = windows.CloseHandle(handle) + t.Fatalf("WaitForSingleObject: %v", err) + } + if processAlive(pid) { + _ = windows.CloseHandle(handle) + t.Fatalf("processAlive(%d) = true after kill; want false (dead)", pid) + } + if err := windows.CloseHandle(handle); err != nil { + t.Errorf("CloseHandle: %v", err) + } + + waitRan = true + if err := cmd.Wait(); err != nil { + var exitErr *exec.ExitError + if !errors.As(err, &exitErr) { + t.Errorf("Wait: %v", err) + } + } +} diff --git a/internal/lockutil/reclaim.go b/internal/lockutil/reclaim.go index 6aaa414c4..e2e804f20 100644 --- a/internal/lockutil/reclaim.go +++ b/internal/lockutil/reclaim.go @@ -11,6 +11,13 @@ import ( // portably on a healthy filesystem. var restoreLockFile = restoreLiveLock +// readSidelinedLock is swappable for the same reason: a healthy filesystem +// does not fail a read of a file this process just renamed, so the +// unreadable-lock path of ReclaimStaleLockRooted needs a seam to exercise. +var readSidelinedLock = func(root *os.Root, name string) ([]byte, error) { + return root.ReadFile(name) +} + // restoreLiveLock puts a lock that turned out to be live back at path after // ReclaimStaleLock moved it aside to inspect it. It first tries a fast, // replacing rename straight from reclaimed to path: a single syscall, which @@ -80,3 +87,95 @@ func ReclaimStaleLock(lockPath, suffix string, isLive func(reclaimedPath string) _ = RemoveLockFile(reclaimed) return true, nil } + +// linkRootedLock is swappable so tests can exercise the hard-link-incapable +// filesystem fallback in restoreRootedLock. +var linkRootedLock = func(root *os.Root, oldname, newname string) error { + return root.Link(oldname, newname) +} + +// ReclaimStaleLockRooted is ReclaimStaleLock for a lock file inside an +// already-opened *os.Root. Every rename/read/remove goes through the root +// handle, so a symlink or reparse point swapped in under lockName after the +// root was opened cannot redirect the operations (the path-based variant +// would re-walk root.Name()+lockName as plain paths). lockName must be a bare +// file name; the root supplies the directory. isLive receives the raw +// sidelined contents and is called only when they were read successfully, so +// it never has to invent a policy for contents it cannot see: an unreadable +// lock is restored here without consulting it. +func ReclaimStaleLockRooted(root *os.Root, lockName, suffix string, isLive func(raw []byte) bool) (bool, error) { + reclaimed := lockName + ".stale." + suffix + if err := root.Rename(lockName, reclaimed); err != nil { + if errors.Is(err, os.ErrNotExist) || isReclaimContended(err) { + return false, nil // another racer already moved/removed it, or it vanished + } + return false, err + } + raw, err := readSidelinedLock(root, reclaimed) + if err != nil { + // Decide here rather than handing nil to isLive. The callback answers + // "do these contents describe a live holder", and a caller can very + // reasonably treat empty or unparseable contents as dead so a holder + // that crashed mid-write is recoverable: kimiidentity's lockHolderAlive + // does exactly that. Passing nil then classified an unreadable lock as + // dead and removed it, which is the opposite of failing closed. A read + // failure is not proof of death, so restore and let the caller wait. + if rerr := restoreRootedLock(root, reclaimed, lockName); rerr != nil { + if errors.Is(rerr, os.ErrExist) { + _ = root.Remove(reclaimed) + return false, nil + } + return false, rerr + } + return false, nil + } + if isLive(raw) { + // Put the live lock back instead of stealing it, and let the caller wait. + if rerr := restoreRootedLock(root, reclaimed, lockName); rerr != nil { + if errors.Is(rerr, os.ErrExist) { + _ = root.Remove(reclaimed) + return false, nil + } + return false, rerr + } + return false, nil + } + if err := root.Remove(reclaimed); err != nil { + return false, err + } + return true, nil +} + +// restoreRootedLock puts a lock that turned out to be live, or that could not +// be read, back at lockName after ReclaimStaleLockRooted moved it aside. It +// tries a no-replace hard link first, so a competing lock created in the gap +// wins (os.ErrExist) instead of being overwritten, then an O_EXCL probe and +// rename for filesystems without hard links without depending on reading contents. +func restoreRootedLock(root *os.Root, reclaimed, lockName string) error { + if err := linkRootedLock(root, reclaimed, lockName); err == nil { + // Restore is complete; the sidelined name is now a redundant link. + // Cleanup is best-effort: a leftover .stale file is harmless, and + // returning a cleanup error here would report the restore as failed + // when it actually succeeded. + _ = root.Remove(reclaimed) + return nil + } else if errors.Is(err, os.ErrExist) { + return err + } + // Hard-link incapable filesystem (FAT, some FUSE/network mounts): probe + // lockName with O_EXCL so we never overwrite a racer. Prefer writing the + // sidelined contents into the probe first so lockName is never observable + // as empty; if the read fails, fall back to the empty probe + rename. + raw, readErr := readSidelinedLock(root, reclaimed) + probe, err := root.OpenFile(lockName, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + return err + } + if readErr == nil && len(raw) > 0 { + if _, werr := probe.Write(raw); werr == nil { + _ = probe.Sync() + } + } + _ = probe.Close() + return root.Rename(reclaimed, lockName) +} diff --git a/internal/lockutil/reclaim_test.go b/internal/lockutil/reclaim_test.go index a4dd992e3..5a1254be9 100644 --- a/internal/lockutil/reclaim_test.go +++ b/internal/lockutil/reclaim_test.go @@ -70,6 +70,217 @@ func TestRestoreLiveLockFastPathOverwritesCompetingFile(t *testing.T) { } } +func TestRestoreRootedLockPreservesCompetingFile(t *testing.T) { + dir := t.TempDir() + root, err := os.OpenRoot(dir) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if err := root.Close(); err != nil { + t.Errorf("root.Close: %v", err) + } + }) + + if err := os.WriteFile(filepath.Join(dir, "lock.stale.tok"), []byte("original-holder"), 0o600); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "lock"), []byte("new-claimant"), 0o600); err != nil { + t.Fatal(err) + } + + err = restoreRootedLock(root, "lock.stale.tok", "lock") + if !errors.Is(err, os.ErrExist) { + t.Fatalf("restoreRootedLock error = %v, want os.ErrExist", err) + } + data, err := os.ReadFile(filepath.Join(dir, "lock")) + if err != nil || string(data) != "new-claimant" { + t.Fatalf("competing lock = %q, err %v; want unchanged content", data, err) + } +} +func TestReclaimStaleLockRooted(t *testing.T) { + dir := t.TempDir() + root, err := os.OpenRoot(dir) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if err := root.Close(); err != nil { + t.Errorf("root.Close: %v", err) + } + }) + lockName := "lock" + dead := func([]byte) bool { return false } + live := func([]byte) bool { return true } + + // A lock the predicate reports dead is reclaimed and removed. + if err := os.WriteFile(filepath.Join(dir, lockName), []byte("crashed-holder"), 0o600); err != nil { + t.Fatal(err) + } + if ok, err := ReclaimStaleLockRooted(root, lockName, "tok-a", dead); err != nil || !ok { + t.Fatalf("a dead lock should be reclaimed (ok=%v err=%v)", ok, err) + } + if _, err := os.Stat(filepath.Join(dir, lockName)); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("reclaimed dead lock should be gone, stat err=%v", err) + } + + // A LIVE lock must be RESTORED intact. + if err := os.WriteFile(filepath.Join(dir, lockName), []byte("live-holder"), 0o600); err != nil { + t.Fatal(err) + } + if ok, err := ReclaimStaleLockRooted(root, lockName, "tok-b", live); err != nil || ok { + t.Fatalf("a live lock must not be reclaimed (ok=%v err=%v)", ok, err) + } + if data, err := os.ReadFile(filepath.Join(dir, lockName)); err != nil || string(data) != "live-holder" { + t.Fatalf("live lock must be left intact, got %q err %v", data, err) + } + + // A missing lock reports no reclaim (nothing to steal). + if ok, err := ReclaimStaleLockRooted(root, lockName, "tok-c", live); err != nil || ok { + t.Fatalf("a missing lock should not report a reclaim (ok=%v err=%v)", ok, err) + } +} + +// TestReclaimStaleLockRootedRestoresUnreadableLock pins that a lock whose +// contents cannot be read is put back rather than deleted. The predicate here +// has the same shape as kimiidentity's lockHolderAlive, which deliberately +// treats empty or unparseable contents as dead so a holder that crashed +// mid-write stays recoverable. Passing it nil for a lock that merely could not +// be read therefore classified a possibly-live lock as dead and removed it, so +// the decision has to be made before the predicate is consulted. +func TestReclaimStaleLockRootedRestoresUnreadableLock(t *testing.T) { + original := readSidelinedLock + readSidelinedLock = func(*os.Root, string) ([]byte, error) { + return nil, errors.New("permission denied") + } + defer func() { readSidelinedLock = original }() + + dir := t.TempDir() + root, err := os.OpenRoot(dir) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if err := root.Close(); err != nil { + t.Errorf("root.Close: %v", err) + } + }) + + lockName := "lock" + lockPath := filepath.Join(dir, lockName) + if err := os.WriteFile(lockPath, []byte("live-holder"), 0o600); err != nil { + t.Fatal(err) + } + // Empty or unreadable contents look dead to this predicate, exactly like + // the real caller's. + deadWhenBlank := func(raw []byte) bool { return len(raw) > 0 } + + ok, err := ReclaimStaleLockRooted(root, lockName, "tok-unreadable", deadWhenBlank) + if err != nil { + t.Fatalf("an unreadable lock should be restored, not error: %v", err) + } + if ok { + t.Fatal("an unreadable lock must not be reported as reclaimed: a read failure is not proof the holder is dead") + } + data, err := os.ReadFile(lockPath) + if err != nil { + t.Fatalf("the unreadable lock must be restored at its original name: %v", err) + } + if string(data) != "live-holder" { + t.Fatalf("restored lock contents = %q, want %q", data, "live-holder") + } + // The sidelined copy must not be left behind. + if entries, _ := os.ReadDir(dir); len(entries) != 1 { + t.Fatalf("expected only the restored lock in the directory, got %d entries", len(entries)) + } +} + +func TestReclaimStaleLockRootedFallbackWhenLinkFails(t *testing.T) { + origLink := linkRootedLock + linkRootedLock = func(*os.Root, string, string) error { + return errors.New("function not implemented") + } + defer func() { linkRootedLock = origLink }() + + dir := t.TempDir() + root, err := os.OpenRoot(dir) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if err := root.Close(); err != nil { + t.Errorf("root.Close: %v", err) + } + }) + + lockName := "lock" + lockPath := filepath.Join(dir, lockName) + if err := os.WriteFile(lockPath, []byte("live-holder"), 0o600); err != nil { + t.Fatal(err) + } + + ok, err := ReclaimStaleLockRooted(root, lockName, "tok-fallback", func(raw []byte) bool { return true }) + if err != nil { + t.Fatalf("unexpected error during restore fallback: %v", err) + } + if ok { + t.Fatal("expected live lock not to be reclaimed") + } + data, err := os.ReadFile(lockPath) + if err != nil || string(data) != "live-holder" { + t.Fatalf("restored lock contents = %q (err=%v), want %q", data, err, "live-holder") + } + if entries, err := os.ReadDir(dir); err != nil || len(entries) != 1 { + t.Fatalf("expected only the restored lock, got %d entries (err=%v)", len(entries), err) + } +} + +func TestReclaimStaleLockRootedFallbackWhenLinkAndReadFail(t *testing.T) { + origLink := linkRootedLock + linkRootedLock = func(*os.Root, string, string) error { + return errors.New("function not implemented") + } + defer func() { linkRootedLock = origLink }() + + origRead := readSidelinedLock + readSidelinedLock = func(*os.Root, string) ([]byte, error) { + return nil, errors.New("permission denied") + } + defer func() { readSidelinedLock = origRead }() + + dir := t.TempDir() + root, err := os.OpenRoot(dir) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if err := root.Close(); err != nil { + t.Errorf("root.Close: %v", err) + } + }) + + lockName := "lock" + lockPath := filepath.Join(dir, lockName) + if err := os.WriteFile(lockPath, []byte("live-holder"), 0o600); err != nil { + t.Fatal(err) + } + + ok, err := ReclaimStaleLockRooted(root, lockName, "tok-unreadable-linkfail", func(raw []byte) bool { return true }) + if err != nil { + t.Fatalf("unexpected error during unreadable restore fallback: %v", err) + } + if ok { + t.Fatal("expected unreadable lock not to be reclaimed") + } + data, err := os.ReadFile(lockPath) + if err != nil || string(data) != "live-holder" { + t.Fatalf("restored lock contents = %q (err=%v), want %q", data, err, "live-holder") + } + if entries, err := os.ReadDir(dir); err != nil || len(entries) != 1 { + t.Fatalf("expected only the restored lock, got %d entries (err=%v)", len(entries), err) + } +} + func TestReclaimStaleLockFailsClosedOnRestoreError(t *testing.T) { // When both the no-replace restore and its copy fallback fail (only // provokable via the seam; a healthy filesystem cannot produce it), the diff --git a/internal/oauth/device.go b/internal/oauth/device.go index 697de057c..038ff6b2c 100644 --- a/internal/oauth/device.go +++ b/internal/oauth/device.go @@ -71,6 +71,7 @@ func RequestDeviceCode(ctx context.Context, client *http.Client, cfg Config, now } request.Header.Set("Content-Type", "application/x-www-form-urlencoded") request.Header.Set("Accept", "application/json") + applyExtraHeaders(request, cfg.ExtraHeaders) response, err := client.Do(request) if err != nil { return DeviceAuth{}, fmt.Errorf("oauth: device authorization request failed: %w", err) @@ -173,6 +174,7 @@ func pollDeviceOnce(ctx context.Context, client *http.Client, cfg Config, device } request.Header.Set("Content-Type", "application/x-www-form-urlencoded") request.Header.Set("Accept", "application/json") + applyExtraHeaders(request, cfg.ExtraHeaders) response, err := client.Do(request) if err != nil { return Token{}, fmt.Errorf("oauth: device token poll failed: %w", err) diff --git a/internal/oauth/device_test.go b/internal/oauth/device_test.go index 9ccb17440..c9786d6b1 100644 --- a/internal/oauth/device_test.go +++ b/internal/oauth/device_test.go @@ -166,3 +166,74 @@ func TestPollDeviceTokenAccessDenied(t *testing.T) { t.Fatalf("err = %v, want access denied", err) } } + +// TestRequestDeviceCodeAppliesExtraHeaders pins that Config.ExtraHeaders land +// on the device-authorization request (Kimi's X-Msh-* identity path). +func TestRequestDeviceCodeAppliesExtraHeaders(t *testing.T) { + var got http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got = r.Header.Clone() + _, _ = w.Write([]byte(`{"device_code":"dc","user_code":"U","verification_uri":"https://x","expires_in":600,"interval":5}`)) + })) + defer server.Close() + cfg := Config{ + ClientID: "c", + DeviceAuthorizationEndpoint: server.URL, + ExtraHeaders: map[string]string{ + "X-Msh-Device-Id": "test-device-id", + "X-Msh-Platform": "kimi_code_cli", + }, + } + if _, err := RequestDeviceCode(context.Background(), server.Client(), cfg, nil); err != nil { + t.Fatalf("RequestDeviceCode: %v", err) + } + if got.Get("X-Msh-Device-Id") != "test-device-id" { + t.Fatalf("X-Msh-Device-Id = %q, want test-device-id", got.Get("X-Msh-Device-Id")) + } + if got.Get("X-Msh-Platform") != "kimi_code_cli" { + t.Fatalf("X-Msh-Platform = %q, want kimi_code_cli", got.Get("X-Msh-Platform")) + } +} + +// TestRequestDeviceCodeWithoutExtraHeadersSendsNone pins that providers with +// no ExtraHeaders do not inject vendor identity headers. +func TestRequestDeviceCodeWithoutExtraHeadersSendsNone(t *testing.T) { + var got http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got = r.Header.Clone() + _, _ = w.Write([]byte(`{"device_code":"dc","user_code":"U","verification_uri":"https://x","expires_in":600,"interval":5}`)) + })) + defer server.Close() + cfg := Config{ClientID: "c", DeviceAuthorizationEndpoint: server.URL} + if _, err := RequestDeviceCode(context.Background(), server.Client(), cfg, nil); err != nil { + t.Fatalf("RequestDeviceCode: %v", err) + } + for _, key := range []string{"X-Msh-Device-Id", "X-Msh-Platform", "X-Msh-Version"} { + if got.Get(key) != "" { + t.Fatalf("unexpected ExtraHeader %s=%q on a provider without ExtraHeaders", key, got.Get(key)) + } + } +} + +// TestPollDeviceTokenAppliesExtraHeaders pins ExtraHeaders on the device-token +// poll request path (same identity headers Kimi requires on every poll). +func TestPollDeviceTokenAppliesExtraHeaders(t *testing.T) { + var got http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got = r.Header.Clone() + _, _ = w.Write([]byte(`{"access_token":"at","token_type":"Bearer","expires_in":3600}`)) + })) + defer server.Close() + cfg := Config{ + ClientID: "c", + TokenEndpoint: server.URL, + ExtraHeaders: map[string]string{"X-Msh-Device-Id": "poll-device-id"}, + } + auth := DeviceAuth{DeviceCode: "dc", Interval: 5 * time.Millisecond, ExpiresAt: time.Now().Add(5 * time.Second)} + if _, err := PollDeviceToken(context.Background(), server.Client(), cfg, auth, nil); err != nil { + t.Fatalf("PollDeviceToken: %v", err) + } + if got.Get("X-Msh-Device-Id") != "poll-device-id" { + t.Fatalf("X-Msh-Device-Id = %q, want poll-device-id", got.Get("X-Msh-Device-Id")) + } +} diff --git a/internal/oauth/flow.go b/internal/oauth/flow.go index 9ff94cccd..09cc68e60 100644 --- a/internal/oauth/flow.go +++ b/internal/oauth/flow.go @@ -154,7 +154,7 @@ func ExchangeCode(ctx context.Context, client *http.Client, cfg Config, code, ve if secret := trimmed(cfg.ClientSecret); secret != "" { form.Set("client_secret", secret) } - return PostToken(ctx, client, cfg.TokenEndpoint, form, Token{Scopes: cfg.Scopes}, now) + return PostToken(ctx, client, cfg.TokenEndpoint, form, Token{Scopes: cfg.Scopes}, cfg.ExtraHeaders, now) } // Refresh exchanges a refresh token for a fresh access token. A response that @@ -181,7 +181,7 @@ func Refresh(ctx context.Context, client *http.Client, cfg Config, current Token // and PostToken only overwrites TokenType when the response supplies one, so // without seeding it here the type would be silently lost across refreshes (L15). base := Token{Scopes: current.Scopes, RefreshToken: refresh, Account: current.Account, IDToken: current.IDToken, TokenType: current.TokenType} - return PostToken(ctx, client, cfg.TokenEndpoint, form, base, now) + return PostToken(ctx, client, cfg.TokenEndpoint, form, base, cfg.ExtraHeaders, now) } // PostToken performs a token-endpoint POST and maps the response onto a Token. @@ -189,7 +189,7 @@ func Refresh(ctx context.Context, client *http.Client, cfg Config, current Token // (e.g. an existing refresh token or scopes) when the response omits them. // Error messages carry only the server's error/error_description — never the raw // body — so token material in an unexpected payload is not leaked. -func PostToken(ctx context.Context, client *http.Client, tokenEndpoint string, form url.Values, base Token, now func() time.Time) (Token, error) { +func PostToken(ctx context.Context, client *http.Client, tokenEndpoint string, form url.Values, base Token, extraHeaders map[string]string, now func() time.Time) (Token, error) { if err := validateTokenEndpoint(tokenEndpoint); err != nil { return Token{}, err } @@ -205,6 +205,7 @@ func PostToken(ctx context.Context, client *http.Client, tokenEndpoint string, f } request.Header.Set("Content-Type", "application/x-www-form-urlencoded") request.Header.Set("Accept", "application/json") + applyExtraHeaders(request, extraHeaders) response, err := client.Do(request) if err != nil { diff --git a/internal/oauth/flow_test.go b/internal/oauth/flow_test.go index 09ef01251..d7312414a 100644 --- a/internal/oauth/flow_test.go +++ b/internal/oauth/flow_test.go @@ -153,7 +153,7 @@ func TestPostTokenRefusesRedirect(t *testing.T) { form.Set("grant_type", "authorization_code") form.Set("code", "the-code") form.Set("client_secret", "shh") - _, err := PostToken(context.Background(), http.DefaultClient, endpoint, form, Token{}, nil) + _, err := PostToken(context.Background(), http.DefaultClient, endpoint, form, Token{}, nil, nil) if !errors.Is(err, ErrUnsafeRedirect) { t.Fatalf("PostToken err = %v, want ErrUnsafeRedirect", err) } @@ -257,7 +257,7 @@ func TestExchangeCodeErrorRedactsBody(t *testing.T) { } func TestPostTokenRefusesInsecureEndpoint(t *testing.T) { - _, err := PostToken(context.Background(), http.DefaultClient, "http://auth.example.com/token", url.Values{}, Token{}, nil) + _, err := PostToken(context.Background(), http.DefaultClient, "http://auth.example.com/token", url.Values{}, Token{}, nil, nil) if !errors.Is(err, ErrInsecureTokenEndpoint) { t.Fatalf("err = %v, want ErrInsecureTokenEndpoint", err) } @@ -299,3 +299,67 @@ func TestRefreshPreservesTokenTypeWhenOmitted(t *testing.T) { t.Fatalf("refresh should carry the existing token_type forward, got %q", tok.TokenType) } } + +// TestExchangeCodeAppliesExtraHeaders pins Config.ExtraHeaders on the +// authorization-code token exchange request. +func TestExchangeCodeAppliesExtraHeaders(t *testing.T) { + var got http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got = r.Header.Clone() + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"access_token":"at","token_type":"Bearer","expires_in":3600}`)) + })) + defer server.Close() + cfg := Config{ + ClientID: "c", + TokenEndpoint: server.URL, + ExtraHeaders: map[string]string{"X-Msh-Device-Id": "exchange-device-id"}, + } + if _, err := ExchangeCode(context.Background(), server.Client(), cfg, "code", "verifier", "http://127.0.0.1/cb", nil); err != nil { + t.Fatalf("ExchangeCode: %v", err) + } + if got.Get("X-Msh-Device-Id") != "exchange-device-id" { + t.Fatalf("X-Msh-Device-Id = %q, want exchange-device-id", got.Get("X-Msh-Device-Id")) + } +} + +// TestRefreshAppliesExtraHeaders pins Config.ExtraHeaders on the refresh +// request path. +func TestRefreshAppliesExtraHeaders(t *testing.T) { + var got http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got = r.Header.Clone() + _, _ = w.Write([]byte(`{"access_token":"new-at","expires_in":3600}`)) + })) + defer server.Close() + cfg := Config{ + ClientID: "c", + TokenEndpoint: server.URL, + ExtraHeaders: map[string]string{"X-Msh-Device-Id": "refresh-device-id"}, + } + if _, err := Refresh(context.Background(), server.Client(), cfg, Token{RefreshToken: "rt"}, nil); err != nil { + t.Fatalf("Refresh: %v", err) + } + if got.Get("X-Msh-Device-Id") != "refresh-device-id" { + t.Fatalf("X-Msh-Device-Id = %q, want refresh-device-id", got.Get("X-Msh-Device-Id")) + } +} + +// TestExchangeCodeWithoutExtraHeadersSendsNone pins that a config without +// ExtraHeaders does not inject vendor identity headers on exchange. +func TestExchangeCodeWithoutExtraHeadersSendsNone(t *testing.T) { + var got http.Header + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got = r.Header.Clone() + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"access_token":"at","token_type":"Bearer","expires_in":3600}`)) + })) + defer server.Close() + cfg := Config{ClientID: "c", TokenEndpoint: server.URL} + if _, err := ExchangeCode(context.Background(), server.Client(), cfg, "code", "verifier", "http://127.0.0.1/cb", nil); err != nil { + t.Fatalf("ExchangeCode: %v", err) + } + if got.Get("X-Msh-Device-Id") != "" { + t.Fatalf("unexpected X-Msh-Device-Id=%q without ExtraHeaders", got.Get("X-Msh-Device-Id")) + } +} diff --git a/internal/oauth/manager.go b/internal/oauth/manager.go index dbdf7817e..457993629 100644 --- a/internal/oauth/manager.go +++ b/internal/oauth/manager.go @@ -35,6 +35,10 @@ type Manager struct { // openBrowser is invoked with the authorization URL for loopback logins. // Tests inject a function that drives the loopback redirect. openBrowser func(authURL string) error + // beforeDeviceCommit is a test-only hook invoked between reading the store + // state and writing the committed device token, so tests can cancel the + // attempt mid-commit and assert the rollback path. + beforeDeviceCommit func() // refreshLocks serializes concurrent refreshes per key so parallel callers // don't each spend the single-use refresh token; the loser reuses the rotated // token. refreshMu guards the map (M7). @@ -126,7 +130,7 @@ func (m *Manager) Login(ctx context.Context, opts LoginOptions) (Status, error) loginCtx, cancel := context.WithTimeout(ctx, timeout) defer cancel() - cfg, err = m.resolveEndpoints(loginCtx, cfg) + cfg, err = m.resolveEndpoints(loginCtx, opts.Provider, cfg) if err != nil { return Status{}, err } @@ -171,7 +175,7 @@ func (m *Manager) PrepareDeviceLogin(ctx context.Context, opts LoginOptions) (De } prepCtx, cancel := context.WithTimeout(ctx, timeout) defer cancel() - cfg, err = m.resolveEndpoints(prepCtx, cfg) + cfg, err = m.resolveEndpoints(prepCtx, opts.Provider, cfg) if err != nil { return DeviceAuth{}, Config{}, err } @@ -199,7 +203,7 @@ func (m *Manager) CompleteDeviceLogin(ctx context.Context, provider string, cfg return Status{}, err } key := ProviderKey(provider) - if err := m.store.Save(key, token); err != nil { + if err := m.store.CommitDeviceToken(key, token, ctx.Err, m.beforeDeviceCommit); err != nil { return Status{}, err } return m.statusFor(key) @@ -207,12 +211,14 @@ func (m *Manager) CompleteDeviceLogin(ctx context.Context, provider string, cfg // resolveEndpoints fills missing authorize/token/device endpoints from issuer // discovery (RFC 8414 then OIDC), leaving any explicitly-pinned endpoint intact. -func (m *Manager) resolveEndpoints(ctx context.Context, cfg Config) (Config, error) { +// ExtraHeaders are derived from the final destinations so header policy cannot +// come from a raw issuer override while requests still use preset URLs. +func (m *Manager) resolveEndpoints(ctx context.Context, name string, cfg Config) (Config, error) { if trimmed(cfg.IssuerURL) == "" { - return cfg, nil + return attachProviderHeaders(name, cfg), nil } if cfg.AuthorizationEndpoint != "" && cfg.TokenEndpoint != "" && cfg.DeviceAuthorizationEndpoint != "" { - return cfg, nil + return attachProviderHeaders(name, cfg), nil } // Discovery is best-effort: a failure is non-fatal because pinned endpoints // may already be sufficient for the chosen flow. Only merge on success (and @@ -241,7 +247,7 @@ func (m *Manager) resolveEndpoints(ctx context.Context, cfg Config) (Config, err cfg.DeviceAuthorizationEndpoint = meta.DeviceAuthorizationEndpoint } } - return cfg, nil + return attachProviderHeaders(name, cfg), nil } // discover tries the OAuth (RFC 8414) well-known path, then the OIDC @@ -404,7 +410,7 @@ func (m *Manager) resolveConfigForKey(ctx context.Context, key string) (Config, // Fill any missing token/authorize/device endpoints from issuer discovery so a // provider configured with only ZERO_OAUTH__ISSUER_URL can still refresh // (refreshAndSave requires the token endpoint). - cfg, err = m.resolveEndpoints(ctx, cfg) + cfg, err = m.resolveEndpoints(ctx, name, cfg) if err != nil { return Config{}, err } diff --git a/internal/oauth/manager_test.go b/internal/oauth/manager_test.go index 219389940..94a460b18 100644 --- a/internal/oauth/manager_test.go +++ b/internal/oauth/manager_test.go @@ -8,6 +8,7 @@ import ( "net/http/httptest" "net/url" "path/filepath" + "sync" "sync/atomic" "testing" "time" @@ -213,7 +214,7 @@ func TestResolveEndpointsRejectsInsecureDiscoveredEndpoint(t *testing.T) { _, _ = io.WriteString(w, `{"issuer":"`+server.URL+`","token_endpoint":"`+server.URL+`/token","authorization_endpoint":"http://evil.example/authorize"}`) }) m := managerFor(t, map[string]string{}, nil) - _, err := m.resolveEndpoints(context.Background(), Config{IssuerURL: server.URL}) + _, err := m.resolveEndpoints(context.Background(), "", Config{IssuerURL: server.URL}) if !errors.Is(err, ErrInsecureTokenEndpoint) { t.Fatalf("resolveEndpoints err = %v, want ErrInsecureTokenEndpoint", err) } @@ -293,3 +294,283 @@ func TestManagerLogout(t *testing.T) { t.Fatal("second logout should report nothing removed") } } + +func TestCompleteDeviceLoginCancelAtCommitRollsBackNewToken(t *testing.T) { + fp := newFakeProvider(t, `{"access_token":"authorized-token"}`) + env := map[string]string{ + "ZERO_OAUTH_DEMO_CLIENT_ID": "client", + "ZERO_OAUTH_DEMO_TOKEN_URL": fp.server.URL + "/token", + } + m := managerFor(t, env, nil) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + m.beforeDeviceCommit = func() { + // Cancel the attempt during commit. + cancel() + } + defer func() { m.beforeDeviceCommit = nil }() + + cfg := Config{ClientID: "client", TokenEndpoint: fp.server.URL + "/token"} + auth := DeviceAuth{DeviceCode: "dc", Interval: 5 * time.Millisecond, ExpiresAt: time.Now().Add(5 * time.Second)} + + _, err := m.CompleteDeviceLogin(ctx, "demo", cfg, auth) + if !errors.Is(err, context.Canceled) { + t.Fatalf("CompleteDeviceLogin err = %v, want context.Canceled", err) + } + + _, ok, err := m.store.Load(ProviderKey("demo")) + if err != nil { + t.Fatalf("Load: %v", err) + } + if ok { + t.Fatal("expected store to have no token after rollback") + } +} + +func TestCompleteDeviceLoginCancelAtCommitRestoresPreviousToken(t *testing.T) { + fp := newFakeProvider(t, `{"access_token":"new-token"}`) + env := map[string]string{ + "ZERO_OAUTH_DEMO_CLIENT_ID": "client", + "ZERO_OAUTH_DEMO_TOKEN_URL": fp.server.URL + "/token", + } + m := managerFor(t, env, nil) + + prevToken := Token{AccessToken: "previous-valid-token", RefreshToken: "prev-rt"} + if err := m.store.Save(ProviderKey("demo"), prevToken); err != nil { + t.Fatalf("seed Save: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + m.beforeDeviceCommit = func() { + cancel() + } + defer func() { m.beforeDeviceCommit = nil }() + + cfg := Config{ClientID: "client", TokenEndpoint: fp.server.URL + "/token"} + auth := DeviceAuth{DeviceCode: "dc", Interval: 5 * time.Millisecond, ExpiresAt: time.Now().Add(5 * time.Second)} + + _, err := m.CompleteDeviceLogin(ctx, "demo", cfg, auth) + if !errors.Is(err, context.Canceled) { + t.Fatalf("CompleteDeviceLogin err = %v, want context.Canceled", err) + } + + stored, ok, err := m.store.Load(ProviderKey("demo")) + if err != nil { + t.Fatalf("Load: %v", err) + } + if !ok { + t.Fatal("expected previous token to be present") + } + if stored.AccessToken != "previous-valid-token" { + t.Fatalf("expected previous token %q restored, got %q", "previous-valid-token", stored.AccessToken) + } +} + +func TestCompleteDeviceLoginTwoManagersSharingStore(t *testing.T) { + fp := newFakeProvider(t, `{"access_token":"concurrent-token"}`) + storePath := filepath.Join(t.TempDir(), "tokens.json") + store, err := NewStore(StoreOptions{FilePath: storePath}) + if err != nil { + t.Fatal(err) + } + + env := map[string]string{ + "ZERO_OAUTH_DEMO_CLIENT_ID": "client", + "ZERO_OAUTH_DEMO_TOKEN_URL": fp.server.URL + "/token", + } + m1, err := NewManager(ManagerOptions{Store: store, Env: env}) + if err != nil { + t.Fatal(err) + } + m2, err := NewManager(ManagerOptions{Store: store, Env: env}) + if err != nil { + t.Fatal(err) + } + + cfg := Config{ClientID: "client", TokenEndpoint: fp.server.URL + "/token"} + auth := DeviceAuth{DeviceCode: "dc", Interval: 5 * time.Millisecond, ExpiresAt: time.Now().Add(5 * time.Second)} + + status, err := m1.CompleteDeviceLogin(context.Background(), "demo", cfg, auth) + if err != nil { + t.Fatalf("m1 CompleteDeviceLogin: %v", err) + } + if !status.HasToken { + t.Fatal("expected status to report HasToken") + } + + tok, ok, err := m2.store.Load(ProviderKey("demo")) + if err != nil || !ok || tok.AccessToken != "concurrent-token" { + t.Fatalf("m2 saw token = %+v, ok = %v, err = %v", tok, ok, err) + } +} + +func TestCompleteDeviceLoginTwoManagersSharingStoreConcurrent(t *testing.T) { + fp1 := newFakeProvider(t, `{"access_token":"token-1"}`) + fp2 := newFakeProvider(t, `{"access_token":"token-2"}`) + storePath := filepath.Join(t.TempDir(), "tokens.json") + // Two independent Store values over the same file exercise the cross-process + // file lock (fileBlob.withLock), not just the in-process Store.mu, which is + // what two concurrent Zero processes actually rely on. + store1, err := NewStore(StoreOptions{FilePath: storePath}) + if err != nil { + t.Fatal(err) + } + store2, err := NewStore(StoreOptions{FilePath: storePath}) + if err != nil { + t.Fatal(err) + } + + m1, err := NewManager(ManagerOptions{Store: store1}) + if err != nil { + t.Fatal(err) + } + m2, err := NewManager(ManagerOptions{Store: store2}) + if err != nil { + t.Fatal(err) + } + + cfg1 := Config{ClientID: "client1", TokenEndpoint: fp1.server.URL + "/token"} + auth1 := DeviceAuth{DeviceCode: "dc1", Interval: 5 * time.Millisecond, ExpiresAt: time.Now().Add(5 * time.Second)} + + cfg2 := Config{ClientID: "client2", TokenEndpoint: fp2.server.URL + "/token"} + auth2 := DeviceAuth{DeviceCode: "dc2", Interval: 5 * time.Millisecond, ExpiresAt: time.Now().Add(5 * time.Second)} + + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + _, _ = m1.CompleteDeviceLogin(context.Background(), "prov1", cfg1, auth1) + }() + go func() { + defer wg.Done() + _, _ = m2.CompleteDeviceLogin(context.Background(), "prov2", cfg2, auth2) + }() + wg.Wait() + + tok1, ok1, err1 := store2.Load(ProviderKey("prov1")) + tok2, ok2, err2 := store2.Load(ProviderKey("prov2")) + if err1 != nil || !ok1 || tok1.AccessToken != "token-1" { + t.Fatalf("prov1: tok=%+v ok=%v err=%v", tok1, ok1, err1) + } + if err2 != nil || !ok2 || tok2.AccessToken != "token-2" { + t.Fatalf("prov2: tok=%+v ok=%v err=%v", tok2, ok2, err2) + } +} + +func TestCompleteDeviceLoginTwoManagersSameKeyConcurrent(t *testing.T) { + fp1 := newFakeProvider(t, `{"access_token":"token-a"}`) + fp2 := newFakeProvider(t, `{"access_token":"token-b"}`) + storePath := filepath.Join(t.TempDir(), "tokens.json") + store1, err := NewStore(StoreOptions{FilePath: storePath}) + if err != nil { + t.Fatal(err) + } + store2, err := NewStore(StoreOptions{FilePath: storePath}) + if err != nil { + t.Fatal(err) + } + m1, err := NewManager(ManagerOptions{Store: store1}) + if err != nil { + t.Fatal(err) + } + m2, err := NewManager(ManagerOptions{Store: store2}) + if err != nil { + t.Fatal(err) + } + + cfg1 := Config{ClientID: "client", TokenEndpoint: fp1.server.URL + "/token"} + cfg2 := Config{ClientID: "client", TokenEndpoint: fp2.server.URL + "/token"} + auth1 := DeviceAuth{DeviceCode: "dc1", Interval: 5 * time.Millisecond, ExpiresAt: time.Now().Add(5 * time.Second)} + auth2 := DeviceAuth{DeviceCode: "dc2", Interval: 5 * time.Millisecond, ExpiresAt: time.Now().Add(5 * time.Second)} + + errs := make(chan error, 2) + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + _, err := m1.CompleteDeviceLogin(context.Background(), "demo", cfg1, auth1) + errs <- err + }() + go func() { + defer wg.Done() + _, err := m2.CompleteDeviceLogin(context.Background(), "demo", cfg2, auth2) + errs <- err + }() + wg.Wait() + close(errs) + for err := range errs { + if err != nil { + t.Fatalf("CompleteDeviceLogin: %v", err) + } + } + + tok, ok, err := store2.Load(ProviderKey("demo")) + if err != nil || !ok { + t.Fatalf("load through second store: tok=%+v ok=%v err=%v", tok, ok, err) + } + if tok.AccessToken != "token-a" && tok.AccessToken != "token-b" { + t.Fatalf("surviving token = %q, want token-a or token-b", tok.AccessToken) + } +} + +func TestPrepareDeviceLoginKimiIssuerOverrideUsesDiscoveredEndpoints(t *testing.T) { + isolateKimiDeviceIDStorage(t) + var discoveryHits, deviceHits atomic.Int32 + var sawIdentity atomic.Bool + mux := http.NewServeMux() + server := httptest.NewServer(mux) + t.Cleanup(server.Close) + mux.HandleFunc("/.well-known/oauth-authorization-server", func(w http.ResponseWriter, _ *http.Request) { + discoveryHits.Add(1) + _, _ = io.WriteString(w, `{"issuer":"`+server.URL+`","token_endpoint":"`+server.URL+`/token","device_authorization_endpoint":"`+server.URL+`/device","authorization_endpoint":"`+server.URL+`/authorize"}`) + }) + mux.HandleFunc("/device", func(w http.ResponseWriter, r *http.Request) { + deviceHits.Add(1) + if r.Header.Get("X-Msh-Device-Id") != "" { + sawIdentity.Store(true) + } + _, _ = io.WriteString(w, `{"device_code":"dc","user_code":"U-1","verification_uri":"https://example/dev","expires_in":600,"interval":1}`) + }) + + store, err := NewStore(StoreOptions{FilePath: filepath.Join(t.TempDir(), "tok.json")}) + if err != nil { + t.Fatal(err) + } + m, err := NewManager(ManagerOptions{ + Store: store, + AllowPresets: true, + Env: map[string]string{ + "ZERO_OAUTH_KIMI_CODE_ISSUER_URL": server.URL, + }, + }) + if err != nil { + t.Fatal(err) + } + + auth, cfg, err := m.PrepareDeviceLogin(context.Background(), LoginOptions{Provider: "kimi-code"}) + if err != nil { + t.Fatalf("PrepareDeviceLogin: %v", err) + } + if auth.UserCode != "U-1" { + t.Fatalf("auth = %+v", auth) + } + if cfg.TokenEndpoint != server.URL+"/token" || cfg.DeviceAuthorizationEndpoint != server.URL+"/device" { + t.Fatalf("discovered endpoints = token %q device %q", cfg.TokenEndpoint, cfg.DeviceAuthorizationEndpoint) + } + if cfg.TokenEndpoint == "https://auth.kimi.com/api/oauth/token" || cfg.DeviceAuthorizationEndpoint == "https://auth.kimi.com/api/oauth/device_authorization" { + t.Fatal("issuer override still used preset Kimi destinations") + } + if len(cfg.ExtraHeaders) != 0 { + t.Fatalf("non-canonical discovered host must not send X-Msh-*: %#v", cfg.ExtraHeaders) + } + if discoveryHits.Load() == 0 || deviceHits.Load() == 0 { + t.Fatalf("discoveryHits=%d deviceHits=%d", discoveryHits.Load(), deviceHits.Load()) + } + if sawIdentity.Load() { + t.Fatal("device request sent X-Msh-* to a non-canonical issuer override") + } +} diff --git a/internal/oauth/oauth.go b/internal/oauth/oauth.go index 4938e7bf3..b55926929 100644 --- a/internal/oauth/oauth.go +++ b/internal/oauth/oauth.go @@ -23,6 +23,7 @@ package oauth import ( "errors" + "net/http" "strings" "time" ) @@ -79,6 +80,14 @@ type Config struct { // ExtraAuthParams are appended to the authorization URL (e.g. login_hint). ExtraAuthParams map[string]string + + // ExtraHeaders are set on every device-authorization, token-poll, code + // exchange, and refresh HTTP request this Config makes. They are derived + // from the resolved destinations (after env overlays and issuer discovery), + // not from the raw override env vars. Most providers need none; Kimi + // Code's backend rejects those requests with 401 unless its vendor-identity + // X-Msh-* headers are present (see kimiExtraHeaders in presets.go). + ExtraHeaders map[string]string } // Errors returned by the engine. Callers can match these with errors.Is. @@ -105,3 +114,13 @@ var ( // trimmed is a tiny helper used across the package. func trimmed(s string) string { return strings.TrimSpace(s) } + +// applyExtraHeaders sets a Config's provider-specific headers (see +// Config.ExtraHeaders) on an outgoing request, after the caller has already +// set the standard Content-Type/Accept pair so a provider entry could +// (in principle) override them. +func applyExtraHeaders(request *http.Request, headers map[string]string) { + for key, value := range headers { + request.Header.Set(key, value) + } +} diff --git a/internal/oauth/presets.go b/internal/oauth/presets.go index 3f8b7fe79..d0394a14c 100644 --- a/internal/oauth/presets.go +++ b/internal/oauth/presets.go @@ -1,8 +1,11 @@ package oauth import ( + "net/url" "os" "strings" + + "github.com/Gitlawb/zero/internal/kimiidentity" ) // envWithPresetsAllowed returns an env map that opts into the baked-in presets @@ -99,6 +102,26 @@ var builtinOAuthPresets = map[string]providerPreset{ Scopes: []string{"openid", "profile", "email", "offline_access", "api.connectors.read", "api.connectors.invoke"}, Flow: FlowLoopback, }, + // Kimi Code uses the same public OAuth client identity the open-source + // kimi-cli ships (github.com/MoonshotAI/kimi-cli, + // src/kimi_cli/auth/oauth.py). The flow is device-code only (RFC 8628) + // against auth.kimi.com — there is no loopback/authorize path — and the + // resulting access token is accepted directly as a bearer on the managed + // coding endpoint https://api.kimi.com/coding/v1 (an OpenAI-compatible + // endpoint). No ID-token claim extraction is needed; the bearer is the + // whole credential. The preset (and catalog descriptor) key is + // "kimi-code", not "kimi": the existing `moonshot` entry already aliases + // "kimi" to itself (its API-key path at api.moonshot.ai), and Get() + // matches an exact descriptor ID before it ever reaches another + // descriptor's aliases, so reusing "kimi" here would silently steal that + // alias out from under existing moonshot profiles. + "kimi-code": { + ClientID: "17e5f671-d194-4dfb-9706-5516cb48c098", + DeviceAuthorizationEndpoint: "https://auth.kimi.com/api/oauth/device_authorization", + TokenEndpoint: "https://auth.kimi.com/api/oauth/token", + Scopes: []string{"openid", "profile", "email", "offline_access"}, + Flow: FlowDevice, + }, } // lookupOAuthPreset returns the baked-in preset for a provider name (if any). @@ -128,3 +151,63 @@ func scopesOrPreset(envScopes string, preset []string) []string { // Copy so a caller appending to cfg.Scopes can't mutate the shared preset slice. return append([]string(nil), preset...) } + +// attachProviderHeaders derives ExtraHeaders from the endpoints this config +// will actually call. Identity headers are minted only when those destinations +// are canonical Kimi hosts. +func attachProviderHeaders(name string, cfg Config) Config { + cfg.ExtraHeaders = providerExtraHeaders(name, cfg.AuthorizationEndpoint, cfg.TokenEndpoint, cfg.DeviceAuthorizationEndpoint, cfg.IssuerURL) + return cfg +} + +// providerExtraHeaders returns the Config.ExtraHeaders a provider's OAuth +// requests need beyond the generic RFC 8628/OAuth2 form bodies this package +// builds. This is a protocol requirement of the provider's OWN backend +// (not tied to whether its preset client_id or an operator-supplied one is in +// use), so unlike the presets above it applies regardless of +// ZERO_OAUTH_ALLOW_PRESETS. endpoints are the resolved destinations, not the +// raw env overlay, so header policy cannot diverge from where requests go. +func providerExtraHeaders(name string, endpoints ...string) map[string]string { + if strings.ToLower(strings.TrimSpace(name)) == "kimi-code" { + for _, ep := range endpoints { + if ep = strings.TrimSpace(ep); ep != "" && !isCanonicalKimiHost(ep) { + return nil + } + } + return kimiExtraHeaders() + } + return nil +} + +// isCanonicalKimiHost reports whether urlStr is an explicitly approved Kimi +// OAuth/API endpoint. Only the hosts this package's preset and managed coding +// path talk to are allowed, and only over https: arbitrary *.kimi.com / +// *.moonshot.cn subdomains and cleartext (or non-http) schemes must not +// receive the persistent device identity. +func isCanonicalKimiHost(urlStr string) bool { + u, err := url.Parse(urlStr) + if err != nil || u.Hostname() == "" { + return false + } + if !strings.EqualFold(u.Scheme, "https") { + return false + } + switch strings.ToLower(u.Hostname()) { + case "auth.kimi.com", "api.kimi.com": + return true + default: + return false + } +} + +// kimiExtraHeaders returns the X-Msh-* vendor-identity headers Kimi Code's +// OAuth/API backend requires on every device-authorization, poll, and refresh +// request — reported to reject all of them with 401 otherwise. They come from +// the shared kimiidentity package so login/refresh and the catalog's +// managed-endpoint completions (kimi-code CustomHeaders) present the SAME +// persistent device identity; values are reverse-engineered from the +// open-source kimi-cli client, not from published documentation, and should +// be confirmed against a real login before this ships. +func kimiExtraHeaders() map[string]string { + return kimiidentity.Headers() +} diff --git a/internal/oauth/presets_test.go b/internal/oauth/presets_test.go index f73f1e102..f25f827aa 100644 --- a/internal/oauth/presets_test.go +++ b/internal/oauth/presets_test.go @@ -179,6 +179,59 @@ func TestResolveConfigHuggingFaceWithEnvClientID(t *testing.T) { } } +// Kimi Code ships a baked-in client_id (the public kimi-cli identity) AND a +// device-code endpoint, so the preset resolves without env. The flow is device +// only (RFC 8628): no loopback/authorize endpoint, no issuer discovery. The +// preset key is "kimi-code", not "kimi": moonshot already aliases "kimi" to +// itself (see TestKimiAliasStillResolvesToMoonshot in the providercatalog +// package's catalog_test.go), so reusing it here would steal that alias from +// existing moonshot profiles. +func TestResolveConfigKimiCodePreset(t *testing.T) { + // isolateKimiDeviceIDStorage lives in providers_test.go (same package). + isolateKimiDeviceIDStorage(t) + r := NewRegistry() + cfg, flow, err := r.ResolveConfig("kimi-code", map[string]string{"ZERO_OAUTH_ALLOW_PRESETS": "1"}) + if err != nil { + t.Fatalf("ResolveConfig(kimi-code): %v", err) + } + if cfg.ClientID != "17e5f671-d194-4dfb-9706-5516cb48c098" { + t.Fatalf("client_id = %q", cfg.ClientID) + } + if cfg.DeviceAuthorizationEndpoint != "https://auth.kimi.com/api/oauth/device_authorization" { + t.Fatalf("device endpoint = %q", cfg.DeviceAuthorizationEndpoint) + } + if cfg.TokenEndpoint != "https://auth.kimi.com/api/oauth/token" { + t.Fatalf("token = %q", cfg.TokenEndpoint) + } + if cfg.AuthorizationEndpoint != "" { + t.Fatalf("authorize endpoint = %q, want empty (Kimi has no loopback flow)", cfg.AuthorizationEndpoint) + } + if flow != FlowDevice { + t.Fatalf("flow = %q, want device (RFC 8628 only)", flow) + } + if len(cfg.Scopes) == 0 { + t.Fatal("preset scopes should be populated") + } + for _, header := range []string{"X-Msh-Platform", "X-Msh-Version", "X-Msh-Device-Name", "X-Msh-Device-Model", "X-Msh-Os-Version", "X-Msh-Device-Id"} { + if cfg.ExtraHeaders[header] == "" { + t.Fatalf("ExtraHeaders[%q] = %q, want a non-empty vendor-identity header (Kimi's backend rejects requests missing these)", header, cfg.ExtraHeaders[header]) + } + } +} + +// A provider with no ExtraHeaders requirement (e.g. xAI) must not pick up +// Kimi's headers or any other provider's. +func TestResolveConfigWithoutExtraHeadersRequirement(t *testing.T) { + r := NewRegistry() + cfg, _, err := r.ResolveConfig("xai", map[string]string{"ZERO_OAUTH_ALLOW_PRESETS": "1"}) + if err != nil { + t.Fatalf("ResolveConfig(xai): %v", err) + } + if len(cfg.ExtraHeaders) != 0 { + t.Fatalf("ExtraHeaders = %#v, want none for a provider with no header requirement", cfg.ExtraHeaders) + } +} + // ChatGPT (Codex) ships a baked-in client_id (the public Codex CLI identity), // so the preset resolves without env. The flow is loopback because the Codex // backend requires a browser; there is no device-code path. diff --git a/internal/oauth/providers.go b/internal/oauth/providers.go index e3b2807a3..98c8da3ea 100644 --- a/internal/oauth/providers.go +++ b/internal/oauth/providers.go @@ -69,15 +69,27 @@ func (r *Registry) ResolveConfig(name string, env map[string]string) (Config, Fl if presetsAllowed(env) { preset, _ = lookupOAuthPreset(name) } + authEnv := strings.TrimSpace(envValue(env, envKey(name, "AUTHORIZE_URL"))) + tokenEnv := strings.TrimSpace(envValue(env, envKey(name, "TOKEN_URL"))) + deviceEnv := strings.TrimSpace(envValue(env, envKey(name, "DEVICE_URL"))) + issuerEnv := strings.TrimSpace(envValue(env, envKey(name, "ISSUER_URL"))) + presetAuth, presetToken, presetDevice := preset.AuthorizationEndpoint, preset.TokenEndpoint, preset.DeviceAuthorizationEndpoint + if issuerEnv != "" { + // An issuer override is a destination change. Keep only endpoints the + // operator also overrode; leftover preset URLs would mix header policy + // from the new issuer with requests that still go to the old host. + presetAuth, presetToken, presetDevice = "", "", "" + } cfg := Config{ ClientID: firstNonEmpty(strings.TrimSpace(envValue(env, envKey(name, "CLIENT_ID"))), preset.ClientID), ClientSecret: firstNonEmpty(strings.TrimSpace(envValue(env, envKey(name, "CLIENT_SECRET"))), preset.ClientSecret), - AuthorizationEndpoint: firstNonEmpty(strings.TrimSpace(envValue(env, envKey(name, "AUTHORIZE_URL"))), preset.AuthorizationEndpoint), - TokenEndpoint: firstNonEmpty(strings.TrimSpace(envValue(env, envKey(name, "TOKEN_URL"))), preset.TokenEndpoint), - DeviceAuthorizationEndpoint: firstNonEmpty(strings.TrimSpace(envValue(env, envKey(name, "DEVICE_URL"))), preset.DeviceAuthorizationEndpoint), - IssuerURL: firstNonEmpty(strings.TrimSpace(envValue(env, envKey(name, "ISSUER_URL"))), preset.IssuerURL), + AuthorizationEndpoint: firstNonEmpty(authEnv, presetAuth), + TokenEndpoint: firstNonEmpty(tokenEnv, presetToken), + DeviceAuthorizationEndpoint: firstNonEmpty(deviceEnv, presetDevice), + IssuerURL: firstNonEmpty(issuerEnv, preset.IssuerURL), Scopes: scopesOrPreset(envValue(env, envKey(name, "SCOPES")), preset.Scopes), } + cfg = attachProviderHeaders(name, cfg) if cfg.ClientID == "" { hint := "" if _, ok := lookupOAuthPreset(name); ok { diff --git a/internal/oauth/providers_test.go b/internal/oauth/providers_test.go index cbffc6b86..774ba2217 100644 --- a/internal/oauth/providers_test.go +++ b/internal/oauth/providers_test.go @@ -5,6 +5,17 @@ import ( "testing" ) +// isolateKimiDeviceIDStorage redirects os.UserConfigDir so kimiidentity never +// writes kimi-device-id under the real user config root. DeviceID is path-keyed, +// so setting these env vars is enough (no separate cache reset). +func isolateKimiDeviceIDStorage(t *testing.T) { + t.Helper() + root := t.TempDir() + t.Setenv("XDG_CONFIG_HOME", root) + t.Setenv("APPDATA", root) + t.Setenv("HOME", root) +} + func TestResolveConfigFromEnv(t *testing.T) { r := NewRegistry() env := map[string]string{ @@ -94,4 +105,139 @@ func TestEnvKey(t *testing.T) { if got := envKey("two.part", "SCOPES"); got != "ZERO_OAUTH_TWO_PART_SCOPES" { t.Fatalf("envKey = %q", got) } + if got := envKey("kimi-code", "TOKEN_URL"); got != "ZERO_OAUTH_KIMI_CODE_TOKEN_URL" { + t.Fatalf("envKey(kimi-code) = %q", got) + } +} + +func TestResolveConfigKimiCodeStripsExtraHeadersOnEndpointOverride(t *testing.T) { + isolateKimiDeviceIDStorage(t) + r := NewRegistry() + // Canonical host (no override): ExtraHeaders has X-Msh-Device-Id + cfgCanonical, _, err := r.ResolveConfig("kimi-code", map[string]string{ + "ZERO_OAUTH_ALLOW_PRESETS": "1", + }) + if err != nil { + t.Fatalf("ResolveConfig: %v", err) + } + if len(cfgCanonical.ExtraHeaders) == 0 || cfgCanonical.ExtraHeaders["X-Msh-Device-Id"] == "" { + t.Fatalf("expected X-Msh-Device-Id on canonical host, got: %v", cfgCanonical.ExtraHeaders) + } + + // Overridden to non-canonical host: ExtraHeaders must be empty + cfgOverride, _, err := r.ResolveConfig("kimi-code", map[string]string{ + "ZERO_OAUTH_ALLOW_PRESETS": "1", + "ZERO_OAUTH_KIMI_CODE_TOKEN_URL": "https://my-custom-proxy.example.com/oauth/token", + }) + if err != nil { + t.Fatalf("ResolveConfig with override: %v", err) + } + if len(cfgOverride.ExtraHeaders) != 0 { + t.Fatalf("expected empty ExtraHeaders on non-canonical host override, got: %v", cfgOverride.ExtraHeaders) + } +} + +func TestIsCanonicalKimiHostAllowlist(t *testing.T) { + approved := []string{ + "https://auth.kimi.com/api/oauth/token", + "https://AUTH.KIMI.COM/path", + "https://api.kimi.com/coding/v1", + "https://api.kimi.com/", + } + for _, ep := range approved { + if !isCanonicalKimiHost(ep) { + t.Fatalf("isCanonicalKimiHost(%q) = false, want true", ep) + } + } + rejected := []string{ + "https://evil.kimi.com/oauth", + "https://kimi.com/oauth", + "https://www.kimi.com/", + "https://auth.moonshot.cn/oauth", + "https://api.moonshot.cn/v1", + "https://attacker.example.com/", + "https://auth.kimi.com.evil.example/", + // Non-HTTPS schemes must never receive the device identity even on an + // approved host (a proxy override could otherwise downgrade to cleartext). + "http://auth.kimi.com/oauth", + "http://api.kimi.com/v1", + "ftp://AUTH.KIMI.COM/oauth", + "//auth.kimi.com/oauth", + "not a url", + "", + } + for _, ep := range rejected { + if isCanonicalKimiHost(ep) { + t.Fatalf("isCanonicalKimiHost(%q) = true, want false", ep) + } + } +} + +// Overriding only to another approved Kimi host must keep identity headers. +func TestResolveConfigKimiCodeKeepsExtraHeadersOnApprovedOverride(t *testing.T) { + isolateKimiDeviceIDStorage(t) + r := NewRegistry() + cfg, _, err := r.ResolveConfig("kimi-code", map[string]string{ + "ZERO_OAUTH_ALLOW_PRESETS": "1", + "ZERO_OAUTH_KIMI_CODE_TOKEN_URL": "https://api.kimi.com/oauth/token", + }) + if err != nil { + t.Fatalf("ResolveConfig: %v", err) + } + if cfg.ExtraHeaders["X-Msh-Device-Id"] == "" { + t.Fatalf("expected ExtraHeaders on approved host override, got: %v", cfg.ExtraHeaders) + } +} + +// A kimi.com subdomain that is not on the allowlist must strip identity headers. +func TestResolveConfigKimiCodeStripsExtraHeadersOnKimiSubdomain(t *testing.T) { + isolateKimiDeviceIDStorage(t) + r := NewRegistry() + cfg, _, err := r.ResolveConfig("kimi-code", map[string]string{ + "ZERO_OAUTH_ALLOW_PRESETS": "1", + "ZERO_OAUTH_KIMI_CODE_TOKEN_URL": "https://evil.kimi.com/oauth/token", + "ZERO_OAUTH_KIMI_CODE_DEVICE_URL": "https://evil.kimi.com/oauth/device", + }) + if err != nil { + t.Fatalf("ResolveConfig: %v", err) + } + if len(cfg.ExtraHeaders) != 0 { + t.Fatalf("expected empty ExtraHeaders on unapproved kimi subdomain, got: %v", cfg.ExtraHeaders) + } +} + +func TestResolveConfigKimiIssuerOnlyClearsPresetEndpoints(t *testing.T) { + isolateKimiDeviceIDStorage(t) + r := NewRegistry() + + canonical, _, err := r.ResolveConfig("kimi-code", map[string]string{ + "ZERO_OAUTH_ALLOW_PRESETS": "1", + "ZERO_OAUTH_KIMI_CODE_ISSUER_URL": "https://auth.kimi.com", + }) + if err != nil { + t.Fatalf("ResolveConfig canonical issuer: %v", err) + } + if canonical.IssuerURL != "https://auth.kimi.com" { + t.Fatalf("IssuerURL = %q", canonical.IssuerURL) + } + if canonical.TokenEndpoint != "" || canonical.DeviceAuthorizationEndpoint != "" { + t.Fatalf("issuer-only override must not keep preset endpoints: token=%q device=%q", canonical.TokenEndpoint, canonical.DeviceAuthorizationEndpoint) + } + if canonical.ExtraHeaders["X-Msh-Device-Id"] == "" { + t.Fatalf("canonical issuer should keep identity headers, got %#v", canonical.ExtraHeaders) + } + + foreign, _, err := r.ResolveConfig("kimi-code", map[string]string{ + "ZERO_OAUTH_ALLOW_PRESETS": "1", + "ZERO_OAUTH_KIMI_CODE_ISSUER_URL": "https://idp.example/oauth", + }) + if err != nil { + t.Fatalf("ResolveConfig foreign issuer: %v", err) + } + if foreign.TokenEndpoint != "" || foreign.DeviceAuthorizationEndpoint != "" { + t.Fatalf("foreign issuer must not keep preset Kimi endpoints: token=%q device=%q", foreign.TokenEndpoint, foreign.DeviceAuthorizationEndpoint) + } + if len(foreign.ExtraHeaders) != 0 { + t.Fatalf("foreign issuer must not send X-Msh-* headers, got %#v", foreign.ExtraHeaders) + } } diff --git a/internal/oauth/store.go b/internal/oauth/store.go index 1bc7f1dc8..42cbf9e41 100644 --- a/internal/oauth/store.go +++ b/internal/oauth/store.go @@ -249,6 +249,49 @@ func (s *Store) Save(key string, token Token) error { }) } +// CommitDeviceToken saves token for key while holding the store-scoped lock +// across snapshot, pre-write cancellation check, write, and rollback on post-write cancellation. +func (s *Store) CommitDeviceToken(key string, token Token, checkCancel func() error, hook func()) error { + if err := ValidateKey(key); err != nil { + return err + } + s.mu.Lock() + defer s.mu.Unlock() + return s.blob.withLock(s.now, func() error { + if checkCancel != nil { + if err := checkCancel(); err != nil { + return err + } + } + state, err := s.readState() + if err != nil { + return err + } + prevToken, prevOK := state.Tokens[key] + if hook != nil { + hook() + } + state.Tokens[key] = token + if err := s.writeState(state); err != nil { + return err + } + if checkCancel != nil { + if err := checkCancel(); err != nil { + if prevOK { + state.Tokens[key] = prevToken + } else { + delete(state.Tokens, key) + } + if rerr := s.writeState(state); rerr != nil { + return errors.Join(err, fmt.Errorf("oauth: rollback device token commit: %w", rerr)) + } + return err + } + } + return nil + }) +} + // Load returns the token for key; the bool is false when none is stored. func (s *Store) Load(key string) (Token, bool, error) { if err := ValidateKey(key); err != nil { diff --git a/internal/providercatalog/catalog.go b/internal/providercatalog/catalog.go index a29f330a5..0b9889253 100644 --- a/internal/providercatalog/catalog.go +++ b/internal/providercatalog/catalog.go @@ -5,6 +5,8 @@ import ( "fmt" "strings" "unicode" + + "github.com/Gitlawb/zero/internal/kimiidentity" ) type Transport string @@ -70,6 +72,19 @@ type Descriptor struct { // OAuthDeviceFlow reports that RFC 8628 device-code login is supported (for // headless / SSH use) in addition to the browser flow. OAuthDeviceFlow bool + // OAuthDeviceOnly reports that device-code is the ONLY OAuth path this + // provider has — there is no browser/loopback authorize endpoint to fall + // back to (Kimi Code). Callers that default a plain Enter/click to the + // browser flow for OAuthDeviceFlow providers must check this first: for a + // device-only provider that generic path has no endpoint to hit at all. + OAuthDeviceOnly bool + + // RuntimeHeaders, if set, lazily produces headers (e.g. a vendor's + // device-identity headers) to attach to CustomHeaders only when + // withRuntimeHeaders is true. Providers that need identity/headers on + // runtime requests (not just OAuth) set this instead of hardcoding an ID + // check inside cloneDescriptor — keeping the clone helper provider-agnostic. + RuntimeHeaders func() map[string]string // Recommended marks a provider surfaced at the top and badged in // catalog-ordered lists and pickers. The recommended descriptors are the first @@ -98,6 +113,13 @@ func RuntimeUnsupportedReason(descriptor Descriptor) string { } } +// IsRuntimeIdentityHeader reports whether a header key is a runtime identity +// header (such as Kimi's X-Msh-* headers) that should be generated dynamically +// rather than persisted in saved configuration profiles. +func IsRuntimeIdentityHeader(key string) bool { + return strings.HasPrefix(strings.ToLower(key), "x-msh-") +} + var descriptors = []Descriptor{ // GitLawb OpenGateway — the recommended default. An OpenAI-compatible gateway // that smart-routes by model id across upstream providers (xiaomi-mimo, @@ -135,6 +157,38 @@ var descriptors = []Descriptor{ d.RequiresAuth = true return oauthProvider(d, false, false) }(), + // Kimi Code (managed OAuth) — the bearer from a Kimi Code device-code OAuth + // login routes to the managed coding endpoint at https://api.kimi.com/coding/v1 + // (NOT api.moonshot.ai/v1, which is the API-key path for the `moonshot` + // catalog entry). Kimi only supports the RFC 8628 device-code flow against + // auth.kimi.com; the access token is accepted directly as a bearer on the + // managed endpoint, so no client spoofing is involved. The baked-in preset + // ships the public Kimi Code client_id and endpoints (off by default, like + // the other presets); env overrides via ZERO_OAUTH_KIMI_CODE_* win. + // + // Descriptor ID is "kimi-code", not "kimi": the `moonshot` entry below + // already aliases "kimi" to itself, and Get() matches an exact descriptor + // ID before it reaches another descriptor's aliases, so reusing "kimi" + // here would silently steal that alias out from under any existing + // moonshot profile (changing its endpoint, default model, and + // MOONSHOT_API_KEY auth without the user asking for it). + // + // Default model is "kimi-for-coding" — the managed endpoint's standard + // tier, available to every Kimi Code member. "kimi-for-coding-highspeed" + // also exists but requires a higher subscription tier (Allegretto+); a + // user who has it can select it explicitly rather than have a fresh + // `zero auth kimi-code` default to a model their plan may not include. + func() Descriptor { + d := openAICompat("kimi-code", "Kimi Code", "https://api.kimi.com/coding/v1", "kimi-for-coding", nil) + d.RequiresAuth = true + d = oauthProvider(d, false, true) + d.OAuthDeviceOnly = true + // Kimi's managed endpoint requires stable vendor-identity headers on + // every request; lazily mint them (with a persistent device ID) only + // when a runtime request actually needs them, not merely on listing. + d.RuntimeHeaders = kimiidentity.Headers + return d + }(), openAICompat("groq", "Groq", "https://api.groq.com/openai/v1", "llama-3.3-70b-versatile", []string{"GROQ_API_KEY"}), openAICompat("deepseek", "DeepSeek", "https://api.deepseek.com/v1", "deepseek-chat", []string{"DEEPSEEK_API_KEY"}), openAICompat("together", "Together AI", "https://api.together.xyz/v1", "meta-llama/Llama-3.3-70B-Instruct-Turbo", []string{"TOGETHER_API_KEY"}), @@ -183,7 +237,8 @@ var descriptors = []Descriptor{ func All() []Descriptor { copied := make([]Descriptor, 0, len(descriptors)) for _, descriptor := range descriptors { - copied = append(copied, cloneDescriptor(descriptor)) + // Listing must not mint Kimi's on-disk device identity. + copied = append(copied, cloneDescriptor(descriptor, false)) } return copied } @@ -192,11 +247,11 @@ func Get(id string) (Descriptor, bool) { normalized := NormalizeID(id) for _, descriptor := range descriptors { if descriptor.ID == normalized { - return cloneDescriptor(descriptor), true + return cloneDescriptor(descriptor, true), true } for _, alias := range descriptor.Aliases { if NormalizeID(alias) == normalized { - return cloneDescriptor(descriptor), true + return cloneDescriptor(descriptor, true), true } } } @@ -308,7 +363,9 @@ func OAuthProviders() []Descriptor { out := []Descriptor{} for _, descriptor := range descriptors { if descriptor.OAuth { - out = append(out, cloneDescriptor(descriptor)) + // Listing only: runtime headers (and the Kimi device-id file they + // mint) are applied when Get/Require builds a real profile. + out = append(out, cloneDescriptor(descriptor, false)) } } return out @@ -363,12 +420,21 @@ func transportDescriptor(id string, name string, transport Transport, baseURL st } } -func cloneDescriptor(descriptor Descriptor) Descriptor { +// cloneDescriptor returns an independent copy of descriptor. When +// withRuntimeHeaders is true, a descriptor with RuntimeHeaders set (e.g. +// kimi-code) also receives its vendor identity headers (including a persistent +// device ID on disk). Listing paths (All, OAuthProviders) pass false so merely +// enumerating providers never mints ~/.config/zero/kimi-device-id for users +// who never touch Kimi; Get/Require pass true so resolve-time profile building +// and completions still present the same identity the OAuth login used. +func cloneDescriptor(descriptor Descriptor, withRuntimeHeaders bool) Descriptor { descriptor.AuthEnvVars = append([]string{}, descriptor.AuthEnvVars...) descriptor.SupportedAPIFormats = append([]APIFormat{}, descriptor.SupportedAPIFormats...) descriptor.Aliases = append([]string{}, descriptor.Aliases...) if descriptor.CustomHeaders != nil { descriptor.CustomHeaders = copyStringMap(descriptor.CustomHeaders) + } else if withRuntimeHeaders && descriptor.RuntimeHeaders != nil { + descriptor.CustomHeaders = descriptor.RuntimeHeaders() } return descriptor } diff --git a/internal/providercatalog/catalog_test.go b/internal/providercatalog/catalog_test.go index 7467bfe7a..40fcbfc68 100644 --- a/internal/providercatalog/catalog_test.go +++ b/internal/providercatalog/catalog_test.go @@ -2,11 +2,36 @@ package providercatalog import ( "errors" + "os" + "path/filepath" "reflect" "strings" "testing" ) +// isolateKimiDeviceIDStorage redirects os.UserConfigDir so Get("kimi-code") +// (which runs RuntimeHeaders → kimiidentity.Headers) never writes +// kimi-device-id under the real user config root. DeviceID is path-keyed, so +// setting these env vars is enough (no separate cache reset). +func isolateKimiDeviceIDStorage(t *testing.T) { + t.Helper() + root := t.TempDir() + t.Setenv("XDG_CONFIG_HOME", root) + t.Setenv("APPDATA", root) + t.Setenv("HOME", root) +} + +// kimiDeviceIDFile returns the path kimiidentity would use under the current +// (possibly test-redirected) os.UserConfigDir. +func kimiDeviceIDFile(t *testing.T) string { + t.Helper() + configDir, err := os.UserConfigDir() + if err != nil { + t.Fatalf("UserConfigDir: %v", err) + } + return filepath.Join(configDir, "zero", "kimi-device-id") +} + var expectedCatalogIDs = []string{ "gitlawb-opengateway", "aimlapi", @@ -19,6 +44,7 @@ var expectedCatalogIDs = []string{ "openrouter", "huggingface", "chatgpt", + "kimi-code", "groq", "deepseek", "together", @@ -201,6 +227,23 @@ func TestAtlasCloudDescriptor(t *testing.T) { } func TestCatalogDescriptorsExposeRequiredDefaults(t *testing.T) { + knownTransports := map[Transport]bool{ + TransportOpenAI: true, + TransportAnthropic: true, + TransportGoogle: true, + TransportBedrock: true, + TransportVertex: true, + TransportOpenAICompatible: true, + TransportAnthropicCompatible: true, + } + knownFormats := map[APIFormat]bool{ + APIFormatOpenAIResponses: true, + APIFormatOpenAIChatCompletions: true, + APIFormatAnthropicMessages: true, + APIFormatGoogleGenerateContent: true, + APIFormatBedrockConverse: true, + APIFormatVertexGenerateContent: true, + } for _, descriptor := range All() { if descriptor.ID == "" { t.Fatal("provider ID is required") @@ -211,7 +254,7 @@ func TestCatalogDescriptorsExposeRequiredDefaults(t *testing.T) { if descriptor.Transport == "" { t.Fatalf("provider %q should expose a transport", descriptor.ID) } - if !ValidTransport(descriptor.Transport) { + if !knownTransports[descriptor.Transport] { t.Fatalf("provider %q has unknown transport %q", descriptor.ID, descriptor.Transport) } if descriptor.DefaultBaseURL == "" { @@ -224,17 +267,11 @@ func TestCatalogDescriptorsExposeRequiredDefaults(t *testing.T) { t.Fatalf("provider %q should expose at least one supported API format", descriptor.ID) } for _, format := range descriptor.SupportedAPIFormats { - if !ValidAPIFormat(format) { + if !knownFormats[format] { t.Fatalf("provider %q has unknown API format %q", descriptor.ID, format) } } } - if ValidTransport("missing") { - t.Fatal("ValidTransport should reject unknown transports") - } - if ValidAPIFormat("missing") { - t.Fatal("ValidAPIFormat should reject unknown API formats") - } } func TestRemoteProvidersDeclareAuthOrExplicitPublicAccess(t *testing.T) { @@ -361,41 +398,6 @@ func TestLookupNormalizesIDsAndAliases(t *testing.T) { } } -func TestListByTransportPreservesCatalogOrder(t *testing.T) { - cases := map[Transport][]string{ - TransportOpenAI: {"openai"}, - TransportAnthropic: {"anthropic"}, - TransportGoogle: {"google"}, - TransportBedrock: {"bedrock"}, - TransportVertex: {"vertex"}, - TransportAnthropicCompat: {"minimax", "minimaxi-cn", "opencode-go-anthropic-compatible", "custom-anthropic-compatible"}, - TransportOpenAICompat: {"gitlawb-opengateway", "aimlapi", "ollama-cloud", "ollama", "lmstudio", "openrouter", "huggingface", "chatgpt", "groq", "deepseek", "together", "fireworks", "dashscope", "moonshot", "atlascloud", "longcat", "nvidia-nim", "mistral", "github", "xai", "venice", "xiaomi-mimo", "bankr", "zai", "zai-cn", "kilocode", "opencode", "opencode-go", "atomic-chat", "chatgpt-proxy", "custom-openai-compatible"}, - } - - for transport, wantIDs := range cases { - descriptors := ListByTransport(transport) - gotIDs := make([]string, 0, len(descriptors)) - for _, descriptor := range descriptors { - if descriptor.Transport != Transport(NormalizeID(string(transport))) { - t.Fatalf("ListByTransport(%q) returned provider %q with transport %q", transport, descriptor.ID, descriptor.Transport) - } - gotIDs = append(gotIDs, descriptor.ID) - } - if !reflect.DeepEqual(gotIDs, wantIDs) { - t.Fatalf("ListByTransport(%q) IDs = %#v, want %#v", transport, gotIDs, wantIDs) - } - } - if descriptors := ListByTransport("missing"); len(descriptors) != 0 { - t.Fatalf("ListByTransport(missing) returned %#v, want empty", descriptors) - } - if gotIDs := descriptorIDs(ListByTransport(TransportOpenAICompatible)); !reflect.DeepEqual(gotIDs, cases[TransportOpenAICompat]) { - t.Fatalf("ListByTransport(openai-compatible alias) IDs = %#v, want %#v", gotIDs, cases[TransportOpenAICompat]) - } - if gotIDs := descriptorIDs(ListByTransport(TransportAnthropicCompatible)); !reflect.DeepEqual(gotIDs, cases[TransportAnthropicCompat]) { - t.Fatalf("ListByTransport(anthropic-compatible alias) IDs = %#v, want %#v", gotIDs, cases[TransportAnthropicCompat]) - } -} - func TestReturnedDescriptorsAreCopies(t *testing.T) { descriptors := All() descriptors[0].ID = "changed" @@ -426,9 +428,53 @@ func TestReturnedDescriptorsAreCopies(t *testing.T) { } } +func TestListByTransportPreservesCatalogOrder(t *testing.T) { + cases := map[Transport][]string{ + TransportOpenAI: {"openai"}, + TransportAnthropic: {"anthropic"}, + TransportGoogle: {"google"}, + TransportBedrock: {"bedrock"}, + TransportVertex: {"vertex"}, + TransportAnthropicCompat: { + "minimax", "minimaxi-cn", "opencode-go-anthropic-compatible", "custom-anthropic-compatible", + }, + TransportOpenAICompat: { + "gitlawb-opengateway", "aimlapi", "ollama-cloud", "ollama", "lmstudio", "openrouter", + "huggingface", "chatgpt", "kimi-code", "groq", "deepseek", "together", "fireworks", "dashscope", + "moonshot", "atlascloud", "longcat", "nvidia-nim", "mistral", "github", "xai", "venice", + "xiaomi-mimo", "bankr", "zai", "zai-cn", "kilocode", "opencode", "opencode-go", "atomic-chat", + "chatgpt-proxy", "custom-openai-compatible", + }, + } + + for transport, wantIDs := range cases { + descriptors := ListByTransport(transport) + gotIDs := make([]string, 0, len(descriptors)) + for _, descriptor := range descriptors { + if descriptor.Transport != Transport(NormalizeID(string(transport))) { + t.Fatalf("ListByTransport(%q) returned provider %q with transport %q", transport, descriptor.ID, descriptor.Transport) + } + gotIDs = append(gotIDs, descriptor.ID) + } + if !reflect.DeepEqual(gotIDs, wantIDs) { + t.Fatalf("ListByTransport(%q) IDs = %#v, want %#v", transport, gotIDs, wantIDs) + } + } + if descriptors := ListByTransport("missing"); len(descriptors) != 0 { + t.Fatalf("ListByTransport(missing) returned %#v, want empty", descriptors) + } + if gotIDs := descriptorIDs(ListByTransport(TransportOpenAICompatible)); !reflect.DeepEqual(gotIDs, cases[TransportOpenAICompat]) { + t.Fatalf("ListByTransport(openai-compatible alias) IDs = %#v, want %#v", gotIDs, cases[TransportOpenAICompat]) + } + if gotIDs := descriptorIDs(ListByTransport(TransportAnthropicCompatible)); !reflect.DeepEqual(gotIDs, cases[TransportAnthropicCompat]) { + t.Fatalf("ListByTransport(anthropic-compatible alias) IDs = %#v, want %#v", gotIDs, cases[TransportAnthropicCompat]) + } +} + func TestOAuthProviderClassification(t *testing.T) { + isolateKimiDeviceIDStorage(t) oauthIDs := descriptorIDs(OAuthProviders()) - if want := []string{"openrouter", "huggingface", "chatgpt", "xai"}; !reflect.DeepEqual(oauthIDs, want) { + if want := []string{"openrouter", "huggingface", "chatgpt", "kimi-code", "xai"}; !reflect.DeepEqual(oauthIDs, want) { t.Fatalf("OAuthProviders() = %#v, want %#v", oauthIDs, want) } if d, _ := Get("openrouter"); !d.OAuthMintsKey { @@ -437,6 +483,9 @@ func TestOAuthProviderClassification(t *testing.T) { if d, _ := Get("xai"); !d.OAuthDeviceFlow { t.Fatal("xai should advertise device-code flow") } + if d, _ := Get("kimi-code"); !d.OAuthDeviceFlow || !d.OAuthDeviceOnly { + t.Fatal("kimi-code should advertise device-only code flow") + } if d, _ := Get("huggingface"); !d.OAuthDeviceFlow { t.Fatal("huggingface should advertise device-code flow") } @@ -445,6 +494,81 @@ func TestOAuthProviderClassification(t *testing.T) { } } +// TestKimiAliasStillResolvesToMoonshot pins the alias-collision fix: moonshot +// already exposes "kimi" as an alias for its API-key endpoint. The Kimi Code +// OAuth descriptor must use a non-conflicting ID ("kimi-code") so resolving +// "kimi" continues to land on moonshot (endpoint, default model, MOONSHOT_API_KEY). +func TestKimiAliasStillResolvesToMoonshot(t *testing.T) { + isolateKimiDeviceIDStorage(t) + d, ok := Get("kimi") + if !ok { + t.Fatal(`Get("kimi") returned false`) + } + if d.ID != "moonshot" { + t.Fatalf(`Get("kimi").ID = %q, want "moonshot" (kimi-code must not steal this alias)`, d.ID) + } + if d.DefaultBaseURL != "https://api.moonshot.ai/v1" { + t.Fatalf(`Get("kimi").DefaultBaseURL = %q, want moonshot API-key endpoint`, d.DefaultBaseURL) + } + if len(d.AuthEnvVars) == 0 || d.AuthEnvVars[0] != "MOONSHOT_API_KEY" { + t.Fatalf(`Get("kimi").AuthEnvVars = %#v, want MOONSHOT_API_KEY`, d.AuthEnvVars) + } + if d.OAuth { + t.Fatal(`Get("kimi") must not be OAuth (that is kimi-code, not the moonshot alias)`) + } + + code, ok := Get("kimi-code") + if !ok { + t.Fatal(`Get("kimi-code") returned false`) + } + if code.ID != "kimi-code" { + t.Fatalf(`Get("kimi-code").ID = %q`, code.ID) + } + if code.DefaultBaseURL != "https://api.kimi.com/coding/v1" { + t.Fatalf(`Get("kimi-code").DefaultBaseURL = %q, want managed coding endpoint`, code.DefaultBaseURL) + } + if !code.OAuth || !code.OAuthDeviceOnly { + t.Fatalf(`Get("kimi-code") oauth flags wrong: OAuth=%v OAuthDeviceOnly=%v`, code.OAuth, code.OAuthDeviceOnly) + } +} + +// TestKimiRuntimeHeadersOnlyOnGet ensures listing providers (All / OAuthProviders) +// does not populate kimi-code's CustomHeaders (which mints a device-id file), +// while Get does so resolve-time request building still gets the vendor headers. +func TestKimiRuntimeHeadersOnlyOnGet(t *testing.T) { + isolateKimiDeviceIDStorage(t) + for _, d := range All() { + if d.ID == "kimi-code" && d.CustomHeaders != nil { + t.Fatalf("All() must not populate kimi-code CustomHeaders: %#v", d.CustomHeaders) + } + } + for _, d := range OAuthProviders() { + if d.ID == "kimi-code" && d.CustomHeaders != nil { + t.Fatalf("OAuthProviders() must not populate kimi-code CustomHeaders: %#v", d.CustomHeaders) + } + } + // Listing must not mint a device id even if a regression invokes + // RuntimeHeaders and discards the returned headers. + devicePath := kimiDeviceIDFile(t) + if _, err := os.Stat(devicePath); err == nil { + t.Fatalf("listing created kimi device id at %s; All/OAuthProviders must not persist identity", devicePath) + } else if !os.IsNotExist(err) { + t.Fatalf("stat kimi device id after listing: %v", err) + } + d, ok := Get("kimi-code") + if !ok { + t.Fatal(`Get("kimi-code") returned false`) + } + for _, header := range []string{"X-Msh-Platform", "X-Msh-Version", "X-Msh-Device-Name", "X-Msh-Device-Model", "X-Msh-Os-Version", "X-Msh-Device-Id"} { + if d.CustomHeaders[header] == "" { + t.Fatalf("Get(kimi-code).CustomHeaders[%q] empty, want vendor-identity header for completions", header) + } + } + if _, err := os.Stat(devicePath); err != nil { + t.Fatalf("Get(kimi-code) should persist device id at %s: %v", devicePath, err) + } +} + func descriptorIDs(descriptors []Descriptor) []string { ids := make([]string, 0, len(descriptors)) for _, descriptor := range descriptors { diff --git a/internal/providercatalog/export_test.go b/internal/providercatalog/export_test.go index a10cfd154..75bab4559 100644 --- a/internal/providercatalog/export_test.go +++ b/internal/providercatalog/export_test.go @@ -9,12 +9,15 @@ func IDs() []string { return ids } +// ListByTransport returns catalog descriptors for transport, preserving catalog +// order. Listing clones only (no RuntimeHeaders) so tests never mint Kimi's +// on-disk device id. func ListByTransport(transport Transport) []Descriptor { normalized := Transport(NormalizeID(string(transport))) items := make([]Descriptor, 0) for _, descriptor := range descriptors { if descriptor.Transport == normalized { - items = append(items, cloneDescriptor(descriptor)) + items = append(items, cloneDescriptor(descriptor, false)) } } return items diff --git a/internal/providercatalog/oauth_test.go b/internal/providercatalog/oauth_test.go index 0733f32ae..1888226c5 100644 --- a/internal/providercatalog/oauth_test.go +++ b/internal/providercatalog/oauth_test.go @@ -4,8 +4,8 @@ import "testing" func TestOAuthProviders(t *testing.T) { providers := OAuthProviders() - if len(providers) != 4 { - t.Fatalf("OAuthProviders() = %d, want 4 (openrouter, xai, huggingface, chatgpt)", len(providers)) + if len(providers) != 5 { + t.Fatalf("OAuthProviders() = %d, want 5 (openrouter, xai, kimi-code, huggingface, chatgpt)", len(providers)) } byID := map[string]Descriptor{} for _, d := range providers { @@ -30,6 +30,10 @@ func TestOAuthProviders(t *testing.T) { if !ok || cg.OAuthMintsKey || cg.OAuthDeviceFlow { t.Fatalf("chatgpt oauth flags wrong: %+v", cg) } + kimi, ok := byID["kimi-code"] + if !ok || kimi.OAuthMintsKey || !kimi.OAuthDeviceFlow || !kimi.OAuthDeviceOnly { + t.Fatalf("kimi-code oauth flags wrong: %+v", kimi) + } } func TestOAuthProvidersReturnsIndependentClones(t *testing.T) { @@ -58,7 +62,7 @@ func TestOAuthProvidersReturnsIndependentClones(t *testing.T) { func TestNonOAuthProvidersNotFlagged(t *testing.T) { for _, d := range All() { switch d.ID { - case "openrouter", "xai", "huggingface", "chatgpt": + case "openrouter", "xai", "kimi-code", "huggingface", "chatgpt": continue } if d.OAuth { diff --git a/internal/tui/model.go b/internal/tui/model.go index 2eb01eeba..20bcad2c9 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -1185,8 +1185,10 @@ func (m model) noBlockingModal() bool { func (m model) quit() (tea.Model, tea.Cmd) { if m.providerWizard != nil { + m.providerWizard.cancelDeviceLogin() m.providerWizard.resetAimlapiOnboard() } + m.setup.cancelDeviceLogin() m.stopPRWatcher() m.stopAllBackgroundTerminalSessions() m.shutdownLSPManager() diff --git a/internal/tui/oauth_device.go b/internal/tui/oauth_device.go index 00bdb557c..e6ec9276b 100644 --- a/internal/tui/oauth_device.go +++ b/internal/tui/oauth_device.go @@ -59,7 +59,13 @@ func oauthDevicePrepare(name string) (oauth.DeviceAuth, oauth.Config, error) { // oauthDeviceComplete polls for the token authorized via oauthDevicePrepare and // stores it under provider: (phase 2). The runtime resolver then attaches // the refreshable token to model calls. -func oauthDeviceComplete(name string, cfg oauth.Config, auth oauth.DeviceAuth) error { +// +// parent must be cancelable by the caller (not context.Background()): the poll +// can run for up to 10 minutes waiting on the user to finish authorizing in +// their browser, and if the caller has since abandoned the flow (e.g. Esc in +// the TUI), canceling parent is the only way to stop this from silently +// completing and persisting a credential the user believed they'd canceled. +func oauthDeviceComplete(parent context.Context, name string, cfg oauth.Config, auth oauth.DeviceAuth) error { store, err := oauth.NewStore(oauth.StoreOptions{}) if err != nil { return err @@ -72,7 +78,7 @@ func oauthDeviceComplete(name string, cfg oauth.Config, auth oauth.DeviceAuth) e if err != nil { return err } - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute) + ctx, cancel := context.WithTimeout(parent, 10*time.Minute) defer cancel() _, err = manager.CompleteDeviceLogin(ctx, name, cfg, auth) return err diff --git a/internal/tui/onboarding.go b/internal/tui/onboarding.go index ed09f3a87..cc6a98e02 100644 --- a/internal/tui/onboarding.go +++ b/internal/tui/onboarding.go @@ -50,23 +50,34 @@ type setupState struct { oauthMode bool oauthPending bool oauthErr string + // oauthAttemptID identifies the current OAuth attempt so a late phase-one + // or poll-result message from an attempt the user already abandoned (Esc, + // then restarted the same provider) cannot be mistaken for the current + // one — providerID alone is not enough since the provider doesn't change + // across attempts. + oauthAttemptID int // Device-code login (RFC 8628) state while an OAuth login is in flight. oauthDevice bool deviceUserCode string deviceVerificationURI string - stage setupStage - err string - baseURL string - name string - apiKey textinput.Model - models []providerWizardModel - modelIndex int - modelQuery string - modelForID string - modelLoad bool - modelErr string - modelSrc string - modelGen uint64 + // deviceLoginCancel cancels the background context backing an in-flight + // device-code poll (setupDevicePollCmd), so abandoning setup (Esc) actually + // stops the poll instead of leaving it to run for up to 10 minutes and + // silently save a credential the user backed out of. + deviceLoginCancel context.CancelFunc + stage setupStage + err string + baseURL string + name string + apiKey textinput.Model + models []providerWizardModel + modelIndex int + modelQuery string + modelForID string + modelLoad bool + modelErr string + modelSrc string + modelGen uint64 // aimlapi holds the shared aimlapi.com onboarding sub-flow while // setup is on setupStageAimlapi. aimlapi *aimlapiOnboardState @@ -93,16 +104,25 @@ type setupOAuthMsg struct { apiKey string tokenLogin bool providerID string + attemptID int err error } // setupOAuthProviderOptions filters the full provider list to the OAuth-capable // ones for the OAuth method path. ChatGPT/Claude are not here — they can't do // real in-app OAuth (use "browse" + a local proxy); see docs/oauth-subscriptions.md. +// +// Uses OAuthProviders() (listing clones) rather than Get(): Get runs +// RuntimeHeaders and would mint ~/.config/zero/kimi-device-id for every user +// who merely paints the "How do you want to connect?" screen. func setupOAuthProviderOptions(all []SetupProviderOption) []SetupProviderOption { + oauthIDs := map[string]struct{}{} + for _, descriptor := range providercatalog.OAuthProviders() { + oauthIDs[descriptor.ID] = struct{}{} + } out := []SetupProviderOption{} for _, option := range all { - if descriptor, ok := providercatalog.Get(option.ID); ok && descriptor.OAuth { + if _, ok := oauthIDs[option.ID]; ok { out = append(out, option) } } @@ -154,10 +174,18 @@ func (m model) handleSetupKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { if m.setup.oauthPending { switch { case keyCtrl(msg, 'c'): + // Stop a device-code poll actually running in the background instead + // of merely quitting the UI — otherwise completing the login later + // in the browser still silently saves the credential. + m.setup.cancelDeviceLogin() return m, tea.Quit case keyIs(msg, tea.KeyEsc): m.setup.oauthPending = false m.setup.oauthDevice = false + // Stop a device-code poll actually running in the background instead + // of merely dismissing the UI — otherwise completing the login later + // in the browser still silently saves the credential. + m.setup.cancelDeviceLogin() } return m, nil } @@ -182,6 +210,7 @@ func (m model) handleSetupKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { } switch { case keyCtrl(msg, 'c'): + m.setup.cancelDeviceLogin() return m, tea.Quit case keyIs(msg, tea.KeyEsc): if m.setup.stage > setupStageWelcome { @@ -194,6 +223,7 @@ func (m model) handleSetupKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { return m, nil } if m.setup.required { + m.setup.cancelDeviceLogin() return m, tea.Quit } return m.exitSetupToChat() @@ -244,6 +274,7 @@ func (m model) handleSetupKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { } switch keyText(msg) { case "q": + m.setup.cancelDeviceLogin() return m, tea.Quit case "k": if m.setup.stage == setupStageProvider { @@ -279,6 +310,7 @@ func (m model) handleSetupKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { return m.advanceSetup() } case "q": + m.setup.cancelDeviceLogin() return m, tea.Quit case "k": switch m.setup.stage { @@ -537,7 +569,7 @@ func (m *model) moveSetupMethod(delta int) { // setupOAuthCmd runs the chosen provider's browser OAuth login off the UI // goroutine for first-run setup. Mirrors the /provider wizard's flow. -func setupOAuthCmd(provider providercatalog.Descriptor) tea.Cmd { +func setupOAuthCmd(provider providercatalog.Descriptor, attemptID int) tea.Cmd { switch { case provider.OAuthMintsKey: return func() tea.Msg { @@ -545,17 +577,17 @@ func setupOAuthCmd(provider providercatalog.Descriptor) tea.Cmd { OpenBrowser: browser.OpenURL, Timeout: 3 * time.Minute, }) - return setupOAuthMsg{apiKey: key, providerID: provider.ID, err: err} + return setupOAuthMsg{apiKey: key, providerID: provider.ID, attemptID: attemptID, err: err} } case provider.ID == "chatgpt": return func() tea.Msg { err := runProviderChatGPTLogin() - return setupOAuthMsg{tokenLogin: true, providerID: provider.ID, err: err} + return setupOAuthMsg{tokenLogin: true, providerID: provider.ID, attemptID: attemptID, err: err} } default: name := provider.ID return func() tea.Msg { - return setupOAuthMsg{tokenLogin: true, providerID: name, err: runProviderTokenLogin(name)} + return setupOAuthMsg{tokenLogin: true, providerID: name, attemptID: attemptID, err: runProviderTokenLogin(name)} } } } @@ -564,6 +596,7 @@ func setupOAuthCmd(provider providercatalog.Descriptor) tea.Cmd { // for first-run setup: the user_code + verification URI to display. type setupOAuthDeviceMsg struct { providerID string + attemptID int userCode string verifyURL string cfg oauth.Config @@ -571,14 +604,15 @@ type setupOAuthDeviceMsg struct { err error } -func setupDevicePrepareCmd(name string) tea.Cmd { +func setupDevicePrepareCmd(name string, attemptID int) tea.Cmd { return func() tea.Msg { auth, cfg, err := oauthDevicePrepare(name) if err != nil { - return setupOAuthDeviceMsg{providerID: name, err: err} + return setupOAuthDeviceMsg{providerID: name, attemptID: attemptID, err: err} } return setupOAuthDeviceMsg{ providerID: name, + attemptID: attemptID, userCode: auth.UserCode, verifyURL: oauthDeviceVerifyTarget(auth), cfg: cfg, @@ -587,10 +621,50 @@ func setupDevicePrepareCmd(name string) tea.Cmd { } } -func setupDevicePollCmd(name string, cfg oauth.Config, auth oauth.DeviceAuth) tea.Cmd { +// setupDevicePollCmd runs phase 2 (poll for the token + store) off the UI +// goroutine and reports completion as a regular OAuth result. ctx must be +// cancelable by the caller so abandoning setup actually stops the poll. +func setupDevicePollCmd(ctx context.Context, name string, attemptID int, cfg oauth.Config, auth oauth.DeviceAuth) tea.Cmd { return func() tea.Msg { - return setupOAuthMsg{tokenLogin: true, providerID: name, err: oauthDeviceComplete(name, cfg, auth)} + return setupOAuthMsg{tokenLogin: true, providerID: name, attemptID: attemptID, err: oauthDeviceComplete(ctx, name, cfg, auth)} + } +} + +// cancelDeviceLogin stops an in-flight device-code poll, if any, and clears +// the stored cancel func. Safe to call even when no poll is running. +func (setup *setupState) cancelDeviceLogin() { + if setup == nil || setup.deviceLoginCancel == nil { + return + } + setup.deviceLoginCancel() + setup.deviceLoginCancel = nil +} + +// beginSetupOAuthAttempt starts a new first-run OAuth attempt and returns its +// id. A new attempt supersedes any poll left over from a previous one and +// bumps oauthAttemptID so a late phase-one/poll-result message tagged with an +// older id is rejected by setupOAuthResultMatches even though the provider +// (and oauthPending) look the same as the new attempt. +func (setup *setupState) beginSetupOAuthAttempt(device bool) int { + setup.cancelDeviceLogin() + setup.oauthAttemptID++ + setup.oauthPending = true + setup.oauthDevice = device + setup.oauthErr = "" + setup.deviceUserCode = "" + setup.deviceVerificationURI = "" + return setup.oauthAttemptID +} + +// setupOAuthResultMatches reports whether an OAuth result message still +// belongs to the in-flight attempt, rejecting a stale message left over from +// an attempt the user abandoned (Esc) and then restarted against the same +// provider. +func (setup *setupState) setupOAuthResultMatches(providerID string, attemptID int) bool { + if setup == nil || !setup.visible || !setup.oauthPending || strings.TrimSpace(providerID) == "" { + return false } + return setup.oauthAttemptID == attemptID } // startSetupDeviceLogin begins the device-code flow for the selected OAuth @@ -599,23 +673,14 @@ func (m model) startSetupDeviceLogin(descriptor providercatalog.Descriptor) (tea if !descriptor.OAuth || !descriptor.OAuthDeviceFlow { return m, nil } - m.setup.oauthPending = true - m.setup.oauthDevice = true - m.setup.oauthErr = "" - m.setup.deviceUserCode = "" - m.setup.deviceVerificationURI = "" - return m, setupDevicePrepareCmd(descriptor.ID) + attemptID := m.setup.beginSetupOAuthAttempt(true) + return m, setupDevicePrepareCmd(descriptor.ID, attemptID) } // applySetupOAuthDeviceCode handles phase 1 of device-code login: show the code, // then start phase 2 (the token poll). On error the redacted message is shown. func (m model) applySetupOAuthDeviceCode(msg setupOAuthDeviceMsg) (tea.Model, tea.Cmd) { - if !m.setup.visible || !m.setup.oauthPending { - return m, nil - } - // Ignore a stale result from a login the user has since replaced with a - // different provider (an in-flight prepare landing after the switch). - if msg.providerID != "" && msg.providerID != m.setupProviderDescriptor().ID { + if !m.setup.setupOAuthResultMatches(msg.providerID, msg.attemptID) { return m, nil } if msg.err != nil { @@ -626,22 +691,26 @@ func (m model) applySetupOAuthDeviceCode(msg setupOAuthDeviceMsg) (tea.Model, te } m.setup.deviceUserCode = msg.userCode m.setup.deviceVerificationURI = msg.verifyURL - return m, setupDevicePollCmd(msg.providerID, msg.cfg, msg.auth) + // The poll runs off the UI goroutine for up to 10 minutes; give it a + // context setup can cancel (Esc) so abandoning the flow actually stops it + // instead of leaving it to complete in the background. m.ctx is the + // parent so quitting zero entirely also unblocks the poll. + ctx, cancel := context.WithCancel(m.ctx) + m.setup.deviceLoginCancel = cancel + return m, setupDevicePollCmd(ctx, msg.providerID, msg.attemptID, msg.cfg, msg.auth) } // applySetupOAuth folds an OAuth login result into the first-run setup: on success // the credential is captured (minted key) or relied upon (stored token) and setup // jumps to model selection; on failure the redacted error is shown. func (m model) applySetupOAuth(msg setupOAuthMsg) (tea.Model, tea.Cmd) { - if !m.setup.visible || !m.setup.oauthPending { - return m, nil - } - // Ignore a stale result for a provider the user has since switched away from, - // so a late login can't capture a credential against the wrong provider. - if msg.providerID != "" && msg.providerID != m.setupProviderDescriptor().ID { + if !m.setup.setupOAuthResultMatches(msg.providerID, msg.attemptID) { return m, nil } m.setup.oauthPending = false + // The poll (if this was a device-code login) is done either way; release + // its context. + m.setup.cancelDeviceLogin() if msg.err != nil { m.setup.oauthErr = redaction.RedactString(msg.err.Error(), redaction.Options{}) return m, nil @@ -784,13 +853,16 @@ func (m model) advanceSetup() (tea.Model, tea.Cmd) { if descriptor.OAuth { // Headless/SSH boxes can't open a browser — use device code there // by default (the user can also force it with "d" from the list). - if descriptor.OAuthDeviceFlow && oauthPreferDeviceFlow() { + // A device-only provider (no loopback/authorization endpoint) + // must take this path on desktops too: the generic manager + // would still run the device flow, but with no UI its + // verification URL and user code are discarded and the + // "browser login" spinner just times out. + if descriptor.OAuthDeviceFlow && (descriptor.OAuthDeviceOnly || oauthPreferDeviceFlow()) { return m.startSetupDeviceLogin(descriptor) } - m.setup.oauthPending = true - m.setup.oauthDevice = false - m.setup.oauthErr = "" - return m, setupOAuthCmd(descriptor) + attemptID := m.setup.beginSetupOAuthAttempt(false) + return m, setupOAuthCmd(descriptor, attemptID) } } if m.setup.stage == setupStageProvider { @@ -1134,6 +1206,7 @@ func (m model) setupNameInputActive() bool { func (m model) handleSetupEndpointKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { switch { case keyCtrl(msg, 'c'): + m.setup.cancelDeviceLogin() return m, tea.Quit case keyIs(msg, tea.KeyEsc) || keyIs(msg, tea.KeyLeft): m.setup.stage = m.previousSetupStage() @@ -1155,6 +1228,7 @@ func (m model) handleSetupEndpointKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { func (m model) handleSetupNameKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { switch { case keyCtrl(msg, 'c'): + m.setup.cancelDeviceLogin() return m, tea.Quit case keyIs(msg, tea.KeyEsc) || keyIs(msg, tea.KeyLeft): m.setup.stage = m.previousSetupStage() @@ -1176,6 +1250,7 @@ func (m model) handleSetupNameKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { func (m model) handleSetupCredentialKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) { switch { case keyCtrl(msg, 'c'): + m.setup.cancelDeviceLogin() return m, tea.Quit case keyIs(msg, tea.KeyEsc) || keyIs(msg, tea.KeyLeft): m.setup.stage = m.previousSetupStage() diff --git a/internal/tui/onboarding_test.go b/internal/tui/onboarding_test.go index 1f70616fb..43257d022 100644 --- a/internal/tui/onboarding_test.go +++ b/internal/tui/onboarding_test.go @@ -5,6 +5,7 @@ import ( "errors" "net/http" "os" + "path/filepath" "strings" "testing" @@ -51,6 +52,84 @@ func TestSetupMethodOptionsDropsOAuthWithoutOAuthProviders(t *testing.T) { } } +// TestSetupMethodAndOAuthProviderScreensDoNotMintKimiDeviceID pins the +// invariant that painting the method screen (and advancing onto the OAuth +// provider list that includes kimi-code) must not create kimi-device-id. +// setupOAuthProviderOptions used to call providercatalog.Get on every option, +// which ran RuntimeHeaders and minted the file for users who never pick Kimi. +func TestSetupMethodAndOAuthProviderScreensDoNotMintKimiDeviceID(t *testing.T) { + root := t.TempDir() + t.Setenv("XDG_CONFIG_HOME", root) + t.Setenv("APPDATA", root) + t.Setenv("HOME", root) + + configDir, err := os.UserConfigDir() + if err != nil { + t.Fatalf("UserConfigDir: %v", err) + } + devicePath := filepath.Join(configDir, "zero", "kimi-device-id") + + // Put non-Kimi OAuth providers first so advancing onto the list does not + // Get(kimi-code) via resetSetupModels for the default selection. kimi-code + // is still present in the list (the regression is enumerating it). + m := newModel(context.Background(), Options{ + Setup: SetupOptions{ + Visible: true, + Providers: []SetupProviderOption{ + {ID: "openai", Name: "OpenAI", EnvVar: "OPENAI_API_KEY", RequiresAuth: true}, + {ID: "openrouter", Name: "OpenRouter", EnvVar: "OPENROUTER_API_KEY", RequiresAuth: true}, + {ID: "xai", Name: "xAI"}, + {ID: "kimi-code", Name: "Kimi Code", RequiresAuth: true}, + }, + }, + }) + m.setup.stage = setupStageMethod + + // Method screen render path calls setupMethodOptions → setupOAuthProviderOptions. + _ = m.setupMethodOptions() + _ = m.setupMethodLines(80) + // The filter alone must not mint (this is the exact regression). + _ = setupOAuthProviderOptions(m.setup.allProviders) + + // Advance onto the OAuth provider list (same filter, now assigned to m.setup.providers). + oauthIdx := -1 + for i, option := range m.setupMethodOptions() { + if option.oauth { + oauthIdx = i + break + } + } + if oauthIdx < 0 { + t.Fatal("expected an OAuth method option when kimi-code is in the setup list") + } + m.setup.selectedMethod = oauthIdx + next, _ := m.advanceSetup() + m = next.(model) + if m.setup.stage != setupStageProvider || !m.setup.oauthMode { + t.Fatalf("stage=%v oauthMode=%v, want provider+oauth", m.setup.stage, m.setup.oauthMode) + } + // Paint the provider list too. + _ = m.setupProviderLines(80, 24) + + if _, err := os.Stat(devicePath); err == nil { + t.Fatalf("method/provider screens minted kimi device id at %s", devicePath) + } else if !os.IsNotExist(err) { + t.Fatalf("stat kimi device id: %v", err) + } + + // Filter still includes kimi-code for users who do want it. + found := false + for _, option := range m.setup.providers { + if option.ID == "kimi-code" { + found = true + break + } + } + if !found { + t.Fatalf("OAuth provider list missing kimi-code: %#v", m.setup.providers) + } +} + func TestAimlapiCheckoutLinkWrapsWithoutTruncation(t *testing.T) { link := "https://checkout.stripe.com/c/pay/cs_test_abcdefghijklmnopqrstuvwxyz0123456789abcdefghijklmnopqrstuvwxyz0123456789#fidkdWxOYHwnPyd1blpxYHZxWjA0" lines := aimlapiLinkLines(link, 32) @@ -2036,6 +2115,154 @@ func TestApplySetupOAuthDeviceCodeShowsCodeAndPolls(t *testing.T) { } } +// TestSetupCtrlCCancelsDeviceLoginPoll regression-tests a bug where Ctrl+C +// during a first-run device-code poll (phase 2) quit the whole program +// without canceling the background context the poll command runs on. Since +// the TUI's parent context is context.Background(), the poll (up to 10 +// minutes) kept running after the process later exited via os.Exit, and if +// the user then finished authorizing in the browser, a completed-in-flight +// write could still land. Ctrl+C must cancel the poll before quitting. +func TestSetupCtrlCCancelsDeviceLoginPoll(t *testing.T) { + t.Setenv("XDG_CONFIG_HOME", t.TempDir()) + t.Setenv("ZERO_OAUTH_TOKENS_PATH", filepath.Join(t.TempDir(), "oauth-tokens.json")) + + m := setupAtOAuthList(t) + for i, p := range m.setup.providers { + if p.ID == "xai" { + m.setup.selected = i + break + } + } + m.setup.oauthPending = true + m.setup.oauthDevice = true + attemptID := m.setup.oauthAttemptID + + res, cmd := m.applySetupOAuthDeviceCode(setupOAuthDeviceMsg{ + providerID: "xai", attemptID: attemptID, userCode: "WXYZ-9", verifyURL: "https://x.ai/device", + }) + m = res.(model) + if cmd == nil { + t.Fatal("device-code msg should start the poll command") + } + if m.setup.deviceLoginCancel == nil { + t.Fatal("starting the poll should store a cancel func on setup") + } + + updated, _ := m.Update(testKeyCtrl('c')) + m = updated.(model) + if m.setup.deviceLoginCancel != nil { + t.Fatal("Ctrl+C should cancel the in-flight device-code poll") + } + + raw := cmd() + msg, ok := raw.(setupOAuthMsg) + if !ok { + t.Fatalf("poll command returned %T, want setupOAuthMsg", raw) + } + if !errors.Is(msg.err, context.Canceled) { + t.Fatalf("poll error = %v, want context.Canceled (Ctrl+C should have canceled the background poll)", msg.err) + } +} + +// TestModelQuitCancelsSetupDeviceLoginPoll regression-tests the direct quit +// path: model.quit() (the path every Ctrl+C-to-exit and "q"-to-exit +// eventually reaches) must cancel the setup (first-run onboarding) device-code +// poll, not just the Ctrl+C handler. Mirrors the provider-wizard quit test in +// provider_wizard_oauth_test.go for the setup flow. +func TestModelQuitCancelsSetupDeviceLoginPoll(t *testing.T) { + t.Setenv("XDG_CONFIG_HOME", t.TempDir()) + t.Setenv("ZERO_OAUTH_TOKENS_PATH", filepath.Join(t.TempDir(), "oauth-tokens.json")) + + m := setupAtOAuthList(t) + for i, p := range m.setup.providers { + if p.ID == "xai" { + m.setup.selected = i + break + } + } + m.setup.oauthPending = true + m.setup.oauthDevice = true + attemptID := m.setup.oauthAttemptID + + res, cmd := m.applySetupOAuthDeviceCode(setupOAuthDeviceMsg{ + providerID: "xai", attemptID: attemptID, userCode: "WXYZ-9", verifyURL: "https://x.ai/device", + }) + m = res.(model) + if cmd == nil { + t.Fatal("device-code msg should start the poll command") + } + if m.setup.deviceLoginCancel == nil { + t.Fatal("starting the poll should store a cancel func on setup") + } + + quit, _ := m.quit() + quitModel := quit.(model) + if quitModel.setup.deviceLoginCancel != nil { + t.Fatal("quit should cancel the in-flight setup device-code poll") + } + + raw := cmd() + msg, ok := raw.(setupOAuthMsg) + if !ok { + t.Fatalf("poll command returned %T, want setupOAuthMsg", raw) + } + if !errors.Is(msg.err, context.Canceled) { + t.Fatalf("poll error = %v, want context.Canceled (quit should have canceled the background poll)", msg.err) + } +} + +// TestSetupStaleDeviceCodeAttemptRejected regression-tests a bug where +// abandoning a device-code login with Esc and immediately restarting it for +// the same provider let a late phase-one result from the FIRST attempt +// overwrite the second attempt's displayed code and start polling an +// authorization the user had already backed out of: setupOAuthDeviceMsg only +// carried providerID, which is identical across both attempts. +func TestSetupStaleDeviceCodeAttemptRejected(t *testing.T) { + m := setupAtOAuthList(t) + for i, p := range m.setup.providers { + if p.ID == "xai" { + m.setup.selected = i + break + } + } + + updated, cmd := m.Update(testKeyText("d")) // attempt 1 + m = updated.(model) + if cmd == nil { + t.Fatal("'d' should return the device-prepare command") + } + staleAttemptID := m.setup.oauthAttemptID + + // Abandon attempt 1 with Esc, then immediately restart against the same + // provider (attempt 2). + updated, _ = m.Update(testKey(tea.KeyEsc)) + m = updated.(model) + if m.setup.oauthPending { + t.Fatal("Esc should abandon the pending device login") + } + updated, cmd = m.Update(testKeyText("d")) // attempt 2 + m = updated.(model) + if cmd == nil { + t.Fatal("restarting should return a new device-prepare command") + } + if m.setup.oauthAttemptID == staleAttemptID { + t.Fatal("restarting the device flow should assign a new attempt id") + } + + // The stale phase-one result from attempt 1 lands after attempt 2 is + // already in flight — same provider, so providerID alone can't reject it. + res, pollCmd := m.applySetupOAuthDeviceCode(setupOAuthDeviceMsg{ + providerID: "xai", attemptID: staleAttemptID, userCode: "STALE-1", verifyURL: "https://x.ai/device/stale", + }) + m = res.(model) + if pollCmd != nil { + t.Fatal("stale attempt's phase-one result must not start a poll") + } + if m.setup.deviceUserCode == "STALE-1" { + t.Fatalf("stale attempt overwrote the current attempt's device code: %+v", m.setup) + } +} + func TestApplySetupOAuthSuccessAdvancesToModel(t *testing.T) { m := newModel(context.Background(), Options{ DiscoverProviderModels: func(ctx context.Context, profile config.ProviderProfile) ([]providermodeldiscovery.Model, error) { @@ -2182,3 +2409,65 @@ func TestCompleteSetupExportsActiveProviderEnv(t *testing.T) { t.Fatalf("%s = %q after setup save, want %q (children would spawn on the stale provider)", config.ActiveProviderEnv, got, next.providerName) } } + +// TestSetupEnterStartsDeviceFlowForDeviceOnlyProvider pins the first-run +// onboarding counterpart of the /provider wizard fix: Kimi Code has no +// loopback/authorize endpoint, so a plain desktop Enter must take the +// device-code path (showing the verification URL and user code) instead of +// the generic browser-login command, whose manager would run the device flow +// with a discarded output writer and leave the spinner to time out. +func TestSetupEnterStartsDeviceFlowForDeviceOnlyProvider(t *testing.T) { + // Isolate config dirs before setupProviderDescriptor → Get("kimi-code") + // so the test never mints ~/.config/zero/kimi-device-id on the host. + root := t.TempDir() + t.Setenv("XDG_CONFIG_HOME", root) + t.Setenv("APPDATA", root) + t.Setenv("HOME", root) + + // Force a "normal desktop with a browser available" environment: + // oauthPreferDeviceFlow() already picks device flow on headless boxes, + // which would mask the bug this test exists to catch. + t.Setenv("ZERO_OAUTH_DEVICE", "") + t.Setenv("SSH_CONNECTION", "") + t.Setenv("SSH_TTY", "") + t.Setenv("DISPLAY", ":0") + t.Setenv("WAYLAND_DISPLAY", "") + + m := newModel(context.Background(), Options{Setup: SetupOptions{ + Visible: true, + Providers: []SetupProviderOption{ + {ID: "kimi-code", Name: "Kimi Code", RequiresAuth: true}, + {ID: "xai", Name: "xAI", DefaultModel: "grok-4", RequiresAuth: true, EnvVar: "XAI_API_KEY"}, + }, + }}) + m.width = 100 + m.height = 30 + m = pressSetupContinueOnce(m) // Welcome → Method + m.setup.selectedMethod = 0 // Sign in with OAuth + updated, _ := m.Update(testKey(tea.KeyEnter)) + m = updated.(model) + + found := false + for i, p := range m.setup.providers { + if p.ID == "kimi-code" { + m.setup.selected = i + found = true + break + } + } + if !found { + t.Fatalf("kimi-code missing from OAuth provider list: %#v", m.setup.providers) + } + if !m.setupProviderDescriptor().OAuthDeviceOnly { + t.Fatal("test fixture assumes kimi-code is OAuthDeviceOnly") + } + + updated, cmd := m.Update(testKey(tea.KeyEnter)) + m = updated.(model) + if !m.setup.oauthPending || !m.setup.oauthDevice { + t.Fatalf("Enter on a device-only provider should start device login (pending=%v device=%v)", m.setup.oauthPending, m.setup.oauthDevice) + } + if cmd == nil { + t.Fatal("Enter on a device-only provider should return the device-prepare command") + } +} diff --git a/internal/tui/provider_wizard.go b/internal/tui/provider_wizard.go index 8f4ee8bb3..e0906fd02 100644 --- a/internal/tui/provider_wizard.go +++ b/internal/tui/provider_wizard.go @@ -9,6 +9,7 @@ import ( "os" "path/filepath" "strings" + "sync/atomic" "time" "unicode" @@ -45,6 +46,9 @@ func (m model) applyProviderWizardOAuth(msg providerWizardOAuthMsg) (model, tea. return m, nil } m.providerWizard.oauthPending = false + // The poll (if this was a device-code login) is done either way; release + // its context. + m.providerWizard.cancelDeviceLogin() if msg.err != nil { m.providerWizard.oauthErr = redaction.ErrorMessage(msg.err, redaction.Options{}) return m, nil @@ -137,7 +141,13 @@ func (m model) applyProviderWizardDeviceCode(msg providerWizardDeviceCodeMsg) (m } m.providerWizard.deviceUserCode = msg.userCode m.providerWizard.deviceVerificationURI = msg.verifyURL - return m, providerWizardDevicePollCmd(msg.providerID, msg.attemptID, msg.cfg, msg.auth) + // The poll runs off the UI goroutine for up to 10 minutes; give it a + // context the wizard can cancel (Esc) so abandoning the flow actually + // stops it instead of leaving it to complete in the background. m.ctx is + // the parent so quitting zero entirely also unblocks the poll. + ctx, cancel := context.WithCancel(m.ctx) + m.providerWizard.deviceLoginCancel = cancel + return m, providerWizardDevicePollCmd(ctx, msg.providerID, msg.attemptID, msg.cfg, msg.auth) } // providerWizardSupportsOAuth reports whether the credential step should offer a @@ -315,10 +325,11 @@ func providerWizardDevicePrepareCmd(name string, attemptID int) tea.Cmd { } // providerWizardDevicePollCmd runs phase 2 (poll for the token + store) off the -// UI goroutine and reports completion as a regular OAuth result. -func providerWizardDevicePollCmd(name string, attemptID int, cfg oauth.Config, auth oauth.DeviceAuth) tea.Cmd { +// UI goroutine and reports completion as a regular OAuth result. ctx must be +// cancelable by the caller so abandoning the wizard actually stops the poll. +func providerWizardDevicePollCmd(ctx context.Context, name string, attemptID int, cfg oauth.Config, auth oauth.DeviceAuth) tea.Cmd { return func() tea.Msg { - return providerWizardOAuthMsg{providerID: name, attemptID: attemptID, tokenLogin: true, err: oauthDeviceComplete(name, cfg, auth)} + return providerWizardOAuthMsg{providerID: name, attemptID: attemptID, tokenLogin: true, err: oauthDeviceComplete(ctx, name, cfg, auth)} } } @@ -376,7 +387,7 @@ func providerWizardMethodOptions() []providerWizardMethodOption { options = append(options, providerWizardMethodOption{ oauth: true, label: "Sign in with OAuth", - subtitle: "One-click browser login, no API key to copy (OpenRouter, xAI, ChatGPT, Hugging Face).", + subtitle: "No API key to copy — one-click browser login (OpenRouter, xAI, ChatGPT, Hugging Face) or device code (Kimi Code).", }) } options = append(options, providerWizardMethodOption{ @@ -439,6 +450,11 @@ type providerWizardState struct { oauthDevice bool deviceUserCode string deviceVerificationURI string + // deviceLoginCancel cancels the background context backing an in-flight + // device-code poll (providerWizardDevicePollCmd), so abandoning the wizard + // (Esc) actually stops the poll instead of leaving it to run for up to 10 + // minutes and silently save a credential the user backed out of. + deviceLoginCancel context.CancelFunc // aimlapi holds the shared aimlapi.com onboarding sub-flow while // the wizard is on providerWizardStepAimlapi. aimlapi *aimlapiOnboardState @@ -499,8 +515,24 @@ func (wizard *providerWizardState) currentProvider() providercatalog.Descriptor return wizard.providers[wizard.selectedProvider] } +// oauthAttemptSeq issues OAuth attempt IDs that stay unique for the lifetime +// of the process. A per-wizard counter is not enough: the wizard is discarded +// wholesale on Esc and rebuilt on the next open, so its counter restarts and +// the second wizard reissues ID 1 for the same provider. A late result from +// the first, canceled attempt then satisfies oauthResultMatches and can +// persist a credential into, or advance, a wizard that never requested it. +var oauthAttemptSeq atomic.Int64 + +// nextOAuthAttemptID returns the next process-wide attempt ID. Never zero, so +// a zero value still reads as "no attempt". +func nextOAuthAttemptID() int { + return int(oauthAttemptSeq.Add(1)) +} + func (wizard *providerWizardState) beginOAuthAttempt(device bool) int { - wizard.oauthAttemptID++ + // A new attempt supersedes any poll left over from a previous one. + wizard.cancelDeviceLogin() + wizard.oauthAttemptID = nextOAuthAttemptID() wizard.oauthPending = true wizard.oauthDevice = device wizard.oauthErr = "" @@ -516,6 +548,16 @@ func (wizard *providerWizardState) oauthResultMatches(providerID string, attempt return wizard.currentProvider().ID == providerID && wizard.oauthAttemptID == attemptID } +// cancelDeviceLogin stops an in-flight device-code poll, if any, and clears +// the stored cancel func. Safe to call even when no poll is running. +func (wizard *providerWizardState) cancelDeviceLogin() { + if wizard == nil || wizard.deviceLoginCancel == nil { + return + } + wizard.deviceLoginCancel() + wizard.deviceLoginCancel = nil +} + // resetAimlapiOnboard drops any in-flight aimlapi.com sub-flow (cancelling a // running top-up stream) when the selected provider changes. func (wizard *providerWizardState) resetAimlapiOnboard() { @@ -591,6 +633,7 @@ func (wizard *providerWizardState) move(delta int) { wizard.modelLoading = false wizard.modelLoadError = "" wizard.oauthPending = false + wizard.cancelDeviceLogin() wizard.oauthErr = "" wizard.resetAimlapiOnboard() wizard.refreshModels() @@ -818,13 +861,19 @@ func (m model) handleProviderWizardKey(msg tea.KeyMsg) (model, tea.Cmd) { // provider list instead of destroying the overlay — and with it the user's // cursor and status — from a step deep in the flow. An in-flight OAuth login // is abandoned the same way; bumping the attempt id makes its late result - // stale so applyProviderWizardOAuth drops it. + // stale so applyProviderWizardOAuth drops it, and cancelDeviceLogin stops a + // device-code poll actually running in the background — otherwise + // completing the login later in the browser still silently saves the + // credential even though the wizard message is discarded. if m.providerWizard.manage && !m.providerWizard.managerStep() && keyIs(msg, tea.KeyEsc) { wizard := m.providerWizard if wizard.oauthPending { wizard.oauthPending = false - wizard.oauthAttemptID++ + // Burn the ID so a result from the attempt being abandoned here + // cannot match once oauthPending goes true again. + wizard.oauthAttemptID = nextOAuthAttemptID() } + wizard.cancelDeviceLogin() wizard.oauthDevice = false wizard.deviceUserCode = "" wizard.deviceVerificationURI = "" @@ -838,9 +887,12 @@ func (m model) handleProviderWizardKey(msg tea.KeyMsg) (model, tea.Cmd) { return m, nil } // While a browser/device OAuth login is in flight, ignore input except Esc, - // which abandons the wizard (the background flow times out and is dropped). + // which abandons the wizard. cancelDeviceLogin stops a device-code poll + // actually running in the background (the browser OAuth flow still relies + // on its own timeout, since it does not persist anything mid-flight). if m.providerWizard.oauthPending { if keyIs(msg, tea.KeyEsc) { + m.providerWizard.cancelDeviceLogin() m.providerWizard = nil } return m, nil @@ -868,10 +920,16 @@ func (m model) handleProviderWizardKey(msg tea.KeyMsg) (model, tea.Cmd) { return m, nil } // On the OAuth provider list, "d" forces device-code login for a device-capable - // provider (xAI) — useful on a desktop when the browser flow won't work. + // provider (xAI) — useful on a desktop when the browser flow won't work. A + // device-ONLY provider (Kimi Code has no loopback/authorize endpoint at + // all) has no browser flow to fall back to, so a plain Enter must also go + // straight to device login: the generic Enter/advanceProviderWizard path + // further below assumes a loopback flow exists and would otherwise hang or + // error against an endpoint that genuinely does not exist. if m.providerWizard.step == providerWizardStepProvider && m.providerWizard.oauthMode && - (keyText(msg) == "d" || keyText(msg) == "D") && - m.providerWizard.currentProvider().OAuthDeviceFlow { + m.providerWizard.currentProvider().OAuthDeviceFlow && + (keyText(msg) == "d" || keyText(msg) == "D" || + (keyIs(msg, tea.KeyEnter) && m.providerWizard.currentProvider().OAuthDeviceOnly)) { return m.startProviderDeviceLogin() } if m.providerWizard.step == providerWizardStepProvider { @@ -1268,7 +1326,12 @@ func (m model) applyProviderWizard() (model, tea.Cmd) { if !preserveExistingCredentialReference { profile = config.SecureProviderProfile(profile, m.userConfigPath) } - if _, err := config.UpsertProvider(m.userConfigPath, profile, true); err != nil { + // Persist without RuntimeHeaders-derived X-Msh-* keys. Those are + // re-attached by applyCatalogDescriptor at resolve time; writing them + // would store hostname + device UUID in config.json for no benefit, and + // would disagree with `zero auth kimi` (EnsureCatalogProvider). + toPersist := stripRuntimeIdentityHeaders(profile) + if _, err := config.UpsertProvider(m.userConfigPath, toPersist, true); err != nil { wizard.err = redaction.RedactString(err.Error(), redaction.Options{ExtraSecretValues: []string{secret, profile.APIKey}}) return m, nil // nothing committed to live state yet } @@ -1354,6 +1417,27 @@ func providerWizardRuntimeProfile(profile config.ProviderProfile) config.Provide return runtimeProfile } +// stripRuntimeIdentityHeaders returns a copy of profile without X-Msh-* headers +// so vendor identity is never written to config.json. In-memory / discovery +// profiles may still carry them; resolve-time applyCatalogDescriptor re-adds +// the fresh RuntimeHeaders values for canonical Kimi endpoints. +func stripRuntimeIdentityHeaders(profile config.ProviderProfile) config.ProviderProfile { + if len(profile.CustomHeaders) == 0 { + return profile + } + out := profile + out.CustomHeaders = maps.Clone(profile.CustomHeaders) + for key := range out.CustomHeaders { + if providercatalog.IsRuntimeIdentityHeader(key) { + delete(out.CustomHeaders, key) + } + } + if len(out.CustomHeaders) == 0 { + out.CustomHeaders = nil + } + return out +} + func (m model) providerWizardOverlay(width int) string { if m.providerWizard == nil { return "" @@ -1463,6 +1547,9 @@ func (wizard *providerWizardState) footerText() string { case providerWizardStepMethod: return "↑/↓ move Enter/→ continue Esc close" case providerWizardStepProvider: + if wizard.oauthMode && wizard.currentProvider().OAuthDeviceOnly { + return "↑/↓ move Enter device code ← back Esc close" + } if wizard.oauthMode && wizard.currentProvider().OAuthDeviceFlow { return "↑/↓ move Enter sign in d device code ← back Esc close" } @@ -2083,12 +2170,25 @@ func providerWizardProfile(provider providercatalog.Descriptor, model string, ap APIFormat: providerWizardAPIFormat(provider), Model: firstProviderDisplayValue(model, provider.DefaultModel), } - // Catalog custom headers (e.g. aimlapi.com's partner attribution) belong to the - // catalog endpoint: the resolver only applies them when the base URL is the - // default, so a profile built against a staging/proxy/custom override must not - // bake them in either — otherwise attribution leaks to an arbitrary host. + // Catalog custom headers (e.g. aimlapi.com's partner attribution, or Kimi + // Code's X-Msh-* vendor-identity headers) belong to the catalog endpoint: + // the resolver only applies them when the base URL is the default, so a + // profile built against a staging/proxy/custom override must not bake + // them in either — otherwise attribution leaks to an arbitrary host. if sameProviderBaseURL(resolvedBaseURL, provider.DefaultBaseURL) { - profile.CustomHeaders = maps.Clone(provider.CustomHeaders) + // `provider` here normally comes from a listing call (OAuthProviders()/ + // All()), which deliberately leaves RuntimeHeaders-backed CustomHeaders + // (Kimi's device-identity headers) unset so merely browsing providers + // doesn't mint Kimi's on-disk device id. Re-resolve through Get(), which + // runs RuntimeHeaders the same way profile-resolve time does (see + // cloneDescriptor / applyCatalogDescriptor), so the wizard's first + // authenticated /models call and the profile it activates immediately + // afterward carry those headers without requiring a restart. + customHeaders := provider.CustomHeaders + if runtimeDescriptor, ok := providercatalog.Get(provider.ID); ok { + customHeaders = runtimeDescriptor.CustomHeaders + } + profile.CustomHeaders = maps.Clone(customHeaders) if providerWizardIsAimlapi(provider) { profile.CustomHeaders = aimlapi.WithResolvedPartnerHeader(profile.CustomHeaders) } diff --git a/internal/tui/provider_wizard_discovery.go b/internal/tui/provider_wizard_discovery.go index b0dcada69..688c443e4 100644 --- a/internal/tui/provider_wizard_discovery.go +++ b/internal/tui/provider_wizard_discovery.go @@ -43,8 +43,12 @@ func (m model) advanceProviderWizard() (model, tea.Cmd) { if m.providerWizard.step == providerWizardStepProvider && m.providerWizard.oauthMode && m.providerWizard.currentProvider().OAuth { provider := m.providerWizard.currentProvider() // Headless/SSH boxes can't open a browser — use device code there by - // default (the user can also force it with "d" from the list). - if provider.OAuthDeviceFlow && oauthPreferDeviceFlow() { + // default (the user can also force it with "d" from the list). A + // device-ONLY provider (Kimi Code has no loopback/authorize endpoint at + // all) must also go straight to device login here: this is the mouse + // double-click activation path, which bypasses the keyboard Enter + // handler's OAuthDeviceOnly check further up in provider_wizard.go. + if provider.OAuthDeviceFlow && (provider.OAuthDeviceOnly || oauthPreferDeviceFlow()) { return m.startProviderDeviceLogin() } attemptID := m.providerWizard.beginOAuthAttempt(false) diff --git a/internal/tui/provider_wizard_oauth_test.go b/internal/tui/provider_wizard_oauth_test.go index bd4ac4d14..a3f7c0ced 100644 --- a/internal/tui/provider_wizard_oauth_test.go +++ b/internal/tui/provider_wizard_oauth_test.go @@ -1,6 +1,7 @@ package tui import ( + "context" "encoding/json" "errors" "os" @@ -8,10 +9,23 @@ import ( "strings" "testing" + tea "charm.land/bubbletea/v2" + "github.com/Gitlawb/zero/internal/config" "github.com/Gitlawb/zero/internal/providercatalog" ) +// isolateKimiDeviceIDStorage redirects os.UserConfigDir so providerWizardProfile +// / providercatalog.Get("kimi-code") never write kimi-device-id under the real +// user config root. DeviceID is path-keyed, so setting these env vars is enough. +func isolateKimiDeviceIDStorage(t *testing.T) { + t.Helper() + root := t.TempDir() + t.Setenv("XDG_CONFIG_HOME", root) + t.Setenv("APPDATA", root) + t.Setenv("HOME", root) +} + // wizardModelAt builds a model whose provider wizard is at step with providerID // selected. func wizardModelAt(t *testing.T, providerID string, step providerWizardStep) model { @@ -87,6 +101,99 @@ func beginTestOAuthAttempt(wizard *providerWizardState, device bool) (string, in return providerID, wizard.beginOAuthAttempt(device) } +// TestProviderWizardManagerEscCancelsDeviceLoginPoll is the manager-mode half +// of TestProviderWizardEscCancelsDeviceLoginPoll. handleProviderWizardKey has +// a separate Esc branch for manage mode that cancels the poll but keeps the +// wizard open, so the non-manager test above never exercised it. The poll must +// still be canceled: leaving it running means finishing the login in the +// browser later silently saves a credential the user backed out of. +func TestProviderWizardManagerEscCancelsDeviceLoginPoll(t *testing.T) { + m := managerTestModel(t) + m.providerWizard = m.newProviderWizard() + m.providerWizard.manage = true + m.providerWizard.selectedMethod = 0 + next, _ := m.advanceProviderWizard() + m = selectWizardOAuthProvider(t, next, "xai") + m.providerWizard.manage = true + if m.providerWizard.managerStep() { + t.Fatal("setup: expected an OAuth step, not a manager list/edit step") + } + providerID, attemptID := beginTestOAuthAttempt(m.providerWizard, true) + + out, cmd := m.applyProviderWizardDeviceCode(providerWizardDeviceCodeMsg{ + providerID: providerID, attemptID: attemptID, userCode: "ABCD-1234", verifyURL: "https://x.ai/device", + }) + if cmd == nil { + t.Fatal("device-code msg should start the poll command") + } + if out.providerWizard.deviceLoginCancel == nil { + t.Fatal("starting the poll should store a cancel func on the wizard") + } + + escaped, _ := out.handleProviderWizardKey(testKey(tea.KeyEsc)) + // Manager mode keeps the wizard; only the login is abandoned. + if escaped.providerWizard == nil { + t.Fatal("Esc in manager mode should keep the wizard open") + } + if escaped.providerWizard.oauthPending { + t.Fatal("Esc should clear the pending login") + } + if escaped.providerWizard.deviceLoginCancel != nil { + t.Fatal("Esc should clear the stored cancel func") + } + + raw := cmd() + msg, ok := raw.(providerWizardOAuthMsg) + if !ok { + t.Fatalf("poll command returned %T, want providerWizardOAuthMsg", raw) + } + if !errors.Is(msg.err, context.Canceled) { + t.Fatalf("poll error = %v, want context.Canceled (manager-mode Esc should have canceled the background poll)", msg.err) + } +} + +// TestProviderWizardAttemptIDsSurviveWizardClose pins that an attempt ID from +// a wizard the user abandoned cannot be reissued by the next one. The wizard +// is discarded wholesale on Esc, so a counter living on it restarts, and a +// late result from the canceled attempt would then satisfy oauthResultMatches +// for a brand new attempt against the same provider, persisting a credential +// the second wizard never asked for. +func TestProviderWizardAttemptIDsSurviveWizardClose(t *testing.T) { + m := mouseTestModel() + m.providerWizard = m.newProviderWizard() + m.providerWizard.selectedMethod = 0 + next, _ := m.advanceProviderWizard() + m = selectWizardOAuthProvider(t, next, "xai") + + firstProvider, firstAttempt := beginTestOAuthAttempt(m.providerWizard, true) + if firstAttempt == 0 { + t.Fatal("first attempt should get a non-zero id") + } + + // Abandon the wizard the way Esc does, then open a fresh one and start the + // same provider again. + m.providerWizard = nil + m.providerWizard = m.newProviderWizard() + m.providerWizard.selectedMethod = 0 + next, _ = m.advanceProviderWizard() + m = selectWizardOAuthProvider(t, next, "xai") + secondProvider, secondAttempt := beginTestOAuthAttempt(m.providerWizard, true) + + if firstProvider != secondProvider { + t.Fatalf("setup: expected the same provider both times, got %q then %q", firstProvider, secondProvider) + } + if secondAttempt == firstAttempt { + t.Fatalf("a reopened wizard reused attempt id %d; a late result from the abandoned attempt would match the new one", firstAttempt) + } + // The stale attempt must not be accepted by the live wizard. + if m.providerWizard.oauthResultMatches(firstProvider, firstAttempt) { + t.Fatal("a result from the abandoned attempt was accepted by the reopened wizard") + } + if !m.providerWizard.oauthResultMatches(secondProvider, secondAttempt) { + t.Fatal("the live attempt should be accepted") + } +} + func TestProviderWizardDeviceShortcutStartsDeviceFlow(t *testing.T) { m := mouseTestModel() m.providerWizard = m.newProviderWizard() @@ -106,6 +213,72 @@ func TestProviderWizardDeviceShortcutStartsDeviceFlow(t *testing.T) { } } +// TestProviderWizardEnterStartsDeviceFlowForDeviceOnlyProvider pins the fix +// for a device-only provider (Kimi Code has no loopback/authorize endpoint at +// all): the generic Enter path assumes a browser flow exists and would +// otherwise hang or error, so Enter must behave exactly like the "d" shortcut +// for a provider with OAuthDeviceOnly set. +func TestProviderWizardEnterStartsDeviceFlowForDeviceOnlyProvider(t *testing.T) { + // oauthPreferDeviceFlow() already defaults to device flow on a headless + // box (no DISPLAY/WAYLAND_DISPLAY, an SSH session, or ZERO_OAUTH_DEVICE + // set) — exactly the environment this test suite runs in — which would + // mask the bug this test exists to catch. Force a "normal desktop with a + // browser available" environment so Enter actually exercises the + // otherwise-loopback-preferring path. + t.Setenv("ZERO_OAUTH_DEVICE", "") + t.Setenv("SSH_CONNECTION", "") + t.Setenv("SSH_TTY", "") + t.Setenv("DISPLAY", ":0") + t.Setenv("WAYLAND_DISPLAY", "") + + m := mouseTestModel() + m.providerWizard = m.newProviderWizard() + m.providerWizard.selectedMethod = 0 + next, _ := m.advanceProviderWizard() // → OAuth list + m = selectWizardOAuthProvider(t, next, "kimi-code") + if !m.providerWizard.currentProvider().OAuthDeviceOnly { + t.Fatal("test fixture assumes kimi-code is OAuthDeviceOnly") + } + + out, cmd := m.handleProviderWizardKey(tea.KeyPressMsg(tea.Key{Code: tea.KeyEnter})) + if !out.providerWizard.oauthPending || !out.providerWizard.oauthDevice { + t.Fatalf("Enter on a device-only provider should start device login (pending=%v device=%v)", out.providerWizard.oauthPending, out.providerWizard.oauthDevice) + } + if cmd == nil { + t.Fatal("Enter on a device-only provider should return the device-prepare command") + } +} + +// TestProviderWizardMouseAdvanceStartsDeviceFlowForDeviceOnlyProvider covers +// double-click activation: mouse.go routes to advanceProviderWizard, which +// bypasses the keyboard handler's OAuthDeviceOnly check. On a desktop (where +// oauthPreferDeviceFlow is false), advance must still start device login for +// a device-only provider so the verification URL/user code are not discarded. +func TestProviderWizardMouseAdvanceStartsDeviceFlowForDeviceOnlyProvider(t *testing.T) { + t.Setenv("ZERO_OAUTH_DEVICE", "") + t.Setenv("SSH_CONNECTION", "") + t.Setenv("SSH_TTY", "") + t.Setenv("DISPLAY", ":0") + t.Setenv("WAYLAND_DISPLAY", "") + + m := mouseTestModel() + m.providerWizard = m.newProviderWizard() + m.providerWizard.selectedMethod = 0 + next, _ := m.advanceProviderWizard() // → OAuth list + m = selectWizardOAuthProvider(t, next, "kimi-code") + if !m.providerWizard.currentProvider().OAuthDeviceOnly { + t.Fatal("test fixture assumes kimi-code is OAuthDeviceOnly") + } + + out, cmd := m.advanceProviderWizard() + if !out.providerWizard.oauthPending || !out.providerWizard.oauthDevice { + t.Fatalf("mouse advance on a device-only provider should start device login (pending=%v device=%v)", out.providerWizard.oauthPending, out.providerWizard.oauthDevice) + } + if cmd == nil { + t.Fatal("mouse advance on a device-only provider should return the device-prepare command") + } +} + func TestProviderWizardDeviceCodeMsgShowsCodeAndPolls(t *testing.T) { m := mouseTestModel() m.providerWizard = m.newProviderWizard() @@ -129,6 +302,100 @@ func TestProviderWizardDeviceCodeMsgShowsCodeAndPolls(t *testing.T) { } } +// TestProviderWizardEscCancelsDeviceLoginPoll regression-tests a bug where +// abandoning the wizard with Esc during a device-code poll (phase 2) never +// canceled the background context the poll command runs on. The wizard +// message was made stale so the UI wouldn't show a stray login, but the +// underlying poll kept running for up to 10 minutes: if the user then +// finished authorizing in their browser, it would still silently succeed and +// save a credential the user believed they had backed out of. Esc must +// actually cancel the context, not just discard the eventual result. +func TestProviderWizardEscCancelsDeviceLoginPoll(t *testing.T) { + // Isolate the oauth token store the poll command touches (see + // managerTestModel), even though the canceled-context path exercised here + // returns before any read/write reaches it. + t.Setenv("XDG_CONFIG_HOME", t.TempDir()) + t.Setenv("ZERO_OAUTH_TOKENS_PATH", filepath.Join(t.TempDir(), "oauth-tokens.json")) + + m := mouseTestModel() + m.providerWizard = m.newProviderWizard() + m.providerWizard.selectedMethod = 0 + next, _ := m.advanceProviderWizard() + m = selectWizardOAuthProvider(t, next, "xai") + providerID, attemptID := beginTestOAuthAttempt(m.providerWizard, true) + + out, cmd := m.applyProviderWizardDeviceCode(providerWizardDeviceCodeMsg{ + providerID: providerID, attemptID: attemptID, userCode: "ABCD-1234", verifyURL: "https://x.ai/device", + }) + if cmd == nil { + t.Fatal("device-code msg should start the poll command") + } + if out.providerWizard.deviceLoginCancel == nil { + t.Fatal("starting the poll should store a cancel func on the wizard") + } + + escaped, _ := out.handleProviderWizardKey(testKey(tea.KeyEsc)) + if escaped.providerWizard != nil { + t.Fatal("Esc while oauthPending should close the wizard") + } + + // The poll command captured the context created for this attempt; run it + // now (after Esc) and confirm CompleteDeviceLogin actually observed + // cancellation instead of running to completion in the background. + raw := cmd() + msg, ok := raw.(providerWizardOAuthMsg) + if !ok { + t.Fatalf("poll command returned %T, want providerWizardOAuthMsg", raw) + } + if !errors.Is(msg.err, context.Canceled) { + t.Fatalf("poll error = %v, want context.Canceled (Esc should have canceled the background poll)", msg.err) + } +} + +// TestModelQuitCancelsProviderWizardDeviceLoginPoll regression-tests a bug +// where model.quit() (the path every Ctrl+C-to-exit and "q"-to-exit +// eventually reaches) only reset the aimlapi.com sub-flow, never the +// provider wizard's device-code poll. Since the TUI runs on +// context.Background(), quitting via Ctrl+C left the poll running in the +// background: authorizing the abandoned login afterward could still save a +// credential the user had just quit to avoid. +func TestModelQuitCancelsProviderWizardDeviceLoginPoll(t *testing.T) { + t.Setenv("XDG_CONFIG_HOME", t.TempDir()) + t.Setenv("ZERO_OAUTH_TOKENS_PATH", filepath.Join(t.TempDir(), "oauth-tokens.json")) + + m := mouseTestModel() + m.providerWizard = m.newProviderWizard() + m.providerWizard.selectedMethod = 0 + next, _ := m.advanceProviderWizard() + m = selectWizardOAuthProvider(t, next, "xai") + providerID, attemptID := beginTestOAuthAttempt(m.providerWizard, true) + + out, cmd := m.applyProviderWizardDeviceCode(providerWizardDeviceCodeMsg{ + providerID: providerID, attemptID: attemptID, userCode: "ABCD-1234", verifyURL: "https://x.ai/device", + }) + if cmd == nil { + t.Fatal("device-code msg should start the poll command") + } + if out.providerWizard.deviceLoginCancel == nil { + t.Fatal("starting the poll should store a cancel func on the wizard") + } + + quit, _ := out.quit() + quitModel := quit.(model) + if quitModel.providerWizard != nil && quitModel.providerWizard.deviceLoginCancel != nil { + t.Fatal("quit should cancel the in-flight device-code poll") + } + + raw := cmd() + msg, ok := raw.(providerWizardOAuthMsg) + if !ok { + t.Fatalf("poll command returned %T, want providerWizardOAuthMsg", raw) + } + if !errors.Is(msg.err, context.Canceled) { + t.Fatalf("poll error = %v, want context.Canceled (quit should have canceled the background poll)", msg.err) + } +} + // A failed OAuth attempt leaves the wizard on the provider list; the error must be // rendered there (not just on the credential step) so a click isn't a silent // no-op, and Hugging Face gets an actionable client_id hint. @@ -449,3 +716,109 @@ func TestAppendOAuthLoginProfileAddsOnceAndRespectsRenames(t *testing.T) { t.Fatalf("unknown provider must not append, got %+v", got) } } + +// TestProviderWizardProfileAppliesKimiRuntimeHeaders regression-tests a bug +// where the /provider wizard built its profile straight from the descriptor +// OAuthProviders() returns. That listing call deliberately omits +// RuntimeHeaders-backed CustomHeaders (Kimi's X-Msh-* vendor-identity +// headers) so merely browsing providers doesn't mint Kimi's on-disk device +// id — but that meant the wizard's first authenticated /models call and the +// profile it activated immediately after finishing were missing those +// headers until zero was restarted. providerWizardProfile must re-resolve +// through providercatalog.Get (which does run RuntimeHeaders) instead of +// using the listing descriptor's CustomHeaders directly. +func TestProviderWizardProfileAppliesKimiRuntimeHeaders(t *testing.T) { + isolateKimiDeviceIDStorage(t) + m := mouseTestModel() + m.providerWizard = m.newProviderWizard() + m.providerWizard.selectedMethod = 0 + next, _ := m.advanceProviderWizard() // → OAuth list + m = selectWizardOAuthProvider(t, next, "kimi-code") + provider := m.providerWizard.currentProvider() + if len(provider.CustomHeaders) != 0 { + t.Fatalf("test fixture assumes the OAuth listing omits runtime headers, got %#v", provider.CustomHeaders) + } + + profile := providerWizardProfile(provider, provider.DefaultModel, "", "", "") + if profile.CustomHeaders["X-Msh-Platform"] != "kimi_code_cli" { + t.Fatalf("providerWizardProfile CustomHeaders[X-Msh-Platform] = %q, want kimi_code_cli", profile.CustomHeaders["X-Msh-Platform"]) + } + if profile.CustomHeaders["X-Msh-Device-Id"] == "" { + t.Fatal("providerWizardProfile should carry Kimi's device id header") + } +} + +// TestProviderWizardProfileOmitsKimiHeadersOnCustomEndpoint is the +// custom-endpoint half of the test above, which only covers the canonical +// Kimi base URL. X-Msh-* carries a persistent device id and the machine +// hostname, so baking it into a profile aimed at an arbitrary staging or proxy +// host would send both somewhere Kimi does not own. +func TestProviderWizardProfileOmitsKimiHeadersOnCustomEndpoint(t *testing.T) { + isolateKimiDeviceIDStorage(t) + m := mouseTestModel() + m.providerWizard = m.newProviderWizard() + m.providerWizard.selectedMethod = 0 + next, _ := m.advanceProviderWizard() + m = selectWizardOAuthProvider(t, next, "kimi-code") + + provider := m.providerWizard.currentProvider() + custom := "https://kimi-proxy.example.test/v1" + if sameProviderBaseURL(custom, provider.DefaultBaseURL) { + t.Fatalf("setup: %q must differ from the canonical endpoint", custom) + } + + profile := providerWizardProfile(provider, provider.DefaultModel, "", custom, "") + + for key := range profile.CustomHeaders { + if strings.HasPrefix(strings.ToLower(key), "x-msh-") { + t.Fatalf("device identity header %q leaked to a custom endpoint: %#v", key, profile.CustomHeaders) + } + } + if profile.BaseURL != custom { + t.Fatalf("profile BaseURL = %q, want the custom endpoint %q", profile.BaseURL, custom) + } +} + +// TestStripRuntimeIdentityHeadersDropsXMshOnly ensures persist strips Kimi +// identity headers while keeping unrelated custom headers (and matching +// `zero auth kimi`, which never writes X-Msh-* into config.json). +func TestStripRuntimeIdentityHeadersDropsXMshOnly(t *testing.T) { + profile := config.ProviderProfile{ + Name: "kimi-code", + CatalogID: "kimi-code", + CustomHeaders: map[string]string{ + "X-Msh-Device-Id": "device-uuid", + "X-Msh-Device-Name": "hostname", + "X-Msh-Platform": "kimi_code_cli", + "X-User-Agent": "keep-me", + }, + } + got := stripRuntimeIdentityHeaders(profile) + for key := range got.CustomHeaders { + if strings.HasPrefix(strings.ToLower(key), "x-msh-") { + t.Fatalf("persist copy still has identity header: %#v", got.CustomHeaders) + } + } + if got.CustomHeaders["X-User-Agent"] != "keep-me" { + t.Fatalf("user header lost: %#v", got.CustomHeaders) + } + // Original in-memory profile must keep runtime headers for discovery. + if profile.CustomHeaders["X-Msh-Device-Id"] != "device-uuid" { + t.Fatalf("strip mutated the in-memory profile: %#v", profile.CustomHeaders) + } +} + +func TestModelQuitCancelsSetupDeviceLogin(t *testing.T) { + m := mouseTestModel() + canceled := false + m.setup.deviceLoginCancel = func() { + canceled = true + } + _, cmd := m.quit() + if !canceled { + t.Fatal("m.quit() did not invoke setup.deviceLoginCancel") + } + if cmd == nil { + t.Fatal("m.quit() should return tea.Quit cmd") + } +}