From 27b69d78e3e7509502965a2cb2a722e732025b11 Mon Sep 17 00:00:00 2001 From: adminturneddevops Date: Sun, 6 Sep 2026 18:36:47 -0400 Subject: [PATCH 1/4] credential implementation --- PLAN-CRED.md | 218 ++++++++ PLAN.md | 60 +- README.md | 30 +- cmd/abox-guest/main.go | 162 +++++- cmd/abox/creds.go | 176 ++++++ cmd/abox/creds_test.go | 128 +++++ cmd/abox/main.go | 76 ++- internal/agent/agent.go | 13 +- internal/agent/agent_test.go | 79 +++ internal/config/config.go | 325 +++++++++-- internal/config/config_test.go | 246 ++++++++- internal/credentials/credentials.go | 22 + internal/credentials/credentials_test.go | 43 ++ internal/credsource/aws.go | 221 ++++++++ internal/credsource/azure.go | 151 +++++ internal/credsource/cloud_test.go | 608 ++++++++++++++++++++ internal/credsource/credsource.go | 116 ++++ internal/credsource/credsource_test.go | 416 ++++++++++++++ internal/credsource/env.go | 28 + internal/credsource/keychain.go | 131 +++++ internal/credsource/resolve.go | 45 ++ internal/credsource/save.go | 47 ++ internal/credsource/vault.go | 130 +++++ internal/guest/brokerclient/client.go | 515 +++++++++++++++++ internal/guest/brokerclient/client_test.go | 540 ++++++++++++++++++ internal/guest/egress/egress.go | 9 +- internal/guest/egress/egress_test.go | 8 +- internal/llmbroker/broker.go | 356 ++++++++++++ internal/llmbroker/broker_test.go | 577 +++++++++++++++++++ internal/mcpauth/oauth.go | 268 +++++++-- internal/mcpauth/oauth_test.go | 210 ++++++- internal/provider/provider.go | 60 +- internal/provider/provider_test.go | 41 +- internal/runtime/runtime.go | 614 ++++++++++++++++++--- internal/runtime/runtime_push_test.go | 445 +++++++++++++++ internal/runtime/runtime_test.go | 4 +- internal/runtime/runtime_turn_test.go | 137 +++++ internal/session/scrub.go | 156 ++++++ internal/session/scrub_test.go | 234 ++++++++ internal/session/session.go | 30 +- internal/session/session_test.go | 46 +- internal/tui/commands.go | 70 ++- internal/tui/commands_test.go | 128 +++++ internal/tui/tui.go | 130 ++++- pkg/abox/abox.go | 119 +++- pkg/abox/abox_test.go | 65 +++ pkg/abox/doc.go | 6 +- protocol/protocol.go | 105 +++- protocol/protocol_test.go | 127 +++++ 49 files changed, 8085 insertions(+), 386 deletions(-) create mode 100644 PLAN-CRED.md create mode 100644 cmd/abox/creds.go create mode 100644 cmd/abox/creds_test.go create mode 100644 internal/credsource/aws.go create mode 100644 internal/credsource/azure.go create mode 100644 internal/credsource/cloud_test.go create mode 100644 internal/credsource/credsource.go create mode 100644 internal/credsource/credsource_test.go create mode 100644 internal/credsource/env.go create mode 100644 internal/credsource/keychain.go create mode 100644 internal/credsource/resolve.go create mode 100644 internal/credsource/save.go create mode 100644 internal/credsource/vault.go create mode 100644 internal/guest/brokerclient/client.go create mode 100644 internal/guest/brokerclient/client_test.go create mode 100644 internal/llmbroker/broker.go create mode 100644 internal/llmbroker/broker_test.go create mode 100644 internal/runtime/runtime_push_test.go create mode 100644 internal/session/scrub.go create mode 100644 internal/session/scrub_test.go create mode 100644 internal/tui/commands_test.go diff --git a/PLAN-CRED.md b/PLAN-CRED.md new file mode 100644 index 0000000..f3b5d73 --- /dev/null +++ b/PLAN-CRED.md @@ -0,0 +1,218 @@ +# ABox Credential Overhaul — Sources, Secrets-at-Rest Removal, Host Provider Broker + +## Context + +At the start of this overhaul, ABox violated the intended host-only LLM credential boundary and stored credential values in session configuration: + +- All secrets lived plaintext in `~/.abox/credentials.env` (internal/credentials/credentials.go). +- `cfg.SecretsFromEnv()` collected **every** provider key + **every** MCP token; every session start dumped the full map into `sessions//guest-config.json` and `config.raw`. Session directories retained plaintext copies, and SDK `SetModel` re-sent all secrets. +- The untrusted guest read the config disk and exported everything into its environment, exposing every key rather than only the selected model credential. +- mcpauth persisted an unused `_REFRESH` token without the client metadata required to refresh it. + +User research (Sept 2026) recommends: host-side credential-source abstraction, resolve only the selected model's credential, never persist resolved values, remove secrets from guest config, and move provider transport behind a host broker. **User approved all three phases**, keychain via `security(1)` subprocess (no cgo — `abox` stays plain `go build`), Vault via `VAULT_ADDR`/`VAULT_TOKEN` KV v2. + +**Approved source set (user decision, Sept 2026):** `env`, `keychain` (macOS), `vault` (HashiCorp Vault KV v2), `azure` (Azure Key Vault), `aws` (AWS Secrets Manager). All cloud stores via stdlib HTTP or a CLI subprocess — no HashiCorp/Azure/AWS SDKs, no cgo. + +**Deferred (documented, not built):** Kubernetes sources, workload identity federation (Azure managed identity, AWS IAM roles — this milestone uses static SP/env credentials only), MCP traffic brokering (guest MCP client keeps its TSI path this milestone; PLAN.md §14.4 is the follow-up that removes MCP tokens from the guest), agentgateway LLM routing (stays "direct base_url" exactly as today — flagged, never claimed enforced, per PLAN.md §14.3). + +## Global decisions + +1. New host-only package `internal/credsource`; `internal/credentials` stays as the credentials.env file store (env-source backend + fallback writer). Import direction: `credsource` may import `config`; `config` never imports `credsource`. +2. Cloud secret stores via stdlib HTTP or CLI subprocess only — no HashiCorp/Azure/AWS SDK dependency, no cgo: Vault = one `GET /v1//data/` with `X-Vault-Token`; Azure Key Vault = stdlib OAuth2 client-credentials token POST plus `GET {vault}/secrets/{name}` Data Plane REST; AWS Secrets Manager = in-package SigV4 over `GetSecretValue` REST with static env credentials. +3. `_REFRESH` write: **delete it** (oauth.go:83). Future work note: persist client_id + refresh token in keychain, implement the refresh grant. +4. One protocol bump, `protocol.Version` 2 → 3, at Phase 3. Phase 2 needs no protocol change: v2 guests already implement `set_model`/`set_mcp_tokens` (cmd/abox-guest/main.go:238-259) and tolerate secretless boot config. Protocol-1 guests cannot run agent sessions from the rewritten secretless config; resume is rejected explicitly rather than reporting a misleading ready state. +5. Phase 3 is **version-gated, not a config mode**: proto ≥ 3 guest binaries have no direct provider transport (broker is the only LLM path); proto == 2 guests get the legacy post-hello secret push + stderr deprecation warning. No `model_transport` knob. +6. Phase 3 prerequisite: before the reader-goroutine demux, the host could not receive guest-initiated frames because `Sandbox.Call` read the connection inline and dropped frames outside its awaited ID. Task 3.1 supplied that demux before broker methods were enabled. +7. Guest context is up to 2 MiB (internal/agent/agent.go:27) but `MaxFrameBytes` is 1 MiB (protocol/protocol.go:14) → broker requests are chunked (mirror of `archive_chunk`), 256 KiB per chunk. + +--- + +## Phase 1 — Host credential Source abstraction + +### 1.1 `internal/credsource` core +New: `internal/credsource/credsource.go` (+test). +```go +type Reference struct{ Source, Name, Field, Version string } +type Value struct { Bytes []byte; Version string; ExpiresAt time.Time; LeaseID string } +func (v *Value) Zero() // best-effort overwrite +func (v Value) String() string // "credsource.Value(redacted)" — defeats accidental %v logging +type Source interface { Resolve(context.Context, Reference) (Value, error); Close() error } +type Resolver struct{ ... } // registers env, keychain (darwin), vault +var ErrNotFound, ErrLocked error +``` +Errors mention only Source/Name, never values. + +### 1.2 env source +New: `internal/credsource/env.go` (+test). Order preserves current semantics: `os.Getenv(ref.Name)` first, then `credentials.Load()` map (reuse credentials.go:22). Remove now-unneeded `credentials.ApplyToEnv()` startup calls at cmd/abox/main.go:61 and pkg/abox/abox.go:59 (keep the one in `mcpLogin`, main.go:296, until 1.7). + +### 1.3 keychain source (security(1), no cgo) +New: `internal/credsource/keychain.go` (+test). Service `abox`, account = `Reference.Name`. +- Get: `/usr/bin/security find-generic-password -s abox -a -w` (password only on stdout; secret never in argv). ASCII-only storage documented (`-w` prints hex for non-ASCII). +- Set: run `/usr/bin/security -i`, write to **stdin**: `add-generic-password -U -s abox -a -X -j "managed by abox"`. `-i` = command-from-stdin mode; `-X` = hex password (avoids argv leak via ps and interactive quoting issues); `-U` upserts. Never use `-w value` (argv leak) or bare `-w` (tty prompt corrupts TUI). +- Delete: `security delete-generic-password -s abox -a `. +- Errors: exit 44 → `ErrNotFound`; stderr `User interaction is not allowed` (locked/headless) → `ErrLocked` with unlock hint; else wrapped stderr (stderr never contains the secret). +- Subprocess behind package var `runSecurity` so tests fake it. `Available()` = darwin + `/usr/bin/security` exists. Export `SetKeychain`/`DeleteKeychain` for TUI/migration. + +### 1.4 vault source +New: `internal/credsource/vault.go` (+test). `VAULT_ADDR`; token from `VAULT_TOKEN` then `~/.vault-token`; optional `VAULT_NAMESPACE` header. `Reference.Name` = KV v2 logical path (`secret/abox/anthropic`), source inserts `/data/`; `?version=` from `Reference.Version`; `Reference.Field` selects key in `data.data` (default `value`); `data.metadata.version` → `Value.Version`. 404 → ErrNotFound, 403 → clear permission error. `http.Client{Timeout: 15s}`. + +### 1.4a azure source (Azure Key Vault) +New: `internal/credsource/azure.go` (+test). `Reference.Name` = Key Vault secret identifier URI (`https://.vault.azure.net/secrets/`); specific version via `Reference.Version` (`/secrets//`), else latest. `Reference.Field` unused (Key Vault secrets are single-value). +- Auth precedence: (1) service principal env — `AZURE_CLIENT_ID` + `AZURE_TENANT_ID` + `AZURE_CLIENT_SECRET` → stdlib OAuth2 client-credentials `POST {AZURE_AUTHORITY_HOST|https://login.microsoftonline.com}//oauth2/v2.0/token` with `scope=https://vault.azure.net/.default`; (2) fallback `az account get-access-token --resource https://vault.azure.net` subprocess (requires `az login`; mirrors keychain's subprocess pattern; behind package var `runAz` for test faking). +- Fetch: `GET {vaultUri}/secrets/{name}[/{version]}?api-version=7.5` with `Authorization: Bearer `; response `{"value": ...}` → `Value.Bytes`. 404 → ErrNotFound, 403 → clear permission error mentioning Key Vault access policy/RBAC. +- Missing SP env **and** `az` unavailable → clear setup error (never mentions secret values). No Azure SDK, no cgo. + +### 1.4b aws source (AWS Secrets Manager) +New: `internal/credsource/aws.go` (+test). `Reference.Name` = secret ID (arbitrary AWS name, may contain `/`). Static env credentials only: `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, optional `AWS_SESSION_TOKEN`; region from `AWS_REGION` then `AWS_DEFAULT_REGION`. Missing creds/region → clear setup error. No AWS SDK, no cgo. +- Fetch: `POST https://secretsmanager..amazonaws.com/` with `X-Amz-Target: secretsmanager.GetSecretValue`, `Content-Type: application/x-amz-json-1.1`, body `{"SecretId": }`, signed with in-package SigV4 (service `secretsmanager`; `x-amz-security-token` included when a session token exists). +- Response: `SecretString` (or base64 `SecretBinary` decoded) → `Value.Bytes`. `Reference.Field` supported: when set, parse `SecretString` as JSON and select the key (cloud-manager convention of multi-key secrets); unset → whole `SecretString`. `ResourceNotFoundException` → ErrNotFound; `AccessDeniedException` → clear permission error. +- SigV4 signing is implemented in-package with `crypto/hmac` (canonical request, SHA256 payload hash, `Authorization: AWS4-HMAC-SHA256 ...`); token/secret never appear in error text or logs. + +### 1.5 config.yaml credential references + back-compat +Modify: internal/config/config.go, providers.go, config_test.go. +```yaml +models: + - name: claude-default + provider: anthropic + model: claude-sonnet-4-20250514 + credential: + source: keychain # env | keychain | vault | azure | aws + name: ANTHROPIC_API_KEY # env: var; keychain: account; vault: KV-v2 path; azure: secret URI; aws: secret ID + field: api_key # vault/aws only (vault default "value"; aws unset = whole SecretString) + version: "4" # vault/azure only (optional) + # credential_env: X # DEPRECATED alias == {source: env, name: X} +mcp_servers: + - name: github + url: https://... + credential: {source: keychain, name: ABOX_MCP_GITHUB_TOKEN} +``` +- `CredentialRef` struct in `config`; `Model.Credential *CredentialRef`, `MCPServer.Credential *CredentialRef`. +- `Model.CredentialReference()`: explicit ref, else `{env, CredentialEnv}`, else `{env, EnvName()}`. New `Model.EnvName()`: CredentialEnv, else canonical provider env from `DefaultProviders()`, else `ABOX_MODEL__KEY` — fills `protocol.GuestModel.CredentialEnv` (ToGuest, config.go:250) so Phase-1 wire format is unchanged. `MCPServer.CredentialReference()` reuses `TokenEnv` (config.go:376). +- Validate: reject both `credential` and `credential_env` set; source ∈ {env, keychain, vault, azure, aws}; env names pass `ValidEnvName`; `field` vault/aws only; `version` vault/azure only; azure `name` must be an `https://…vault.azure.net/secrets/…` (or other region suffix) URI. +- **Delete `SecretsFromEnv`** (config.go:281-300) + its test. Replace `Model.CredentialPresent` (config.go:396; sole caller tui.go:444) with resolver-backed presence check so keychain/vault keys don't render "missing". Presence is resolved **once when the picker opens** (cached per session), never in the render path — a `security` subprocess or Vault HTTP call per frame would freeze the TUI. +- `credsource.FromConfig(config.CredentialRef) Reference` glue. + +### 1.6 Resolve only what's selected; fix SDK SetModel +New: `internal/credsource/resolve.go` (+test). +```go +// Selected model's credential keyed by model.EnvName() + one token per enabled +// MCP server keyed by TokenEnv. Missing model credential -> error naming the +// reference. Missing MCP token -> skipped (guest MCP degrades gracefully). +func ResolveSelected(ctx, *Resolver, config.File, config.Model) (map[string]string, error) +``` +Uses `cfg.ResolvedMCPServers()` (config.go:307) — offline resolves no MCP tokens. Call sites: cmd/abox/main.go:124 and pkg/abox/abox.go:104 replace `cfg.SecretsFromEnv()`; pkg/abox/abox.go:246 `SetModel` resolves **only** the new model's credential. Zero intermediate `Value`s after copying (best-effort, documented). + +### 1.7 TUI keychain-by-default, migration, mcpauth +Modify: internal/tui/commands.go (:47, :56), tui.go (:335-389, :444), internal/mcpauth/oauth.go (:57-86), cmd/abox/main.go. +- `applyProviderKey`: try `SetKeychain`; on success upsert the model's `credential: {keychain, ...}` ref and `cfg.Save()` (config.go:354); status "key saved to macOS keychain (service abox)". On ErrLocked/unavailable, fall back to `credentials.Save` and replace any stale explicit cloud/keychain reference with `credential: {source: env, ...}`. `applyMCPKey` mirrors this selected-source update. +- `mcpauth.LoginNamed`: same keychain-preferred writer; **delete the `_REFRESH` write**. +- New CLI `abox creds migrate` (dispatched like `mcp`, main.go:35): move credentials.env entries to keychain (including MCP OAuth tokens; re-login is the fallback for expired ones), update matching config refs, drop `*_REFRESH` keys, rewrite credentials.env to a comment (kept, 0600). No silent startup migration — env source keeps working indefinitely. + +### Phase 1 verification +`go build ./...` && `CGO_ENABLED=0 GOOS=linux GOARCH=arm64 go build ./cmd/abox-guest` && `go test ./protocol ./internal/... ./pkg/...` && `golangci-lint run ./...`. +Tests in existing style (single-purpose funcs, `t.Setenv("HOME"/"ABOX_HOME", t.TempDir())`): env precedence; keychain fake `runSecurity` asserting no secret in argv, exit-44 → ErrNotFound, locked → ErrLocked; vault httptest asserting path/version/token header/field, `~/.vault-token` fallback; azure httptest asserting token POST form + secret GET api-version + version path + fake `runAz` fallback + 404/403 mapping; aws httptest asserting SigV4 Authorization header shape, X-Amz-Target, SecretId body, session-token header, JSON `field` selection, 404/403 mapping; config alias mapping + both-set rejection + YAML round-trip; ResolveSelected selectivity + offline; credentials.env 0600 mode assertion; oauth_test asserts no `*_REFRESH` persisted. + +--- + +## Phase 2 — Remove secrets at rest + +### 2.1 Stop writing secrets to guest-config.json / config.raw +Modify: session.go:175-199 (`WriteGuestConfig` drops secrets param, never sets `GuestConfig.Secrets`), runtime.go:55-75 (`Prepare` drops secrets param), call sites main.go:124 / abox.go:104, tests session_test.go:13, runtime_test.go. `protocol.GuestConfig.Secrets` stays in the struct, commented Deprecated (old images still parse; old hosts still work). + +### 2.2 Post-hello secret push +New: `Sandbox.PushSecrets(ctx, model, secrets)` — for protocol 2, `set_model` carries the selected model credential, then `set_mcp_tokens` carries MCP tokens; for protocol 3, `set_model` carries metadata only and only MCP tokens are pushed. It reuses SetModel/SetMCPTokens (runtime.go:352-361). Ordering: `runtime.Start` returns only after `waitHello` (runtime.go:183); callers push immediately after Start, **before** TransferArchive/first UserTurn (main.go:141-154, abox.go:124-138). Resume paths push too because `config.raw` is rewritten secretless. Protocol-1 agent startup and resume are rejected; only a diagnostic no-key probe may use an old image. + +### 2.3 Scrub existing session dirs +New: `internal/session/scrub.go` (+test). `ScrubSecrets(root)` — surgical, never deletes sessions: +1. Walk `sessions//` (pattern from `LatestForRepo`, session.go:94). +2. `guest-config.json`: unmarshal to `map[string]json.RawMessage` (preserves unknown fields); no `"secrets"` key → skip (idempotent); else delete key, tmp+rename write, 0600. +3. `config.raw`: chmod 0600 → write scrubbed JSON zero-padded to 1 MiB → chmod 0400. Extract the pad-write body of `writeConfigDisk` (runtime.go:84-95) into shared `session.WritePaddedConfig` and have runtime reuse it (runtime already imports session). +4. root.raw / transcript.json / console.log / session.json untouched. +Invoke once per startup (main.go `run()` and pkg/abox `open()` after `config.Load()`); stderr "abox: scrubbed plaintext secrets from N old session(s)" when N > 0. An incomplete scrub is returned as a startup error by both CLI and SDK rather than allowing a session to proceed with uncertain secret removal. + +### 2.4 Old-image compatibility (decision) +Protocol-2 images supported via the 2.2 push; `make image` required only for protocol < 2. Rationale: v2 guests already implement both push handlers and tolerate secretless boot config. + +### Phase 2 verification +Build/test/lint gates as Phase 1. session_test: no `"secrets"` key + 0600 mode. runtime_test: keep 0400 resume assertion (runtime_test.go:61); config.raw contains no secret marker. scrub_test: legacy fixture scrubbed, modes preserved, other fields byte-identical, second run no-op, root.raw untouched. net.Pipe fake-guest test (runtime_turn_test.go style): PushSecrets emits set_model then set_mcp_tokens before any user_turn; protocol-1 agent startup/resume is refused. Manual: `make build && make image`, run a session, `strings ~/.abox/sessions//config.raw | grep -c API_KEY` → 0. + +--- + +## Phase 3 — Host provider broker + +### 3.1 Sandbox reader-goroutine demux (prerequisite; land as its own PR) +Modify: runtime.go (Call :226, userTurnLocked :270, waitHello :190, Stop :383). After hello, a single `readLoop()` goroutine owns the conn and routes: response frames → per-call channel keyed by ID (subsumes inline loops incl. late-cancel logic tested at runtime_turn_test.go:63-256); `agent_event` → active turn's channel; guest-originated frames (ID prefix `g-`, Method set) → `GuestCallHandler`: +```go +type GuestCallHandler interface { + Handle(ctx context.Context, method string, params json.RawMessage, + notify func(method string, params any) error) (any, *protocol.Error) +} +``` +Unknown guest methods → typed error frame. **Each guest call is dispatched on its own goroutine** — never inline on readLoop (a `provider_send` Last=true otherwise blocks the loop for the whole provider stream, deadlocking `provider_cancel` and every other frame); `notify` and reply frames serialize through `writeMu`. All existing turn/cancel tests must stay green. + +### 3.2 Protocol types, Version = 3 +Modify: protocol/protocol.go (+test). Methods: `provider_open`, `provider_send`, `provider_cancel` (guest→host); `provider_event` (host→guest push). +```go +const Version = 3 +const MaxProviderChunk = 256 << 10 +const MaxProviderToolArgs = 512 << 10 +const MaxProviderStreams = 2 + +type ProviderOpenParams struct{ Model string; Rich bool } // configured alias ONLY — no URL/headers/credential names +type ProviderOpenResult struct{ StreamID string } +type ProviderSendParams struct{ StreamID string; Data []byte; Last bool } // chunks of marshaled ProviderRequest +type ProviderRequest struct{ Messages []ProviderMessage; Tools []ProviderToolSchema } +type ProviderCancelParams struct{ StreamID string } +type ProviderEventParams struct{ StreamID, Type, Text, ToolID, ToolName, ToolArgs string; Usage *UsageInfo; StopReason, Err string } +``` +`ProviderMessage`/`ProviderToolSchema` mirror `provider.Message`/`ToolSchema` (protocol imports no internal packages). Bounds host-side: ≤ MaxProviderStreams open, reassembled request ≤ 4 MiB, tool args truncated at MaxProviderToolArgs (error event beyond), 5-min idle stream timeout, cancel path on every stream. Broker methods accepted only post-hello from proto ≥ 3 guests. `SetModelParams.Secrets` documented unused for LLM at proto ≥ 3. + +### 3.3 Provider transport host-callable +Modify: internal/provider/provider.go. Inject credential + client: `Stream(ctx, model, key string, client *http.Client, ...)` / same for `StreamWithUsage` — deletes `os.Getenv(model.CredentialEnv)` reads (:48, :72) and the `egress.Client()` default (:19). SSE loops (:164-229, :310-374) reused verbatim; `provider.Event` maps 1:1 onto ProviderEventParams. Host client `Timeout: 5m`; request URL built exclusively from cfg — guest input contributes only the alias. + +### 3.4 Host LLM broker +New: `internal/llmbroker/broker.go` (+test); wire as `sb.OnGuestCall` in cmd/abox/main.go + pkg/abox. `Broker{cfg, resolver}`: +- `provider_open`: reject offline; `cfg.ModelNamed` lookup (config.go:384), unknown alias → error; allocate stream. +- `provider_send`: budget-checked append; on Last, unmarshal, **resolve credential at call time** via `Resolve(model.CredentialReference())`, call `provider.StreamWithUsage`, fan events into `notify("provider_event", ...)`; zero Value after request build; terminal done/error + cleanup. +- `provider_cancel`: cancel HTTP context. +- agentgateway mode: LLM keeps configured base_url exactly as today (current code never routed LLM through the gateway either); gateway LLM adapter is follow-up, fail-closed noted for `enforcement: required` once it exists. +- Logs stream lifecycle only (alias, status, byte counts) — never headers/bodies. + +### 3.5 Guest broker client + agent hook +New: `internal/guest/brokerclient/client.go` (+test). Modify: cmd/abox-guest/main.go, internal/agent/agent.go. +- `agent.Loop` gains `Stream func(...) (<-chan provider.Event, error)`; `Turn` (agent.go:69-104, call sites :73/:75) uses it instead of calling provider directly; event-consumption loop untouched. +- brokerclient: `provider_open` (IDs `g-1…`), chunked `provider_send` (mirror TransferArchive, runtime.go:363-381), `provider_event` → `provider.Event`, `provider_cancel` on ctx cancel. +- Guest read loop (main.go:93-127): route response frames (empty Method) → brokerclient pending map; `provider_event` → stream dispatch. `set_model` ignores Secrets for the model; boot `applySecrets(cfg.Secrets)` (main.go:46) deleted in proto-3 guest; MCP tokens still via `set_mcp_tokens`. +- Mixed-binary fail-fast: guest learns the host's protocol in the hello exchange; if host < 3, refuse model calls with "host binary too old; run make build" (an old host's inline `Call` loop silently discards guest-initiated frames, so without this a proto-3 image against an old `abox` hangs — exactly the state after `make image` without `make build`). +- Hardening: drop the three provider hosts from `defaultAllowed` in internal/guest/egress/egress.go:16-20 — proto-3 guests need TSI egress only for configured MCP URLs. + +### 3.6 Host stops pushing LLM secrets to proto-3 guests +Modify: PushSecrets/SetModel (runtime.go:352-361), abox.go:246, tui.go:379-380. Proto ≥ 3 → `set_model` carries model only; proto == 2 → legacy full push + stderr deprecation ("run make image"). On proto ≥ 3, `/provider` keys stay host-side entirely. + +### Phase 3 verification +Gates as before + `make build && make image` smoke (`abox --probe-vm`, then a real turn). protocol_test: round-trip new types; oversized provider_send rejected. runtime_turn_test additions: guest `provider_open` mid-`user_turn` — both complete (the demux test); unknown guest method → typed error; all existing tests green. broker_test: httptest SSE provider asserting Authorization/x-api-key built host-side; unknown-alias reject; offline reject; cancel aborts HTTP; chunk budget; credential re-resolved per call (rotate fixture). agent_test: Turn with fake Stream (no HTTP). brokerclient net.Pipe test: chunking, reassembly, cancel. + +--- + +## Documentation updates +- PLAN.md §4.1 and §13.4 distinguish protocol-3 host-only LLM credentials from guest-held MCP tokens and the protocol-2 legacy push. The §2 model-traffic decision and README security story describe the implemented protocol-3 host broker while retaining compatibility `base_url` metadata in the guest config. + +## Risks / notes +1. Demux refactor (3.1) touches every RPC path incl. cancel edge cases — own PR, existing suite green before broker methods. +2. Keychain headless/SSH (`abox exec`, locked keychain) → ErrLocked with guidance; use env-source refs in CI. Document. +3. MCP tokens still enter the guest this milestone — accepted; §14.4 brokering is the follow-up. +4. Zeroing is best-effort (Go GC copies); stated in package docs. +5. Old `abox` binary resuming a scrubbed session re-writes secrets into config.raw (old Prepare); next new-binary start re-scrubs. Mixed-binary users only. +6. Tool args > 512 KiB → error event (today's ceiling was 1 MiB); named const, acceptable. + +## Integration status + +- Protocol 3 rich turns request usage through the guest broker client and return accumulated usage and stop reason through the SDK `TurnResult`. +- Protocol 3 provider HTTPS and LLM authentication are host-brokered. LLM credential values are absent from the guest and session config; compatibility model metadata, including `base_url` and the credential environment-variable name, remains on the guest config disk but is not trusted for protocol-3 routing. +- Startup credential resolution is partial: successfully resolved MCP tokens are pushed even if the selected model credential is missing. Missing optional MCP tokens are skipped; other source failures are reported after the partial push. Interactive CLI may continue after reporting the error; headless CLI and SDK startup return it. +- Protocol-1 resume is rejected after the host rewrites `config.raw` without secrets. Protocol 2 remains the legacy secret-push path; protocol 3 keeps LLM credentials host-side. +- MCP tokens still enter guest memory through `set_mcp_tokens`; MCP traffic and credential brokering remain follow-up work. +- Session scrubbing reports aggregate per-session errors and aborts CLI/SDK startup if any legacy session could not be scrubbed. +- Azure and AWS credential-source authentication is limited to static host credentials (or an existing Azure CLI login); managed/workload identity is deferred. +- `llmbroker.Broker` still has no config-update API. The SDK and TUI therefore install a newly constructed broker after an idle model change so subsequent streams use current aliases, base URLs, and credential references. A broker-owned atomic `UpdateConfig` API would remove direct handler replacement and better define concurrent SDK `SetModel` behavior. +- Isolation remains **Planned** until the named hardware tests pass. diff --git a/PLAN.md b/PLAN.md index 5b56a03..8d9b2e6 100644 --- a/PLAN.md +++ b/PLAN.md @@ -50,8 +50,8 @@ The current plan makes these decisions: | Initial guest | ARM64 Linux | | Initial microVM backend | libkrun over Apple Hypervisor.framework | | Runtime integration | Dedicated `abox-vmm` Go helper with a narrow cgo boundary | -| Guest network | No NIC and no libkrun TSI in milestone one | -| Model traffic | Guest agent calls providers; TSI inet for HTTPS only | +| Guest network | No guest NIC. TSI inet for allowlisted MCP HTTPS only | +| Model traffic | Host provider broker. TSI inet is MCP-only | | Providers | OpenAI, Anthropic, and Grok through xAI | | Repository state | Clean Git worktree only | | Host workspace sharing | Prohibited | @@ -60,7 +60,7 @@ The current plan makes these decisions: | TUI framework | Bubble Tea v2, Bubbles, and Lip Gloss v2 | | TUI style | Full-screen near-black interface with restrained status colors | | agentgateway | Optional adapter; never required for basic operation | -| Connectivity broker | Host-owned, typed, endpoint-bound package and MCP broker | +| Connectivity broker | LLM broker is in. Package/MCP broker remains Planned | | Package-manager compatibility | Origin rewrite to a guest loopback adapter, not HTTP(S) proxy | | Instruction loading | Supervisor reads the captured host snapshot and host configuration | | Repo instruction authority | Repo text cannot change policy, limits, connectivity, or tools | @@ -182,9 +182,11 @@ The host-side `abox` process owns: - Audit records - Patch review and confirmed import -Provider credentials may be entered on the host (`/provider`) and are -copied into the guest agent so the model client runs inside the -microVM. The host must not run the agent loop or call provider APIs. +Provider credentials are entered on the host (`/provider`) and resolved +from env, macOS keychain, Vault, Azure Key Vault, or AWS Secrets Manager. +They are never written to session dirs or the guest disk. The host broker +calls the provider API; the host must not run the agent loop. MCP tokens +still enter the guest. The host supervisor must remain small. It must not contain an arbitrary shell execution path, generated-code runner, or generic guest-to-host file service. @@ -220,8 +222,8 @@ The `abox-guest` worker and everything it starts are untrusted. The design assumes the guest can become fully compromised, including guest root and the guest kernel. -The guest owns the agent: the prompt, the model client, tools, and -everything the model starts. +The guest owns the agent: the prompt, tools, and everything the model +starts. Provider HTTPS is host-brokered. The guest owns all effectful tools: @@ -237,7 +239,8 @@ The guest owns all effectful tools: - Applications started by the agent The guest receives no model-provider credentials, host home-directory access, -cloud credentials, SSH keys, Docker socket, or read-write host mount. +cloud credentials, SSH keys, Docker socket, or read-write host mount. MCP +tokens still enter the guest. ### 4.4 External Services @@ -682,10 +685,11 @@ Host may call: Guest may call only: -- `FetchPackage` -- MCP stream methods defined in section 14.4 +- `provider_open`, `provider_send`, `provider_cancel` (host LLM broker) - Readiness and bounded log or status notifications +`FetchPackage` and MCP stream methods (section 14.4) remain Planned. + The guest must not invoke host tool, import, shell, or arbitrary-fetch methods. Phase 3 tests both directions. @@ -820,14 +824,14 @@ defense in depth but is not the primary boundary. ## 12. Agent Loop -The host supervisor owns the model interaction loop: +The guest owns the model interaction loop. The host broker performs provider HTTPS: 1. Receive the user's prompt from the TUI or `abox exec`. 2. Build the model request using configured instructions, the five ABox tool schemas, and any approved discovered MCP tool schemas. -3. Stream model output into normalized host events. +3. Send the request through the host provider broker and stream events back. 4. When the model requests a tool, validate the tool name and arguments. -5. Send a typed tool request to `abox-guest` over RPC. +5. Run the tool in the guest. 6. Stream or collect the bounded guest result. 7. Display activity and result status in the TUI. 8. Return the result to the same provider conversation. @@ -1011,11 +1015,11 @@ blocks while preserving the assistant content needed for subsequent turns. ### 13.4 Credentials -- Credentials remain only in host memory. -- Credentials are resolved by the host credential source from environment - variables or the operating system credential store. -- Credentials are never written to session logs. -- Credentials are never copied into the guest. +- LLM credentials remain only in host memory. +- Sources: env, macOS keychain, Vault KV v2, Azure Key Vault, AWS Secrets Manager. +- Credentials are never written to session logs or `config.raw`. +- LLM credentials are never copied into the guest. +- MCP tokens still enter the guest until section 14.4. - Configuration stores credential references, not secret values. The credential source is distinct from the connectivity broker. It resolves @@ -1062,15 +1066,16 @@ Connectivity is independent from the guest runtime isolation profile. - The trusted host supervisor may contact explicitly configured model-provider endpoints. -- The connectivity broker may contact exact configured remote MCP endpoints on - behalf of the guest MCP client. -- The connectivity broker may fetch from exact configured package indexes on - behalf of guest package tooling. -- The guest remains without a NIC and without TSI. +- The guest MCP client may contact configured MCP endpoints over TSI inet. +- Package-index fetch remains Planned (section 14.4). +- The guest remains without a NIC. - Direct mode does not imply unrestricted guest egress. ### 14.3 `agentgateway` +The LLM gateway adapter is Planned. Today `agentgateway` mode applies to MCP +endpoints; the host broker dials each model's `base_url`. + - ABox is a standalone client of a pre-existing agentgateway endpoint. - ABox does not install a local gateway, Kubernetes CRDs, Helm charts, or an agentgateway control plane. @@ -1117,8 +1122,9 @@ enforced. ### 14.4 Connectivity Broker Contract -The first milestone includes a typed, allowlisted host broker for configured -package indexes and remote MCP servers. The broker is implemented by +This section is Planned for MCP and package indexes. The LLM provider broker +is already in. The first milestone includes a typed, allowlisted host broker +for configured package indexes and remote MCP servers. The broker is implemented by `internal/connectivity` inside the trusted supervisor and does not run as a separate daemon. @@ -1185,7 +1191,7 @@ The contract enforces: bounded and cancellable. - The broker is not a TCP, CONNECT, SOCKS, DNS, or general HTTP forwarder. - Provider, gateway, package-index, MCP, and host credentials remain on the - host and are never returned to the guest. + host and are never returned to the guest. MCP tokens are the current exception. - In offline mode, all remote broker methods are refused. - With required agentgateway enforcement, the broker may open only the configured agentgateway endpoint and never a direct backend or package-index diff --git a/README.md b/README.md index 1b35096..5363a93 100644 --- a/README.md +++ b/README.md @@ -137,7 +137,7 @@ Clone of (1) for that run. This is /dev/vda → /. Repo, guest Git, agent writes ABox does not boot (1). It copies (1) → (2), then the microVM uses (2). --resume skips the copy and boots the existing (2). 3. Config disk — sessions//config.raw -~1 MiB, read-only /dev/vdb. Session id, model, keys. Not cloned from the golden image, not an OS. It lives inside of the directory where your sandbox harness session lives. +~1 MiB, read-only /dev/vdb. Session id, model. Not cloned from the golden image, not an OS. It lives inside of the directory where your sandbox harness session lives. The VM boots **only** the session clone, not the golden file. Destroy a session directory and that run’s guest files are gone; the golden image stays clean for the next `abox`. `make image-update` patches `/usr/local/bin/abox-guest` on an existing golden disk; `make image` rebuilds the golden disk from scratch. @@ -239,10 +239,10 @@ abox exec --prompt "list the repository files" ## LLM Integration -ABox is an LLM **client/harness**. The model loop/context is not on the host (your ABox instance/harness running on your computer). Prompts, streaming, tool calls, and provider HTTPS all run inside `abox-guest` in the microVM. The host TUI forwards your text over vsock (`user_turn`) and renders `agent_event` frames. That is the same isolation idea as MCP: the sandbox is the trust boundary for anything the model sees or starts. +ABox is an LLM **client/harness**. The model loop/context is not on the host (your ABox instance/harness running on your computer). Prompts, streaming, and tool calls run inside `abox-guest` in the microVM. Provider HTTPS is brokered by the host. The host TUI forwards your text over vsock (`user_turn`) and renders `agent_event` frames. That is the same isolation idea as MCP: the sandbox is the trust boundary for anything the model sees or starts. ```go -func Stream(ctx context.Context, model config.Model, messages []Message, tools []ToolSchema) (<-chan Event, error) +func Stream(ctx context.Context, model config.Model, key string, client *http.Client, messages []Message, tools []ToolSchema) (<-chan Event, error) ``` `Stream` talks to one configured profile. xAI and OpenAI use Chat Completions (`/chat/completions`). Anthropic uses Messages (`/v1/messages`). Provider-side shell, code execution, and file tools stay off. The model only sees ABox’s five guest tools plus any MCP tools discovered in the guest. @@ -250,7 +250,17 @@ func Stream(ctx context.Context, model config.Model, messages []Message, tools [ ![](img/prov1.png) ![](img/prov2.png) -Config lives at `~/.abox/config.yaml`. Keys are **not** stored in that file. `/provider` in the TUI writes `~/.abox/credentials.env` (mode 0600) and copies the value onto the sealed guest `config.raw` disk so the microVM can dial the API. Same as direct-mode MCP: the token has to live in the guest because the guest makes the HTTPS call. +Config lives at `~/.abox/config.yaml`. Keys are **not** stored in that file. Credential sources: `env`, `keychain` (macOS), `vault`, `azure`, `aws`. `/provider` in the TUI saves to the macOS keychain first (service `abox`), falling back to `~/.abox/credentials.env` (mode 0600). LLM keys stay on the host. MCP tokens still go to the guest because the guest makes those HTTPS calls. + +```yaml +credential: + source: keychain # env | keychain | vault | azure | aws + name: ANTHROPIC_API_KEY # env var, keychain account, vault path, Azure secret URI, or AWS secret id + # field: value # vault/aws only + # version: "4" # vault/azure only +``` + +`credential_env: XAI_API_KEY` is the same as `{source: env, name: XAI_API_KEY}`. Vault needs `VAULT_ADDR` + `VAULT_TOKEN`. Azure needs `AZURE_CLIENT_ID` / `AZURE_TENANT_ID` / `AZURE_CLIENT_SECRET` (or `az login`). AWS needs `AWS_ACCESS_KEY_ID` + `AWS_SECRET_ACCESS_KEY`. `abox creds migrate` moves existing `credentials.env` entries into the keychain. Default profiles: @@ -277,9 +287,9 @@ models: Pick one in the TUI with `/provider`, or pass `--model grok-default` (and the other profile names) on `abox` / `abox exec`. Missing `XAI_API_KEY` / `OPENAI_API_KEY` / `ANTHROPIC_API_KEY` fails the turn, not VM boot (`abox --probe-vm` still works). -Guest egress is allowlisted: `api.x.ai`, `api.openai.com`, `api.anthropic.com` on HTTPS `:443` only. Those sockets leave via libkrun TSI inet (no guest NIC). Isolation is still **Planned**. A compromised guest can read the key on `config.raw`; the allowlist is ABox’s Go dialer, not a VMM guarantee. +Guest egress is allowlisted for configured MCP origins on HTTPS `:443` via libkrun TSI inet (no guest NIC). Provider HTTPS is host-brokered, so those hosts are not on the guest allowlist. Isolation is still **Planned**. The allowlist is ABox’s Go dialer, not a VMM guarantee. -LLM traffic does **not** take the MCP `connectivity.mode` path. Direct vs agentgateway today applies to MCP servers. The model client always hits the provider `base_url` above. +LLM traffic does **not** take the MCP `connectivity.mode` path. Direct vs agentgateway today applies to MCP servers. The host broker hits the provider `base_url` above. ## MCP Integration @@ -412,14 +422,14 @@ Because of the above, Go or Rust are naturally great languages. Because I like G ## What is not done yet -Compaction, checkpoint/rollback/fork, stdio MCP, host broker, and resource +Compaction, checkpoint/rollback/fork, stdio MCP, host MCP/package broker, and resource acceptance. See `PLAN.md`. Streamable HTTP MCP is in; isolation stays Planned. ## Security Do not describe this build as verified isolation. The device plan is -allowlisted (no guest NIC, no host-path virtio-fs, TSI flags zero). Claims -stay Planned until the hardware suite in `PLAN.md` §21.4 passes. +allowlisted (no guest NIC, no host-path virtio-fs). TSI inet is MCP HTTPS only. +Claims stay Planned until the hardware suite in `PLAN.md` §21.4 passes. ## Whats Currently In Place ``` @@ -436,7 +446,7 @@ stay Planned until the hardware suite in `PLAN.md` §21.4 passes. ├──────────────────┼──────────────────────────────────────────────────────────────────────┤ │ Providers │ Grok/OpenAI (chat completions) + Anthropic Messages. /provider keys. │ ├──────────────────┼──────────────────────────────────────────────────────────────────────┤ -│ LLM egress │ Allowlist: api.x.ai, api.openai.com, api.anthropic.com via TSI inet │ +│ LLM egress │ Host broker dials providers. Guest TSI inet is MCP-only │ ├──────────────────┼──────────────────────────────────────────────────────────────────────┤ │ MCP │ Guest Streamable HTTP client; direct URLs or exclusive agentgateway │ ├──────────────────┼──────────────────────────────────────────────────────────────────────┤ diff --git a/cmd/abox-guest/main.go b/cmd/abox-guest/main.go index 7e19b6a..56c7fd8 100644 --- a/cmd/abox-guest/main.go +++ b/cmd/abox-guest/main.go @@ -20,6 +20,7 @@ import ( "github.com/AdminTurnedDevOps/ABox/internal/agent" "github.com/AdminTurnedDevOps/ABox/internal/config" + "github.com/AdminTurnedDevOps/ABox/internal/guest/brokerclient" "github.com/AdminTurnedDevOps/ABox/internal/guest/egress" guestmcp "github.com/AdminTurnedDevOps/ABox/internal/guest/mcp" "github.com/AdminTurnedDevOps/ABox/internal/guest/tools" @@ -43,7 +44,6 @@ func run() error { if err != nil { return err } - applySecrets(cfg.Secrets) repo := tools.Repo{Root: cfg.RepoDir} if err := os.MkdirAll(repo.Root, 0o755); err != nil { return err @@ -58,7 +58,14 @@ func run() error { fmt.Fprintf(os.Stderr, "abox-guest: mcp: %v\n", err) } defer mcpMgr.Close() - loop := &agent.Loop{Model: config.ModelFromGuest(cfg.Model), Repo: repo, MCP: mcpMgr, ContextFile: agent.DefaultContextFile} + bclient := brokerclient.New() + loop := &agent.Loop{ + Model: config.ModelFromGuest(cfg.Model), + Repo: repo, + MCP: mcpMgr, + ContextFile: agent.DefaultContextFile, + Stream: bclient.Stream, + } if err := loop.LoadContext(); err != nil { fmt.Fprintf(os.Stderr, "abox-guest: context: %v\n", err) } @@ -86,9 +93,26 @@ func run() error { if ack.Error != nil { return ack.Error } + if len(ack.Result) > 0 { + var ackRes protocol.HelloResult + if err := json.Unmarshal(ack.Result, &ackRes); err == nil { + bclient.SetHostProtocol(ackRes.Protocol) + } + } w := &connWriter{c: conn} + bclient.AttachContext(w.writeContext) turns := &turnTracker{} + shutdownHandled := false + defer bclient.Close(errors.New("guest connection closed")) + defer func() { + if shutdownHandled { + return + } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + turns.cancelAllAndWait(ctx) + }() var archive bytes.Buffer for { frame, err := protocol.ReadFrame(conn) @@ -98,13 +122,16 @@ func run() error { } return err } + if bclient.HandleFrame(frame) { + continue + } switch frame.Method { case "user_turn": if !turns.start(frame.ID) { _ = w.write(protocol.Frame{V: protocol.Version, ID: frame.ID, Error: &protocol.Error{Code: "guest", Message: "turn already in progress"}}) continue } - go runTurn(w, turns, loop, frame) + go runTurn(w, turns, loop, bclient, frame) case "cancel_turn": p, e := protocol.DecodeParams[protocol.CancelTurnParams](frame.Params) if e != nil { @@ -114,34 +141,84 @@ func run() error { turns.cancel(p.ID) ok, _ := protocol.EncodeParams(map[string]bool{"ok": true}) _ = w.write(protocol.Frame{V: protocol.Version, ID: frame.ID, Result: ok}) + case "shutdown": + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + stopped := turns.cancelAllAndWait(ctx) + cancel() + shutdownHandled = true + if !stopped { + return fmt.Errorf("active turn did not stop before shutdown timeout") + } + _ = loop.SaveContext() + ok, _ := protocol.EncodeParams(map[string]bool{"ok": true}) + if err := w.write(protocol.Frame{V: protocol.Version, ID: frame.ID, Result: ok}); err != nil { + return err + } + return nil default: resp := handle(loop, repo, mcpMgr, &archive, frame) if err := w.write(resp); err != nil { return err } - if frame.Method == "shutdown" { - turns.cancelAll() - return nil - } } } } type connWriter struct { - mu sync.Mutex - c net.Conn + once sync.Once + gate chan struct{} + c net.Conn } func (w *connWriter) write(f protocol.Frame) error { - w.mu.Lock() - defer w.mu.Unlock() - return protocol.WriteFrame(w.c, f) + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + return w.writeContext(ctx, f) +} + +func (w *connWriter) writeContext(ctx context.Context, f protocol.Frame) error { + w.once.Do(func() { + w.gate = make(chan struct{}, 1) + w.gate <- struct{}{} + }) + select { + case <-ctx.Done(): + return ctx.Err() + case <-w.gate: + } + defer func() { w.gate <- struct{}{} }() + done := make(chan struct{}) + watchDone := make(chan struct{}) + go func() { + select { + case <-done: + case <-ctx.Done(): + select { + case <-done: + default: + _ = w.c.Close() + } + } + close(watchDone) + }() + err := protocol.WriteFrame(w.c, f) + close(done) + <-watchDone + if ctxErr := ctx.Err(); ctxErr != nil { + return ctxErr + } + if err != nil { + return err + } + return nil } type turnTracker struct { - mu sync.Mutex - active string - stop context.CancelFunc + mu sync.Mutex + active string + stop context.CancelFunc + canceled bool + finished chan struct{} } func (t *turnTracker) start(id string) bool { @@ -151,30 +228,55 @@ func (t *turnTracker) start(id string) bool { return false } t.active = id + t.canceled = false + t.finished = make(chan struct{}) return true } func (t *turnTracker) setCancel(id string, cancel context.CancelFunc) { t.mu.Lock() - defer t.mu.Unlock() + shouldCancel := false if t.active == id { t.stop = cancel + shouldCancel = t.canceled + } + t.mu.Unlock() + if shouldCancel { + cancel() } } func (t *turnTracker) cancel(id string) { t.mu.Lock() - defer t.mu.Unlock() - if t.active == id && t.stop != nil { - t.stop() + var stop context.CancelFunc + if t.active == id { + t.canceled = true + stop = t.stop + } + t.mu.Unlock() + if stop != nil { + stop() } } -func (t *turnTracker) cancelAll() { +func (t *turnTracker) cancelAllAndWait(ctx context.Context) bool { t.mu.Lock() - defer t.mu.Unlock() - if t.stop != nil { - t.stop() + if t.active == "" { + t.mu.Unlock() + return true + } + t.canceled = true + stop := t.stop + finished := t.finished + t.mu.Unlock() + if stop != nil { + stop() + } + select { + case <-finished: + return true + case <-ctx.Done(): + return false } } @@ -182,12 +284,18 @@ func (t *turnTracker) done(id string) { t.mu.Lock() defer t.mu.Unlock() if t.active == id { + finished := t.finished t.active = "" t.stop = nil + t.canceled = false + t.finished = nil + if finished != nil { + close(finished) + } } } -func runTurn(w *connWriter, turns *turnTracker, loop *agent.Loop, req protocol.Frame) { +func runTurn(w *connWriter, turns *turnTracker, loop *agent.Loop, bclient *brokerclient.Client, req protocol.Frame) { defer turns.done(req.ID) p, err := protocol.DecodeParams[protocol.UserTurnParams](req.Params) if err != nil { @@ -204,6 +312,11 @@ func runTurn(w *connWriter, turns *turnTracker, loop *agent.Loop, req protocol.F turns.setCancel(req.ID, cancel) loop.MaxTurns = p.MaxTurns loop.Rich = p.RichEvents + if p.RichEvents { + loop.Stream = bclient.StreamWithUsage + } else { + loop.Stream = bclient.Stream + } loop.OnEvent = func(ev protocol.AgentEvent) { raw, _ := protocol.EncodeParams(ev) _ = w.write(protocol.Frame{ID: req.ID, Method: "agent_event", Params: raw}) @@ -254,7 +367,6 @@ func handle(loop *agent.Loop, repo tools.Repo, mcpMgr *guestmcp.Manager, archive err = e break } - applySecrets(p.Secrets) loop.Model = config.ModelFromGuest(p.Model) out.Result, _ = protocol.EncodeParams(map[string]bool{"ok": true}) case "archive_chunk": diff --git a/cmd/abox/creds.go b/cmd/abox/creds.go new file mode 100644 index 0000000..9ce51e3 --- /dev/null +++ b/cmd/abox/creds.go @@ -0,0 +1,176 @@ +package main + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + "time" + + "github.com/AdminTurnedDevOps/ABox/internal/config" + "github.com/AdminTurnedDevOps/ABox/internal/credentials" + "github.com/AdminTurnedDevOps/ABox/internal/credsource" +) + +var ( + migrationKeychainAvailable = credsource.KeychainAvailable + migrationSetKeychain = credsource.SetKeychain +) + +func runCreds(args []string) error { + if len(args) == 0 { + return fmt.Errorf("usage: abox creds migrate (move credentials.env entries to the macOS keychain)") + } + switch args[0] { + case "migrate": + return credsMigrate() + default: + return fmt.Errorf("unknown creds command %q (try: abox creds migrate)", args[0]) + } +} + +func credsMigrate() error { + if !migrationKeychainAvailable() { + return fmt.Errorf("macOS keychain unavailable (this command needs /usr/bin/security on darwin); the env credential source keeps working") + } + cfg, _, err := config.Load() + if err != nil { + return err + } + creds, err := credentials.Load() + if err != nil { + return err + } + if len(creds) == 0 { + fmt.Println("no credentials to migrate") + return nil + } + + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + + remaining := make(map[string]string, len(creds)) + for name, value := range creds { + remaining[name] = value + } + var migrated, failed, dropped int + configChanged := false + legacyRefresh := legacyMCPRefreshEntries(cfg) + for _, name := range sortedCredNames(creds) { + if _, drop := legacyRefresh[name]; drop { + fmt.Printf("dropping %s (refresh tokens are no longer stored; re-login when the access token expires)\n", name) + delete(remaining, name) + dropped++ + continue + } + if err := migrationSetKeychain(ctx, name, []byte(creds[name])); err != nil { + fmt.Printf("skipped %s: %v (entry stays in the file; set it manually or re-login)\n", name, err) + failed++ + continue + } + delete(remaining, name) + configChanged = upsertCredentialRefs(&cfg, name) || configChanged + migrated++ + } + if configChanged { + if err := cfg.Save(); err != nil { + return fmt.Errorf("keychain writes succeeded but config update failed: %w", err) + } + } + if err := rewriteCredentialFile(remaining); err != nil { + return err + } + if failed > 0 { + fmt.Printf("migrated %d, dropped %d refresh token(s), skipped %d (only skipped entries remain; fix them and run abox creds migrate again)\n", migrated, dropped, failed) + return nil + } + fmt.Printf("migrated %d credential(s) to the macOS keychain (service %s), dropped %d refresh token(s)\n", + migrated, credsource.KeychainService, dropped) + fmt.Printf("rewrote %s (kept, mode 0600)\n", credentials.Path()) + return nil +} + +func legacyMCPRefreshEntries(cfg config.File) map[string]struct{} { + drops := make(map[string]struct{}, len(cfg.MCPServers)) + protected := make(map[string]struct{}, len(cfg.Models)+len(cfg.MCPServers)) + for _, model := range cfg.Models { + protected[model.CredentialReference().Name] = struct{}{} + } + for _, server := range cfg.MCPServers { + ref := server.CredentialReference() + protected[ref.Name] = struct{}{} + drops[config.TokenEnv(server)+"_REFRESH"] = struct{}{} + if ref.Source == "env" { + drops[ref.Name+"_REFRESH"] = struct{}{} + } + } + for name := range protected { + delete(drops, name) + } + return drops +} + +func upsertCredentialRefs(cfg *config.File, name string) bool { + changed := false + for i := range cfg.Models { + ref := cfg.Models[i].CredentialReference() + if ref.Source == "env" && ref.Name == name { + cfg.Models[i].CredentialEnv = "" + cfg.Models[i].Credential = &config.CredentialRef{Source: "keychain", Name: name} + changed = true + } + } + for i := range cfg.MCPServers { + ref := cfg.MCPServers[i].CredentialReference() + if ref.Source == "env" && ref.Name == name { + cfg.MCPServers[i].CredentialEnv = "" + cfg.MCPServers[i].Credential = &config.CredentialRef{Source: "keychain", Name: name} + changed = true + } + } + return changed +} + +func rewriteCredentialFile(creds map[string]string) error { + var body strings.Builder + body.WriteString("# ABox credentials. Mode 0600. Do not commit.\n") + if len(creds) == 0 { + body.WriteString("# Credentials migrated to the macOS keychain; env fallback remains supported.\n") + } else { + for _, name := range sortedCredNames(creds) { + body.WriteString(name) + body.WriteByte('=') + body.WriteString(creds[name]) + body.WriteByte('\n') + } + } + path := credentials.Path() + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + return fmt.Errorf("rewrite credentials: %w", err) + } + tmp := path + ".tmp" + if err := os.WriteFile(tmp, []byte(body.String()), 0o600); err != nil { + return fmt.Errorf("rewrite credentials: %w", err) + } + if err := os.Rename(tmp, path); err != nil { + return fmt.Errorf("rewrite credentials: %w", err) + } + if err := os.Chmod(path, 0o600); err != nil { + return fmt.Errorf("rewrite credentials: %w", err) + } + return nil +} + +func sortedCredNames(m map[string]string) []string { + out := make([]string, 0, len(m)) + for k := range m { + out = append(out, k) + } + for i := 1; i < len(out); i++ { + for j := i; j > 0 && out[j] < out[j-1]; j-- { + out[j], out[j-1] = out[j-1], out[j] + } + } + return out +} diff --git a/cmd/abox/creds_test.go b/cmd/abox/creds_test.go new file mode 100644 index 0000000..18b6d26 --- /dev/null +++ b/cmd/abox/creds_test.go @@ -0,0 +1,128 @@ +package main + +import ( + "context" + "errors" + "testing" + + "github.com/AdminTurnedDevOps/ABox/internal/config" + "github.com/AdminTurnedDevOps/ABox/internal/credentials" +) + +func TestCredsMigrateRemovesSuccessAndRefreshButKeepsFailure(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + t.Setenv("ABOX_HOME", "") + envRef := config.CredentialRef{Source: "env", Name: "CUSTOM_SOURCE"} + cfg := config.Defaults() + cfg.Models = []config.Model{ + {Name: "custom", Provider: "other", Credential: &envRef}, + {Name: "failed", Provider: "other", CredentialEnv: "FAILED_KEY"}, + {Name: "refresh-model", Provider: "other", CredentialEnv: "MODEL_REFRESH"}, + } + cfg.MCPServers = []config.MCPServer{ + {Name: "refresh-mcp", URL: "https://mcp.example/api", CredentialEnv: "MCP_CRED_REFRESH"}, + {Name: "normal-mcp", URL: "https://normal.example/api", CredentialEnv: "MCP_TOKEN"}, + } + if err := cfg.Save(); err != nil { + t.Fatal(err) + } + for name, value := range map[string]string{ + "CUSTOM_SOURCE": "custom-value", + "FAILED_KEY": "failed-value", + "MODEL_REFRESH": "model-value", + "MCP_CRED_REFRESH": "mcp-value", + "MCP_CRED_REFRESH_REFRESH": "legacy-refresh-value", + "MCP_TOKEN_REFRESH": "legacy-normal-refresh-value", + } { + if err := credentials.Save(name, value); err != nil { + t.Fatal(err) + } + } + + origAvailable := migrationKeychainAvailable + origSet := migrationSetKeychain + migrationKeychainAvailable = func() bool { return true } + called := map[string]bool{} + migrationSetKeychain = func(_ context.Context, name string, _ []byte) error { + called[name] = true + if name == "FAILED_KEY" || name == "MODEL_REFRESH" || name == "MCP_CRED_REFRESH" { + return errors.New("write failed") + } + return nil + } + t.Cleanup(func() { + migrationKeychainAvailable = origAvailable + migrationSetKeychain = origSet + }) + + if err := credsMigrate(); err != nil { + t.Fatal(err) + } + remaining, err := credentials.Load() + if err != nil { + t.Fatal(err) + } + if len(remaining) != 3 || remaining["FAILED_KEY"] != "failed-value" || remaining["MODEL_REFRESH"] != "model-value" || remaining["MCP_CRED_REFRESH"] != "mcp-value" { + t.Fatalf("remaining credentials %#v", remaining) + } + if called["MCP_CRED_REFRESH_REFRESH"] || called["MCP_TOKEN_REFRESH"] { + t.Fatal("known legacy refresh token was sent to the keychain") + } + if !called["MODEL_REFRESH"] || !called["MCP_CRED_REFRESH"] { + t.Fatalf("configured refresh-suffixed credentials were dropped: calls %#v", called) + } + savedCfg, _, err := config.Load() + if err != nil { + t.Fatal(err) + } + if got := savedCfg.Models[0].CredentialReference(); got != (config.CredentialRef{Source: "keychain", Name: "CUSTOM_SOURCE"}) { + t.Fatalf("custom reference %#v", got) + } + if got := savedCfg.Models[1].CredentialReference(); got != (config.CredentialRef{Source: "env", Name: "FAILED_KEY"}) { + t.Fatalf("failed reference %#v", got) + } + if got := savedCfg.Models[2].CredentialReference(); got != (config.CredentialRef{Source: "env", Name: "MODEL_REFRESH"}) { + t.Fatalf("refresh model reference %#v", got) + } + if got := savedCfg.MCPServers[0].CredentialReference(); got != (config.CredentialRef{Source: "env", Name: "MCP_CRED_REFRESH"}) { + t.Fatalf("refresh mcp reference %#v", got) + } +} + +func TestUpsertCredentialRefsUsesEffectiveEnvReference(t *testing.T) { + vault := config.CredentialRef{Source: "vault", Name: "secret/abox/vault"} + azure := config.CredentialRef{Source: "azure", Name: "https://testkv.vault.azure.net/secrets/azure"} + aws := config.CredentialRef{Source: "aws", Name: "prod/aws"} + customEnv := config.CredentialRef{Source: "env", Name: "CUSTOM_ENV"} + customMCPEnv := config.CredentialRef{Source: "env", Name: "CUSTOM_MCP_ENV"} + cfg := config.File{ + Models: []config.Model{ + {Name: "vault", Provider: "other", Credential: &vault}, + {Name: "azure", Provider: "other", Credential: &azure}, + {Name: "aws", Provider: "other", Credential: &aws}, + {Name: "custom", Provider: "other", Credential: &customEnv}, + }, + MCPServers: []config.MCPServer{{Name: "custom", URL: "https://mcp.example/api", Credential: &customMCPEnv}}, + } + + for i := 0; i < 3; i++ { + if upsertCredentialRefs(&cfg, cfg.Models[i].EnvName()) { + t.Fatalf("cloud reference %d was changed", i) + } + } + if cfg.Models[0].CredentialReference() != vault || cfg.Models[1].CredentialReference() != azure || cfg.Models[2].CredentialReference() != aws { + t.Fatalf("cloud references changed: %#v", cfg.Models) + } + if !upsertCredentialRefs(&cfg, "CUSTOM_ENV") { + t.Fatal("custom env reference was not changed") + } + if got := cfg.Models[3].CredentialReference(); got != (config.CredentialRef{Source: "keychain", Name: "CUSTOM_ENV"}) { + t.Fatalf("custom env reference %#v", got) + } + if !upsertCredentialRefs(&cfg, "CUSTOM_MCP_ENV") { + t.Fatal("custom MCP env reference was not changed") + } + if got := cfg.MCPServers[0].CredentialReference(); got != (config.CredentialRef{Source: "keychain", Name: "CUSTOM_MCP_ENV"}) { + t.Fatalf("custom MCP env reference %#v", got) + } +} diff --git a/cmd/abox/main.go b/cmd/abox/main.go index ef4d8cb..9d9d36f 100644 --- a/cmd/abox/main.go +++ b/cmd/abox/main.go @@ -5,8 +5,10 @@ package main import ( "context" "encoding/json" + "errors" "flag" "fmt" + "log" "os" "os/signal" "path/filepath" @@ -16,6 +18,8 @@ import ( "github.com/AdminTurnedDevOps/ABox/internal/agent" "github.com/AdminTurnedDevOps/ABox/internal/config" "github.com/AdminTurnedDevOps/ABox/internal/credentials" + "github.com/AdminTurnedDevOps/ABox/internal/credsource" + "github.com/AdminTurnedDevOps/ABox/internal/llmbroker" "github.com/AdminTurnedDevOps/ABox/internal/mcpauth" "github.com/AdminTurnedDevOps/ABox/internal/repository" "github.com/AdminTurnedDevOps/ABox/internal/runtime" @@ -35,6 +39,9 @@ func run() error { if len(os.Args) > 1 && os.Args[1] == "mcp" { return runMCP(os.Args[2:]) } + if len(os.Args) > 1 && os.Args[1] == "creds" { + return runCreds(os.Args[2:]) + } fs := flag.NewFlagSet("abox", flag.ContinueOnError) execFlag := fs.Bool("exec", false, "headless driver") prompt := fs.String("prompt", "", "prompt for exec mode") @@ -58,9 +65,12 @@ func run() error { if err != nil { return err } - if err := credentials.ApplyToEnv(); err != nil { + if err := scrubLegacySessions(); err != nil { return err } + resolver := credsource.NewResolver() + defer resolver.Close() + sel, ok := cfg.ModelNamed(*modelName) if !ok { return fmt.Errorf("no model profile %q (config %s)", *modelName, cfgPath) @@ -121,7 +131,7 @@ func run() error { if err != nil { return err } - if err := runtime.Prepare(sess, image, sel, cfg.SecretsFromEnv(), mcpServers, *resume); err != nil { + if err := runtime.Prepare(sess, image, sel, mcpServers, *resume); err != nil { if execMode { return err } @@ -139,9 +149,26 @@ func run() error { fmt.Fprintf(os.Stderr, "abox: vm start: %v\n", err) vmState = "failed" } else { + if started.GuestProtocol < 2 { + if *resume { + started.Stop() + return fmt.Errorf("cannot resume protocol-1 session %s after secretless config rewrite; rebuild the guest image and start a new session", sess.ID) + } + if !*probeVM { + started.Stop() + return fmt.Errorf("protocol-1 guest cannot use the secretless config; rebuild the guest image") + } + } sb = started vmState = "ready" defer sb.Stop() + sb.OnGuestCall = brokerForMode(cfg, resolver, execMode) + if err := pushSecrets(sb, cfg, resolver, sel); err != nil { + if execMode { + return err + } + fmt.Fprintf(os.Stderr, "abox: %v\n", err) + } if !*resume { archive, err := repository.ArchiveHEAD(snap.Root) if err != nil { @@ -171,14 +198,49 @@ func run() error { if execMode { return runExec(sb, *prompt) } - var log []string + var transcript []string if *resume { - log = resumeLog(sess, sb) - if len(log) > 0 { - _ = session.WriteTranscript(sess.TranscriptPath(), log) + transcript = resumeLog(sess, sb) + if len(transcript) > 0 { + _ = session.WriteTranscript(sess.TranscriptPath(), transcript) } } - return tui.Run(cfg, sel, sb, vmState, log, sess.TranscriptPath()) + return tui.Run(cfg, sel, sb, vmState, transcript, resolver, sess.TranscriptPath()) +} + +// Headless logs stream lifecycle; the TUI stays quiet so logs never paint into the UI. +func brokerForMode(cfg config.File, resolver *credsource.Resolver, execMode bool) *llmbroker.Broker { + b := llmbroker.New(cfg, resolver) + if execMode { + b.SetLogger(log.Printf) + } + return b +} + +func pushSecrets(sb *runtime.Sandbox, cfg config.File, resolver *credsource.Resolver, sel config.Model) error { + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + secrets, resolveErr := credsource.ResolveSelected(ctx, resolver, cfg, sel) + pushErr := sb.PushSecrets(ctx, sel, secrets) + var errs []error + if resolveErr != nil { + errs = append(errs, fmt.Errorf("resolve credentials: %w", resolveErr)) + } + if pushErr != nil { + errs = append(errs, fmt.Errorf("push resolved credentials: %w", pushErr)) + } + return errors.Join(errs...) +} + +func scrubLegacySessions() error { + n, err := session.ScrubSecretsEverywhere() + if n > 0 { + fmt.Fprintf(os.Stderr, "abox: scrubbed plaintext secrets from %d old session(s)\n", n) + } + if err != nil { + return fmt.Errorf("legacy session scrub incomplete; affected sessions may still contain plaintext secrets: %w", err) + } + return nil } func resumeLog(sess *session.Session, sb *runtime.Sandbox) []string { diff --git a/internal/agent/agent.go b/internal/agent/agent.go index 3bcf225..c005e78 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -36,6 +36,8 @@ type Loop struct { OnEvent func(protocol.AgentEvent) MaxTurns int Rich bool + + Stream func(ctx context.Context, model config.Model, messages []provider.Message, tools []provider.ToolSchema) (<-chan provider.Event, error) } func BuiltinTools() []provider.ToolSchema { @@ -67,13 +69,12 @@ func (l *Loop) Turn(ctx context.Context, user string) error { var usage protocol.UsageInfo stopReason := "" for i := 0; i < limit; i++ { - var events <-chan provider.Event - var err error - if l.Rich { - events, err = provider.StreamWithUsage(ctx, l.Model, l.Messages, l.allTools()) - } else { - events, err = provider.Stream(ctx, l.Model, l.Messages, l.allTools()) + if l.Stream == nil { + err := fmt.Errorf("no stream function configured for the agent loop") + l.emit(protocol.AgentEvent{Kind: "error", Err: err.Error()}) + return err } + events, err := l.Stream(ctx, l.Model, l.Messages, l.allTools()) if err != nil { l.emit(protocol.AgentEvent{Kind: "error", Err: err.Error()}) _ = l.SaveContext() diff --git a/internal/agent/agent_test.go b/internal/agent/agent_test.go index 70fc1da..af198c8 100644 --- a/internal/agent/agent_test.go +++ b/internal/agent/agent_test.go @@ -8,9 +8,11 @@ import ( "strings" "testing" + "github.com/AdminTurnedDevOps/ABox/internal/config" "github.com/AdminTurnedDevOps/ABox/internal/guest/mcp" "github.com/AdminTurnedDevOps/ABox/internal/guest/tools" "github.com/AdminTurnedDevOps/ABox/internal/provider" + "github.com/AdminTurnedDevOps/ABox/protocol" ) func TestMaxTurnsStops(t *testing.T) { @@ -196,3 +198,80 @@ func TestAllToolsIncludesMCP(t *testing.T) { t.Fatalf("got %q", tools[5].Name) } } + +// fakeStream returns a canned event stream; used to exercise Turn without HTTP. +type fakeStream struct { + calls int + events []provider.Event +} + +func (f *fakeStream) stream(_ context.Context, _ config.Model, _ []provider.Message, _ []provider.ToolSchema) (<-chan provider.Event, error) { + f.calls++ + out := make(chan provider.Event, len(f.events)) + for _, ev := range f.events { + out <- ev + } + close(out) + return out, nil +} + +func TestTurnUsesInjectedStream(t *testing.T) { + fs := &fakeStream{events: []provider.Event{ + {Type: "text", Text: "hello"}, + {Type: "done"}, + }} + l := &Loop{ + Repo: tools.Repo{Root: t.TempDir()}, + Stream: fs.stream, + } + var got []string + l.OnEvent = func(ev protocol.AgentEvent) { got = append(got, ev.Kind) } + if err := l.Turn(context.Background(), "hi"); err != nil { + t.Fatal(err) + } + if fs.calls != 1 { + t.Fatalf("stream calls %d", fs.calls) + } + if len(l.Messages) != 2 || l.Messages[1].Content != "hello" { + t.Fatalf("messages %+v", l.Messages) + } + if len(got) < 2 { + t.Fatalf("events %v", got) + } +} + +func TestTurnEmitsRichUsageResult(t *testing.T) { + fs := &fakeStream{events: []provider.Event{ + {Type: "text", Text: "hello"}, + {Type: "usage", Usage: &protocol.UsageInfo{InputTokens: 7, OutputTokens: 3}, StopReason: "end_turn"}, + {Type: "done"}, + }} + l := &Loop{ + Repo: tools.Repo{Root: t.TempDir()}, + Rich: true, + Stream: fs.stream, + } + var result protocol.AgentEvent + l.OnEvent = func(ev protocol.AgentEvent) { + if ev.Kind == "result" { + result = ev + } + } + if err := l.Turn(context.Background(), "hi"); err != nil { + t.Fatal(err) + } + if result.Usage == nil || result.Usage.InputTokens != 7 || result.Usage.OutputTokens != 3 { + t.Fatalf("result usage %+v", result.Usage) + } + if result.StopReason != "end_turn" { + t.Fatalf("stop reason %q", result.StopReason) + } +} + +func TestTurnWithoutStreamErrors(t *testing.T) { + l := &Loop{Repo: tools.Repo{Root: t.TempDir()}} + err := l.Turn(context.Background(), "hi") + if err == nil || !strings.Contains(err.Error(), "stream function") { + t.Fatalf("got %v", err) + } +} diff --git a/internal/config/config.go b/internal/config/config.go index 1c7d31e..33333ca 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -25,11 +25,32 @@ type File struct { } type Model struct { - Name string `yaml:"name"` - Provider string `yaml:"provider"` - Model string `yaml:"model"` - CredentialEnv string `yaml:"credential_env"` - BaseURL string `yaml:"base_url,omitempty"` + Name string `yaml:"name"` + Provider string `yaml:"provider"` + Model string `yaml:"model"` + CredentialEnv string `yaml:"credential_env,omitempty"` // DEPRECATED: alias for credential {source: env, name: X} + Credential *CredentialRef `yaml:"credential,omitempty"` + BaseURL string `yaml:"base_url,omitempty"` +} + +type CredentialRef struct { + Source string `yaml:"source"` + Name string `yaml:"name"` + Field string `yaml:"field,omitempty"` // vault/aws only + Version string `yaml:"version,omitempty"` // vault/azure only +} + +type AzureCloud struct { + KeyVaultDNSSuffix string + AuthorityHost string + VaultResource string +} + +var azureClouds = []AzureCloud{ + {KeyVaultDNSSuffix: "vault.azure.net", AuthorityHost: "https://login.microsoftonline.com", VaultResource: "https://vault.azure.net"}, + {KeyVaultDNSSuffix: "vault.usgovcloudapi.net", AuthorityHost: "https://login.microsoftonline.us", VaultResource: "https://vault.usgovcloudapi.net"}, + {KeyVaultDNSSuffix: "vault.azure.cn", AuthorityHost: "https://login.chinacloudapi.cn", VaultResource: "https://vault.azure.cn"}, + {KeyVaultDNSSuffix: "vault.microsoftazure.de", AuthorityHost: "https://login.microsoftonline.de", VaultResource: "https://vault.microsoftazure.de"}, } type Connectivity struct { @@ -38,12 +59,13 @@ type Connectivity struct { } type MCPServer struct { - Name string `yaml:"name"` - URL string `yaml:"url"` - CredentialEnv string `yaml:"credential_env,omitempty"` - ClientID string `yaml:"client_id,omitempty"` - Scopes []string `yaml:"scopes,omitempty"` - ToolAllowlist []string `yaml:"tool_allowlist,omitempty"` + Name string `yaml:"name"` + URL string `yaml:"url"` + CredentialEnv string `yaml:"credential_env,omitempty"` // DEPRECATED: alias for credential {source: env, name: X} + Credential *CredentialRef `yaml:"credential,omitempty"` + ClientID string `yaml:"client_id,omitempty"` + Scopes []string `yaml:"scopes,omitempty"` + ToolAllowlist []string `yaml:"tool_allowlist,omitempty"` } var mcpNameRE = regexp.MustCompile(`^[a-z0-9-]+$`) @@ -126,11 +148,10 @@ func EnsureLayout() error { } func seedFromLegacy() error { - home := homeDir() - if home == "" { + legacy := LegacyAppSupportDir() + if legacy == "" { return nil } - legacy := filepath.Join(home, "Library", "Application Support", "ABox") for _, name := range []string{"config.yaml", "credentials.env"} { dst := filepath.Join(Dir(), name) if exists(dst) { @@ -148,6 +169,35 @@ func seedFromLegacy() error { return fmt.Errorf("seed %s: %w", name, err) } } + return scrubLegacyAppSupportCredentials(legacy) +} + +// LegacyAppSupportDir is the pre-~/.abox macOS location. Dir() never returns +// it; it is only used to seed a missing ~/.abox and to scrub leftover secrets. +func LegacyAppSupportDir() string { + home := homeDir() + if home == "" { + return "" + } + return filepath.Join(home, "Library", "Application Support", "ABox") +} + +func scrubLegacyAppSupportCredentials(legacy string) error { + path := filepath.Join(legacy, "credentials.env") + if !exists(path) { + return nil + } + body := []byte("# ABox credentials. Mode 0600. Do not commit.\n# Leftover Application Support copy; credentials now live under ~/.abox or the macOS keychain.\n") + tmp := path + ".tmp" + if err := os.WriteFile(tmp, body, 0o600); err != nil { + return fmt.Errorf("scrub legacy credentials: %w", err) + } + if err := os.Rename(tmp, path); err != nil { + return fmt.Errorf("scrub legacy credentials: %w", err) + } + if err := os.Chmod(path, 0o600); err != nil { + return fmt.Errorf("scrub legacy credentials: %w", err) + } return nil } @@ -180,6 +230,24 @@ func (c File) Validate() error { } seen[s.Name] = struct{}{} } + for i, m := range c.Models { + if err := m.validate(); err != nil { + return fmt.Errorf("models[%d]: %w", i, err) + } + } + destinations := make(map[string]string, len(c.Models)+len(c.MCPServers)) + for _, m := range c.Models { + if _, exists := destinations[m.EnvName()]; !exists { + destinations[m.EnvName()] = fmt.Sprintf("model %q", m.Name) + } + } + for _, s := range c.MCPServers { + dest := TokenEnv(s) + if previous, exists := destinations[dest]; exists { + return fmt.Errorf("credential destination env %q is shared by %s and mcp server %q", dest, previous, s.Name) + } + destinations[dest] = fmt.Sprintf("mcp server %q", s.Name) + } if c.Runtime.Isolation != "" && c.Runtime.Isolation != "microvm" { return fmt.Errorf("isolation must be microvm") } @@ -199,12 +267,182 @@ func (s MCPServer) validate() error { if err := validateHTTPSURL("url", s.URL); err != nil { return err } + if s.CredentialEnv != "" && s.Credential != nil { + return fmt.Errorf("set either credential or credential_env, not both") + } if s.CredentialEnv != "" && !ValidEnvName(s.CredentialEnv) { return fmt.Errorf("invalid credential_env %q", s.CredentialEnv) } + if s.Credential != nil { + if err := s.Credential.validate(); err != nil { + return fmt.Errorf("credential: %w", err) + } + } + return nil +} + +func (m Model) validate() error { + if m.CredentialEnv != "" && m.Credential != nil { + return fmt.Errorf("set either credential or credential_env, not both") + } + if m.CredentialEnv != "" && !ValidEnvName(m.CredentialEnv) { + return fmt.Errorf("invalid credential_env %q", m.CredentialEnv) + } + if m.Credential != nil { + if err := m.Credential.validate(); err != nil { + return fmt.Errorf("credential: %w", err) + } + } + return nil +} + +var credentialSources = map[string]struct{}{ + "env": {}, + "keychain": {}, + "vault": {}, + "azure": {}, + "aws": {}, +} + +func (c CredentialRef) validate() error { + if _, ok := credentialSources[c.Source]; !ok { + return fmt.Errorf("unknown source %q (want env, keychain, vault, azure, or aws)", c.Source) + } + if strings.TrimSpace(c.Name) == "" { + return fmt.Errorf("name is required") + } + switch c.Source { + case "env", "keychain": + if c.Field != "" { + return fmt.Errorf("field is not supported for source %q", c.Source) + } + if c.Version != "" { + return fmt.Errorf("version is not supported for source %q", c.Source) + } + if !ValidEnvName(c.Name) { + return fmt.Errorf("invalid %s credential name %q", c.Source, c.Name) + } + case "vault": + if c.Version != "" && !isNumeric(c.Version) { + return fmt.Errorf("version %q must be numeric for vault", c.Version) + } + case "aws": + if c.Version != "" { + return fmt.Errorf("version is not supported for source %q", c.Source) + } + case "azure": + if c.Field != "" { + return fmt.Errorf("field is not supported for source %q", c.Source) + } + if _, _, _, _, err := ParseAzureSecretReference(c.Name, c.Version); err != nil { + return err + } + } return nil } +// ParseAzureSecretReference validates and canonicalizes an Azure Key Vault +// secret identifier. Only first-party data-plane DNS suffixes are accepted so +// callers can safely attach an Azure bearer token to the returned URI. +func ParseAzureSecretReference(name, requestedVersion string) (uri, secretName, version string, cloud AzureCloud, err error) { + raw := strings.TrimSpace(name) + u, parseErr := url.Parse(raw) + if parseErr != nil || u.Scheme != "https" || u.Opaque != "" || u.Host == "" || u.User != nil || u.Port() != "" || u.RawQuery != "" || u.Fragment != "" || u.ForceQuery { + err = fmt.Errorf("azure credential name must be an https secret URI for Azure Key Vault like https://vault.vault.azure.net/secrets/name") + return + } + host := strings.ToLower(u.Hostname()) + for _, candidate := range azureClouds { + suffix := "." + candidate.KeyVaultDNSSuffix + if strings.HasSuffix(host, suffix) { + vaultName := strings.TrimSuffix(host, suffix) + if strings.Contains(vaultName, ".") || !validAzureVaultName(vaultName) { + break + } + cloud = candidate + break + } + } + if cloud.KeyVaultDNSSuffix == "" { + err = fmt.Errorf("azure credential host %q is not a supported Azure Key Vault endpoint", u.Hostname()) + return + } + if u.RawPath != "" { + err = fmt.Errorf("azure credential path must not contain escaped characters") + return + } + parts := strings.Split(strings.TrimPrefix(u.Path, "/"), "/") + if len(parts) < 2 || len(parts) > 3 || parts[0] != "secrets" || !validAzureSecretPart(parts[1]) { + err = fmt.Errorf("azure credential name must point at /secrets/ with an optional version") + return + } + secretName = parts[1] + if len(parts) == 3 { + if !validAzureSecretPart(parts[2]) { + err = fmt.Errorf("azure credential URI has an invalid secret version") + return + } + version = parts[2] + } else if requestedVersion != "" { + if !validAzureSecretPart(requestedVersion) { + err = fmt.Errorf("azure credential version %q is invalid", requestedVersion) + return + } + version = requestedVersion + } + u.Scheme = "https" + u.Host = host + u.Path = "/secrets/" + secretName + if version != "" { + u.Path += "/" + version + } + uri = u.String() + return +} + +func validAzureVaultName(name string) bool { + if len(name) < 3 || len(name) > 24 || name[0] < 'a' || name[0] > 'z' { + return false + } + if last := name[len(name)-1]; !asciiAlphaNumeric(last) { + return false + } + for i := 1; i < len(name)-1; i++ { + if !asciiAlphaNumeric(name[i]) && name[i] != '-' { + return false + } + } + return true +} + +func validAzureSecretPart(part string) bool { + if len(part) == 0 || len(part) > 127 { + return false + } + for i := range part { + if !asciiAlphaNumeric(part[i]) && part[i] != '-' { + return false + } + } + return true +} + +func asciiAlphaNumeric(b byte) bool { + return b >= 'a' && b <= 'z' || b >= 'A' && b <= 'Z' || b >= '0' && b <= '9' +} + +func isNumeric(s string) bool { + if s == "" { + return false + } + for _, r := range s { + if r < '0' || r > '9' { + return false + } + } + return true +} + func validateHTTPSURL(field, raw string) error { if strings.TrimSpace(raw) == "" { return fmt.Errorf("%s is required", field) @@ -247,12 +485,42 @@ func ValidEnvName(name string) bool { return true } +func (m Model) EnvName() string { + if m.CredentialEnv != "" { + return m.CredentialEnv + } + for _, p := range DefaultProviders() { + if p.Provider == m.Provider { + return p.Env + } + } + n := strings.ToUpper(strings.ReplaceAll(m.Name, "-", "_")) + return "ABOX_MODEL_" + n + "_KEY" +} + +func (m Model) CredentialReference() CredentialRef { + if m.Credential != nil { + return *m.Credential + } + if m.CredentialEnv != "" { + return CredentialRef{Source: "env", Name: m.CredentialEnv} + } + return CredentialRef{Source: "env", Name: m.EnvName()} +} + +func (s MCPServer) CredentialReference() CredentialRef { + if s.Credential != nil { + return *s.Credential + } + return CredentialRef{Source: "env", Name: TokenEnv(s)} +} + func (m Model) ToGuest() protocol.GuestModel { return protocol.GuestModel{ Name: m.Name, Provider: m.Provider, Model: m.Model, - CredentialEnv: m.CredentialEnv, + CredentialEnv: m.EnvName(), BaseURL: m.BaseURL, } } @@ -278,26 +546,6 @@ func (r Resources) Resolved() (vcpu, ram int) { return } -func (c File) SecretsFromEnv() map[string]string { - out := map[string]string{} - for _, env := range ProviderCredentialEnvs() { - if v := os.Getenv(env); v != "" { - out[env] = v - } - } - servers, err := c.ResolvedMCPServers() - if err != nil { - return out - } - for _, s := range servers { - name := TokenEnv(s) - if v := os.Getenv(name); v != "" { - out[name] = v - } - } - return out -} - func GuestImagePath() string { return filepath.Join(ImageDir(), GuestImageName) } @@ -393,13 +641,6 @@ func (c File) ModelNamed(name string) (Model, bool) { return Model{}, false } -func (m Model) CredentialPresent() bool { - if m.CredentialEnv == "" { - return false - } - return strings.TrimSpace(os.Getenv(m.CredentialEnv)) != "" -} - func Path() string { return filepath.Join(Dir(), "config.yaml") } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index f5f3a76..c109699 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -3,6 +3,7 @@ package config import ( "os" "path/filepath" + "strings" "testing" ) @@ -40,18 +41,176 @@ func TestDefaultProvidersDriveDefaults(t *testing.T) { } } -func TestSecretsFromEnv(t *testing.T) { - t.Setenv("XAI_API_KEY", "xk") - t.Setenv("OPENAI_API_KEY", "") +func TestModelEnvName(t *testing.T) { + if got := (Model{Name: "grok-default", Provider: "xai"}).EnvName(); got != "XAI_API_KEY" { + t.Fatalf("canonical provider env: %q", got) + } + if got := (Model{Name: "custom", CredentialEnv: "MY_KEY"}).EnvName(); got != "MY_KEY" { + t.Fatalf("explicit env: %q", got) + } + if got := (Model{Name: "my-model", Provider: "other"}).EnvName(); got != "ABOX_MODEL_MY_MODEL_KEY" { + t.Fatalf("derived env: %q", got) + } +} + +func TestModelCredentialReference(t *testing.T) { + if got := (Model{CredentialEnv: "X"}).CredentialReference(); got != (CredentialRef{Source: "env", Name: "X"}) { + t.Fatalf("alias: %+v", got) + } + if got := (Model{Name: "g", Provider: "xai"}).CredentialReference(); got != (CredentialRef{Source: "env", Name: "XAI_API_KEY"}) { + t.Fatalf("canonical: %+v", got) + } + explicit := CredentialRef{Source: "vault", Name: "secret/abox/x", Field: "api_key"} + if got := (Model{Credential: &explicit}).CredentialReference(); got != explicit { + t.Fatalf("explicit: %+v", got) + } +} + +func TestMCPServerCredentialReference(t *testing.T) { + if got := (MCPServer{Name: "gh"}).CredentialReference(); got != (CredentialRef{Source: "env", Name: "ABOX_MCP_GH_TOKEN"}) { + t.Fatalf("derived: %+v", got) + } + if got := (MCPServer{Name: "gh", CredentialEnv: "GH"}).CredentialReference(); got != (CredentialRef{Source: "env", Name: "GH"}) { + t.Fatalf("alias: %+v", got) + } +} + +func TestValidateRejectsCredentialAndCredentialEnvBoth(t *testing.T) { c := Defaults() - c.MCPServers = []MCPServer{{Name: "gh", URL: "https://api.githubcopilot.com/mcp/", CredentialEnv: "ABOX_MCP_GH_TOKEN"}} - t.Setenv("ABOX_MCP_GH_TOKEN", "mt") - got := c.SecretsFromEnv() - if got["XAI_API_KEY"] != "xk" || got["ABOX_MCP_GH_TOKEN"] != "mt" { - t.Fatalf("%#v", got) + c.Models[0].Credential = &CredentialRef{Source: "env", Name: "X"} + c.Models[0].CredentialEnv = "X" + if err := c.Validate(); err == nil { + t.Fatal("expected both-set rejection") + } + c2 := Defaults() + c2.MCPServers = []MCPServer{{ + Name: "gh", URL: "https://api.githubcopilot.com/mcp/", + CredentialEnv: "GH", Credential: &CredentialRef{Source: "env", Name: "GH"}, + }} + if err := c2.Validate(); err == nil { + t.Fatal("expected mcp both-set rejection") } - if _, ok := got["OPENAI_API_KEY"]; ok { - t.Fatalf("empty openai key leaked: %#v", got) +} + +func TestValidateRejectsUnknownCredentialSource(t *testing.T) { + c := Defaults() + c.Models[0].CredentialEnv = "" + c.Models[0].Credential = &CredentialRef{Source: "kube", Name: "x"} + if err := c.Validate(); err == nil || !strings.Contains(err.Error(), "unknown source") { + t.Fatalf("got %v", err) + } +} + +func TestValidateFieldVersionRules(t *testing.T) { + c := Defaults() + c.Models[0].CredentialEnv = "" + c.Models[0].Credential = &CredentialRef{Source: "keychain", Name: "K", Field: "f"} + if err := c.Validate(); err == nil || !strings.Contains(err.Error(), "field is not supported") { + t.Fatalf("got %v", err) + } + c = Defaults() + c.Models[0].CredentialEnv = "" + c.Models[0].Credential = &CredentialRef{Source: "aws", Name: "prod/x", Version: "1"} + if err := c.Validate(); err == nil || !strings.Contains(err.Error(), "version is not supported") { + t.Fatalf("got %v", err) + } + c = Defaults() + c.Models[0].CredentialEnv = "" + c.Models[0].Credential = &CredentialRef{Source: "vault", Name: "secret/a/b", Field: "api_key", Version: "3"} + if err := c.Validate(); err != nil { + t.Fatalf("vault with field+version should pass: %v", err) + } + c = Defaults() + c.Models[0].CredentialEnv = "" + c.Models[0].Credential = &CredentialRef{Source: "azure", Name: "https://testkv.vault.azure.net/secrets/x", Version: "v1"} + if err := c.Validate(); err != nil { + t.Fatalf("azure with version should pass: %v", err) + } +} + +func TestAzureCredentialRejectsUntrustedAndAmbiguousURLs(t *testing.T) { + for _, raw := range []string{ + "https://testkv.vault.azure.net.attacker.example/secrets/x", + "https://testkv.vault.azure.net@attacker.example/secrets/x", + "https://testkv.vault.azure.net:443/secrets/x", + "https://testkv.vault.azure.net/secrets/x?redirect=https://attacker.example", + "https://testkv.vault.azure.net/secrets/x/one/too-many", + "https://testkv.vault.azure.net/secrets/x%2Fversion", + } { + c := Defaults() + c.Models[0].CredentialEnv = "" + c.Models[0].Credential = &CredentialRef{Source: "azure", Name: raw} + if err := c.Validate(); err == nil { + t.Fatalf("accepted %q", raw) + } + } +} + +func TestAzureCredentialSupportsSovereignClouds(t *testing.T) { + for _, raw := range []string{ + "https://testkv.vault.azure.net/secrets/x", + "https://testkv.vault.usgovcloudapi.net/secrets/x", + "https://testkv.vault.azure.cn/secrets/x", + "https://testkv.vault.microsoftazure.de/secrets/x", + } { + if _, _, _, cloud, err := ParseAzureSecretReference(raw, "version-1"); err != nil { + t.Fatalf("%s: %v", raw, err) + } else if cloud.AuthorityHost == "" || cloud.VaultResource == "" { + t.Fatalf("%s: incomplete cloud %#v", raw, cloud) + } + } +} + +func TestValidateRejectsCredentialDestinationCollisions(t *testing.T) { + c := Defaults() + c.MCPServers = []MCPServer{ + {Name: "one", URL: "https://one.example/mcp", CredentialEnv: "SHARED_TOKEN"}, + {Name: "two", URL: "https://two.example/mcp", CredentialEnv: "SHARED_TOKEN"}, + } + if err := c.Validate(); err == nil || !strings.Contains(err.Error(), "credential destination env") { + t.Fatalf("mcp collision: %v", err) + } + + c = Defaults() + c.MCPServers = []MCPServer{{Name: "model", URL: "https://mcp.example/api", CredentialEnv: "XAI_API_KEY"}} + if err := c.Validate(); err == nil || !strings.Contains(err.Error(), "credential destination env") { + t.Fatalf("model collision: %v", err) + } +} + +func TestValidateRejectsUnsafeKeychainAccount(t *testing.T) { + c := Defaults() + c.Models[0].CredentialEnv = "" + c.Models[0].Credential = &CredentialRef{Source: "keychain", Name: "safe-name; delete"} + if err := c.Validate(); err == nil || !strings.Contains(err.Error(), "invalid keychain credential name") { + t.Fatalf("got %v", err) + } +} + +func TestCredentialYAMLRoundTrip(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("ABOX_HOME", "") + c := Defaults() + c.Models[0].CredentialEnv = "" + c.Models[0].Credential = &CredentialRef{Source: "vault", Name: "secret/abox/grok", Field: "api_key", Version: "4"} + c.MCPServers = []MCPServer{{ + Name: "gh", URL: "https://api.githubcopilot.com/mcp/", + Credential: &CredentialRef{Source: "keychain", Name: "ABOX_MCP_GH_TOKEN"}, + }} + if err := c.Save(); err != nil { + t.Fatal(err) + } + got, _, err := Load() + if err != nil { + t.Fatal(err) + } + want := c.Models[0].Credential + if got.Models[0].Credential == nil || *got.Models[0].Credential != *want { + t.Fatalf("model credential: %+v", got.Models[0].Credential) + } + if got.MCPServers[0].Credential == nil || got.MCPServers[0].Credential.Source != "keychain" { + t.Fatalf("mcp credential: %+v", got.MCPServers[0].Credential) } } @@ -128,6 +287,73 @@ func TestEnsureLayoutSeedsLegacyConfig(t *testing.T) { } } +func TestEnsureLayoutScrubsLegacyCredentialsEvenWhenDestExists(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("ABOX_HOME", "") + legacy := filepath.Join(home, "Library", "Application Support", "ABox") + if err := os.MkdirAll(legacy, 0o700); err != nil { + t.Fatal(err) + } + if err := os.MkdirAll(filepath.Join(home, ".abox"), 0o700); err != nil { + t.Fatal(err) + } + modern := []byte("# ABox credentials. Mode 0600. Do not commit.\nXAI_API_KEY=modern-file-value\n") + if err := os.WriteFile(filepath.Join(home, ".abox", "credentials.env"), modern, 0o600); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(legacy, "credentials.env"), []byte("XAI_API_KEY=legacy-file-value\n"), 0o600); err != nil { + t.Fatal(err) + } + if err := EnsureLayout(); err != nil { + t.Fatal(err) + } + gotModern, err := os.ReadFile(filepath.Join(home, ".abox", "credentials.env")) + if err != nil { + t.Fatal(err) + } + if string(gotModern) != string(modern) { + t.Fatalf("modern credentials rewritten: %q", gotModern) + } + gotLegacy, err := os.ReadFile(filepath.Join(legacy, "credentials.env")) + if err != nil { + t.Fatal(err) + } + if strings.Contains(string(gotLegacy), "legacy-file-value") || strings.Contains(string(gotLegacy), "XAI_API_KEY=") { + t.Fatalf("legacy credentials not scrubbed: %q", gotLegacy) + } +} + +func TestEnsureLayoutSeedsThenScrubsLegacyCredentials(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("ABOX_HOME", "") + legacy := filepath.Join(home, "Library", "Application Support", "ABox") + if err := os.MkdirAll(legacy, 0o700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(legacy, "credentials.env"), []byte("XAI_API_KEY=legacy-file-value\n"), 0o600); err != nil { + t.Fatal(err) + } + if err := EnsureLayout(); err != nil { + t.Fatal(err) + } + gotModern, err := os.ReadFile(filepath.Join(home, ".abox", "credentials.env")) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(gotModern), "legacy-file-value") { + t.Fatalf("missing seeded credentials: %q", gotModern) + } + gotLegacy, err := os.ReadFile(filepath.Join(legacy, "credentials.env")) + if err != nil { + t.Fatal(err) + } + if strings.Contains(string(gotLegacy), "legacy-file-value") { + t.Fatalf("source credentials not scrubbed: %q", gotLegacy) + } +} + func TestDirRespectsABOX_HOME(t *testing.T) { override := t.TempDir() t.Setenv("ABOX_HOME", override) diff --git a/internal/credentials/credentials.go b/internal/credentials/credentials.go index e463eef..4cd0068 100644 --- a/internal/credentials/credentials.go +++ b/internal/credentials/credentials.go @@ -64,6 +64,28 @@ func Save(envName, value string) error { return err } cur[envName] = value + return writeAll(cur) +} + +func Delete(envName string) error { + if envName == "" { + return fmt.Errorf("empty credential name") + } + if !config.ValidEnvName(envName) { + return fmt.Errorf("invalid credential name %q", envName) + } + cur, err := Load() + if err != nil { + return err + } + if _, ok := cur[envName]; !ok { + return nil + } + delete(cur, envName) + return writeAll(cur) +} + +func writeAll(cur map[string]string) error { if err := os.MkdirAll(config.AppSupportDir(), 0o700); err != nil { return err } diff --git a/internal/credentials/credentials_test.go b/internal/credentials/credentials_test.go index aaf6e1a..7c9bf75 100644 --- a/internal/credentials/credentials_test.go +++ b/internal/credentials/credentials_test.go @@ -52,3 +52,46 @@ func TestSavePreservesMCPTokens(t *testing.T) { t.Fatalf("llm key: %#v", got) } } + +func TestDeleteRemovesOneKey(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + if err := Save("XAI_API_KEY", "keep-me"); err != nil { + t.Fatal(err) + } + if err := Save("ABOX_MCP_GITHUB_TOKEN", "drop-me"); err != nil { + t.Fatal(err) + } + if err := Delete("ABOX_MCP_GITHUB_TOKEN"); err != nil { + t.Fatal(err) + } + got, err := Load() + if err != nil { + t.Fatal(err) + } + if got["XAI_API_KEY"] != "keep-me" { + t.Fatalf("kept %#v", got) + } + if _, ok := got["ABOX_MCP_GITHUB_TOKEN"]; ok { + t.Fatalf("deleted key still present: %#v", got) + } +} + +func TestDeleteMissingIsNoop(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + if err := Delete("XAI_API_KEY"); err != nil { + t.Fatal(err) + } + if err := Save("XAI_API_KEY", "keep-me"); err != nil { + t.Fatal(err) + } + if err := Delete("OPENAI_API_KEY"); err != nil { + t.Fatal(err) + } + got, err := Load() + if err != nil { + t.Fatal(err) + } + if got["XAI_API_KEY"] != "keep-me" { + t.Fatalf("got %#v", got) + } +} diff --git a/internal/credsource/aws.go b/internal/credsource/aws.go new file mode 100644 index 0000000..6328505 --- /dev/null +++ b/internal/credsource/aws.go @@ -0,0 +1,221 @@ +package credsource + +import ( + "bytes" + "context" + "crypto/hmac" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "os" + "strings" + "time" +) + +type awsSource struct{} + +const ( + awsRequestTimeout = 15 * time.Second + awsAlgorithm = "AWS4-HMAC-SHA256" + awsService = "secretsmanager" +) + +var newAWSClient = func() *http.Client { + return &http.Client{Timeout: awsRequestTimeout} +} + +func (awsSource) Resolve(ctx context.Context, ref Reference) (Value, error) { + accessKey := strings.TrimSpace(os.Getenv("AWS_ACCESS_KEY_ID")) + secretKey := os.Getenv("AWS_SECRET_ACCESS_KEY") + sessionToken := strings.TrimSpace(os.Getenv("AWS_SESSION_TOKEN")) + region := strings.TrimSpace(os.Getenv("AWS_REGION")) + if region == "" { + region = strings.TrimSpace(os.Getenv("AWS_DEFAULT_REGION")) + } + if accessKey == "" || secretKey == "" { + return Value{}, fmt.Errorf("%w: aws source requires AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY", ErrLocked) + } + if region == "" { + return Value{}, fmt.Errorf("aws source requires AWS_REGION or AWS_DEFAULT_REGION") + } + host := fmt.Sprintf("%s.%s.amazonaws.com", awsService, region) + reqURL := "https://" + host + "/" + if endpoint := strings.TrimSpace(os.Getenv("AWS_ENDPOINT_URL")); endpoint != "" { + u, err := url.Parse(strings.TrimRight(endpoint, "/") + "/") + if err != nil || u.Host == "" { + return Value{}, fmt.Errorf("invalid AWS_ENDPOINT_URL") + } + host = u.Host + reqURL = u.String() + } + body, err := json.Marshal(map[string]string{"SecretId": ref.Name}) + if err != nil { + return Value{}, err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, reqURL, bytes.NewReader(body)) + if err != nil { + return Value{}, err + } + signAWSRequest(req, host, region, accessKey, secretKey, sessionToken, body) + + resp, err := newAWSClient().Do(req) + if err != nil { + return Value{}, fmt.Errorf("aws secrets manager request for %s: %w", ref.Name, err) + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<20)) + if resp.StatusCode >= 300 { + return Value{}, awsError(ref.Name, resp.StatusCode, respBody) + } + var parsed struct { + SecretString string `json:"SecretString"` + SecretBinary string `json:"SecretBinary"` + } + if err := json.Unmarshal(respBody, &parsed); err != nil { + return Value{}, fmt.Errorf("aws secrets manager %s: malformed response", ref.Name) + } + if parsed.SecretString != "" { + return awsFieldBytes(ref.Name, ref.Field, parsed.SecretString) + } + if parsed.SecretBinary != "" { + raw, err := base64.StdEncoding.DecodeString(parsed.SecretBinary) + if err != nil { + return Value{}, fmt.Errorf("aws secrets manager %s: malformed SecretBinary", ref.Name) + } + return Value{Bytes: raw}, nil + } + return Value{}, fmt.Errorf("%w: aws secrets manager %s", ErrNotFound, ref.Name) +} + +func (awsSource) Close() error { return nil } + +func awsFieldBytes(name, field, secretString string) (Value, error) { + if field == "" { + return Value{Bytes: []byte(secretString)}, nil + } + var obj map[string]any + if err := decodeJSONUseNumber([]byte(secretString), &obj); err != nil { + return Value{}, fmt.Errorf("aws secrets manager %s: SecretString is not a JSON object, cannot select field %q", name, field) + } + raw, ok := obj[field] + if !ok { + return Value{}, fmt.Errorf("%w: aws secrets manager %s has no field %q", ErrNotFound, name, field) + } + return Value{Bytes: vaultFieldBytes(raw)}, nil +} + +func awsError(name string, status int, body []byte) error { + var parsed struct { + Type string `json:"__type"` + Message string `json:"message"` + } + _ = json.Unmarshal(body, &parsed) + kind := parsed.Type + if i := strings.LastIndex(kind, "#"); i >= 0 { + kind = kind[i+1:] + } + switch { + case kind == "ResourceNotFoundException" || status == http.StatusNotFound: + return fmt.Errorf("%w: aws secrets manager %s", ErrNotFound, name) + case kind == "AccessDeniedException" || status == http.StatusForbidden: + return fmt.Errorf("aws secrets manager %s: access denied (check the IAM policy for secretsmanager:GetSecretValue)", name) + case kind == "UnrecognizedClientException": + return fmt.Errorf("aws secrets manager %s: invalid credentials (check AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY)", name) + default: + msg := strings.TrimSpace(parsed.Message) + if msg == "" { + msg = strings.TrimSpace(string(body)) + } + return fmt.Errorf("aws secrets manager %s: %d %s", name, status, msg) + } +} + +func signAWSRequest(req *http.Request, host, region, accessKey, secretKey, sessionToken string, body []byte) { + now := time.Now().UTC() + amzDate := now.Format("20060102T150405Z") + dateStamp := now.Format("20060102") + + payloadHash := hex.EncodeToString(sum256(body)) + + req.Header.Set("Content-Type", "application/x-amz-json-1.1") + req.Header.Set("X-Amz-Date", amzDate) + req.Header.Set("X-Amz-Target", awsService+".GetSecretValue") + if sessionToken != "" { + req.Header.Set("X-Amz-Security-Token", sessionToken) + } + req.Host = host + + type hdr struct{ name, value string } + headers := []hdr{ + {"content-type", req.Header.Get("Content-Type")}, + {"host", host}, + {"x-amz-date", amzDate}, + } + if sessionToken != "" { + headers = append(headers, hdr{"x-amz-security-token", sessionToken}) + } + headers = append(headers, hdr{"x-amz-target", req.Header.Get("X-Amz-Target")}) + + var canonHeaders strings.Builder + var signedNames strings.Builder + for i, h := range headers { + canonHeaders.WriteString(h.name) + canonHeaders.WriteString(":") + canonHeaders.WriteString(strings.TrimSpace(h.value)) + canonHeaders.WriteString("\n") + if i > 0 { + signedNames.WriteString(";") + } + signedNames.WriteString(h.name) + } + signedHeaders := signedNames.String() + + canonicalRequest := strings.Join([]string{ + "POST", + "/", + "", + canonHeaders.String(), + signedHeaders, + payloadHash, + }, "\n") + + credentialScope := strings.Join([]string{dateStamp, region, awsService, "aws4_request"}, "/") + stringToSign := strings.Join([]string{ + awsAlgorithm, + amzDate, + credentialScope, + hex.EncodeToString(sum256([]byte(canonicalRequest))), + }, "\n") + + signature := hex.EncodeToString(hmacSHA256( + signingKey(secretKey, dateStamp, region, awsService), + []byte(stringToSign), + )) + req.Header.Set("Authorization", fmt.Sprintf( + "%s Credential=%s/%s, SignedHeaders=%s, Signature=%s", + awsAlgorithm, accessKey, credentialScope, signedHeaders, signature, + )) +} + +func signingKey(secret, dateStamp, region, service string) []byte { + kDate := hmacSHA256([]byte("AWS4"+secret), []byte(dateStamp)) + kRegion := hmacSHA256(kDate, []byte(region)) + kService := hmacSHA256(kRegion, []byte(service)) + return hmacSHA256(kService, []byte("aws4_request")) +} + +func hmacSHA256(key, data []byte) []byte { + h := hmac.New(sha256.New, key) + h.Write(data) + return h.Sum(nil) +} + +func sum256(data []byte) []byte { + s := sha256.Sum256(data) + return s[:] +} diff --git a/internal/credsource/azure.go b/internal/credsource/azure.go new file mode 100644 index 0000000..d7360f2 --- /dev/null +++ b/internal/credsource/azure.go @@ -0,0 +1,151 @@ +package credsource + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "os" + "os/exec" + "strings" + "time" + + "github.com/AdminTurnedDevOps/ABox/internal/config" +) + +type azureSource struct{} + +const ( + azureAPIVersion = "7.5" + azureTokenTimeout = 15 * time.Second +) + +var runAz = func(ctx context.Context, args []string) (stdout string, err error) { + cmd := exec.CommandContext(ctx, "az", args...) + out, err := cmd.Output() + return string(out), err +} + +var newAzureClient = func() *http.Client { + return &http.Client{Timeout: 15 * time.Second} +} + +var azAvailable = func() bool { + _, err := exec.LookPath("az") + return err == nil +} + +func (azureSource) Resolve(ctx context.Context, ref Reference) (Value, error) { + secretURI, name, version, cloud, err := config.ParseAzureSecretReference(ref.Name, ref.Version) + if err != nil { + return Value{}, err + } + token, err := azureToken(ctx, cloud) + if err != nil { + return Value{}, err + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, secretURI+"?api-version="+azureAPIVersion, nil) + if err != nil { + return Value{}, err + } + req.Header.Set("Authorization", "Bearer "+token) + resp, err := newAzureClient().Do(req) + if err != nil { + return Value{}, fmt.Errorf("azure key vault request for %s: %w", name, err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + switch { + case resp.StatusCode == http.StatusNotFound: + return Value{}, fmt.Errorf("%w: azure key vault %s", ErrNotFound, name) + case resp.StatusCode == http.StatusForbidden: + return Value{}, fmt.Errorf("azure key vault %s: permission denied (check the key vault access policy or RBAC role)", name) + case resp.StatusCode >= 300: + return Value{}, fmt.Errorf("azure key vault %s: %s", name, resp.Status) + } + var parsed struct { + Value string `json:"value"` + ID string `json:"id"` + } + if err := json.Unmarshal(body, &parsed); err != nil { + return Value{}, fmt.Errorf("azure key vault %s: malformed response", name) + } + if parsed.Value == "" { + return Value{}, fmt.Errorf("%w: azure key vault %s is empty", ErrNotFound, name) + } + if version == "" { + if parts := strings.Split(parsed.ID, "/"); len(parts) > 0 { + version = parts[len(parts)-1] + } + } + return Value{Bytes: []byte(parsed.Value), Version: version}, nil +} + +func (azureSource) Close() error { return nil } + +func azureToken(ctx context.Context, cloud config.AzureCloud) (string, error) { + clientID := strings.TrimSpace(os.Getenv("AZURE_CLIENT_ID")) + tenantID := strings.TrimSpace(os.Getenv("AZURE_TENANT_ID")) + clientSecret := os.Getenv("AZURE_CLIENT_SECRET") + if clientID != "" && tenantID != "" && clientSecret != "" { + authority := cloud.AuthorityHost + if configured := strings.TrimRight(strings.TrimSpace(os.Getenv("AZURE_AUTHORITY_HOST")), "/"); configured != "" && !strings.EqualFold(configured, authority) { + return "", fmt.Errorf("AZURE_AUTHORITY_HOST %q does not match Key Vault cloud authority %q", configured, authority) + } + tokenCtx, cancel := context.WithTimeout(ctx, azureTokenTimeout) + defer cancel() + form := strings.NewReader(strings.Join([]string{ + "grant_type=client_credentials", + "client_id=" + url.QueryEscape(clientID), + "client_secret=" + url.QueryEscape(clientSecret), + "scope=" + url.QueryEscape(cloud.VaultResource+"/.default"), + }, "&")) + req, err := http.NewRequestWithContext(tokenCtx, http.MethodPost, + authority+"/"+tenantID+"/oauth2/v2.0/token", form) + if err != nil { + return "", err + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + resp, err := newAzureClient().Do(req) + if err != nil { + return "", fmt.Errorf("azure token request: %w", err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if resp.StatusCode >= 300 { + return "", fmt.Errorf("azure token request: %s (check AZURE_CLIENT_ID/AZURE_TENANT_ID/AZURE_CLIENT_SECRET)", resp.Status) + } + var parsed struct { + AccessToken string `json:"access_token"` + } + if err := json.Unmarshal(body, &parsed); err != nil || parsed.AccessToken == "" { + return "", fmt.Errorf("azure token response: malformed") + } + return parsed.AccessToken, nil + } + if !azAvailable() { + return "", fmt.Errorf("%w: azure source needs AZURE_CLIENT_ID+AZURE_TENANT_ID+AZURE_CLIENT_SECRET or the Azure CLI (az login)", ErrLocked) + } + cmdCtx, cancel := context.WithTimeout(ctx, 30*time.Second) + defer cancel() + out, err := runAz(cmdCtx, []string{"account", "get-access-token", "--resource", cloud.VaultResource, "--output", "json"}) + if err != nil { + return "", fmt.Errorf("az account get-access-token failed (run: az login): %w", err) + } + var parsed struct { + AccessToken string `json:"accessToken"` + Token string `json:"token"` + } + if err := json.Unmarshal([]byte(out), &parsed); err != nil { + return "", fmt.Errorf("az token output: malformed") + } + if parsed.AccessToken != "" { + return parsed.AccessToken, nil + } + if parsed.Token != "" { + return parsed.Token, nil + } + return "", fmt.Errorf("az token output: no access token") +} diff --git a/internal/credsource/cloud_test.go b/internal/credsource/cloud_test.go new file mode 100644 index 0000000..84dd2b4 --- /dev/null +++ b/internal/credsource/cloud_test.go @@ -0,0 +1,608 @@ +package credsource + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "strings" + "sync" + "testing" + + "github.com/AdminTurnedDevOps/ABox/internal/config" + "github.com/AdminTurnedDevOps/ABox/internal/credentials" +) + +func newVaultServer(t *testing.T, h http.HandlerFunc) *httptest.Server { + t.Helper() + srv := httptest.NewServer(h) + t.Cleanup(srv.Close) + return srv +} + +// --- Azure Key Vault --- + +func newAzureTLSServer(t *testing.T, h http.HandlerFunc) *httptest.Server { + t.Helper() + srv := httptest.NewTLSServer(h) + t.Cleanup(srv.Close) + target, err := url.Parse(srv.URL) + if err != nil { + t.Fatal(err) + } + orig := newAzureClient + newAzureClient = func() *http.Client { + client := srv.Client() + transport := client.Transport + client.Transport = roundTripFunc(func(req *http.Request) (*http.Response, error) { + clone := req.Clone(req.Context()) + clone.URL.Scheme = target.Scheme + clone.URL.Host = target.Host + return transport.RoundTrip(clone) + }) + return client + } + t.Cleanup(func() { newAzureClient = orig }) + return srv +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) } + +func fakeAzAvailable(t *testing.T) { + t.Helper() + orig := azAvailable + azAvailable = func() bool { return true } + t.Cleanup(func() { azAvailable = orig }) +} + +func TestAzureResolveServicePrincipal(t *testing.T) { + var gotForm string + var gotScope string + var mu sync.Mutex + newAzureTLSServer(t, func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/secrets/anthropic" { + if r.Header.Get("Authorization") != "Bearer sp-token" { + t.Errorf("missing bearer") + } + if r.URL.Query().Get("api-version") != "7.5" { + t.Errorf("api-version %q", r.URL.Query().Get("api-version")) + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "value": "azure-secret", + "id": "https://testkv.vault.azure.net/secrets/anthropic/abc123", + }) + return + } + mu.Lock() + defer mu.Unlock() + gotForm = r.URL.Path + _ = r.ParseForm() + gotScope = r.Form.Get("scope") + _ = json.NewEncoder(w).Encode(map[string]any{"access_token": "sp-token"}) + }) + t.Setenv("AZURE_CLIENT_ID", "client") + t.Setenv("AZURE_TENANT_ID", "tenant") + t.Setenv("AZURE_CLIENT_SECRET", "sp-secret") + t.Setenv("AZURE_AUTHORITY_HOST", "") + + v, err := testResolver().Resolve(context.Background(), Reference{ + Source: "azure", + Name: "https://testkv.vault.azure.net/secrets/anthropic", + }) + if err != nil { + t.Fatal(err) + } + mu.Lock() + path := gotForm + scope := gotScope + mu.Unlock() + if path != "/tenant/oauth2/v2.0/token" { + t.Fatalf("token path %q", path) + } + if scope != "https://vault.azure.net/.default" { + t.Fatalf("scope %q", scope) + } + if string(v.Bytes) != "azure-secret" || v.Version != "abc123" { + t.Fatalf("got %q version %q", v.Bytes, v.Version) + } +} + +func TestAzureResolveAzCLIFallback(t *testing.T) { + fakeAzAvailable(t) + t.Setenv("AZURE_CLIENT_ID", "") + t.Setenv("AZURE_TENANT_ID", "") + t.Setenv("AZURE_CLIENT_SECRET", "") + orig := runAz + runAz = func(_ context.Context, args []string) (string, error) { + if len(args) != 6 || args[0] != "account" || args[1] != "get-access-token" || args[2] != "--resource" || args[3] != "https://vault.azure.net" || args[4] != "--output" || args[5] != "json" { + return "", fmt.Errorf("unexpected az args %v", args) + } + return `{"accessToken":"cli-token"}`, nil + } + t.Cleanup(func() { runAz = orig }) + + newAzureTLSServer(t, func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "Bearer cli-token" { + t.Errorf("missing cli bearer") + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "value": "cli-secret", + "id": "https://testkv.vault.azure.net/secrets/anthropic/v2", + }) + }) + v, err := testResolver().Resolve(context.Background(), Reference{ + Source: "azure", + Name: "https://testkv.vault.azure.net/secrets/anthropic", + Version: "v2", + }) + if err != nil { + t.Fatal(err) + } + if string(v.Bytes) != "cli-secret" { + t.Fatalf("got %q", v.Bytes) + } + if v.Version != "v2" { + t.Fatalf("version %q", v.Version) + } +} + +func TestAzureNotFound(t *testing.T) { + fakeAzAvailable(t) + t.Setenv("AZURE_CLIENT_ID", "") + t.Setenv("AZURE_TENANT_ID", "") + t.Setenv("AZURE_CLIENT_SECRET", "") + orig := runAz + runAz = func(_ context.Context, _ []string) (string, error) { + return `{"accessToken":"t"}`, nil + } + t.Cleanup(func() { runAz = orig }) + newAzureTLSServer(t, func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + }) + _, err := testResolver().Resolve(context.Background(), Reference{Source: "azure", Name: "https://testkv.vault.azure.net/secrets/anthropic"}) + if !errors.Is(err, ErrNotFound) { + t.Fatalf("got %v", err) + } +} + +func TestAzureForbidden(t *testing.T) { + fakeAzAvailable(t) + t.Setenv("AZURE_CLIENT_ID", "") + t.Setenv("AZURE_TENANT_ID", "") + t.Setenv("AZURE_CLIENT_SECRET", "") + orig := runAz + runAz = func(_ context.Context, _ []string) (string, error) { + return `{"accessToken":"t"}`, nil + } + t.Cleanup(func() { runAz = orig }) + newAzureTLSServer(t, func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusForbidden) + }) + _, err := testResolver().Resolve(context.Background(), Reference{Source: "azure", Name: "https://testkv.vault.azure.net/secrets/anthropic"}) + if err == nil || !strings.Contains(err.Error(), "permission denied") { + t.Fatalf("got %v", err) + } +} + +func TestAzureRejectsNonURI(t *testing.T) { + _, err := testResolver().Resolve(context.Background(), Reference{Source: "azure", Name: "not-a-uri"}) + if err == nil || !strings.Contains(err.Error(), "https secret URI") { + t.Fatalf("got %v", err) + } +} + +func TestAzureRejectsAttackerHostBeforeFetchingToken(t *testing.T) { + orig := runAz + runAz = func(_ context.Context, _ []string) (string, error) { + t.Fatal("attacker URI reached token acquisition") + return "", nil + } + t.Cleanup(func() { runAz = orig }) + t.Setenv("AZURE_CLIENT_ID", "") + t.Setenv("AZURE_TENANT_ID", "") + t.Setenv("AZURE_CLIENT_SECRET", "") + _, err := testResolver().Resolve(context.Background(), Reference{ + Source: "azure", + Name: "https://testkv.vault.azure.net.attacker.example/secrets/x", + }) + if err == nil || !strings.Contains(err.Error(), "not a supported Azure Key Vault endpoint") { + t.Fatalf("got %v", err) + } +} + +func TestAzureRejectsMismatchedAuthorityBeforeSendingClientSecret(t *testing.T) { + t.Setenv("AZURE_CLIENT_ID", "client") + t.Setenv("AZURE_TENANT_ID", "tenant") + t.Setenv("AZURE_CLIENT_SECRET", "client-secret") + t.Setenv("AZURE_AUTHORITY_HOST", "https://attacker.example") + _, _, _, cloud, err := config.ParseAzureSecretReference("https://testkv.vault.azure.net/secrets/x", "") + if err != nil { + t.Fatal(err) + } + if _, err := azureToken(context.Background(), cloud); err == nil || !strings.Contains(err.Error(), "does not match") { + t.Fatalf("got %v", err) + } +} + +func TestAzureSovereignCLIResources(t *testing.T) { + fakeAzAvailable(t) + t.Setenv("AZURE_CLIENT_ID", "") + t.Setenv("AZURE_TENANT_ID", "") + t.Setenv("AZURE_CLIENT_SECRET", "") + for _, tc := range []struct { + name string + host string + authority string + resource string + }{ + {name: "usgov", host: "testkv.vault.usgovcloudapi.net", authority: "https://login.microsoftonline.us", resource: "https://vault.usgovcloudapi.net"}, + {name: "china", host: "testkv.vault.azure.cn", authority: "https://login.chinacloudapi.cn", resource: "https://vault.azure.cn"}, + {name: "germany", host: "testkv.vault.microsoftazure.de", authority: "https://login.microsoftonline.de", resource: "https://vault.microsoftazure.de"}, + } { + t.Run(tc.name, func(t *testing.T) { + _, _, _, cloud, err := config.ParseAzureSecretReference("https://"+tc.host+"/secrets/x", "") + if err != nil { + t.Fatal(err) + } + if cloud.AuthorityHost != tc.authority { + t.Fatalf("authority %q", cloud.AuthorityHost) + } + orig := runAz + runAz = func(_ context.Context, args []string) (string, error) { + if len(args) != 6 || args[3] != tc.resource || args[4] != "--output" || args[5] != "json" { + t.Fatalf("unexpected az args %v", args) + } + return `{"accessToken":"cli-token"}`, nil + } + t.Cleanup(func() { runAz = orig }) + if _, err := azureToken(context.Background(), cloud); err != nil { + t.Fatal(err) + } + }) + } +} + +func TestAzureNoCredentials(t *testing.T) { + t.Setenv("AZURE_CLIENT_ID", "") + t.Setenv("AZURE_TENANT_ID", "") + t.Setenv("AZURE_CLIENT_SECRET", "") + origAvail := azAvailable + azAvailable = func() bool { return false } + t.Cleanup(func() { azAvailable = origAvail }) + _, err := testResolver().Resolve(context.Background(), Reference{Source: "azure", Name: "https://testkv.vault.azure.net/secrets/x"}) + if !errors.Is(err, ErrLocked) { + t.Fatalf("got %v", err) + } +} + +// --- AWS Secrets Manager --- + +func newAWSServer(t *testing.T, status int, body string, check func(r *http.Request, sawAuth string)) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + auth := r.Header.Get("Authorization") + if r.Header.Get("X-Amz-Target") != "secretsmanager.GetSecretValue" { + t.Errorf("target %q", r.Header.Get("X-Amz-Target")) + } + if r.Header.Get("Content-Type") != "application/x-amz-json-1.1" { + t.Errorf("content type %q", r.Header.Get("Content-Type")) + } + if check != nil { + check(r, auth) + } + w.WriteHeader(status) + _, _ = w.Write([]byte(body)) + })) + t.Cleanup(srv.Close) + t.Setenv("AWS_ENDPOINT_URL", srv.URL) + return srv +} + +func TestAWSResolveSigV4(t *testing.T) { + var body map[string]string + var mu sync.Mutex + newAWSServer(t, 200, `{"SecretString":"aws-secret"}`, func(r *http.Request, auth string) { + mu.Lock() + defer mu.Unlock() + _ = json.NewDecoder(r.Body).Decode(&body) + if !strings.HasPrefix(auth, "AWS4-HMAC-SHA256 Credential=AKID-TEST/") { + t.Errorf("authorization %q", auth) + } + if !strings.Contains(auth, "SignedHeaders=content-type;host;x-amz-date;x-amz-target") { + t.Errorf("signed headers missing: %q", auth) + } + if strings.Contains(auth, "aws-secret-key") { + t.Errorf("secret key leaked into Authorization") + } + }) + t.Setenv("AWS_ACCESS_KEY_ID", "AKID-TEST") + t.Setenv("AWS_SECRET_ACCESS_KEY", "aws-secret-key") + t.Setenv("AWS_SESSION_TOKEN", "") + t.Setenv("AWS_REGION", "us-east-1") + + v, err := testResolver().Resolve(context.Background(), Reference{Source: "aws", Name: "prod/anthropic"}) + if err != nil { + t.Fatal(err) + } + if string(v.Bytes) != "aws-secret" { + t.Fatalf("got %q", v.Bytes) + } + mu.Lock() + defer mu.Unlock() + if body["SecretId"] != "prod/anthropic" { + t.Fatalf("body %v", body) + } +} + +func TestAWSResolveSessionTokenAndField(t *testing.T) { + var sawAuth string + var mu sync.Mutex + newAWSServer(t, 200, `{"SecretString":"{\"api_key\":\"nested\",\"other\":\"x\"}"}`, func(r *http.Request, auth string) { + mu.Lock() + defer mu.Unlock() + sawAuth = auth + }) + t.Setenv("AWS_ACCESS_KEY_ID", "AKID-TEST") + t.Setenv("AWS_SECRET_ACCESS_KEY", "aws-secret-key") + t.Setenv("AWS_SESSION_TOKEN", "session-tok") + t.Setenv("AWS_REGION", "eu-west-1") + + v, err := testResolver().Resolve(context.Background(), Reference{Source: "aws", Name: "multikey", Field: "api_key"}) + if err != nil { + t.Fatal(err) + } + if string(v.Bytes) != "nested" { + t.Fatalf("got %q", v.Bytes) + } + mu.Lock() + auth := sawAuth + mu.Unlock() + if !strings.Contains(auth, "x-amz-security-token") { + t.Fatalf("session token not signed: %q", auth) + } + if !strings.Contains(auth, "eu-west-1") { + t.Fatalf("region not in scope: %q", auth) + } +} + +func TestAWSNumericFieldPreservesPrecision(t *testing.T) { + const large = "9007199254740993123456789" + newAWSServer(t, 200, `{"SecretString":"{\"large\":`+large+`}"}`, nil) + t.Setenv("AWS_ACCESS_KEY_ID", "AKID-TEST") + t.Setenv("AWS_SECRET_ACCESS_KEY", "aws-secret-key") + t.Setenv("AWS_SESSION_TOKEN", "") + t.Setenv("AWS_REGION", "us-east-1") + + v, err := testResolver().Resolve(context.Background(), Reference{Source: "aws", Name: "numeric", Field: "large"}) + if err != nil { + t.Fatal(err) + } + if string(v.Bytes) != large { + t.Fatalf("got %q", v.Bytes) + } +} + +func TestAWSNotFound(t *testing.T) { + newAWSServer(t, 400, `{"__type":"com.amazonaws.secretsmanager#ResourceNotFoundException"}`, nil) + t.Setenv("AWS_ACCESS_KEY_ID", "AKID-TEST") + t.Setenv("AWS_SECRET_ACCESS_KEY", "k") + t.Setenv("AWS_SESSION_TOKEN", "") + t.Setenv("AWS_REGION", "us-east-1") + _, err := testResolver().Resolve(context.Background(), Reference{Source: "aws", Name: "missing"}) + if !errors.Is(err, ErrNotFound) { + t.Fatalf("got %v", err) + } +} + +func TestAWSAccessDenied(t *testing.T) { + newAWSServer(t, 400, `{"__type":"AccessDeniedException","message":"no"}`, nil) + t.Setenv("AWS_ACCESS_KEY_ID", "AKID-TEST") + t.Setenv("AWS_SECRET_ACCESS_KEY", "k") + t.Setenv("AWS_SESSION_TOKEN", "") + t.Setenv("AWS_REGION", "us-east-1") + _, err := testResolver().Resolve(context.Background(), Reference{Source: "aws", Name: "nope"}) + if err == nil || !strings.Contains(err.Error(), "access denied") { + t.Fatalf("got %v", err) + } +} + +func TestAWSMissingCredentials(t *testing.T) { + t.Setenv("AWS_ACCESS_KEY_ID", "") + t.Setenv("AWS_SECRET_ACCESS_KEY", "") + _, err := testResolver().Resolve(context.Background(), Reference{Source: "aws", Name: "x"}) + if !errors.Is(err, ErrLocked) { + t.Fatalf("got %v", err) + } +} + +func TestAWSMissingRegion(t *testing.T) { + t.Setenv("AWS_ACCESS_KEY_ID", "AKID") + t.Setenv("AWS_SECRET_ACCESS_KEY", "k") + t.Setenv("AWS_REGION", "") + t.Setenv("AWS_DEFAULT_REGION", "") + _, err := testResolver().Resolve(context.Background(), Reference{Source: "aws", Name: "x"}) + if err == nil || !strings.Contains(err.Error(), "AWS_REGION") { + t.Fatalf("got %v", err) + } +} + +// --- ResolveSelected --- + +type mapSource struct { + mu sync.Mutex + vals map[string]string + errs map[string]error + calls []string +} + +func (m *mapSource) Resolve(_ context.Context, ref Reference) (Value, error) { + m.mu.Lock() + defer m.mu.Unlock() + m.calls = append(m.calls, ref.Source+"/"+ref.Name) + if err := m.errs[ref.Source+"/"+ref.Name]; err != nil { + return Value{}, err + } + if v, ok := m.vals[ref.Source+"/"+ref.Name]; ok { + return Value{Bytes: []byte(v)}, nil + } + return Value{}, ErrNotFound +} + +func (m *mapSource) Close() error { return nil } + +func TestResolveSelectedOnlySelectedModel(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + if err := credentials.Save("XAI_API_KEY", "xk"); err != nil { + t.Fatal(err) + } + // OPENAI_API_KEY is present too; it must NOT be resolved or returned. + if err := credentials.Save("OPENAI_API_KEY", "ok"); err != nil { + t.Fatal(err) + } + + r := testResolver() + cfg := config.Defaults() + cfg.Connectivity.Mode = "direct" + cfg.MCPServers = []config.MCPServer{ + {Name: "gh", URL: "https://api.githubcopilot.com/mcp/", CredentialEnv: "ABOX_MCP_GH_TOKEN"}, + } + model := config.Model{Name: "grok-default", Provider: "xai", CredentialEnv: "XAI_API_KEY"} + t.Setenv("ABOX_MCP_GH_TOKEN", "mtok") + + got, err := ResolveSelected(context.Background(), r, cfg, model) + if err != nil { + t.Fatal(err) + } + if got["XAI_API_KEY"] != "xk" { + t.Fatalf("model key missing: %#v", got) + } + if got["ABOX_MCP_GH_TOKEN"] != "mtok" { + t.Fatalf("mcp token missing: %#v", got) + } + if _, ok := got["OPENAI_API_KEY"]; ok { + t.Fatalf("unselected model key leaked: %#v", got) + } +} + +func TestResolveSelectedMissingModelError(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + t.Setenv("ABOX_TEST_NEVER_KEY", "") + r := testResolver() + cfg := config.Defaults() + model := config.Model{Name: "m", Provider: "x", CredentialEnv: "ABOX_TEST_NEVER_KEY"} + _, err := ResolveSelected(context.Background(), r, cfg, model) + if err == nil || !strings.Contains(err.Error(), "ABOX_TEST_NEVER_KEY") { + t.Fatalf("got %v", err) + } +} + +func TestResolveSelectedReturnsMCPTokenWithModelError(t *testing.T) { + r := NewResolver() + source := &mapSource{vals: map[string]string{"env/MCP_TOKEN": "mcp-value"}} + r.Register("env", source) + cfg := config.Defaults() + cfg.Connectivity.Mode = "direct" + cfg.MCPServers = []config.MCPServer{{Name: "svc", URL: "https://mcp.example/api", CredentialEnv: "MCP_TOKEN"}} + model := config.Model{Name: "missing", Provider: "other", CredentialEnv: "MODEL_TOKEN"} + got, err := ResolveSelected(context.Background(), r, cfg, model) + if err == nil || !strings.Contains(err.Error(), "MODEL_TOKEN") { + t.Fatalf("got error %v", err) + } + if got["MCP_TOKEN"] != "mcp-value" { + t.Fatalf("partial secrets %#v", got) + } +} + +func TestResolveSelectedReportsMCPBackendErrorAndContinues(t *testing.T) { + backendErr := errors.New("backend unavailable") + r := NewResolver() + source := &mapSource{ + vals: map[string]string{"env/MODEL_TOKEN": "model", "env/GOOD_TOKEN": "good"}, + errs: map[string]error{"env/BAD_TOKEN": backendErr}, + } + r.Register("env", source) + cfg := config.Defaults() + cfg.MCPServers = []config.MCPServer{ + {Name: "bad", URL: "https://bad.example/mcp", CredentialEnv: "BAD_TOKEN"}, + {Name: "good", URL: "https://good.example/mcp", CredentialEnv: "GOOD_TOKEN"}, + } + model := config.Model{Name: "selected", Provider: "other", CredentialEnv: "MODEL_TOKEN"} + got, err := ResolveSelected(context.Background(), r, cfg, model) + if !errors.Is(err, backendErr) || !strings.Contains(err.Error(), `mcp server "bad"`) { + t.Fatalf("got error %v", err) + } + if got["MODEL_TOKEN"] != "model" || got["GOOD_TOKEN"] != "good" { + t.Fatalf("partial secrets %#v", got) + } +} + +func TestResolveSelectedMissingMCPTokenSkipped(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + if err := credentials.Save("XAI_API_KEY", "xk"); err != nil { + t.Fatal(err) + } + r := testResolver() + cfg := config.Defaults() + cfg.Connectivity.Mode = "direct" + cfg.MCPServers = []config.MCPServer{ + {Name: "gh", URL: "https://api.githubcopilot.com/mcp/", CredentialEnv: "ABOX_MCP_MISSING_TOKEN"}, + } + t.Setenv("ABOX_MCP_MISSING_TOKEN", "") + model := config.Model{Name: "grok-default", Provider: "xai", CredentialEnv: "XAI_API_KEY"} + got, err := ResolveSelected(context.Background(), r, cfg, model) + if err != nil { + t.Fatal(err) + } + if got["XAI_API_KEY"] != "xk" { + t.Fatalf("%#v", got) + } + if _, ok := got["ABOX_MCP_MISSING_TOKEN"]; ok { + t.Fatalf("missing token not skipped: %#v", got) + } +} + +func TestResolveSelectedOfflineSkipsMCP(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + if err := credentials.Save("XAI_API_KEY", "xk"); err != nil { + t.Fatal(err) + } + r := testResolver() + cfg := config.Defaults() + cfg.Connectivity.Mode = "offline" + cfg.MCPServers = []config.MCPServer{ + {Name: "gh", URL: "https://api.githubcopilot.com/mcp/", CredentialEnv: "ABOX_MCP_GH_TOKEN"}, + } + model := config.Model{Name: "grok-default", Provider: "xai", CredentialEnv: "XAI_API_KEY"} + got, err := ResolveSelected(context.Background(), r, cfg, model) + if err != nil { + t.Fatal(err) + } + if _, ok := got["ABOX_MCP_GH_TOKEN"]; ok { + t.Fatalf("offline resolved an mcp token: %#v", got) + } +} + +func TestCredentialsFileMode(t *testing.T) { + dir := t.TempDir() + t.Setenv("HOME", dir) + if err := credentials.Save("XAI_API_KEY", "k"); err != nil { + t.Fatal(err) + } + st, err := os.Stat(credentials.Path()) + if err != nil { + t.Fatal(err) + } + if st.Mode().Perm() != 0o600 { + t.Fatalf("perm %o", st.Mode().Perm()) + } + if !strings.HasSuffix(filepath.ToSlash(credentials.Path()), "/.abox/credentials.env") { + t.Fatalf("path %q", credentials.Path()) + } +} diff --git a/internal/credsource/credsource.go b/internal/credsource/credsource.go new file mode 100644 index 0000000..4743a8f --- /dev/null +++ b/internal/credsource/credsource.go @@ -0,0 +1,116 @@ +// Package credsource resolves host-held credentials. Guest packages must not +// import it. Zeroing Values is best-effort: Go's GC copies memory. +package credsource + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + "time" + + "github.com/AdminTurnedDevOps/ABox/internal/config" +) + +type Reference struct { + Source string + Name string + Field string + Version string +} + +type Value struct { + Bytes []byte + Version string + ExpiresAt time.Time + LeaseID string +} + +func (v *Value) Zero() { + if v == nil || v.Bytes == nil { + return + } + for i := range v.Bytes { + v.Bytes[i] = 0 + } + v.Bytes = nil +} + +func (v Value) String() string { return "credsource.Value(redacted)" } + +func (v Value) GoString() string { return "credsource.Value(redacted)" } + +// Source implementations must never include secret values in errors. +type Source interface { + Resolve(ctx context.Context, ref Reference) (Value, error) + Close() error +} + +var ( + ErrNotFound = errors.New("credential not found") + ErrLocked = errors.New("credential store locked or unavailable") +) + +type Resolver struct { + mu sync.Mutex + sources map[string]Source +} + +func NewResolver() *Resolver { + r := &Resolver{sources: map[string]Source{}} + r.Register("env", envSource{}) + r.Register("keychain", keychainSource{}) + r.Register("vault", vaultSource{}) + r.Register("azure", azureSource{}) + r.Register("aws", awsSource{}) + return r +} + +func (r *Resolver) Register(name string, s Source) { + r.mu.Lock() + defer r.mu.Unlock() + r.sources[name] = s +} + +func (r *Resolver) Resolve(ctx context.Context, ref Reference) (Value, error) { + r.mu.Lock() + s, ok := r.sources[ref.Source] + r.mu.Unlock() + if !ok { + return Value{}, fmt.Errorf("unknown credential source %q for %q", ref.Source, ref.Name) + } + if strings.TrimSpace(ref.Name) == "" { + return Value{}, fmt.Errorf("empty credential name for source %q", ref.Source) + } + return s.Resolve(ctx, ref) +} + +func (r *Resolver) Close() error { + r.mu.Lock() + sources := make([]Source, 0, len(r.sources)) + for _, s := range r.sources { + sources = append(sources, s) + } + r.mu.Unlock() + var first error + for _, s := range sources { + if err := s.Close(); err != nil && first == nil { + first = err + } + } + return first +} + +func Present(ctx context.Context, r *Resolver, ref Reference) bool { + v, err := r.Resolve(ctx, ref) + if err != nil { + return false + } + v.Zero() + return true +} + +func FromConfig(c config.CredentialRef) Reference { + return Reference{Source: c.Source, Name: c.Name, Field: c.Field, Version: c.Version} +} diff --git a/internal/credsource/credsource_test.go b/internal/credsource/credsource_test.go new file mode 100644 index 0000000..845f3fd --- /dev/null +++ b/internal/credsource/credsource_test.go @@ -0,0 +1,416 @@ +package credsource + +import ( + "context" + "errors" + "fmt" + "net/http" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/AdminTurnedDevOps/ABox/internal/credentials" +) + +func testResolver() *Resolver { + return NewResolver() +} + +type fakeExitError int + +func (e fakeExitError) Error() string { return fmt.Sprintf("exit status %d", int(e)) } +func (e fakeExitError) ExitCode() int { return int(e) } + +func TestValueStringRedacts(t *testing.T) { + v := Value{Bytes: []byte("super-secret")} + if got := v.String(); got != "credsource.Value(redacted)" { + t.Fatalf("String()=%q", got) + } + if got := strings.TrimSpace(strings.ReplaceAll(strings.ToLower("credsource.Value(redacted)"), " ", "")); strings.Contains(got, "super-secret") { + t.Fatal("String leaked secret") + } + if v.Bytes == nil { + t.Fatal("String must not consume the value") + } +} + +func TestValueZeroOverwrites(t *testing.T) { + v := Value{Bytes: []byte("super-secret")} + v.Zero() + if v.Bytes != nil { + t.Fatalf("Bytes not cleared: %q", string(v.Bytes)) + } +} + +func TestResolverUnknownSource(t *testing.T) { + _, err := testResolver().Resolve(context.Background(), Reference{Source: "nope", Name: "X"}) + if err == nil || !strings.Contains(err.Error(), `unknown credential source "nope"`) { + t.Fatalf("got %v", err) + } +} + +func TestResolverEmptyName(t *testing.T) { + _, err := testResolver().Resolve(context.Background(), Reference{Source: "env", Name: " "}) + if err == nil || !strings.Contains(err.Error(), "empty credential name") { + t.Fatalf("got %v", err) + } +} + +func TestEnvSourceProcessEnvWins(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + t.Setenv("ABOX_TEST_ENV_KEY", "from-process") + if err := credentials.Save("ABOX_TEST_ENV_KEY", "from-file"); err != nil { + t.Fatal(err) + } + v, err := testResolver().Resolve(context.Background(), Reference{Source: "env", Name: "ABOX_TEST_ENV_KEY"}) + if err != nil { + t.Fatal(err) + } + if string(v.Bytes) != "from-process" { + t.Fatalf("got %q", v.Bytes) + } +} + +func TestEnvSourceFallsBackToFile(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + t.Setenv("ABOX_TEST_ENV_KEY2", "") + if err := credentials.Save("ABOX_TEST_ENV_KEY2", "from-file"); err != nil { + t.Fatal(err) + } + v, err := testResolver().Resolve(context.Background(), Reference{Source: "env", Name: "ABOX_TEST_ENV_KEY2"}) + if err != nil { + t.Fatal(err) + } + if string(v.Bytes) != "from-file" { + t.Fatalf("got %q", v.Bytes) + } +} + +func TestEnvSourceNotFound(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + t.Setenv("ABOX_TEST_ENV_MISSING", "") + _, err := testResolver().Resolve(context.Background(), Reference{Source: "env", Name: "ABOX_TEST_ENV_MISSING"}) + if !errors.Is(err, ErrNotFound) { + t.Fatalf("got %v", err) + } +} + +func TestKeychainResolveReadsPasswordOnly(t *testing.T) { + var argv []string + orig := runSecurity + runSecurity = func(_ context.Context, args []string, stdin string) (string, string, error) { + argv = args + return "kchain-value\n", "", nil + } + t.Cleanup(func() { runSecurity = orig }) + + v, err := testResolver().Resolve(context.Background(), Reference{Source: "keychain", Name: "ANTHROPIC_API_KEY"}) + if err != nil { + t.Fatal(err) + } + if string(v.Bytes) != "kchain-value" { + t.Fatalf("got %q", v.Bytes) + } + want := []string{"find-generic-password", "-s", "abox", "-a", "ANTHROPIC_API_KEY", "-w"} + if len(argv) != len(want) { + t.Fatalf("argv=%v", argv) + } + for i := range want { + if argv[i] != want[i] { + t.Fatalf("argv=%v", argv) + } + } + for _, a := range argv { + if strings.Contains(a, "kchain-value") { + t.Fatalf("secret value leaked into argv: %v", argv) + } + } +} + +func TestKeychainResolveNotFoundExit44(t *testing.T) { + orig := runSecurity + runSecurity = func(_ context.Context, args []string, _ string) (string, string, error) { + return "", "could not be found", fakeExitError(44) + } + t.Cleanup(func() { runSecurity = orig }) + + _, err := testResolver().Resolve(context.Background(), Reference{Source: "keychain", Name: "NOPE"}) + if !errors.Is(err, ErrNotFound) { + t.Fatalf("got %v", err) + } +} + +func TestKeychainResolveLocked(t *testing.T) { + orig := runSecurity + runSecurity = func(_ context.Context, _ []string, _ string) (string, string, error) { + return "", "security: SecKeychainSearchCopyNext(): User interaction is not allowed.", fakeExitError(1) + } + t.Cleanup(func() { runSecurity = orig }) + + _, err := testResolver().Resolve(context.Background(), Reference{Source: "keychain", Name: "X"}) + if !errors.Is(err, ErrLocked) { + t.Fatalf("got %v", err) + } +} + +func TestKeychainSetUsesStdinHexNoArgvLeak(t *testing.T) { + var argv, stdin, stderr []string + orig := runSecurity + runSecurity = func(_ context.Context, args []string, in string) (string, string, error) { + argv = args + stdin = append(stdin, in) + return "", "", nil + } + t.Cleanup(func() { runSecurity = orig }) + + if err := SetKeychain(context.Background(), "ANTHROPIC_API_KEY", []byte("plain-key")); err != nil { + t.Fatal(err) + } + if len(argv) != 1 || argv[0] != "-i" { + t.Fatalf("argv=%v, want only [\"-i\"]", argv) + } + if len(stdin) != 1 { + t.Fatalf("stdin writes: %d", len(stdin)) + } + cmd := stdin[0] + if strings.Contains(cmd, "plain-key") { + t.Fatal("plaintext secret in security stdin command") + } + if !strings.Contains(cmd, "add-generic-password -U -s abox -a ANTHROPIC_API_KEY -X") { + t.Fatalf("stdin command: %q", cmd) + } + _ = stderr +} + +func TestKeychainRejectsCommandInputAccountName(t *testing.T) { + called := false + orig := runSecurity + runSecurity = func(_ context.Context, _ []string, _ string) (string, string, error) { + called = true + return "", "", nil + } + t.Cleanup(func() { runSecurity = orig }) + + err := SetKeychain(context.Background(), "SAFE_NAME\n delete-generic-password", []byte("value")) + if err == nil || !strings.Contains(err.Error(), "invalid keychain account name") { + t.Fatalf("got %v", err) + } + if called { + t.Fatal("security invoked for invalid account name") + } +} + +func TestKeychainCommandErrorWrapsCause(t *testing.T) { + cause := errors.New("security failed") + orig := runSecurity + runSecurity = func(_ context.Context, _ []string, _ string) (string, string, error) { + return "", "diagnostic", cause + } + t.Cleanup(func() { runSecurity = orig }) + + err := SetKeychain(context.Background(), "SAFE_NAME", []byte("value")) + if !errors.Is(err, cause) || !strings.Contains(err.Error(), "diagnostic") { + t.Fatalf("got %v", err) + } +} + +func TestKeychainMissingToolIsUnavailableAndWrapsCause(t *testing.T) { + cause := &os.PathError{Op: "fork/exec", Path: "/usr/bin/security", Err: os.ErrNotExist} + orig := runSecurity + runSecurity = func(_ context.Context, _ []string, _ string) (string, string, error) { + return "", "", cause + } + t.Cleanup(func() { runSecurity = orig }) + + _, err := testResolver().Resolve(context.Background(), Reference{Source: "keychain", Name: "SAFE_NAME"}) + if !errors.Is(err, ErrLocked) || !errors.Is(err, os.ErrNotExist) { + t.Fatalf("got %v", err) + } +} + +func TestSecurityToolAvailable(t *testing.T) { + if !securityToolAvailable("darwin", 0o755) { + t.Fatal("executable regular file should be available on darwin") + } + if securityToolAvailable("linux", 0o755) || securityToolAvailable("darwin", os.ModeDir|0o755) || securityToolAvailable("darwin", 0o644) { + t.Fatal("unsupported OS, directory, or non-executable file reported available") + } +} + +func TestSavePreferredReportsKeychainSource(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + origEnabled := KeychainEnabled + origSecurity := runSecurity + KeychainEnabled = func() bool { return true } + runSecurity = func(_ context.Context, args []string, _ string) (string, string, error) { + if len(args) != 1 || args[0] != "-i" { + t.Fatalf("args %v", args) + } + return "", "", nil + } + t.Cleanup(func() { + KeychainEnabled = origEnabled + runSecurity = origSecurity + }) + + result, err := SavePreferred(context.Background(), "SAFE_NAME", "value") + if err != nil { + t.Fatal(err) + } + if result.Source != "keychain" || !result.Keychain { + t.Fatalf("result %#v", result) + } +} + +func TestSavePreferredRemovesLeftoverFileEntry(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + if err := credentials.Save("SAFE_NAME", "old-file-value"); err != nil { + t.Fatal(err) + } + origEnabled := KeychainEnabled + origSecurity := runSecurity + KeychainEnabled = func() bool { return true } + runSecurity = func(_ context.Context, args []string, _ string) (string, string, error) { + return "", "", nil + } + t.Cleanup(func() { + KeychainEnabled = origEnabled + runSecurity = origSecurity + }) + + if _, err := SavePreferred(context.Background(), "SAFE_NAME", "new-keychain-value"); err != nil { + t.Fatal(err) + } + got, err := credentials.Load() + if err != nil { + t.Fatal(err) + } + if _, ok := got["SAFE_NAME"]; ok { + t.Fatalf("leftover file entry %#v", got) + } +} + +func TestKeychainDelete(t *testing.T) { + var argv []string + orig := runSecurity + runSecurity = func(_ context.Context, args []string, _ string) (string, string, error) { + argv = args + return "", "", nil + } + t.Cleanup(func() { runSecurity = orig }) + + if err := DeleteKeychain(context.Background(), "ANTHROPIC_API_KEY"); err != nil { + t.Fatal(err) + } + want := []string{"delete-generic-password", "-s", "abox", "-a", "ANTHROPIC_API_KEY"} + if len(argv) != len(want) { + t.Fatalf("argv=%v", argv) + } + for i := range want { + if argv[i] != want[i] { + t.Fatalf("argv=%v", argv) + } + } +} + +func TestVaultResolvePathAndToken(t *testing.T) { + srv := newVaultServer(t, func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/secret/data/abox/anthropic" { + t.Errorf("path %s", r.URL.Path) + } + if r.Header.Get("X-Vault-Token") != "vault-token" { + t.Errorf("token header missing") + } + if r.URL.Query().Get("version") != "" { + t.Errorf("unexpected version %q", r.URL.Query().Get("version")) + } + _, _ = w.Write([]byte(`{"data":{"data":{"value":"vv"},"metadata":{"version":7}}}`)) + }) + t.Setenv("VAULT_ADDR", srv.URL) + t.Setenv("VAULT_TOKEN", "vault-token") + v, err := testResolver().Resolve(context.Background(), Reference{Source: "vault", Name: "secret/abox/anthropic"}) + if err != nil { + t.Fatal(err) + } + if string(v.Bytes) != "vv" || v.Version != "7" { + t.Fatalf("got %q version %q", v.Bytes, v.Version) + } +} + +func TestVaultResolveVersionAndField(t *testing.T) { + srv := newVaultServer(t, func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("version") != "3" { + t.Errorf("version %q", r.URL.Query().Get("version")) + } + _, _ = w.Write([]byte(`{"data":{"data":{"api_key":"fieldval","other":"x"},"metadata":{"version":3}}}`)) + }) + t.Setenv("VAULT_ADDR", srv.URL) + t.Setenv("VAULT_TOKEN", "vault-token") + v, err := testResolver().Resolve(context.Background(), Reference{Source: "vault", Name: "secret/abox/anthropic", Field: "api_key", Version: "3"}) + if err != nil { + t.Fatal(err) + } + if string(v.Bytes) != "fieldval" { + t.Fatalf("got %q", v.Bytes) + } +} + +func TestVaultNumericFieldPreservesPrecision(t *testing.T) { + const large = "9007199254740993123456789" + srv := newVaultServer(t, func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(`{"data":{"data":{"large":` + large + `},"metadata":{"version":9}}}`)) + }) + t.Setenv("VAULT_ADDR", srv.URL) + t.Setenv("VAULT_TOKEN", "vault-token") + v, err := testResolver().Resolve(context.Background(), Reference{Source: "vault", Name: "secret/abox/number", Field: "large"}) + if err != nil { + t.Fatal(err) + } + if string(v.Bytes) != large { + t.Fatalf("got %q", v.Bytes) + } +} + +func TestVaultTokenFileFallback(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("VAULT_TOKEN", "") + if err := os.WriteFile(filepath.Join(home, ".vault-token"), []byte("file-token"), 0o600); err != nil { + t.Fatal(err) + } + srv := newVaultServer(t, func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("X-Vault-Token") != "file-token" { + t.Errorf("token from file missing") + } + _, _ = w.Write([]byte(`{"data":{"data":{"value":"ok"},"metadata":{"version":1}}}`)) + }) + t.Setenv("VAULT_ADDR", srv.URL) + v, err := testResolver().Resolve(context.Background(), Reference{Source: "vault", Name: "secret/abox/anthropic"}) + if err != nil { + t.Fatal(err) + } + if string(v.Bytes) != "ok" { + t.Fatalf("got %q", v.Bytes) + } +} + +func TestVaultNotFound(t *testing.T) { + srv := newVaultServer(t, func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + }) + t.Setenv("VAULT_ADDR", srv.URL) + t.Setenv("VAULT_TOKEN", "vault-token") + _, err := testResolver().Resolve(context.Background(), Reference{Source: "vault", Name: "secret/abox/anthropic"}) + if !errors.Is(err, ErrNotFound) { + t.Fatalf("got %v", err) + } +} + +func TestVaultMissingAddr(t *testing.T) { + t.Setenv("VAULT_ADDR", "") + _, err := testResolver().Resolve(context.Background(), Reference{Source: "vault", Name: "secret/abox/anthropic"}) + if err == nil || !strings.Contains(err.Error(), "VAULT_ADDR") { + t.Fatalf("got %v", err) + } +} diff --git a/internal/credsource/env.go b/internal/credsource/env.go new file mode 100644 index 0000000..ebaa9c6 --- /dev/null +++ b/internal/credsource/env.go @@ -0,0 +1,28 @@ +package credsource + +import ( + "context" + "fmt" + "os" + "strings" + + "github.com/AdminTurnedDevOps/ABox/internal/credentials" +) + +type envSource struct{} + +func (envSource) Resolve(_ context.Context, ref Reference) (Value, error) { + if v := strings.TrimSpace(os.Getenv(ref.Name)); v != "" { + return Value{Bytes: []byte(v)}, nil + } + cur, err := credentials.Load() + if err != nil { + return Value{}, fmt.Errorf("load %s: %w", credentials.Path(), err) + } + if v := strings.TrimSpace(cur[ref.Name]); v != "" { + return Value{Bytes: []byte(v)}, nil + } + return Value{}, fmt.Errorf("%w: env %s", ErrNotFound, ref.Name) +} + +func (envSource) Close() error { return nil } diff --git a/internal/credsource/keychain.go b/internal/credsource/keychain.go new file mode 100644 index 0000000..59e3e9f --- /dev/null +++ b/internal/credsource/keychain.go @@ -0,0 +1,131 @@ +package credsource + +import ( + "context" + "encoding/hex" + "errors" + "fmt" + "os" + "os/exec" + "runtime" + "strings" + + "github.com/AdminTurnedDevOps/ABox/internal/config" +) + +// keychainSource shells out to /usr/bin/security (no cgo). Values are ASCII: +// `security find-generic-password -w` hex-encodes anything else. +type keychainSource struct{} + +const KeychainService = "abox" + +var runSecurity = func(ctx context.Context, args []string, stdin string) (stdout, stderr string, err error) { + cmd := exec.CommandContext(ctx, "/usr/bin/security", args...) + if stdin != "" { + cmd.Stdin = strings.NewReader(stdin) + } + var outBuf, errBuf strings.Builder + cmd.Stdout = &outBuf + cmd.Stderr = &errBuf + err = cmd.Run() + return outBuf.String(), errBuf.String(), err +} + +func KeychainAvailable() bool { + info, err := os.Stat("/usr/bin/security") + if err != nil { + return false + } + return securityToolAvailable(runtime.GOOS, info.Mode()) +} + +func securityToolAvailable(goos string, mode os.FileMode) bool { + return goos == "darwin" && mode.IsRegular() && mode.Perm()&0o111 != 0 +} + +func (keychainSource) Resolve(ctx context.Context, ref Reference) (Value, error) { + if !config.ValidEnvName(ref.Name) { + return Value{}, fmt.Errorf("invalid keychain account name %q", ref.Name) + } + stdout, stderr, err := runSecurity(ctx, []string{ + "find-generic-password", "-s", KeychainService, "-a", ref.Name, "-w", + }, "") + if err == nil { + v := strings.TrimSpace(stdout) + if v == "" { + return Value{}, fmt.Errorf("%w: keychain %s is empty", ErrNotFound, ref.Name) + } + return Value{Bytes: []byte(v)}, nil + } + if exitCode(err) == 44 { + return Value{}, fmt.Errorf("%w: keychain %s: %w", ErrNotFound, ref.Name, err) + } + if strings.Contains(stderr, "User interaction is not allowed") { + return Value{}, fmt.Errorf("%w: keychain locked while reading %s (unlock the login keychain or use credential source env): %w", ErrLocked, ref.Name, err) + } + return Value{}, keychainCommandError(ctx, "read", ref.Name, stderr, err) +} + +func (keychainSource) Close() error { return nil } + +func exitCode(err error) int { + var coder interface{ ExitCode() int } + if errors.As(err, &coder) { + return coder.ExitCode() + } + return -1 +} + +// SetKeychain writes via `security -i` with `-X` hex on stdin so the secret +// never appears in argv. Names must be env-var syntax: -i parses a command language. +func SetKeychain(ctx context.Context, name string, value []byte) error { + if !config.ValidEnvName(name) { + return fmt.Errorf("invalid keychain account name %q", name) + } + cmdStr := fmt.Sprintf("add-generic-password -U -s %s -a %s -X %s -j \"managed by abox\"\n", + KeychainService, name, hex.EncodeToString(value)) + _, stderr, err := runSecurity(ctx, []string{"-i"}, cmdStr) + if err != nil { + if strings.Contains(stderr, "User interaction is not allowed") { + return fmt.Errorf("%w: keychain locked while writing %s (unlock the login keychain or use credential source env): %w", ErrLocked, name, err) + } + return keychainCommandError(ctx, "write", name, stderr, err) + } + return nil +} + +func DeleteKeychain(ctx context.Context, name string) error { + if !config.ValidEnvName(name) { + return fmt.Errorf("invalid keychain account name %q", name) + } + _, stderr, err := runSecurity(ctx, []string{ + "delete-generic-password", "-s", KeychainService, "-a", name, + }, "") + if err != nil { + if exitCode(err) == 44 { + return fmt.Errorf("%w: keychain %s: %w", ErrNotFound, name, err) + } + if strings.Contains(stderr, "User interaction is not allowed") { + return fmt.Errorf("%w: keychain locked while deleting %s: %w", ErrLocked, name, err) + } + return keychainCommandError(ctx, "delete", name, stderr, err) + } + return nil +} + +func keychainCommandError(ctx context.Context, action, name, stderr string, err error) error { + if ctxErr := ctx.Err(); ctxErr != nil { + return fmt.Errorf("keychain %s %s: %w", action, name, ctxErr) + } + detail := strings.TrimSpace(stderr) + if errors.Is(err, os.ErrNotExist) || errors.Is(err, os.ErrPermission) { + if detail != "" { + return fmt.Errorf("%w: keychain %s %s: %s: %w", ErrLocked, action, name, detail, err) + } + return fmt.Errorf("%w: keychain %s %s: %w", ErrLocked, action, name, err) + } + if detail != "" { + return fmt.Errorf("keychain %s %s: %s: %w", action, name, detail, err) + } + return fmt.Errorf("keychain %s %s: %w", action, name, err) +} diff --git a/internal/credsource/resolve.go b/internal/credsource/resolve.go new file mode 100644 index 0000000..364fc0e --- /dev/null +++ b/internal/credsource/resolve.go @@ -0,0 +1,45 @@ +package credsource + +import ( + "context" + "errors" + "fmt" + + "github.com/AdminTurnedDevOps/ABox/internal/config" +) + +// ResolveSelected returns the selected model credential plus enabled MCP +// tokens. A missing model credential is an error; a missing MCP token is skipped. +func ResolveSelected(ctx context.Context, r *Resolver, cfg config.File, model config.Model) (map[string]string, error) { + out := map[string]string{} + var resolveErrs []error + ref := model.CredentialReference() + val, err := r.Resolve(ctx, FromConfig(ref)) + if err != nil { + val.Zero() + resolveErrs = append(resolveErrs, fmt.Errorf("credential for model %q (%s %s): %w", model.Name, ref.Source, ref.Name, err)) + } else { + out[model.EnvName()] = string(val.Bytes) + val.Zero() + } + + servers, err := cfg.ResolvedMCPServers() + if err != nil { + resolveErrs = append(resolveErrs, fmt.Errorf("resolve mcp servers: %w", err)) + return out, errors.Join(resolveErrs...) + } + for _, srv := range servers { + sref := srv.CredentialReference() + tok, err := r.Resolve(ctx, FromConfig(sref)) + if err != nil { + tok.Zero() + if !errors.Is(err, ErrNotFound) { + resolveErrs = append(resolveErrs, fmt.Errorf("credential for mcp server %q (%s %s): %w", srv.Name, sref.Source, sref.Name, err)) + } + continue + } + out[config.TokenEnv(srv)] = string(tok.Bytes) + tok.Zero() + } + return out, errors.Join(resolveErrs...) +} diff --git a/internal/credsource/save.go b/internal/credsource/save.go new file mode 100644 index 0000000..30a0f9e --- /dev/null +++ b/internal/credsource/save.go @@ -0,0 +1,47 @@ +package credsource + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/AdminTurnedDevOps/ABox/internal/credentials" +) + +// KeychainEnabled is swapped to false in tests: the macOS keychain is +// machine-wide, so HOME-scoped temp dirs do not isolate it. +var KeychainEnabled = KeychainAvailable + +type SaveResult struct { + Source string + Keychain bool + Note string +} + +func SavePreferred(ctx context.Context, envName, value string) (SaveResult, error) { + ctx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + if KeychainEnabled() { + if err := SetKeychain(ctx, envName, []byte(value)); err != nil { + if !errors.Is(err, ErrLocked) { + return SaveResult{}, fmt.Errorf("keychain: %w", err) + } + if err := credentials.Save(envName, value); err != nil { + return SaveResult{}, err + } + credentials.SetEnv(envName, value) + return SaveResult{Source: "env", Note: "keychain locked; saved to " + credentials.Path()}, nil + } + if err := credentials.Delete(envName); err != nil { + return SaveResult{}, fmt.Errorf("keychain saved %s but leftover file entry could not be removed: %w", envName, err) + } + credentials.SetEnv(envName, value) + return SaveResult{Source: "keychain", Keychain: true, Note: "key saved to macOS keychain (service abox)"}, nil + } + if err := credentials.Save(envName, value); err != nil { + return SaveResult{}, err + } + credentials.SetEnv(envName, value) + return SaveResult{Source: "env", Note: "key saved to " + credentials.Path()}, nil +} diff --git a/internal/credsource/vault.go b/internal/credsource/vault.go new file mode 100644 index 0000000..1457681 --- /dev/null +++ b/internal/credsource/vault.go @@ -0,0 +1,130 @@ +package credsource + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "time" +) + +type vaultSource struct{} + +const vaultRequestTimeout = 15 * time.Second + +func (vaultSource) Resolve(ctx context.Context, ref Reference) (Value, error) { + addr := strings.TrimRight(strings.TrimSpace(os.Getenv("VAULT_ADDR")), "/") + if addr == "" { + return Value{}, fmt.Errorf("vault source requires VAULT_ADDR") + } + token := strings.TrimSpace(os.Getenv("VAULT_TOKEN")) + if token == "" { + if home, err := os.UserHomeDir(); err == nil { + if b, err := os.ReadFile(filepath.Join(home, ".vault-token")); err == nil { + token = strings.TrimSpace(string(b)) + } + } + } + if token == "" { + return Value{}, fmt.Errorf("vault source requires VAULT_TOKEN or ~/.vault-token") + } + mount, rest, found := strings.Cut(strings.Trim(ref.Name, "/"), "/") + if !found || mount == "" || rest == "" { + return Value{}, fmt.Errorf("vault reference %q must be a KV v2 path like secret/abox/name", ref.Name) + } + url := fmt.Sprintf("%s/v1/%s/data/%s", addr, mount, rest) + if ref.Version != "" { + url += "?version=" + strings.TrimPrefix(ref.Version, "?") + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return Value{}, err + } + req.Header.Set("X-Vault-Token", token) + if ns := strings.TrimSpace(os.Getenv("VAULT_NAMESPACE")); ns != "" { + req.Header.Set("X-Vault-Namespace", ns) + } + client := &http.Client{Timeout: vaultRequestTimeout} + resp, err := client.Do(req) + if err != nil { + return Value{}, fmt.Errorf("vault request %s: %w", mount+"/"+rest, err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + switch { + case resp.StatusCode == http.StatusNotFound: + return Value{}, fmt.Errorf("%w: vault %s", ErrNotFound, ref.Name) + case resp.StatusCode == http.StatusForbidden: + return Value{}, fmt.Errorf("vault %s: permission denied (check the token's policy)", ref.Name) + case resp.StatusCode >= 300: + return Value{}, fmt.Errorf("vault %s: %s %s", ref.Name, resp.Status, vaultErrMessage(body)) + } + var parsed struct { + Data struct { + Data map[string]any `json:"data"` + Metadata struct { + Version json.Number `json:"version"` + } `json:"metadata"` + } `json:"data"` + } + if err := decodeJSONUseNumber(body, &parsed); err != nil { + return Value{}, fmt.Errorf("vault %s: malformed response", ref.Name) + } + field := ref.Field + if field == "" { + field = "value" + } + raw, ok := parsed.Data.Data[field] + if !ok { + return Value{}, fmt.Errorf("%w: vault %s has no field %q", ErrNotFound, ref.Name, field) + } + return Value{Bytes: vaultFieldBytes(raw), Version: parsed.Data.Metadata.Version.String()}, nil +} + +func decodeJSONUseNumber(data []byte, dst any) error { + dec := json.NewDecoder(bytes.NewReader(data)) + dec.UseNumber() + if err := dec.Decode(dst); err != nil { + return err + } + var extra any + if err := dec.Decode(&extra); err != io.EOF { + if err == nil { + return fmt.Errorf("multiple JSON values") + } + return err + } + return nil +} + +func (vaultSource) Close() error { return nil } + +func vaultFieldBytes(raw any) []byte { + switch v := raw.(type) { + case string: + return []byte(v) + case json.Number: + return []byte(v.String()) + default: + b, err := json.Marshal(raw) + if err != nil { + return nil + } + return b + } +} + +func vaultErrMessage(body []byte) string { + var parsed struct { + Errors []string `json:"errors"` + } + if err := json.Unmarshal(body, &parsed); err == nil && len(parsed.Errors) > 0 { + return strings.Join(parsed.Errors, "; ") + } + return strings.TrimSpace(string(body)) +} diff --git a/internal/guest/brokerclient/client.go b/internal/guest/brokerclient/client.go new file mode 100644 index 0000000..5cdc20d --- /dev/null +++ b/internal/guest/brokerclient/client.go @@ -0,0 +1,515 @@ +// Package brokerclient is the guest-side client for the host provider broker. +// No credential, base URL, or header exists in the guest. +package brokerclient + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "sync" + "time" + + "github.com/AdminTurnedDevOps/ABox/internal/config" + "github.com/AdminTurnedDevOps/ABox/internal/provider" + "github.com/AdminTurnedDevOps/ABox/protocol" +) + +// ErrHostTooOld: an old host silently drops guest-initiated frames, so a +// proto-3 guest would hang without this check. +var ErrHostTooOld = errors.New("host binary too old; run make build") + +var ErrEventQueueOverflow = errors.New("provider event queue budget exceeded") + +const ( + eventQueueMaxCount = 256 + eventQueueMaxBytes = 4 << 20 + canceledCallGrace = 5 * time.Second +) + +type Client struct { + writeMu sync.Mutex + write func(context.Context, protocol.Frame) error + + mu sync.Mutex + nextID int + hostProto int + closedErr error + pending map[string]*pendingCall + streams map[string]*clientStream +} + +type pendingCall struct { + once sync.Once + ready chan struct{} + frame protocol.Frame +} + +func newPendingCall() *pendingCall { + return &pendingCall{ready: make(chan struct{})} +} + +func (p *pendingCall) complete(frame protocol.Frame) { + p.once.Do(func() { + p.frame = frame + close(p.ready) + }) +} + +type clientStream struct { + out chan provider.Event + notify chan struct{} + abort chan struct{} + settled chan struct{} + + mu sync.Mutex + queue []queuedEvent + queueBytes int + terminalErr error + sealed bool + abortOnce sync.Once + settledOnce sync.Once +} + +type queuedEvent struct { + event provider.Event + size int +} + +func newClientStream() *clientStream { + s := &clientStream{ + out: make(chan provider.Event), notify: make(chan struct{}, 1), + abort: make(chan struct{}), settled: make(chan struct{}), + } + go s.pump() + return s +} + +func (s *clientStream) enqueue(ev provider.Event, terminal bool, size int) error { + s.mu.Lock() + if s.sealed { + s.mu.Unlock() + return errStreamSealed + } + if size < 0 || size > protocol.MaxProviderEvent || len(s.queue) >= eventQueueMaxCount || s.queueBytes+size > eventQueueMaxBytes { + s.mu.Unlock() + return ErrEventQueueOverflow + } + s.queue = append(s.queue, queuedEvent{event: ev, size: size}) + s.queueBytes += size + if terminal { + s.sealed = true + } + s.mu.Unlock() + select { + case s.notify <- struct{}{}: + default: + } + return nil +} + +var errStreamSealed = errors.New("provider stream sealed") + +func (s *clientStream) abortNow() { + s.abortOnce.Do(func() { + s.mu.Lock() + s.queue = nil + s.queueBytes = 0 + s.terminalErr = nil + s.sealed = true + s.mu.Unlock() + close(s.abort) + }) +} + +func (s *clientStream) fail(err error) { + if err == nil { + err = errors.New("broker connection closed") + } + s.mu.Lock() + if s.sealed { + s.mu.Unlock() + return + } + s.sealed = true + s.terminalErr = err + s.mu.Unlock() + select { + case s.notify <- struct{}{}: + default: + } +} + +func (s *clientStream) pump() { + defer func() { + close(s.out) + s.settledOnce.Do(func() { close(s.settled) }) + }() + for { + s.mu.Lock() + if len(s.queue) > 0 { + item := s.queue[0] + s.queue[0] = queuedEvent{} + s.queue = s.queue[1:] + s.queueBytes -= item.size + s.mu.Unlock() + select { + case s.out <- item.event: + case <-s.abort: + return + } + continue + } + if s.terminalErr != nil { + err := s.terminalErr + s.terminalErr = nil + s.mu.Unlock() + select { + case s.out <- provider.Event{Type: "error", Err: err}: + case <-s.abort: + return + } + continue + } + sealed := s.sealed + s.mu.Unlock() + if sealed { + return + } + select { + case <-s.notify: + case <-s.abort: + return + } + } +} + +func New() *Client { + return &Client{ + pending: map[string]*pendingCall{}, + streams: map[string]*clientStream{}, + } +} + +func (c *Client) Attach(write func(protocol.Frame) error) { + if write == nil { + c.AttachContext(nil) + return + } + c.AttachContext(func(_ context.Context, frame protocol.Frame) error { return write(frame) }) +} + +func (c *Client) AttachContext(write func(context.Context, protocol.Frame) error) { + c.writeMu.Lock() + c.write = write + c.writeMu.Unlock() +} + +func (c *Client) SetHostProtocol(p int) { + c.mu.Lock() + c.hostProto = p + c.mu.Unlock() +} + +func (c *Client) hostProtocol() int { + c.mu.Lock() + defer c.mu.Unlock() + return c.hostProto +} + +func (c *Client) HandleFrame(f protocol.Frame) bool { + if f.Method == "provider_event" { + c.dispatchEvent(f) + return true + } + if f.Method == "" && strings.HasPrefix(f.ID, "g-") { + c.mu.Lock() + call := c.pending[f.ID] + if call != nil { + call.complete(f) + } + c.mu.Unlock() + return true + } + return false +} + +func (c *Client) call(ctx context.Context, method string, params any) (protocol.Frame, error) { + c.writeMu.Lock() + write := c.write + c.writeMu.Unlock() + if write == nil { + return protocol.Frame{}, fmt.Errorf("broker client not attached") + } + raw, err := protocol.EncodeParams(params) + if err != nil { + return protocol.Frame{}, err + } + c.mu.Lock() + if c.closedErr != nil { + err := c.closedErr + c.mu.Unlock() + return protocol.Frame{}, err + } + c.nextID++ + id := fmt.Sprintf("g-%d", c.nextID) + call := newPendingCall() + c.pending[id] = call + c.mu.Unlock() + forget := func() { + c.mu.Lock() + if c.pending[id] == call { + delete(c.pending, id) + } + c.mu.Unlock() + } + if err := write(ctx, protocol.Frame{V: protocol.Version, ID: id, Method: method, Params: raw}); err != nil { + forget() + c.Close(err) + return protocol.Frame{}, err + } + select { + case <-ctx.Done(): + select { + case <-call.ready: + forget() + return call.frame, nil + default: + go c.reapCanceledCall(id, call, method) + return protocol.Frame{}, ctx.Err() + } + case <-call.ready: + forget() + return call.frame, nil + } +} + +// A canceled provider_open can still succeed on the host; cancel that stream +// so it does not occupy a slot until idle timeout. +func (c *Client) reapCanceledCall(id string, call *pendingCall, method string) { + timer := time.NewTimer(canceledCallGrace) + defer timer.Stop() + select { + case <-call.ready: + case <-timer.C: + c.forgetPending(id, call) + return + } + c.forgetPending(id, call) + if method == "provider_open" { + c.cancelOpenedFromFrame(call.frame) + } +} + +func (c *Client) forgetPending(id string, call *pendingCall) { + c.mu.Lock() + if c.pending[id] == call { + delete(c.pending, id) + } + c.mu.Unlock() +} + +func (c *Client) cancelOpenedFromFrame(frame protocol.Frame) { + if frame.Error != nil { + return + } + var openRes protocol.ProviderOpenResult + if json.Unmarshal(frame.Result, &openRes) != nil || openRes.StreamID == "" { + return + } + c.mu.Lock() + closed := c.closedErr != nil + c.mu.Unlock() + if closed { + return + } + c.cancelHost(openRes.StreamID) +} + +func (c *Client) Stream(ctx context.Context, model config.Model, messages []provider.Message, tools []provider.ToolSchema) (<-chan provider.Event, error) { + return c.stream(ctx, model, messages, tools, false) +} + +func (c *Client) StreamWithUsage(ctx context.Context, model config.Model, messages []provider.Message, tools []provider.ToolSchema) (<-chan provider.Event, error) { + return c.stream(ctx, model, messages, tools, true) +} + +func (c *Client) stream(ctx context.Context, model config.Model, messages []provider.Message, tools []provider.ToolSchema, rich bool) (<-chan provider.Event, error) { + if c.hostProtocol() < 3 { + return nil, fmt.Errorf("%w (host speaks protocol %d)", ErrHostTooOld, c.hostProtocol()) + } + openCtx, openCancel := context.WithTimeout(ctx, time.Minute) + defer openCancel() + + open, err := c.call(openCtx, "provider_open", protocol.ProviderOpenParams{Model: model.Name, Rich: rich}) + if err != nil { + return nil, fmt.Errorf("provider_open: %w", err) + } + if open.Error != nil { + return nil, open.Error + } + var openRes protocol.ProviderOpenResult + decodeErr := json.Unmarshal(open.Result, &openRes) + if decodeErr != nil || openRes.StreamID == "" { + if openRes.StreamID != "" { + go c.cancelHost(openRes.StreamID) + } + return nil, fmt.Errorf("provider_open: malformed result") + } + streamID := openRes.StreamID + cancelOpened := func() { + go c.cancelHost(streamID) + } + + req := protocol.ProviderRequest{ + Messages: make([]protocol.ProviderMessage, len(messages)), + } + for i, m := range messages { + req.Messages[i] = protocol.ProviderMessage{ + Role: m.Role, Content: m.Content, ToolID: m.ToolID, + ToolName: m.ToolName, ToolArgs: m.ToolArgs, ToolResult: m.ToolResult, + } + } + for _, t := range tools { + req.Tools = append(req.Tools, protocol.ProviderToolSchema{ + Name: t.Name, Description: t.Description, Parameters: t.Parameters, + }) + } + body, err := json.Marshal(req) + if err != nil { + cancelOpened() + return nil, err + } + if len(body) > protocol.MaxProviderRequest { + cancelOpened() + return nil, fmt.Errorf("provider request too large") + } + + // Host may start streaming as soon as it processes the Last chunk. + stream := newClientStream() + c.mu.Lock() + if c.closedErr != nil { + err := c.closedErr + c.mu.Unlock() + stream.abortNow() + cancelOpened() + return nil, err + } + c.streams[streamID] = stream + c.mu.Unlock() + + const chunk = protocol.MaxProviderChunk + for off := 0; off < len(body); off += chunk { + end := off + chunk + if end > len(body) { + end = len(body) + } + sendCtx, cancelSend := context.WithTimeout(ctx, time.Minute) + resp, err := c.call(sendCtx, "provider_send", protocol.ProviderSendParams{ + StreamID: streamID, Data: body[off:end], Last: end == len(body), + }) + cancelSend() + if err != nil { + c.abortStream(streamID, stream) + cancelOpened() + return nil, fmt.Errorf("provider_send: %w", err) + } + if resp.Error != nil { + c.abortStream(streamID, stream) + cancelOpened() + return nil, resp.Error + } + } + + // Also wait on settled so this goroutine does not leak when ctx never cancels. + go func() { + select { + case <-ctx.Done(): + c.abortStream(streamID, stream) + c.cancelHost(streamID) + case <-stream.settled: + } + }() + return stream.out, nil +} + +func (c *Client) abortStream(id string, stream *clientStream) { + c.mu.Lock() + if c.streams[id] == stream { + delete(c.streams, id) + } + c.mu.Unlock() + stream.abortNow() +} + +func (c *Client) cancelHost(streamID string) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _, _ = c.call(ctx, "provider_cancel", protocol.ProviderCancelParams{StreamID: streamID}) +} + +func (c *Client) dispatchEvent(f protocol.Frame) { + var p protocol.ProviderEventParams + if err := json.Unmarshal(f.Params, &p); err != nil { + c.Close(fmt.Errorf("malformed provider event: %w", err)) + return + } + ev := provider.Event{ + Type: p.Type, Text: p.Text, + ToolID: p.ToolID, ToolName: p.ToolName, ToolArgs: p.ToolArgs, + Usage: p.Usage, StopReason: p.StopReason, + } + if p.Err != "" { + ev.Err = errors.New(p.Err) + } else if p.Type == "error" { + ev.Err = errors.New("provider error") + } + terminal := p.Type == "done" || p.Type == "error" + c.mu.Lock() + stream := c.streams[p.StreamID] + if stream == nil { + c.mu.Unlock() + return + } + if err := stream.enqueue(ev, terminal, len(f.Params)); err != nil { + if errors.Is(err, ErrEventQueueOverflow) && c.streams[p.StreamID] == stream { + delete(c.streams, p.StreamID) + c.mu.Unlock() + stream.fail(ErrEventQueueOverflow) + go c.cancelHost(p.StreamID) + return + } + c.mu.Unlock() + return + } + if terminal { + if c.streams[p.StreamID] == stream { + delete(c.streams, p.StreamID) + } + } + c.mu.Unlock() +} + +func (c *Client) Close(err error) { + if err == nil { + err = errors.New("broker connection closed") + } + c.mu.Lock() + if c.closedErr != nil { + c.mu.Unlock() + return + } + c.closedErr = err + pending := c.pending + streams := c.streams + c.pending = map[string]*pendingCall{} + c.streams = map[string]*clientStream{} + c.mu.Unlock() + for _, call := range pending { + call.complete(protocol.Frame{Error: &protocol.Error{Code: "connection", Message: err.Error()}}) + } + for _, stream := range streams { + stream.fail(err) + } +} diff --git a/internal/guest/brokerclient/client_test.go b/internal/guest/brokerclient/client_test.go new file mode 100644 index 0000000..03caeae --- /dev/null +++ b/internal/guest/brokerclient/client_test.go @@ -0,0 +1,540 @@ +package brokerclient + +import ( + "context" + "encoding/json" + "errors" + "net" + "strings" + "sync" + "testing" + "time" + + "github.com/AdminTurnedDevOps/ABox/internal/config" + "github.com/AdminTurnedDevOps/ABox/internal/provider" + "github.com/AdminTurnedDevOps/ABox/protocol" +) + +type hostStub struct { + t *testing.T + conn net.Conn + writeMu sync.Mutex + + mu sync.Mutex + requests []json.RawMessage // reassembled requests + lastSeen []string + openReqs []protocol.ProviderOpenParams + + cancelHook func(streamID string) + sendErr *protocol.Error +} + +func (h *hostStub) write(f protocol.Frame) error { + h.writeMu.Lock() + defer h.writeMu.Unlock() + return protocol.WriteFrame(h.conn, f) +} + +func (h *hostStub) serve() { + var currentStream string + var reassembled []byte + for { + frame, err := protocol.ReadFrame(h.conn) + if err != nil { + return + } + switch frame.Method { + case "provider_open": + p, _ := protocol.DecodeParams[protocol.ProviderOpenParams](frame.Params) + h.mu.Lock() + h.openReqs = append(h.openReqs, p) + h.mu.Unlock() + currentStream = "s1" + res, _ := protocol.EncodeParams(protocol.ProviderOpenResult{StreamID: currentStream}) + _ = h.write(protocol.Frame{ID: frame.ID, Result: res}) + case "provider_send": + p, _ := protocol.DecodeParams[protocol.ProviderSendParams](frame.Params) + h.mu.Lock() + sendErr := h.sendErr + h.mu.Unlock() + if sendErr != nil { + _ = h.write(protocol.Frame{ID: frame.ID, Error: sendErr}) + continue + } + reassembled = append(reassembled, p.Data...) + if p.Last { + h.mu.Lock() + h.lastSeen = append(h.lastSeen, "last") + h.requests = append(h.requests, append([]byte(nil), reassembled...)) + h.mu.Unlock() + reassembled = nil + } + ok, _ := protocol.EncodeParams(map[string]bool{"ok": true}) + _ = h.write(protocol.Frame{ID: frame.ID, Result: ok}) + case "provider_cancel": + p, _ := protocol.DecodeParams[protocol.ProviderCancelParams](frame.Params) + h.mu.Lock() + if h.cancelHook != nil { + h.cancelHook(p.StreamID) + } + h.mu.Unlock() + ok, _ := protocol.EncodeParams(map[string]bool{"ok": true}) + _ = h.write(protocol.Frame{ID: frame.ID, Result: ok}) + } + } +} + +func newClientPair(t *testing.T, hostProto int) (*Client, *hostStub) { + t.Helper() + host, guest := net.Pipe() + t.Cleanup(func() { host.Close(); guest.Close() }) + stub := &hostStub{t: t, conn: host} + go stub.serve() + c := New() + c.SetHostProtocol(hostProto) + c.Attach(func(f protocol.Frame) error { + return protocol.WriteFrame(guest, f) + }) + go func() { + for { + frame, err := protocol.ReadFrame(guest) + if err != nil { + return + } + c.HandleFrame(frame) + } + }() + return c, stub +} + +func TestStreamChunkingAndReassembly(t *testing.T) { + c, stub := newClientPair(t, 3) + big := strings.Repeat("x", protocol.MaxProviderChunk+512) // forces two chunks + msgs := []provider.Message{{Role: "user", Content: big}} + tools := []provider.ToolSchema{{Name: "list_files", Description: "d", Parameters: map[string]any{"type": "object"}}} + + events, err := c.Stream(context.Background(), config.Model{Name: "grok-default", Provider: "xai"}, msgs, tools) + if err != nil { + t.Fatal(err) + } + pushText(t, stub, "s1", "hello ") + pushText(t, stub, "s1", "world") + pushDone(t, stub, "s1") + + var text string + for ev := range events { + if ev.Type == "text" { + text += ev.Text + } + } + if text != "hello world" { + t.Fatalf("text %q", text) + } + stub.mu.Lock() + defer stub.mu.Unlock() + if len(stub.requests) != 1 { + t.Fatalf("reassembled requests %d", len(stub.requests)) + } + var req protocol.ProviderRequest + if err := json.Unmarshal(stub.requests[0], &req); err != nil { + t.Fatal(err) + } + if len(req.Messages) != 1 || req.Messages[0].Content != big { + t.Fatalf("message lost in chunking") + } + if len(req.Tools) != 1 || req.Tools[0].Name != "list_files" { + t.Fatalf("tools lost in chunking") + } + if len(stub.openReqs) != 1 || stub.openReqs[0].Model != "grok-default" { + t.Fatalf("open params %+v", stub.openReqs) + } +} + +func TestStreamWithUsageForwardsRichOpen(t *testing.T) { + c, stub := newClientPair(t, 3) + events, err := c.StreamWithUsage(context.Background(), config.Model{Name: "grok-default"}, nil, nil) + if err != nil { + t.Fatal(err) + } + pushDone(t, stub, "s1") + for range events { + } + stub.mu.Lock() + defer stub.mu.Unlock() + if len(stub.openReqs) != 1 || !stub.openReqs[0].Rich { + t.Fatalf("open params %+v", stub.openReqs) + } +} + +func TestStreamDoesNotDropBufferedEvents(t *testing.T) { + c, stub := newClientPair(t, 3) + events, err := c.Stream(context.Background(), config.Model{Name: "grok-default"}, nil, nil) + if err != nil { + t.Fatal(err) + } + const count = 200 + for i := 0; i < count; i++ { + pushText(t, stub, "s1", "x") + } + pushDone(t, stub, "s1") + got := 0 + for ev := range events { + if ev.Type == "text" { + got++ + } + } + if got != count { + t.Fatalf("got %d events, want %d", got, count) + } +} + +func TestClientStreamQueueBudgets(t *testing.T) { + s := &clientStream{notify: make(chan struct{}, 1)} + for i := 0; i < eventQueueMaxCount; i++ { + if err := s.enqueue(provider.Event{Type: "text"}, false, 1); err != nil { + t.Fatalf("enqueue %d: %v", i, err) + } + } + if err := s.enqueue(provider.Event{Type: "text"}, false, 1); !errors.Is(err, ErrEventQueueOverflow) { + t.Fatalf("count overflow: %v", err) + } + + s = &clientStream{notify: make(chan struct{}, 1)} + for i := 0; i < eventQueueMaxBytes/protocol.MaxProviderEvent; i++ { + if err := s.enqueue(provider.Event{Type: "text"}, false, protocol.MaxProviderEvent); err != nil { + t.Fatalf("byte enqueue %d: %v", i, err) + } + } + if err := s.enqueue(provider.Event{Type: "text"}, false, 1); !errors.Is(err, ErrEventQueueOverflow) { + t.Fatalf("byte overflow: %v", err) + } +} + +func TestStreamQueueOverflowFailsAndCancelsHost(t *testing.T) { + c, stub := newClientPair(t, 3) + canceled := make(chan string, 1) + stub.mu.Lock() + stub.cancelHook = func(id string) { canceled <- id } + stub.mu.Unlock() + events, err := c.Stream(context.Background(), config.Model{Name: "grok-default"}, nil, nil) + if err != nil { + t.Fatal(err) + } + for i := 0; i < eventQueueMaxCount+2; i++ { + p, _ := json.Marshal(protocol.ProviderEventParams{StreamID: "s1", Type: "text", Text: "x"}) + c.HandleFrame(protocol.Frame{Method: "provider_event", Params: p}) + } + + var last provider.Event + for ev := range events { + last = ev + } + if last.Type != "error" || !errors.Is(last.Err, ErrEventQueueOverflow) { + t.Fatalf("terminal event %+v", last) + } + select { + case id := <-canceled: + if id != "s1" { + t.Fatalf("canceled stream %q", id) + } + case <-time.After(2 * time.Second): + t.Fatal("overflow did not cancel host stream") + } +} + +func TestStreamRequiresHostProtocol3(t *testing.T) { + c, _ := newClientPair(t, 2) + _, err := c.Stream(context.Background(), config.Model{Name: "g"}, nil, nil) + if !errors.Is(err, ErrHostTooOld) { + t.Fatalf("got %v", err) + } +} + +func TestStreamCancelForwardsProviderCancel(t *testing.T) { + c, stub := newClientPair(t, 3) + ctx, cancel := context.WithCancel(context.Background()) + canceled := make(chan struct{}, 1) + stub.cancelHook = func(string) { canceled <- struct{}{} } + + events, err := c.Stream(ctx, config.Model{Name: "grok-default"}, []provider.Message{{Role: "user", Content: "q"}}, nil) + if err != nil { + t.Fatal(err) + } + pushText(t, stub, "s1", "partial") + cancel() + select { + case <-canceled: + case <-time.After(2 * time.Second): + t.Fatal("provider_cancel never forwarded") + } + closed := make(chan struct{}) + go func() { + for range events { + } + close(closed) + }() + select { + case <-closed: + case <-time.After(2 * time.Second): + t.Fatal("local stream did not close on cancellation") + } +} + +func TestStreamCancelClosesLocallyWhenHostWriteFails(t *testing.T) { + c, _ := newClientPair(t, 3) + ctx, cancel := context.WithCancel(context.Background()) + events, err := c.Stream(ctx, config.Model{Name: "grok-default"}, nil, nil) + if err != nil { + t.Fatal(err) + } + c.Attach(func(protocol.Frame) error { return errors.New("host unavailable") }) + cancel() + + closed := make(chan struct{}) + go func() { + for range events { + } + close(closed) + }() + select { + case <-closed: + case <-time.After(2 * time.Second): + t.Fatal("local stream remained open after cancellation") + } +} + +func TestCloseIsIdempotentAndClosesStreams(t *testing.T) { + c, _ := newClientPair(t, 3) + events, err := c.Stream(context.Background(), config.Model{Name: "grok-default"}, nil, nil) + if err != nil { + t.Fatal(err) + } + c.Close(errors.New("disconnected")) + c.Close(errors.New("again")) + select { + case ev, ok := <-events: + if !ok || ev.Type != "error" || ev.Err == nil || !strings.Contains(ev.Err.Error(), "disconnected") { + t.Fatalf("terminal event %+v, open=%v", ev, ok) + } + case <-time.After(2 * time.Second): + t.Fatal("connection error was not delivered") + } + select { + case _, ok := <-events: + if ok { + t.Fatal("event channel remained open after terminal error") + } + case <-time.After(2 * time.Second): + t.Fatal("event channel did not close") + } +} + +func TestClosePreservesQueuedEvents(t *testing.T) { + c, _ := newClientPair(t, 3) + events, err := c.Stream(context.Background(), config.Model{Name: "grok-default"}, nil, nil) + if err != nil { + t.Fatal(err) + } + p, _ := json.Marshal(protocol.ProviderEventParams{StreamID: "s1", Type: "text", Text: "before disconnect"}) + c.HandleFrame(protocol.Frame{Method: "provider_event", Params: p}) + c.Close(errors.New("disconnected")) + + var got []provider.Event + for ev := range events { + got = append(got, ev) + } + if len(got) != 2 || got[0].Text != "before disconnect" || got[1].Type != "error" || got[1].Err == nil { + t.Fatalf("events %+v", got) + } +} + +func TestStreamHostOpenErrorPropagates(t *testing.T) { + host, guest := net.Pipe() + t.Cleanup(func() { host.Close(); guest.Close() }) + go func() { + frame, err := protocol.ReadFrame(host) + if err != nil { + return + } + _ = protocol.WriteFrame(host, protocol.Frame{ + ID: frame.ID, Error: &protocol.Error{Code: "host", Message: "unknown model profile \"nope\""}, + }) + }() + c := New() + c.SetHostProtocol(3) + c.Attach(func(f protocol.Frame) error { return protocol.WriteFrame(guest, f) }) + go func() { + for { + frame, err := protocol.ReadFrame(guest) + if err != nil { + return + } + c.HandleFrame(frame) + } + }() + _, err := c.Stream(context.Background(), config.Model{Name: "nope"}, nil, nil) + if err == nil || !strings.Contains(err.Error(), "unknown model profile") { + t.Fatalf("got %v", err) + } +} + +func TestCanceledOpenCancelsLateHostStream(t *testing.T) { + host, guest := net.Pipe() + t.Cleanup(func() { host.Close(); guest.Close() }) + + c := New() + c.SetHostProtocol(3) + c.Attach(func(f protocol.Frame) error { + return protocol.WriteFrame(guest, f) + }) + go func() { + for { + frame, err := protocol.ReadFrame(guest) + if err != nil { + return + } + c.HandleFrame(frame) + } + }() + + openSeen := make(chan struct{}) + releaseOpen := make(chan struct{}) + canceled := make(chan string, 1) + go func() { + for { + frame, err := protocol.ReadFrame(host) + if err != nil { + return + } + switch frame.Method { + case "provider_open": + close(openSeen) + <-releaseOpen + res, _ := protocol.EncodeParams(protocol.ProviderOpenResult{StreamID: "s-late"}) + _ = protocol.WriteFrame(host, protocol.Frame{ID: frame.ID, Result: res}) + case "provider_cancel": + p, _ := protocol.DecodeParams[protocol.ProviderCancelParams](frame.Params) + canceled <- p.StreamID + ok, _ := protocol.EncodeParams(map[string]bool{"ok": true}) + _ = protocol.WriteFrame(host, protocol.Frame{ID: frame.ID, Result: ok}) + default: + ok, _ := protocol.EncodeParams(map[string]bool{"ok": true}) + _ = protocol.WriteFrame(host, protocol.Frame{ID: frame.ID, Result: ok}) + } + } + }() + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + _, err := c.Stream(ctx, config.Model{Name: "g"}, nil, nil) + done <- err + }() + select { + case <-openSeen: + case <-time.After(2 * time.Second): + t.Fatal("provider_open was not written") + } + cancel() + select { + case err := <-done: + if !errors.Is(err, context.Canceled) { + t.Fatalf("got %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("canceled open did not return") + } + close(releaseOpen) + select { + case id := <-canceled: + if id != "s-late" { + t.Fatalf("canceled stream %q", id) + } + case <-time.After(2 * time.Second): + t.Fatal("late provider_open was not canceled") + } +} + +func TestStreamSendErrorCancelsOpenedHostStream(t *testing.T) { + c, stub := newClientPair(t, 3) + canceled := make(chan string, 1) + stub.mu.Lock() + stub.sendErr = &protocol.Error{Code: "host", Message: "send rejected"} + stub.cancelHook = func(id string) { canceled <- id } + stub.mu.Unlock() + _, err := c.Stream(context.Background(), config.Model{Name: "grok-default"}, nil, nil) + if err == nil || !strings.Contains(err.Error(), "send rejected") { + t.Fatalf("got %v", err) + } + select { + case id := <-canceled: + if id != "s1" { + t.Fatalf("canceled stream %q", id) + } + case <-time.After(2 * time.Second): + t.Fatal("opened host stream was not canceled") + } +} + +func TestCloseWakesPendingCall(t *testing.T) { + c := New() + written := make(chan struct{}) + c.Attach(func(protocol.Frame) error { + close(written) + return nil + }) + done := make(chan protocol.Frame, 1) + go func() { + frame, _ := c.call(context.Background(), "provider_open", protocol.ProviderOpenParams{Model: "g"}) + done <- frame + }() + <-written + c.Close(errors.New("disconnected")) + select { + case frame := <-done: + if frame.Error == nil || frame.Error.Code != "connection" { + t.Fatalf("frame %+v", frame) + } + case <-time.After(2 * time.Second): + t.Fatal("pending call was not woken") + } +} + +func TestCallPassesContextToBlockedWriter(t *testing.T) { + c := New() + c.AttachContext(func(ctx context.Context, _ protocol.Frame) error { + <-ctx.Done() + return ctx.Err() + }) + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + done := make(chan error, 1) + go func() { + _, err := c.call(ctx, "provider_cancel", protocol.ProviderCancelParams{StreamID: "s1"}) + done <- err + }() + select { + case err := <-done: + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("got %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("blocked writer ignored context") + } +} + +func pushText(t *testing.T, stub *hostStub, streamID, text string) { + t.Helper() + p, _ := json.Marshal(protocol.ProviderEventParams{StreamID: streamID, Type: "text", Text: text}) + if err := stub.write(protocol.Frame{V: protocol.Version, ID: "h-push", Method: "provider_event", Params: p}); err != nil { + t.Fatal(err) + } +} + +func pushDone(t *testing.T, stub *hostStub, streamID string) { + t.Helper() + p, _ := json.Marshal(protocol.ProviderEventParams{StreamID: streamID, Type: "done"}) + if err := stub.write(protocol.Frame{V: protocol.Version, ID: "h-push", Method: "provider_event", Params: p}); err != nil { + t.Fatal(err) + } +} diff --git a/internal/guest/egress/egress.go b/internal/guest/egress/egress.go index 1d8faa3..6c3305f 100644 --- a/internal/guest/egress/egress.go +++ b/internal/guest/egress/egress.go @@ -1,4 +1,4 @@ -// Package egress allows the guest agent to reach only configured LLM APIs. +// Package egress allowlists guest outbound hosts. package egress import ( @@ -13,11 +13,8 @@ import ( "time" ) -var defaultAllowed = map[string]struct{}{ - "api.x.ai": {}, - "api.openai.com": {}, - "api.anthropic.com": {}, -} +// Empty: protocol-3 guests only allow MCP origins added at boot. +var defaultAllowed = map[string]struct{}{} var ( allowedMu sync.Mutex diff --git a/internal/guest/egress/egress_test.go b/internal/guest/egress/egress_test.go index f68a083..12be3fd 100644 --- a/internal/guest/egress/egress_test.go +++ b/internal/guest/egress/egress_test.go @@ -6,8 +6,12 @@ import ( ) func TestAllowedHosts(t *testing.T) { - if !Allowed("api.x.ai") || !Allowed("api.openai.com") || !Allowed("api.anthropic.com") { - t.Fatal("expected LLM hosts allowed") + t.Cleanup(ResetForTest) + ResetForTest() + // Protocol-3 guests broker LLM traffic through the host: no provider + // host is allowed by default, only configured MCP origins via Allow. + if Allowed("api.x.ai") || Allowed("api.openai.com") || Allowed("api.anthropic.com") { + t.Fatal("expected no LLM hosts allowed by default") } if Allowed("example.com") || Allowed("169.254.169.254") { t.Fatal("expected arbitrary hosts denied") diff --git a/internal/llmbroker/broker.go b/internal/llmbroker/broker.go new file mode 100644 index 0000000..9098292 --- /dev/null +++ b/internal/llmbroker/broker.go @@ -0,0 +1,356 @@ +// Package llmbroker is the host provider broker. The guest sends a model +// alias; this package resolves credentials and dials the provider. The guest +// never sees a credential, base URL, or HTTP header. +package llmbroker + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + "sync" + "time" + + "github.com/AdminTurnedDevOps/ABox/internal/config" + "github.com/AdminTurnedDevOps/ABox/internal/credsource" + "github.com/AdminTurnedDevOps/ABox/internal/provider" + "github.com/AdminTurnedDevOps/ABox/protocol" +) + +const idleTimeout = 5 * time.Minute + +type Broker struct { + cfg config.File + resolver *credsource.Resolver + client *http.Client + logf func(format string, args ...any) + + mu sync.Mutex + next int + streams map[string]*stream +} + +type streamState uint8 + +const ( + streamReceiving streamState = iota + streamStarted + streamFinished +) + +type stream struct { + id string + model config.Model + rich bool + mu sync.Mutex + state streamState + buf bytes.Buffer + ctx context.Context + cancel context.CancelFunc + idleTimer *time.Timer + bytesIn int +} + +func New(cfg config.File, resolver *credsource.Resolver) *Broker { + return &Broker{ + cfg: cfg, + resolver: resolver, + client: &http.Client{Timeout: 5 * time.Minute}, + streams: map[string]*stream{}, + } +} + +func (b *Broker) WithHTTPClient(c *http.Client) *Broker { + if c != nil { + b.client = c + } + return b +} + +func (b *Broker) SetLogger(f func(format string, args ...any)) { + b.logf = f +} + +func (b *Broker) log(format string, args ...any) { + if b.logf != nil { + b.logf(format, args...) + } +} + +func (b *Broker) Handle(ctx context.Context, method string, params json.RawMessage, notify func(method string, params any) error) (any, *protocol.Error) { + switch method { + case "provider_open": + return b.open(ctx, params) + case "provider_send": + return b.send(params, notify) + case "provider_cancel": + return b.cancel(params) + default: + return nil, &protocol.Error{Code: "host", Message: "unknown broker method " + method} + } +} + +func (b *Broker) open(parent context.Context, raw json.RawMessage) (any, *protocol.Error) { + p, err := protocol.DecodeParams[protocol.ProviderOpenParams](raw) + if err != nil { + return nil, &protocol.Error{Code: "host", Message: err.Error()} + } + if b.cfg.Connectivity.Mode == "offline" { + return nil, &protocol.Error{Code: "host", Message: "offline mode: provider access is disabled"} + } + if strings.TrimSpace(p.Model) == "" { + return nil, &protocol.Error{Code: "host", Message: "model alias required"} + } + model, ok := b.cfg.ModelNamed(p.Model) + if !ok { + return nil, &protocol.Error{Code: "host", Message: fmt.Sprintf("unknown model profile %q", p.Model)} + } + b.mu.Lock() + defer b.mu.Unlock() + if len(b.streams) >= protocol.MaxProviderStreams { + return nil, &protocol.Error{Code: "host", Message: "too many open provider streams"} + } + b.next++ + id := fmt.Sprintf("s%d", b.next) + ctx, cancel := context.WithCancel(parent) + st := &stream{id: id, model: model, rich: p.Rich, state: streamReceiving, ctx: ctx, cancel: cancel} + st.idleTimer = time.AfterFunc(idleTimeout, cancel) + b.streams[id] = st + b.log("provider stream %s opened (model=%s rich=%v)", id, p.Model, p.Rich) + go func() { + <-ctx.Done() + b.finish(st) + }() + return protocol.ProviderOpenResult{StreamID: id}, nil +} + +func (b *Broker) send(raw json.RawMessage, notify func(string, any) error) (any, *protocol.Error) { + p, err := protocol.DecodeParams[protocol.ProviderSendParams](raw) + if err != nil { + return nil, &protocol.Error{Code: "host", Message: err.Error()} + } + if p.StreamID == "" { + return nil, &protocol.Error{Code: "host", Message: "provider stream id required"} + } + b.mu.Lock() + st, ok := b.streams[p.StreamID] + b.mu.Unlock() + if !ok { + return nil, &protocol.Error{Code: "host", Message: "unknown provider stream"} + } + st.mu.Lock() + if st.state != streamReceiving { + st.mu.Unlock() + return nil, &protocol.Error{Code: "host", Message: "provider stream already started"} + } + if len(p.Data) > protocol.MaxProviderChunk { + st.mu.Unlock() + b.finish(st) + return nil, &protocol.Error{Code: "host", Message: "provider chunk too large"} + } + if st.buf.Len()+len(p.Data) > protocol.MaxProviderRequest { + st.mu.Unlock() + b.finish(st) + return nil, &protocol.Error{Code: "host", Message: "provider request too large"} + } + _, _ = st.buf.Write(p.Data) + st.bytesIn += len(p.Data) + if st.idleTimer != nil { + st.idleTimer.Reset(idleTimeout) + } + if !p.Last { + st.mu.Unlock() + return map[string]bool{"ok": true}, nil + } + st.state = streamStarted + body := append([]byte(nil), st.buf.Bytes()...) + st.buf = bytes.Buffer{} + st.mu.Unlock() + return b.start(st, body, notify) +} + +func (b *Broker) start(st *stream, body []byte, notify func(string, any) error) (any, *protocol.Error) { + var req protocol.ProviderRequest + if err := json.Unmarshal(body, &req); err != nil { + b.finish(st) + return nil, &protocol.Error{Code: "host", Message: "malformed provider request"} + } + body = nil + if len(req.Messages) > protocol.MaxProviderMessages { + b.finish(st) + return nil, &protocol.Error{Code: "host", Message: "too many provider messages"} + } + if len(req.Tools) > protocol.MaxProviderTools { + b.finish(st) + return nil, &protocol.Error{Code: "host", Message: "too many provider tools"} + } + for _, m := range req.Messages { + if len(m.ToolArgs) > protocol.MaxProviderToolArgs { + b.finish(st) + return nil, &protocol.Error{Code: "host", Message: "tool args too large"} + } + } + msgs := make([]provider.Message, len(req.Messages)) + for i, m := range req.Messages { + msgs[i] = provider.Message{ + Role: m.Role, Content: m.Content, ToolID: m.ToolID, + ToolName: m.ToolName, ToolArgs: m.ToolArgs, ToolResult: m.ToolResult, + } + } + tools := make([]provider.ToolSchema, len(req.Tools)) + for i, t := range req.Tools { + tools[i] = provider.ToolSchema{Name: t.Name, Description: t.Description, Parameters: t.Parameters} + } + if b.resolver == nil { + b.finish(st) + return nil, &protocol.Error{Code: "host", Message: "no credential resolver configured"} + } + ref := st.model.CredentialReference() + val, err := b.resolver.Resolve(st.ctx, credsource.FromConfig(ref)) + if err != nil { + b.finish(st) + return nil, &protocol.Error{Code: "host", Message: fmt.Sprintf("credential for model %q (%s %s): %v", st.model.Name, ref.Source, ref.Name, err)} + } + key := string(val.Bytes) + val.Zero() + + st.mu.Lock() + if st.state != streamStarted || st.ctx.Err() != nil { + st.mu.Unlock() + b.finish(st) + return nil, &protocol.Error{Code: "canceled", Message: "provider stream canceled"} + } + if st.idleTimer == nil { + st.idleTimer = time.AfterFunc(idleTimeout, st.cancel) + } else { + st.idleTimer.Reset(idleTimeout) + } + st.mu.Unlock() + go b.pump(st.ctx, st, key, msgs, tools, notify) + return map[string]bool{"ok": true}, nil +} + +func (b *Broker) pump(ctx context.Context, st *stream, key string, msgs []provider.Message, tools []provider.ToolSchema, notify func(string, any) error) { + var events <-chan provider.Event + var err error + if st.rich { + events, err = provider.StreamWithUsage(ctx, st.model, key, b.client, msgs, tools) + } else { + events, err = provider.Stream(ctx, st.model, key, b.client, msgs, tools) + } + if err != nil { + b.notify(notify, protocol.ProviderEventParams{StreamID: st.id, Type: "error", Err: err.Error()}) + b.finish(st) + return + } + b.forwardEvents(st, events, notify) +} + +func (b *Broker) forwardEvents(st *stream, events <-chan provider.Event, notify func(string, any) error) { + terminalSent := false + eventCount := 0 + draining := false + for ev := range events { + if draining { + continue + } + eventCount++ + if eventCount > protocol.MaxProviderEvents { + st.cancel() + b.notify(notify, protocol.ProviderEventParams{StreamID: st.id, Type: "error", Err: "too many provider events"}) + terminalSent = true + draining = true + continue + } + st.mu.Lock() + if st.state == streamStarted && st.idleTimer != nil { + st.idleTimer.Reset(idleTimeout) + } + st.mu.Unlock() + p := protocol.ProviderEventParams{ + StreamID: st.id, Type: ev.Type, Text: ev.Text, + ToolID: ev.ToolID, ToolName: ev.ToolName, ToolArgs: ev.ToolArgs, + Usage: ev.Usage, StopReason: ev.StopReason, + } + if ev.Err != nil { + p.Err = ev.Err.Error() + } + terminal := ev.Type == "done" || ev.Type == "error" + raw, marshalErr := json.Marshal(p) + if marshalErr != nil || len(raw) > protocol.MaxProviderEvent || len(p.ToolArgs) > protocol.MaxProviderToolArgs { + st.cancel() + b.notify(notify, protocol.ProviderEventParams{StreamID: st.id, Type: "error", Err: "provider event too large"}) + terminalSent = true + draining = true + continue + } + if err := b.notify(notify, p); err != nil { + terminalSent = true + draining = true + st.cancel() + continue + } + if terminal { + terminalSent = true + draining = true + st.cancel() + } + } + if !terminalSent { + b.notify(notify, protocol.ProviderEventParams{StreamID: st.id, Type: "done"}) + } + b.finish(st) +} + +func (b *Broker) notify(notify func(string, any) error, p protocol.ProviderEventParams) error { + if notify == nil { + return fmt.Errorf("provider event notifier unavailable") + } + return notify("provider_event", p) +} + +func (b *Broker) cancel(raw json.RawMessage) (any, *protocol.Error) { + p, err := protocol.DecodeParams[protocol.ProviderCancelParams](raw) + if err != nil { + return nil, &protocol.Error{Code: "host", Message: err.Error()} + } + b.mu.Lock() + st := b.streams[p.StreamID] + b.mu.Unlock() + if st != nil { + b.log("provider stream %s canceled", st.id) + st.mu.Lock() + receiving := st.state == streamReceiving + st.mu.Unlock() + st.cancel() + if receiving { + b.finish(st) + } + } + return map[string]bool{"ok": true}, nil +} + +func (b *Broker) finish(st *stream) { + st.mu.Lock() + if st.state == streamFinished { + st.mu.Unlock() + return + } + st.state = streamFinished + st.buf = bytes.Buffer{} + if st.idleTimer != nil { + st.idleTimer.Stop() + st.idleTimer = nil + } + bytesIn := st.bytesIn + st.mu.Unlock() + st.cancel() + b.mu.Lock() + if b.streams[st.id] == st { + delete(b.streams, st.id) + } + b.mu.Unlock() + b.log("provider stream %s closed (bytes_in=%d)", st.id, bytesIn) +} diff --git a/internal/llmbroker/broker_test.go b/internal/llmbroker/broker_test.go new file mode 100644 index 0000000..06322cb --- /dev/null +++ b/internal/llmbroker/broker_test.go @@ -0,0 +1,577 @@ +package llmbroker + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/AdminTurnedDevOps/ABox/internal/config" + "github.com/AdminTurnedDevOps/ABox/internal/credsource" + "github.com/AdminTurnedDevOps/ABox/internal/provider" + "github.com/AdminTurnedDevOps/ABox/protocol" +) + +type rotatingSource struct { + mu sync.Mutex + vals []string + calls int + refs []string +} + +func (r *rotatingSource) Resolve(_ context.Context, ref credsource.Reference) (credsource.Value, error) { + r.mu.Lock() + defer r.mu.Unlock() + r.refs = append(r.refs, ref.Source+"/"+ref.Name) + if r.calls >= len(r.vals) { + return credsource.Value{}, credsource.ErrNotFound + } + v := r.vals[r.calls] + r.calls++ + return credsource.Value{Bytes: []byte(v)}, nil +} + +func (r *rotatingSource) Close() error { return nil } + +func (r *rotatingSource) count() int { + r.mu.Lock() + defer r.mu.Unlock() + return r.calls +} + +type notifyRecorder struct { + mu sync.Mutex + events []protocol.ProviderEventParams + waiters []chan struct{} +} + +func (n *notifyRecorder) notify(method string, params any) error { + if method != "provider_event" { + return nil + } + raw, err := json.Marshal(params) + if err != nil { + return err + } + var p protocol.ProviderEventParams + _ = json.Unmarshal(raw, &p) + n.mu.Lock() + n.events = append(n.events, p) + waiters := n.waiters + n.waiters = nil + n.mu.Unlock() + for _, w := range waiters { + close(w) + } + return nil +} + +func (n *notifyRecorder) waitFor(t *testing.T, timeout time.Duration, pred func(protocol.ProviderEventParams) bool) protocol.ProviderEventParams { + t.Helper() + deadline := time.After(timeout) + for { + n.mu.Lock() + for _, ev := range n.events { + if pred(ev) { + n.mu.Unlock() + return ev + } + } + w := make(chan struct{}) + n.waiters = append(n.waiters, w) + n.mu.Unlock() + select { + case <-w: + case <-deadline: + t.Fatal("timed out waiting for provider_event") + } + } +} + +func sseProvider(t *testing.T, provider string, stream func(w http.ResponseWriter, auth, body string)) (srv *httptest.Server, authSeen func() []string) { + t.Helper() + var mu sync.Mutex + var auths []string + srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + auth := r.Header.Get("Authorization") + if auth == "" { + auth = r.Header.Get("x-api-key") + } + body, _ := io.ReadAll(r.Body) + mu.Lock() + auths = append(auths, auth) + mu.Unlock() + stream(w, auth, string(body)) + })) + t.Cleanup(srv.Close) + return srv, func() []string { mu.Lock(); defer mu.Unlock(); return append([]string(nil), auths...) } +} + +func openStream(t *testing.T, b *Broker, alias string) string { + t.Helper() + res, perr := b.Handle(context.Background(), "provider_open", mustJSON(t, protocol.ProviderOpenParams{Model: alias}), nil) + if perr != nil { + t.Fatalf("open: %v", perr) + } + raw, err := json.Marshal(res) + if err != nil { + t.Fatal(err) + } + var openRes protocol.ProviderOpenResult + _ = json.Unmarshal(raw, &openRes) + return openRes.StreamID +} + +func mustJSON(t *testing.T, v any) json.RawMessage { + t.Helper() + raw, err := json.Marshal(v) + if err != nil { + t.Fatal(err) + } + return raw +} + +func sendChunk(t *testing.T, b *Broker, rec *notifyRecorder, streamID string, data []byte, last bool) (any, *protocol.Error) { + t.Helper() + return b.Handle(context.Background(), "provider_send", mustJSON(t, protocol.ProviderSendParams{ + StreamID: streamID, Data: data, Last: last, + }), rec.notify) +} + +func defaultCfg() config.File { + cfg := config.Defaults() + cfg.Connectivity.Mode = "direct" + return cfg +} + +func TestBrokerStreamsOpenAITextHostAuth(t *testing.T) { + srv, authSeen := sseProvider(t, "openai", func(w http.ResponseWriter, auth, body string) { + if !strings.HasPrefix(auth, "Bearer k1") { + t.Errorf("host-side auth missing: %q", auth) + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, strings.Join([]string{ + `data: {"choices":[{"delta":{"content":"hi"}}]}`, + `data: [DONE]`, + "", + }, "\n")) + }) + cfg := defaultCfg() + cfg.Models[1].BaseURL = srv.URL // openai-default + r := credsource.NewResolver() + r.Register("rot", &rotatingSource{vals: []string{"k1"}}) + cfg.Models[1].CredentialEnv = "" + cfg.Models[1].Credential = &config.CredentialRef{Source: "rot", Name: "openai"} + b := New(cfg, r).WithHTTPClient(srv.Client()) + + rec := ¬ifyRecorder{} + id := openStream(t, b, "openai-default") + req := protocol.ProviderRequest{Messages: []protocol.ProviderMessage{{Role: "user", Content: "q"}}} + raw, _ := json.Marshal(req) + if _, perr := sendChunk(t, b, rec, id, raw, true); perr != nil { + t.Fatalf("send: %v", perr) + } + ev := rec.waitFor(t, 2*time.Second, func(p protocol.ProviderEventParams) bool { return p.Text == "hi" && p.StreamID == id }) + if ev.Type != "text" { + t.Fatalf("event %+v", ev) + } + rec.waitFor(t, 2*time.Second, func(p protocol.ProviderEventParams) bool { return p.Type == "done" }) + if got := authSeen(); len(got) != 1 { + t.Fatalf("provider calls %d", len(got)) + } +} + +func TestBrokerAnthropicKeyHeader(t *testing.T) { + srv, authSeen := sseProvider(t, "anthropic", func(w http.ResponseWriter, auth, body string) { + if auth != "k-ant" { + t.Errorf("x-api-key %q", auth) + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, strings.Join([]string{ + `data: {"type":"content_block_delta","delta":{"text":"yo"}}`, + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":1}}`, + "", + }, "\n")) + }) + cfg := defaultCfg() + cfg.Models[2].BaseURL = srv.URL // claude-default + r := credsource.NewResolver() + r.Register("rot", &rotatingSource{vals: []string{"k-ant"}}) + cfg.Models[2].CredentialEnv = "" + cfg.Models[2].Credential = &config.CredentialRef{Source: "rot", Name: "anthropic"} + b := New(cfg, r).WithHTTPClient(srv.Client()) + + rec := ¬ifyRecorder{} + id := openStream(t, b, "claude-default") + req := protocol.ProviderRequest{Messages: []protocol.ProviderMessage{{Role: "user", Content: "q"}}} + raw, _ := json.Marshal(req) + if _, perr := sendChunk(t, b, rec, id, raw, true); perr != nil { + t.Fatalf("send: %v", perr) + } + rec.waitFor(t, 2*time.Second, func(p protocol.ProviderEventParams) bool { return p.Text == "yo" }) + rec.waitFor(t, 2*time.Second, func(p protocol.ProviderEventParams) bool { return p.Type == "done" }) + if got := authSeen(); len(got) != 1 || got[0] != "k-ant" { + t.Fatalf("auth %v", got) + } +} + +func TestBrokerUnknownAliasRejected(t *testing.T) { + b := New(defaultCfg(), credsource.NewResolver()) + _, perr := b.Handle(context.Background(), "provider_open", mustJSON(t, providerOpenJSON("nope")), nil) + if perr == nil || !strings.Contains(perr.Message, "unknown model profile") { + t.Fatalf("got %+v", perr) + } +} + +func providerOpenJSON(name string) protocol.ProviderOpenParams { + return protocol.ProviderOpenParams{Model: name} +} + +func TestBrokerOfflineRejected(t *testing.T) { + cfg := defaultCfg() + cfg.Connectivity.Mode = "offline" + b := New(cfg, credsource.NewResolver()) + _, perr := b.Handle(context.Background(), "provider_open", mustJSON(t, providerOpenJSON("openai-default")), nil) + if perr == nil || !strings.Contains(perr.Message, "offline") { + t.Fatalf("got %+v", perr) + } +} + +func TestBrokerUnknownStreamRejected(t *testing.T) { + b := New(defaultCfg(), credsource.NewResolver()) + _, perr := b.Handle(context.Background(), "provider_send", mustJSON(t, protocol.ProviderSendParams{StreamID: "s99", Last: true}), nil) + if perr == nil || !strings.Contains(perr.Message, "unknown provider stream") { + t.Fatalf("got %+v", perr) + } +} + +func TestBrokerChunkBudgetRejected(t *testing.T) { + cfg := defaultCfg() + r := credsource.NewResolver() + b := New(cfg, r) + id := openStream(t, b, "openai-default") + remaining := protocol.MaxProviderRequest - 4 + for remaining > 0 { + n := protocol.MaxProviderChunk + if n > remaining { + n = remaining + } + if _, perr := sendChunk(t, b, ¬ifyRecorder{}, id, make([]byte, n), false); perr != nil { + t.Fatalf("chunk: %v", perr) + } + remaining -= n + } + rec := ¬ifyRecorder{} + _, perr := sendChunk(t, b, rec, id, []byte("overflow"), false) + if perr == nil || !strings.Contains(perr.Message, "too large") { + t.Fatalf("got %+v", perr) + } +} + +func TestBrokerRejectsOversizedChunkAndCleansStream(t *testing.T) { + b := New(defaultCfg(), credsource.NewResolver()) + id := openStream(t, b, "openai-default") + _, perr := sendChunk(t, b, ¬ifyRecorder{}, id, make([]byte, protocol.MaxProviderChunk+1), false) + if perr == nil || !strings.Contains(perr.Message, "chunk too large") { + t.Fatalf("got %+v", perr) + } + _, perr = sendChunk(t, b, ¬ifyRecorder{}, id, []byte(`{}`), true) + if perr == nil || !strings.Contains(perr.Message, "unknown provider stream") { + t.Fatalf("stream was not removed: %+v", perr) + } +} + +func TestBrokerLastStartsProviderExactlyOnce(t *testing.T) { + srv, authSeen := sseProvider(t, "openai", func(w http.ResponseWriter, auth, body string) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "data: [DONE]\n\n") + }) + cfg := defaultCfg() + cfg.Models[1].BaseURL = srv.URL + r := credsource.NewResolver() + r.Register("rot", &rotatingSource{vals: []string{"key"}}) + cfg.Models[1].CredentialEnv = "" + cfg.Models[1].Credential = &config.CredentialRef{Source: "rot", Name: "openai"} + b := New(cfg, r).WithHTTPClient(srv.Client()) + rec := ¬ifyRecorder{} + id := openStream(t, b, "openai-default") + body, _ := json.Marshal(protocol.ProviderRequest{Messages: []protocol.ProviderMessage{{Role: "user", Content: "q"}}}) + if _, perr := sendChunk(t, b, rec, id, body, true); perr != nil { + t.Fatalf("first Last: %v", perr) + } + if _, perr := sendChunk(t, b, rec, id, nil, true); perr == nil { + t.Fatal("repeated Last was accepted") + } + rec.waitFor(t, 2*time.Second, func(p protocol.ProviderEventParams) bool { return p.Type == "done" }) + if got := authSeen(); len(got) != 1 { + t.Fatalf("provider calls %d", len(got)) + } +} + +func TestBrokerRejectsDataAfterStart(t *testing.T) { + srv, _ := sseProvider(t, "openai", func(w http.ResponseWriter, auth, body string) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "data: [DONE]\n\n") + }) + cfg := defaultCfg() + cfg.Models[1].BaseURL = srv.URL + r := credsource.NewResolver() + r.Register("rot", &rotatingSource{vals: []string{"key"}}) + cfg.Models[1].CredentialEnv = "" + cfg.Models[1].Credential = &config.CredentialRef{Source: "rot", Name: "openai"} + b := New(cfg, r).WithHTTPClient(srv.Client()) + id := openStream(t, b, "openai-default") + body, _ := json.Marshal(protocol.ProviderRequest{}) + if _, perr := sendChunk(t, b, ¬ifyRecorder{}, id, body, true); perr != nil { + t.Fatal(perr) + } + if _, perr := sendChunk(t, b, ¬ifyRecorder{}, id, []byte("more"), false); perr == nil { + t.Fatal("data after start was accepted") + } +} + +func TestBrokerRequestCountBounds(t *testing.T) { + b := New(defaultCfg(), credsource.NewResolver()) + id := openStream(t, b, "openai-default") + req := protocol.ProviderRequest{Messages: make([]protocol.ProviderMessage, protocol.MaxProviderMessages+1)} + body, _ := json.Marshal(req) + _, perr := sendChunk(t, b, ¬ifyRecorder{}, id, body, true) + if perr == nil || !strings.Contains(perr.Message, "too many provider messages") { + t.Fatalf("got %+v", perr) + } +} + +func TestBrokerToolCountBounds(t *testing.T) { + b := New(defaultCfg(), credsource.NewResolver()) + id := openStream(t, b, "openai-default") + req := protocol.ProviderRequest{Tools: make([]protocol.ProviderToolSchema, protocol.MaxProviderTools+1)} + body, _ := json.Marshal(req) + _, perr := sendChunk(t, b, ¬ifyRecorder{}, id, body, true) + if perr == nil || !strings.Contains(perr.Message, "too many provider tools") { + t.Fatalf("got %+v", perr) + } +} + +func TestBrokerRejectsOversizedProviderEvent(t *testing.T) { + srv, _ := sseProvider(t, "openai", func(w http.ResponseWriter, auth, body string) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"`+strings.Repeat("x", protocol.MaxProviderEvent)+`"}}]}`+"\n\n") + }) + cfg := defaultCfg() + cfg.Models[1].BaseURL = srv.URL + r := credsource.NewResolver() + r.Register("rot", &rotatingSource{vals: []string{"key"}}) + cfg.Models[1].CredentialEnv = "" + cfg.Models[1].Credential = &config.CredentialRef{Source: "rot", Name: "openai"} + b := New(cfg, r).WithHTTPClient(srv.Client()) + rec := ¬ifyRecorder{} + id := openStream(t, b, "openai-default") + body, _ := json.Marshal(protocol.ProviderRequest{}) + if _, perr := sendChunk(t, b, rec, id, body, true); perr != nil { + t.Fatal(perr) + } + ev := rec.waitFor(t, 2*time.Second, func(p protocol.ProviderEventParams) bool { return p.Type == "error" }) + if !strings.Contains(ev.Err, "event too large") { + t.Fatalf("event %+v", ev) + } +} + +func TestBrokerDrainsProviderAfterNotifyFailure(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + st := &stream{id: "s1", state: streamStarted, ctx: ctx, cancel: cancel} + b := &Broker{streams: map[string]*stream{"s1": st}} + events := make(chan provider.Event) + producerDone := make(chan struct{}) + go func() { + defer close(producerDone) + defer close(events) + for i := 0; i < 100; i++ { + events <- provider.Event{Type: "text", Text: "x"} + } + }() + forwardDone := make(chan struct{}) + go func() { + b.forwardEvents(st, events, func(string, any) error { return errors.New("connection closed") }) + close(forwardDone) + }() + select { + case <-producerDone: + case <-time.After(2 * time.Second): + t.Fatal("provider producer remained blocked after notify failure") + } + select { + case <-forwardDone: + case <-time.After(2 * time.Second): + t.Fatal("broker did not finish draining provider events") + } +} + +func TestBrokerOpenContextClosesReceivingStream(t *testing.T) { + b := New(defaultCfg(), credsource.NewResolver()) + ctx, cancel := context.WithCancel(context.Background()) + res, perr := b.Handle(ctx, "provider_open", mustJSON(t, protocol.ProviderOpenParams{Model: "openai-default"}), nil) + if perr != nil { + t.Fatal(perr) + } + openRaw, _ := json.Marshal(res) + var open protocol.ProviderOpenResult + _ = json.Unmarshal(openRaw, &open) + cancel() + deadline := time.Now().Add(time.Second) + for { + b.mu.Lock() + _, exists := b.streams[open.StreamID] + b.mu.Unlock() + if !exists { + break + } + if time.Now().After(deadline) { + t.Fatal("stream survived its connection context") + } + time.Sleep(time.Millisecond) + } +} + +func TestBrokerCredentialResolvedPerCall(t *testing.T) { + srv, authSeen := sseProvider(t, "openai", func(w http.ResponseWriter, auth, body string) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "data: [DONE]\n\n") + }) + cfg := defaultCfg() + cfg.Models[1].BaseURL = srv.URL + rot := &rotatingSource{vals: []string{"key-one", "key-two"}} + r := credsource.NewResolver() + r.Register("rot", rot) + cfg.Models[1].CredentialEnv = "" + cfg.Models[1].Credential = &config.CredentialRef{Source: "rot", Name: "openai"} + b := New(cfg, r).WithHTTPClient(srv.Client()) + + for i := 0; i < 2; i++ { + rec := ¬ifyRecorder{} + id := openStream(t, b, "openai-default") + raw, _ := json.Marshal(protocol.ProviderRequest{Messages: []protocol.ProviderMessage{{Role: "user", Content: "q"}}}) + if _, perr := sendChunk(t, b, rec, id, raw, true); perr != nil { + t.Fatalf("send %d: %v", i, perr) + } + rec.waitFor(t, 2*time.Second, func(p protocol.ProviderEventParams) bool { return p.Type == "done" }) + } + if rot.count() != 2 { + t.Fatalf("resolver calls %d", rot.count()) + } + got := authSeen() + if len(got) != 2 { + t.Fatalf("provider calls %d", len(got)) + } + if !strings.HasPrefix(got[0], "Bearer key-one") || !strings.HasPrefix(got[1], "Bearer key-two") { + t.Fatalf("rotated credentials not used per call: %v", got) + } +} + +func TestBrokerMissingCredentialTypedError(t *testing.T) { + cfg := defaultCfg() + r := credsource.NewResolver() + r.Register("rot", &rotatingSource{vals: nil}) + cfg.Models[1].CredentialEnv = "" + cfg.Models[1].Credential = &config.CredentialRef{Source: "rot", Name: "openai"} + b := New(cfg, r) + id := openStream(t, b, "openai-default") + raw, _ := json.Marshal(protocol.ProviderRequest{}) + _, perr := sendChunk(t, b, ¬ifyRecorder{}, id, raw, true) + if perr == nil || !strings.Contains(perr.Message, "credential for model") { + t.Fatalf("got %+v", perr) + } +} + +func TestBrokerCancelAbortsHTTP(t *testing.T) { + requestCtx := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"slow\"}}]}\n\n") + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + <-r.Context().Done() // hold the stream open until cancel + close(requestCtx) + })) + t.Cleanup(srv.Close) + + cfg := defaultCfg() + cfg.Models[1].BaseURL = srv.URL + r := credsource.NewResolver() + r.Register("rot", &rotatingSource{vals: []string{"k"}}) + cfg.Models[1].CredentialEnv = "" + cfg.Models[1].Credential = &config.CredentialRef{Source: "rot", Name: "openai"} + b := New(cfg, r).WithHTTPClient(srv.Client()) + + rec := ¬ifyRecorder{} + id := openStream(t, b, "openai-default") + raw, _ := json.Marshal(protocol.ProviderRequest{Messages: []protocol.ProviderMessage{{Role: "user", Content: "q"}}}) + if _, perr := sendChunk(t, b, rec, id, raw, true); perr != nil { + t.Fatalf("send: %v", perr) + } + rec.waitFor(t, 2*time.Second, func(p protocol.ProviderEventParams) bool { return p.Text == "slow" }) + if _, perr := b.Handle(context.Background(), "provider_cancel", mustJSON(t, protocol.ProviderCancelParams{StreamID: id}), nil); perr != nil { + t.Fatalf("cancel: %v", perr) + } + select { + case <-requestCtx: + case <-time.After(2 * time.Second): + t.Fatal("cancel did not abort the provider HTTP request") + } + if _, perr := b.Handle(context.Background(), "provider_cancel", mustJSON(t, protocol.ProviderCancelParams{StreamID: id}), nil); perr != nil { + t.Fatalf("second cancel: %v", perr) + } +} + +func TestBrokerToolArgsBound(t *testing.T) { + cfg := defaultCfg() + r := credsource.NewResolver() + b := New(cfg, r) + id := openStream(t, b, "openai-default") + req := protocol.ProviderRequest{Messages: []protocol.ProviderMessage{{ + Role: "assistant", ToolName: "x", ToolArgs: strings.Repeat("a", protocol.MaxProviderToolArgs+1), + }}} + raw, _ := json.Marshal(req) + var perr *protocol.Error + for off := 0; off < len(raw); off += protocol.MaxProviderChunk { + end := min(off+protocol.MaxProviderChunk, len(raw)) + _, perr = sendChunk(t, b, ¬ifyRecorder{}, id, raw[off:end], end == len(raw)) + if perr != nil { + break + } + } + if perr == nil || !strings.Contains(perr.Message, "tool args too large") { + t.Fatalf("got %+v", perr) + } +} + +func TestBrokerUnknownMethodTypedError(t *testing.T) { + b := New(defaultCfg(), credsource.NewResolver()) + _, perr := b.Handle(context.Background(), "fetch_url", []byte(`{}`), nil) + if perr == nil || !strings.Contains(perr.Message, "unknown broker method") { + t.Fatalf("got %+v", perr) + } +} + +func TestBrokerMaxConcurrentStreams(t *testing.T) { + cfg := defaultCfg() + b := New(cfg, credsource.NewResolver()) + first := openStream(t, b, "openai-default") + second := openStream(t, b, "grok-default") + _, perr := b.Handle(context.Background(), "provider_open", mustJSON(t, providerOpenJSON("claude-default")), nil) + if perr == nil || !strings.Contains(perr.Message, "too many open provider streams") { + t.Fatalf("got %+v", perr) + } + _ = first + _ = second +} + +var _ = errors.New +var _ = fmt.Sprintf diff --git a/internal/mcpauth/oauth.go b/internal/mcpauth/oauth.go index c07cc8d..b044047 100644 --- a/internal/mcpauth/oauth.go +++ b/internal/mcpauth/oauth.go @@ -19,7 +19,7 @@ import ( "time" "github.com/AdminTurnedDevOps/ABox/internal/config" - "github.com/AdminTurnedDevOps/ABox/internal/credentials" + "github.com/AdminTurnedDevOps/ABox/internal/credsource" ) type Result struct { @@ -54,6 +54,32 @@ type tokenResp struct { Error string `json:"error"` } +var savePreferred = credsource.SavePreferred + +var ( + oauthURLValidator = validatePublicOAuthURL + lookupOAuthHost = net.DefaultResolver.LookupIPAddr + blockedOAuthNets = []*net.IPNet{ + mustCIDR("0.0.0.0/8"), + mustCIDR("100.64.0.0/10"), + mustCIDR("192.0.0.0/24"), + mustCIDR("192.0.2.0/24"), + mustCIDR("198.18.0.0/15"), + mustCIDR("198.51.100.0/24"), + mustCIDR("203.0.113.0/24"), + mustCIDR("240.0.0.0/4"), + mustCIDR("2001:db8::/32"), + } +) + +func mustCIDR(raw string) *net.IPNet { + _, network, err := net.ParseCIDR(raw) + if err != nil { + panic(err) + } + return network +} + func LoginNamed(ctx context.Context, cfg config.File, name string) error { srv, err := serverNamed(cfg, name) if err != nil { @@ -65,8 +91,15 @@ func LoginNamed(ctx context.Context, cfg config.File, name string) error { if val == "" { return fmt.Errorf("set %s or omit credential_env to use OAuth", srv.CredentialEnv) } - credentials.SetEnv(env, val) - return credentials.Save(env, val) + res, err := savePreferred(ctx, env, val) + if err != nil { + return err + } + if err := persistCredentialReference(cfg, srv.Name, env, res.Source); err != nil { + return err + } + fmt.Printf("mcp %s token saved (%s)\n", srv.Name, res.Note) + return nil } res, err := Login(ctx, srv, Options{}) if err != nil { @@ -75,16 +108,35 @@ func LoginNamed(ctx context.Context, cfg config.File, name string) error { if res.AccessToken == "" { return nil } - if err := credentials.Save(env, res.AccessToken); err != nil { + saved, err := savePreferred(ctx, env, res.AccessToken) + if err != nil { return err } - credentials.SetEnv(env, res.AccessToken) - if res.RefreshToken != "" { - _ = credentials.Save(env+"_REFRESH", res.RefreshToken) + if err := persistCredentialReference(cfg, srv.Name, env, saved.Source); err != nil { + return err } + fmt.Printf("mcp %s token saved (%s)\n", srv.Name, saved.Note) return nil } +func persistCredentialReference(cfg config.File, serverName, credentialName, source string) error { + if source != "env" && source != "keychain" { + return fmt.Errorf("mcp %s token saved to unknown credential source %q", serverName, source) + } + for i := range cfg.MCPServers { + if cfg.MCPServers[i].Name != serverName { + continue + } + cfg.MCPServers[i].CredentialEnv = "" + cfg.MCPServers[i].Credential = &config.CredentialRef{Source: source, Name: credentialName} + if err := cfg.Save(); err != nil { + return fmt.Errorf("save mcp %s credential reference: %w", serverName, err) + } + return nil + } + return fmt.Errorf("unknown mcp server %q", serverName) +} + func serverNamed(cfg config.File, name string) (config.MCPServer, error) { for _, s := range cfg.MCPServers { if s.Name == name { @@ -99,7 +151,12 @@ func Login(ctx context.Context, srv config.MCPServer, opts Options) (Result, err if client == nil { client = &http.Client{Timeout: 30 * time.Second} } - status, hdr, err := probeMCP(ctx, client, srv.URL) + client = withoutRedirects(client) + mcpURL, err := oauthURLValidator(ctx, srv.URL) + if err != nil { + return Result{}, fmt.Errorf("mcp URL: %w", err) + } + status, hdr, err := probeMCP(ctx, client, mcpURL.String()) if err != nil { return Result{}, err } @@ -111,17 +168,25 @@ func Login(ctx context.Context, srv config.MCPServer, opts Options) (Result, err } metaURL := resourceMetadataURL(hdr) if metaURL == "" { - metaURL, err = discoverPRM(ctx, client, srv.URL) + metaURL, err = discoverPRM(ctx, client, mcpURL.String()) if err != nil { return Result{}, err } } - prm, err := fetchPRM(ctx, client, metaURL) + metadataURL, err := oauthURLValidator(ctx, metaURL) + if err != nil { + return Result{}, fmt.Errorf("protected resource metadata URL: %w", err) + } + if !sameOrigin(mcpURL, metadataURL) { + return Result{}, fmt.Errorf("protected resource metadata URL must use the configured MCP origin %s", oauthOrigin(mcpURL)) + } + prm, err := fetchPRM(ctx, client, metadataURL.String()) if err != nil { return Result{}, err } - if len(prm.AuthorizationServers) == 0 { - return Result{}, fmt.Errorf("protected resource metadata has no authorization_servers") + resource, err := validateProtectedResourceMetadata(ctx, mcpURL, prm) + if err != nil { + return Result{}, err } as, err := fetchAS(ctx, client, prm.AuthorizationServers[0]) if err != nil { @@ -158,7 +223,7 @@ func Login(ctx context.Context, srv config.MCPServer, opts Options) (Result, err if scope == "" { scope = strings.Join(prm.ScopesSupported, " ") } - authURL := authorizeURL(as.AuthorizationEndpoint, clientID, redir, challenge, state, resourceParam(srv.URL, prm.Resource), scope) + authURL := authorizeURL(as.AuthorizationEndpoint, clientID, redir, challenge, state, resource, scope) open := opts.OpenURL if open == nil { open = openBrowser @@ -179,7 +244,119 @@ func Login(ctx context.Context, srv config.MCPServer, opts Options) (Result, err case <-time.After(5 * time.Minute): return Result{}, fmt.Errorf("oauth timed out waiting for browser callback") } - return exchangeCode(ctx, client, as.TokenEndpoint, clientID, redir, code, verifier, resourceParam(srv.URL, prm.Resource)) + return exchangeCode(ctx, client, as.TokenEndpoint, clientID, redir, code, verifier, resource) +} + +func withoutRedirects(client *http.Client) *http.Client { + clone := *client + clone.CheckRedirect = func(_ *http.Request, _ []*http.Request) error { + return http.ErrUseLastResponse + } + return &clone +} + +func validatePublicOAuthURL(ctx context.Context, raw string) (*url.URL, error) { + u, err := url.Parse(strings.TrimSpace(raw)) + if err != nil || u.Scheme != "https" || u.Opaque != "" || u.Host == "" || u.User != nil || u.Fragment != "" { + return nil, fmt.Errorf("must be an https URL without userinfo or fragment") + } + host := strings.ToLower(u.Hostname()) + if host == "" || strings.HasSuffix(host, ".") { + return nil, fmt.Errorf("has an invalid host") + } + if ip := net.ParseIP(host); ip != nil { + if !publicOAuthIP(ip) { + return nil, fmt.Errorf("host %q is not a public address", host) + } + return u, nil + } + if !strings.Contains(host, ".") || host == "localhost" || strings.HasSuffix(host, ".localhost") || strings.HasSuffix(host, ".local") || strings.HasSuffix(host, ".internal") || strings.HasSuffix(host, ".home.arpa") { + return nil, fmt.Errorf("host %q is not a public DNS name", host) + } + lookupCtx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + addrs, err := lookupOAuthHost(lookupCtx, host) + if err != nil { + return nil, fmt.Errorf("resolve host %q: %w", host, err) + } + if len(addrs) == 0 { + return nil, fmt.Errorf("host %q has no addresses", host) + } + for _, addr := range addrs { + if !publicOAuthIP(addr.IP) { + return nil, fmt.Errorf("host %q resolves to non-public address %s", host, addr.IP) + } + } + return u, nil +} + +func publicOAuthIP(ip net.IP) bool { + if !ip.IsGlobalUnicast() || ip.IsPrivate() || ip.IsLoopback() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsUnspecified() { + return false + } + for _, network := range blockedOAuthNets { + if network.Contains(ip) { + return false + } + } + return true +} + +func oauthOrigin(u *url.URL) string { + scheme := strings.ToLower(u.Scheme) + host := strings.ToLower(u.Hostname()) + port := u.Port() + defaultPort := "443" + if scheme == "http" { + defaultPort = "80" + } + if port != "" && port != defaultPort { + return scheme + "://" + net.JoinHostPort(host, port) + } + if strings.Contains(host, ":") { + host = "[" + host + "]" + } + return scheme + "://" + host +} + +func sameOrigin(a, b *url.URL) bool { + return strings.EqualFold(a.Scheme, b.Scheme) && strings.EqualFold(a.Hostname(), b.Hostname()) && effectivePort(a) == effectivePort(b) +} + +func effectivePort(u *url.URL) string { + if port := u.Port(); port != "" { + return port + } + if strings.EqualFold(u.Scheme, "https") { + return "443" + } + return "80" +} + +func validateProtectedResourceMetadata(ctx context.Context, mcpURL *url.URL, prm prmDoc) (string, error) { + if strings.TrimSpace(prm.Resource) == "" { + return "", fmt.Errorf("protected resource metadata has no resource identifier") + } + resourceURL, err := oauthURLValidator(ctx, prm.Resource) + if err != nil { + return "", fmt.Errorf("protected resource metadata resource: %w", err) + } + if !sameOrigin(mcpURL, resourceURL) { + return "", fmt.Errorf("protected resource metadata resource must use the configured MCP origin %s", oauthOrigin(mcpURL)) + } + if len(prm.AuthorizationServers) == 0 { + return "", fmt.Errorf("protected resource metadata has no authorization_servers") + } + for _, raw := range prm.AuthorizationServers { + issuer, err := oauthURLValidator(ctx, raw) + if err != nil { + return "", fmt.Errorf("authorization server %q: %w", raw, err) + } + if issuer.RawQuery != "" { + return "", fmt.Errorf("authorization server issuer must not contain a query") + } + } + return strings.TrimRight(resourceURL.String(), "/"), nil } func probeMCP(ctx context.Context, client *http.Client, rawURL string) (int, http.Header, error) { @@ -262,9 +439,12 @@ func fetchPRM(ctx context.Context, client *http.Client, raw string) (prmDoc, err func fetchAS(ctx context.Context, client *http.Client, issuer string) (asDoc, error) { issuer = strings.TrimRight(issuer, "/") - u, err := url.Parse(issuer) + u, err := oauthURLValidator(ctx, issuer) if err != nil { - return asDoc{}, err + return asDoc{}, fmt.Errorf("authorization server issuer: %w", err) + } + if u.RawQuery != "" { + return asDoc{}, fmt.Errorf("authorization server issuer must not contain a query") } var candidates []string if u.Path != "" && u.Path != "/" { @@ -287,10 +467,12 @@ func fetchAS(ctx context.Context, client *http.Client, issuer string) (asDoc, er last = err continue } - if doc.AuthorizationEndpoint != "" && doc.TokenEndpoint != "" { + if err := validateAuthorizationServerMetadata(ctx, u, doc); err == nil { return doc, nil + } else { + last = fmt.Errorf("%s: %w", c, err) + continue } - last = fmt.Errorf("%s: missing endpoints", c) } if last == nil { last = fmt.Errorf("authorization server metadata not found") @@ -298,6 +480,36 @@ func fetchAS(ctx context.Context, client *http.Client, issuer string) (asDoc, er return asDoc{}, last } +func validateAuthorizationServerMetadata(ctx context.Context, issuer *url.URL, doc asDoc) error { + if strings.TrimRight(doc.Issuer, "/") != strings.TrimRight(issuer.String(), "/") { + return fmt.Errorf("authorization server metadata issuer %q does not match %q", doc.Issuer, issuer.String()) + } + for _, endpoint := range []struct { + name string + raw string + required bool + }{ + {name: "authorization_endpoint", raw: doc.AuthorizationEndpoint, required: true}, + {name: "token_endpoint", raw: doc.TokenEndpoint, required: true}, + {name: "registration_endpoint", raw: doc.RegistrationEndpoint}, + } { + if endpoint.raw == "" { + if endpoint.required { + return fmt.Errorf("authorization server metadata has no %s", endpoint.name) + } + continue + } + u, err := oauthURLValidator(ctx, endpoint.raw) + if err != nil { + return fmt.Errorf("%s: %w", endpoint.name, err) + } + if !sameOrigin(issuer, u) { + return fmt.Errorf("%s must use authorization server origin %s", endpoint.name, oauthOrigin(issuer)) + } + } + return nil +} + func supportsS256(methods []string) bool { for _, m := range methods { if strings.EqualFold(m, "S256") { @@ -402,12 +614,15 @@ func exchangeCode(ctx context.Context, client *http.Client, tokenURL, clientID, } defer resp.Body.Close() b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if resp.StatusCode >= 300 { + return Result{}, fmt.Errorf("token endpoint: %s", resp.Status) + } var tr tokenResp if err := json.Unmarshal(b, &tr); err != nil { return Result{}, fmt.Errorf("token json: %w", err) } - if resp.StatusCode >= 300 || tr.AccessToken == "" { - return Result{}, fmt.Errorf("token %s: %s", resp.Status, b) + if tr.AccessToken == "" { + return Result{}, fmt.Errorf("token endpoint response has no access token") } return Result{AccessToken: tr.AccessToken, RefreshToken: tr.RefreshToken, TokenType: tr.TokenType}, nil } @@ -448,19 +663,6 @@ func randomHex(n int) (string, error) { return hex.EncodeToString(b), nil } -func resourceParam(mcpURL, prmResource string) string { - if prmResource != "" { - return strings.TrimRight(prmResource, "/") - } - u, err := url.Parse(mcpURL) - if err != nil { - return mcpURL - } - u.RawQuery = "" - u.Fragment = "" - return strings.TrimRight(u.String(), "/") -} - func openBrowser(raw string) error { return exec.Command("open", raw).Start() } diff --git a/internal/mcpauth/oauth_test.go b/internal/mcpauth/oauth_test.go index cba4af4..7bcbfe3 100644 --- a/internal/mcpauth/oauth_test.go +++ b/internal/mcpauth/oauth_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "io" + "net" "net/http" "net/http/httptest" "net/url" @@ -11,9 +12,28 @@ import ( "testing" "github.com/AdminTurnedDevOps/ABox/internal/config" + "github.com/AdminTurnedDevOps/ABox/internal/credentials" + "github.com/AdminTurnedDevOps/ABox/internal/credsource" ) +func allowLocalOAuthURLs(t *testing.T) { + t.Helper() + orig := oauthURLValidator + oauthURLValidator = func(_ context.Context, raw string) (*url.URL, error) { + u, err := url.Parse(raw) + if err != nil { + return nil, err + } + if u.Host == "" { + return nil, url.InvalidHostError(raw) + } + return u, nil + } + t.Cleanup(func() { oauthURLValidator = orig }) +} + func TestLoginUnauthenticated(t *testing.T) { + allowLocalOAuthURLs(t) ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{}}`)) @@ -29,18 +49,12 @@ func TestLoginUnauthenticated(t *testing.T) { } func TestLoginOAuthCodeExchange(t *testing.T) { + allowLocalOAuthURLs(t) var authorizeURL string mux := http.NewServeMux() as := httptest.NewServer(nil) mcp := httptest.NewServer(nil) - mux.HandleFunc("/.well-known/oauth-protected-resource", func(w http.ResponseWriter, r *http.Request) { - _ = json.NewEncoder(w).Encode(map[string]any{ - "resource": mcp.URL, - "authorization_servers": []string{as.URL}, - "scopes_supported": []string{"mcp"}, - }) - }) mux.HandleFunc("/.well-known/oauth-authorization-server", func(w http.ResponseWriter, r *http.Request) { _ = json.NewEncoder(w).Encode(map[string]any{ "issuer": as.URL, @@ -77,8 +91,15 @@ func TestLoginOAuthCodeExchange(t *testing.T) { as.Config.Handler = mux mcpMux := http.NewServeMux() + mcpMux.HandleFunc("/.well-known/oauth-protected-resource", func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(map[string]any{ + "resource": mcp.URL, + "authorization_servers": []string{as.URL}, + "scopes_supported": []string{"mcp"}, + }) + }) mcpMux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("WWW-Authenticate", `Bearer resource_metadata="`+as.URL+`/.well-known/oauth-protected-resource"`) + w.Header().Set("WWW-Authenticate", `Bearer resource_metadata="`+mcp.URL+`/.well-known/oauth-protected-resource"`) http.Error(w, "auth", http.StatusUnauthorized) }) mcp.Config.Handler = mcpMux @@ -107,3 +128,176 @@ func TestLoginOAuthCodeExchange(t *testing.T) { t.Fatalf("got %#v authorize=%s", res, authorizeURL) } } + +func TestValidatePublicOAuthURLRejectsPrivateTargets(t *testing.T) { + for _, raw := range []string{ + "http://public.example/oauth", + "https://localhost/oauth", + "https://127.0.0.1/oauth", + "https://10.0.0.1/oauth", + "https://169.254.169.254/oauth", + "https://[::1]/oauth", + "https://[fe80::1]/oauth", + } { + if _, err := validatePublicOAuthURL(context.Background(), raw); err == nil { + t.Fatalf("accepted %q", raw) + } + } +} + +func TestValidatePublicOAuthURLRejectsDNSResolvingPrivate(t *testing.T) { + orig := lookupOAuthHost + lookupOAuthHost = func(_ context.Context, host string) ([]net.IPAddr, error) { + if host != "public.example" { + t.Fatalf("host %q", host) + } + return []net.IPAddr{{IP: net.ParseIP("192.168.1.10")}}, nil + } + t.Cleanup(func() { lookupOAuthHost = orig }) + if _, err := validatePublicOAuthURL(context.Background(), "https://public.example/oauth"); err == nil || !strings.Contains(err.Error(), "non-public") { + t.Fatalf("got %v", err) + } +} + +func TestLoginRejectsCrossOriginProtectedResourceMetadata(t *testing.T) { + allowLocalOAuthURLs(t) + metadataRequests := 0 + attacker := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + metadataRequests++ + _ = json.NewEncoder(w).Encode(map[string]any{}) + })) + defer attacker.Close() + mcp := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("WWW-Authenticate", `Bearer resource_metadata="`+attacker.URL+`/metadata"`) + http.Error(w, "auth", http.StatusUnauthorized) + })) + defer mcp.Close() + + _, err := Login(context.Background(), config.MCPServer{Name: "svc", URL: mcp.URL}, Options{HTTPClient: http.DefaultClient}) + if err == nil || !strings.Contains(err.Error(), "configured MCP origin") { + t.Fatalf("got %v", err) + } + if metadataRequests != 0 { + t.Fatal("cross-origin metadata endpoint was requested") + } +} + +func TestProtectedResourceMetadataRejectsResourceMismatch(t *testing.T) { + allowLocalOAuthURLs(t) + mcpURL, _ := url.Parse("http://mcp.example:8443/mcp") + _, err := validateProtectedResourceMetadata(context.Background(), mcpURL, prmDoc{ + Resource: "http://attacker.example/resource", + AuthorizationServers: []string{"http://auth.example"}, + }) + if err == nil || !strings.Contains(err.Error(), "configured MCP origin") { + t.Fatalf("got %v", err) + } +} + +func TestAuthorizationMetadataRejectsAttackerTokenEndpoint(t *testing.T) { + allowLocalOAuthURLs(t) + issuer, _ := url.Parse("http://auth.example:8443") + err := validateAuthorizationServerMetadata(context.Background(), issuer, asDoc{ + Issuer: issuer.String(), + AuthorizationEndpoint: issuer.String() + "/authorize", + TokenEndpoint: "http://attacker.example/token", + }) + if err == nil || !strings.Contains(err.Error(), "authorization server origin") { + t.Fatalf("got %v", err) + } +} + +func TestTokenExchangeDoesNotFollowRedirect(t *testing.T) { + attackerRequests := 0 + attacker := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + attackerRequests++ + _ = json.NewEncoder(w).Encode(map[string]any{"access_token": "attacker-token"}) + })) + defer attacker.Close() + tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, attacker.URL+"/token", http.StatusTemporaryRedirect) + })) + defer tokenServer.Close() + + _, err := exchangeCode(context.Background(), withoutRedirects(http.DefaultClient), tokenServer.URL, "client", "http://127.0.0.1/callback", "code", "verifier", "https://mcp.example") + if err == nil { + t.Fatal("expected redirect response to fail") + } + if attackerRequests != 0 { + t.Fatal("token request followed attacker redirect") + } +} + +func TestLoginNamedPersistsNoRefreshToken(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("ABOX_HOME", "") + // Never touch the real macOS keychain from tests: force the file-store + // fallback path of the keychain-preferred writer. + origKC := credsource.KeychainEnabled + credsource.KeychainEnabled = func() bool { return false } + t.Cleanup(func() { credsource.KeychainEnabled = origKC }) + + cfg := config.Defaults() + cfg.MCPServers = []config.MCPServer{{ + Name: "gh", URL: "https://api.githubcopilot.com/mcp/", + CredentialEnv: "ABOX_MCP_GH_PAT", + }} + t.Setenv("ABOX_MCP_GH_PAT", "pat-value") + if err := LoginNamed(context.Background(), cfg, "gh"); err != nil { + t.Fatal(err) + } + creds, err := credentials.Load() + if err != nil { + t.Fatal(err) + } + if creds["ABOX_MCP_GH_PAT"] != "pat-value" { + t.Fatalf("token not saved: %#v", creds) + } + for name := range creds { + if strings.HasSuffix(name, "_REFRESH") { + t.Fatalf("refresh token persisted: %s", name) + } + } + savedCfg, _, err := config.Load() + if err != nil { + t.Fatal(err) + } + ref := savedCfg.MCPServers[0].CredentialReference() + if ref.Source != "env" || ref.Name != "ABOX_MCP_GH_PAT" { + t.Fatalf("credential reference %#v", ref) + } +} + +func TestLoginNamedPersistsKeychainReference(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("ABOX_HOME", "") + t.Setenv("CUSTOM_MCP_TOKEN", "pat-value") + origSave := savePreferred + savePreferred = func(_ context.Context, name, value string) (credsource.SaveResult, error) { + if name != "CUSTOM_MCP_TOKEN" || value != "pat-value" { + t.Fatalf("save %q=%q", name, value) + } + return credsource.SaveResult{Source: "keychain", Keychain: true, Note: "keychain"}, nil + } + t.Cleanup(func() { savePreferred = origSave }) + + cfg := config.Defaults() + cfg.MCPServers = []config.MCPServer{{ + Name: "gh", + URL: "https://api.githubcopilot.com/mcp/", + CredentialEnv: "CUSTOM_MCP_TOKEN", + }} + if err := LoginNamed(context.Background(), cfg, "gh"); err != nil { + t.Fatal(err) + } + savedCfg, _, err := config.Load() + if err != nil { + t.Fatal(err) + } + server := savedCfg.MCPServers[0] + if server.CredentialEnv != "" || server.Credential == nil || *server.Credential != (config.CredentialRef{Source: "keychain", Name: "CUSTOM_MCP_TOKEN"}) { + t.Fatalf("saved server %#v", server) + } +} diff --git a/internal/provider/provider.go b/internal/provider/provider.go index fae65c6..53731aa 100644 --- a/internal/provider/provider.go +++ b/internal/provider/provider.go @@ -8,16 +8,13 @@ import ( "fmt" "io" "net/http" - "os" "strings" + "time" "github.com/AdminTurnedDevOps/ABox/internal/config" - "github.com/AdminTurnedDevOps/ABox/internal/guest/egress" "github.com/AdminTurnedDevOps/ABox/protocol" ) -var newHTTPClient = func() *http.Client { return egress.Client() } - type Event struct { Type string Text string @@ -44,34 +41,25 @@ type ToolSchema struct { Parameters map[string]any } -func Stream(ctx context.Context, model config.Model, messages []Message, tools []ToolSchema) (<-chan Event, error) { - key := strings.TrimSpace(os.Getenv(model.CredentialEnv)) - if key == "" { +// Stream never reads the environment or builds its own client. +func Stream(ctx context.Context, model config.Model, key string, client *http.Client, messages []Message, tools []ToolSchema) (<-chan Event, error) { + return stream(ctx, model, key, client, messages, tools, false) +} + +func StreamWithUsage(ctx context.Context, model config.Model, key string, client *http.Client, messages []Message, tools []ToolSchema) (<-chan Event, error) { + includeUsage := model.Provider != "xai" + return stream(ctx, model, key, client, messages, tools, includeUsage) +} + +func stream(ctx context.Context, model config.Model, key string, client *http.Client, messages []Message, tools []ToolSchema, includeUsage bool) (<-chan Event, error) { + if strings.TrimSpace(key) == "" { return nil, fmt.Errorf("missing credential %s", model.CredentialEnv) } - base := strings.TrimRight(model.BaseURL, "/") - if base == "" { - switch model.Provider { - case "xai": - base = "https://api.x.ai/v1" - case "openai": - base = "https://api.openai.com/v1" - case "anthropic": - return streamAnthropic(ctx, model, key, messages, tools) - default: - return nil, fmt.Errorf("unsupported provider %q", model.Provider) - } + if client == nil { + client = &http.Client{Timeout: 5 * time.Minute} } if model.Provider == "anthropic" { - return streamAnthropic(ctx, model, key, messages, tools) - } - return streamOpenAICompat(ctx, base, key, model.Model, messages, tools, false) -} - -func StreamWithUsage(ctx context.Context, model config.Model, messages []Message, tools []ToolSchema) (<-chan Event, error) { - key := strings.TrimSpace(os.Getenv(model.CredentialEnv)) - if key == "" { - return nil, fmt.Errorf("missing credential %s", model.CredentialEnv) + return streamAnthropic(ctx, model, key, client, messages, tools) } base := strings.TrimRight(model.BaseURL, "/") if base == "" { @@ -80,20 +68,14 @@ func StreamWithUsage(ctx context.Context, model config.Model, messages []Message base = "https://api.x.ai/v1" case "openai": base = "https://api.openai.com/v1" - case "anthropic": - return streamAnthropic(ctx, model, key, messages, tools) default: return nil, fmt.Errorf("unsupported provider %q", model.Provider) } } - if model.Provider == "anthropic" { - return streamAnthropic(ctx, model, key, messages, tools) - } - includeUsage := model.Provider != "xai" - return streamOpenAICompat(ctx, base, key, model.Model, messages, tools, includeUsage) + return streamOpenAICompat(ctx, base, key, model.Model, client, messages, tools, includeUsage) } -func streamOpenAICompat(ctx context.Context, base, key, model string, messages []Message, tools []ToolSchema, includeUsage bool) (<-chan Event, error) { +func streamOpenAICompat(ctx context.Context, base, key, model string, client *http.Client, messages []Message, tools []ToolSchema, includeUsage bool) (<-chan Event, error) { var oaiMsgs []map[string]any for _, m := range messages { switch { @@ -152,7 +134,7 @@ func streamOpenAICompat(ctx context.Context, base, key, model string, messages [ } req.Header.Set("Authorization", "Bearer "+key) req.Header.Set("Content-Type", "application/json") - resp, err := newHTTPClient().Do(req) + resp, err := client.Do(req) if err != nil { return nil, err } @@ -229,7 +211,7 @@ func streamOpenAICompat(ctx context.Context, base, key, model string, messages [ return out, nil } -func streamAnthropic(ctx context.Context, model config.Model, key string, messages []Message, tools []ToolSchema) (<-chan Event, error) { +func streamAnthropic(ctx context.Context, model config.Model, key string, client *http.Client, messages []Message, tools []ToolSchema) (<-chan Event, error) { base := strings.TrimRight(model.BaseURL, "/") if base == "" { base = "https://api.anthropic.com" @@ -298,7 +280,7 @@ func streamAnthropic(ctx context.Context, model config.Model, key string, messag req.Header.Set("x-api-key", key) req.Header.Set("anthropic-version", "2023-06-01") req.Header.Set("Content-Type", "application/json") - resp, err := newHTTPClient().Do(req) + resp, err := client.Do(req) if err != nil { return nil, err } diff --git a/internal/provider/provider_test.go b/internal/provider/provider_test.go index 46c9d0b..ec9b1f2 100644 --- a/internal/provider/provider_test.go +++ b/internal/provider/provider_test.go @@ -5,7 +5,6 @@ import ( "io" "net/http" "net/http/httptest" - "os" "strings" "testing" @@ -13,7 +12,6 @@ import ( ) func TestStreamOpenAITextAndUsage(t *testing.T) { - t.Setenv("TEST_KEY", "k") body := strings.Join([]string{ `data: {"choices":[{"delta":{"content":"hi"}}]}`, `data: {"choices":[],"usage":{"prompt_tokens":3,"completion_tokens":2}}`, @@ -24,6 +22,9 @@ func TestStreamOpenAITextAndUsage(t *testing.T) { if r.URL.Path != "/chat/completions" { t.Errorf("path %s", r.URL.Path) } + if r.Header.Get("Authorization") != "Bearer test-key" { + t.Errorf("authorization %q", r.Header.Get("Authorization")) + } raw, _ := io.ReadAll(r.Body) if !strings.Contains(string(raw), "include_usage") { t.Error("expected stream_options include_usage") @@ -32,13 +33,10 @@ func TestStreamOpenAITextAndUsage(t *testing.T) { _, _ = io.WriteString(w, body) })) t.Cleanup(srv.Close) - old := newHTTPClient - newHTTPClient = func() *http.Client { return srv.Client() } - t.Cleanup(func() { newHTTPClient = old }) ch, err := StreamWithUsage(context.Background(), config.Model{ Provider: "openai", Model: "gpt", CredentialEnv: "TEST_KEY", BaseURL: srv.URL, - }, []Message{{Role: "user", Content: "q"}}, nil) + }, "test-key", srv.Client(), []Message{{Role: "user", Content: "q"}}, nil) if err != nil { t.Fatal(err) } @@ -57,32 +55,26 @@ func TestStreamOpenAITextAndUsage(t *testing.T) { } func TestStreamOpenAIErrorStatus(t *testing.T) { - t.Setenv("TEST_KEY", "k") srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Error(w, "nope", http.StatusBadRequest) })) t.Cleanup(srv.Close) - old := newHTTPClient - newHTTPClient = func() *http.Client { return srv.Client() } - t.Cleanup(func() { newHTTPClient = old }) _, err := Stream(context.Background(), config.Model{ Provider: "openai", Model: "gpt", CredentialEnv: "TEST_KEY", BaseURL: srv.URL, - }, nil, nil) + }, "test-key", srv.Client(), nil, nil) if err == nil || !strings.Contains(err.Error(), "400") { t.Fatalf("got %v", err) } } func TestStreamMissingKey(t *testing.T) { - os.Unsetenv("MISSING_ABOX_KEY") - _, err := Stream(context.Background(), config.Model{CredentialEnv: "MISSING_ABOX_KEY"}, nil, nil) - if err == nil { - t.Fatal("expected missing credential") + _, err := Stream(context.Background(), config.Model{CredentialEnv: "MISSING_ABOX_KEY"}, "", nil, nil, nil) + if err == nil || !strings.Contains(err.Error(), "missing credential") { + t.Fatalf("got %v", err) } } -func TestStreamAnthropicText(t *testing.T) { - t.Setenv("TEST_KEY", "k") +func TestStreamAnthropicTextAndKeyHeader(t *testing.T) { body := strings.Join([]string{ `data: {"type":"message_start","message":{"usage":{"input_tokens":4}}}`, `data: {"type":"content_block_delta","delta":{"text":"yo"}}`, @@ -90,16 +82,16 @@ func TestStreamAnthropicText(t *testing.T) { "", }, "\n") srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("x-api-key") != "test-key" { + t.Errorf("x-api-key %q", r.Header.Get("x-api-key")) + } w.Header().Set("Content-Type", "text/event-stream") _, _ = io.WriteString(w, body) })) t.Cleanup(srv.Close) - old := newHTTPClient - newHTTPClient = func() *http.Client { return srv.Client() } - t.Cleanup(func() { newHTTPClient = old }) ch, err := Stream(context.Background(), config.Model{ Provider: "anthropic", Model: "claude", CredentialEnv: "TEST_KEY", BaseURL: srv.URL, - }, []Message{{Role: "user", Content: "q"}}, nil) + }, "test-key", srv.Client(), []Message{{Role: "user", Content: "q"}}, nil) if err != nil { t.Fatal(err) } @@ -116,3 +108,10 @@ func TestStreamAnthropicText(t *testing.T) { t.Fatalf("text=%v usage=%v", text, usage) } } + +func TestStreamUnsupportedProvider(t *testing.T) { + _, err := Stream(context.Background(), config.Model{Provider: "nope"}, "k", &http.Client{}, nil, nil) + if err == nil || !strings.Contains(err.Error(), "unsupported provider") { + t.Fatalf("got %v", err) + } +} diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index 1ca644d..19951a1 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -10,6 +10,8 @@ import ( "os" "os/exec" "path/filepath" + "strconv" + "strings" "sync" "time" @@ -25,6 +27,15 @@ var ErrGuestTooOld = errors.New("guest protocol too old") const cancelResponseTimeout = 5 * time.Second +const ( + callQueueFrames = 1 + callQueueBytes = protocol.MaxFrameBytes + turnQueueFrames = 256 + turnQueueBytes = 8 << 20 + writeTimeout = 30 * time.Second + shutdownTimeout = 3 * time.Second +) + type TurnOptions struct { MaxTurns int TimeoutSec int @@ -41,18 +52,163 @@ type TurnOutcome struct { StopReason string } +// Handle must not block the connection read loop; stream via notify instead. +type GuestCallHandler interface { + Handle(ctx context.Context, method string, params json.RawMessage, + notify func(method string, params any) error) (any, *protocol.Error) +} + type Sandbox struct { Sess *session.Session History []protocol.HistoryLine GuestProtocol int cmd *exec.Cmd conn net.Conn - mu sync.Mutex - writeMu sync.Mutex - nextID int + OnGuestCall GuestCallHandler + + writeOnce sync.Once + writeGate chan struct{} + turnMu sync.Mutex + + mu sync.Mutex + nextID int + calls map[string]*frameQueue + activeTurn string + turnQ *frameQueue + lifeCtx context.Context + lifeCancel context.CancelFunc + guestSlots chan struct{} + cancelSlot chan struct{} + busyQueue chan protocol.Frame + + readOnce sync.Once + failOnce sync.Once + readDone chan struct{} +} + +var ( + errQueueClosed = errors.New("frame queue closed") + errQueueFull = errors.New("frame queue budget exceeded") +) + +// frameQueue gives the connection read loop bounded, nonblocking delivery. +// Exceeding either budget fails the owning call/turn instead of silently +// dropping a response or blocking cancellation traffic. +type frameQueue struct { + mu sync.Mutex + frames []queuedFrame + bytes int + maxN int + maxB int + wake chan struct{} + closed bool + err error +} + +type queuedFrame struct { + frame protocol.Frame + size int +} + +func newFrameQueue(maxFrames, maxBytes int) *frameQueue { + return &frameQueue{maxN: maxFrames, maxB: maxBytes, wake: make(chan struct{}, 1)} } -func Prepare(sess *session.Session, imagePath string, model config.Model, secrets map[string]string, mcpServers []config.MCPServer, resume bool) error { +func (q *frameQueue) push(frame protocol.Frame) error { + raw, err := json.Marshal(frame) + if err != nil { + return err + } + q.mu.Lock() + if q.closed { + q.mu.Unlock() + return errQueueClosed + } + if len(q.frames) >= q.maxN || q.bytes+len(raw) > q.maxB { + q.mu.Unlock() + return errQueueFull + } + q.frames = append(q.frames, queuedFrame{frame: frame, size: len(raw)}) + q.bytes += len(raw) + q.mu.Unlock() + select { + case q.wake <- struct{}{}: + default: + } + return nil +} + +func (q *frameQueue) abort(err error) { + q.mu.Lock() + if !q.closed { + q.frames = nil + q.bytes = 0 + q.closed = true + q.err = err + } + q.mu.Unlock() + select { + case q.wake <- struct{}{}: + default: + } +} + +func (q *frameQueue) close(err error) { + q.mu.Lock() + if !q.closed { + q.closed = true + q.err = err + } + q.mu.Unlock() + select { + case q.wake <- struct{}{}: + default: + } +} + +func (q *frameQueue) pop(ctx context.Context) (protocol.Frame, bool, error) { + for { + q.mu.Lock() + if len(q.frames) > 0 { + item := q.frames[0] + q.frames[0] = queuedFrame{} + q.frames = q.frames[1:] + q.bytes -= item.size + q.mu.Unlock() + return item.frame, true, nil + } + if q.closed { + err := q.err + q.mu.Unlock() + return protocol.Frame{}, false, err + } + q.mu.Unlock() + select { + case <-ctx.Done(): + // Prefer a frame that became ready with the cancellation signal. + // This avoids racing a real terminal response at the deadline. + q.mu.Lock() + if len(q.frames) > 0 { + item := q.frames[0] + q.frames[0] = queuedFrame{} + q.frames = q.frames[1:] + q.bytes -= item.size + q.mu.Unlock() + return item.frame, true, nil + } + if q.closed { + err := q.err + q.mu.Unlock() + return protocol.Frame{}, false, err + } + q.mu.Unlock() + return protocol.Frame{}, false, ctx.Err() + case <-q.wake: + } + } +} + +func Prepare(sess *session.Session, imagePath string, model config.Model, mcpServers []config.MCPServer, resume bool) error { if imagePath == "" { imagePath = config.GuestImagePath() } @@ -68,31 +224,14 @@ func Prepare(sess *session.Session, imagePath string, model config.Model, secret return fmt.Errorf("clone session disk: %w", err) } } - if err := sess.WriteGuestConfig(model, secrets, mcpServers); err != nil { + if err := sess.WriteGuestConfig(model, mcpServers); err != nil { return err } - return writeConfigDisk(sess) -} - -func writeConfigDisk(sess *session.Session) error { data, err := os.ReadFile(sess.GuestConfigJSON()) if err != nil { return err } - // Resume rewrites config.raw; the previous run left it mode 0400. - _ = os.Chmod(sess.ConfigDisk(), 0o600) - f, err := os.OpenFile(sess.ConfigDisk(), os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600) - if err != nil { - return err - } - defer f.Close() - if _, err := f.Write(data); err != nil { - return err - } - if _, err := f.Write(make([]byte, 1<<20-len(data))); err != nil { - return err - } - return os.Chmod(sess.ConfigDisk(), 0o400) + return session.WritePaddedConfig(sess.ConfigDisk(), data) } func Start(ctx context.Context, sess *session.Session, vmmPath string, vcpu int, ram int) (*Sandbox, error) { @@ -179,7 +318,7 @@ func Start(ctx context.Context, sess *session.Session, vmmPath string, vcpu int, conn = a.c } - sb := &Sandbox{Sess: sess, cmd: cmd, conn: conn} + sb := &Sandbox{Sess: sess, cmd: cmd, conn: conn, calls: map[string]*frameQueue{}} if err := sb.waitHello(ctx); err != nil { sb.Stop() return nil, err @@ -203,7 +342,7 @@ func (s *Sandbox) waitHello(ctx context.Context) error { if hello.SessionID != s.Sess.ID || hello.Capability != s.Sess.Capability { return fmt.Errorf("guest capability mismatch") } - ok, _ := protocol.EncodeParams(protocol.HelloResult{Accepted: true}) + ok, _ := protocol.EncodeParams(protocol.HelloResult{Accepted: true, Protocol: protocol.Version}) if err := protocol.WriteFrame(s.conn, protocol.Frame{ID: frame.ID, Result: ok}); err != nil { return err } @@ -217,37 +356,269 @@ func (s *Sandbox) waitHello(ctx context.Context) error { return nil } +func (s *Sandbox) startReading() { + s.readOnce.Do(func() { + s.readDone = make(chan struct{}) + s.lifeCtx, s.lifeCancel = context.WithCancel(context.Background()) + s.guestSlots = make(chan struct{}, protocol.MaxGuestCalls-1) // reserve one slot for cancel + s.cancelSlot = make(chan struct{}, 1) + s.busyQueue = make(chan protocol.Frame, protocol.MaxGuestCalls) + go s.writeBusyReplies() + go s.readLoop() + }) +} + +func (s *Sandbox) readLoop() { + for { + frame, err := protocol.ReadFrame(s.conn) + if err != nil { + s.failAll(err) + return + } + if frame.Method != "" { + switch { + case strings.HasPrefix(frame.ID, "g-"): + slots := s.guestSlots + if frame.Method == "provider_cancel" { + slots = s.cancelSlot + } + select { + case slots <- struct{}{}: + go func(frame protocol.Frame, slots chan struct{}) { + defer func() { <-slots }() + s.dispatchGuestCall(frame) + }(frame, slots) + default: + select { + case s.busyQueue <- frame: + default: + s.failConnection(fmt.Errorf("guest busy reply backpressure exceeded")) + return + } + } + case frame.Method == "agent_event": + s.mu.Lock() + q := s.turnQ + active := s.activeTurn + if q != nil && frame.ID == active { + if err := q.push(frame); err != nil && !errors.Is(err, errQueueClosed) { + queueErr := fmt.Errorf("turn frame queue overflow: %w", err) + q.abort(queueErr) + s.mu.Unlock() + s.failConnection(queueErr) + return + } + } + s.mu.Unlock() + default: + } + continue + } + s.mu.Lock() + q, known := s.calls[frame.ID] + turnFrame := false + if !known && frame.ID == s.activeTurn { + q = s.turnQ + turnFrame = true + } + if q == nil { + s.mu.Unlock() + continue + } + if err := q.push(frame); err != nil && !errors.Is(err, errQueueClosed) { + queueName := "call" + if turnFrame { + queueName = "turn" + } + queueErr := fmt.Errorf("%s frame queue overflow: %w", queueName, err) + q.abort(queueErr) + s.mu.Unlock() + s.failConnection(queueErr) + return + } + s.mu.Unlock() + } +} + +func (s *Sandbox) failConnection(err error) { + _ = s.conn.Close() + s.failAll(err) +} + +func (s *Sandbox) failAll(err error) { + s.failOnce.Do(func() { + if err == nil { + err = errors.New("guest connection closed") + } + if s.lifeCancel != nil { + s.lifeCancel() + } + s.mu.Lock() + pending := s.calls + s.calls = map[string]*frameQueue{} + turnQ := s.turnQ + s.turnQ = nil + s.activeTurn = "" + s.mu.Unlock() + for _, q := range pending { + q.close(err) + } + if turnQ != nil { + turnQ.close(err) + } + close(s.readDone) + }) +} + +func (s *Sandbox) dispatchGuestCall(frame protocol.Frame) { + out := protocol.Frame{V: protocol.Version, ID: frame.ID} + if s.GuestProtocol < 3 { + out.Error = &protocol.Error{Code: "host", Message: fmt.Sprintf("provider broker requires protocol 3, guest speaks %d", s.GuestProtocol)} + if err := s.writeFrame(out); err != nil { + s.failConnection(err) + } + return + } + if s.OnGuestCall == nil { + out.Error = &protocol.Error{Code: "host", Message: "guest calls not supported by this host"} + if err := s.writeFrame(out); err != nil { + s.failConnection(err) + } + return + } + notify := func(method string, params any) error { + raw, err := protocol.EncodeParams(params) + if err != nil { + return err + } + s.mu.Lock() + s.nextID++ + n := s.nextID + s.mu.Unlock() + if err := s.writeFrame(protocol.Frame{V: protocol.Version, ID: fmt.Sprintf("h-%d", n), Method: method, Params: raw}); err != nil { + s.failConnection(err) + return err + } + return nil + } + res, perr := s.OnGuestCall.Handle(s.lifeCtx, frame.Method, frame.Params, notify) + switch { + case perr != nil: + out.Error = perr + case res != nil: + raw, err := protocol.EncodeParams(res) + if err != nil { + out.Error = &protocol.Error{Code: "host", Message: err.Error()} + } else { + out.Result = raw + } + default: + out.Result = []byte(`{"ok":true}`) + } + if err := s.writeFrame(out); err != nil { + s.failConnection(err) + } +} + +func (s *Sandbox) writeBusyReplies() { + for { + select { + case <-s.lifeCtx.Done(): + return + case frame := <-s.busyQueue: + if err := s.writeFrame(protocol.Frame{V: protocol.Version, ID: frame.ID, Error: &protocol.Error{ + Code: "busy", Message: "too many concurrent guest calls", + }}); err != nil { + s.failConnection(err) + return + } + } + } +} + func (s *Sandbox) writeFrame(f protocol.Frame) error { - s.writeMu.Lock() - defer s.writeMu.Unlock() - return protocol.WriteFrame(s.conn, f) + return s.writeFrameContext(context.Background(), f) } -func (s *Sandbox) Call(ctx context.Context, method string, params any, result any) error { +func (s *Sandbox) writeFrameContext(parent context.Context, f protocol.Frame) error { + ctx, cancel := context.WithTimeout(parent, writeTimeout) + defer cancel() + s.writeOnce.Do(func() { + s.writeGate = make(chan struct{}, 1) + s.writeGate <- struct{}{} + }) + select { + case <-ctx.Done(): + return ctx.Err() + case <-s.writeGate: + } + defer func() { s.writeGate <- struct{}{} }() + done := make(chan struct{}) + watchDone := make(chan struct{}) + go func() { + select { + case <-done: + case <-ctx.Done(): + select { + case <-done: + default: + _ = s.conn.Close() + } + } + close(watchDone) + }() + err := protocol.WriteFrame(s.conn, f) + close(done) + <-watchDone + if ctxErr := ctx.Err(); ctxErr != nil { + return ctxErr + } + if err != nil { + return err + } + return nil +} + +func (s *Sandbox) dropCall(id string) { s.mu.Lock() - defer s.mu.Unlock() - s.nextID++ - id := fmt.Sprintf("%d", s.nextID) + q := s.calls[id] + delete(s.calls, id) + s.mu.Unlock() + if q != nil { + q.close(context.Canceled) + } +} + +func (s *Sandbox) Call(ctx context.Context, method string, params any, result any) error { + s.startReading() raw, err := protocol.EncodeParams(params) if err != nil { return err } - if deadline, ok := ctx.Deadline(); ok { - _ = s.conn.SetDeadline(deadline) - defer s.conn.SetDeadline(time.Time{}) + s.mu.Lock() + if s.lifeCtx.Err() != nil { + s.mu.Unlock() + return errors.New("guest connection closed") + } + if s.calls == nil { + s.calls = map[string]*frameQueue{} } - if err := s.writeFrame(protocol.Frame{ID: id, Method: method, Params: raw}); err != nil { + s.nextID++ + id := strconv.Itoa(s.nextID) + q := newFrameQueue(callQueueFrames, callQueueBytes) + s.calls[id] = q + s.mu.Unlock() + if err := s.writeFrameContext(ctx, protocol.Frame{V: protocol.Version, ID: id, Method: method, Params: raw}); err != nil { + s.failConnection(err) return err } - var frame protocol.Frame - for { - frame, err = protocol.ReadFrame(s.conn) + frame, ok, err := q.pop(ctx) + s.dropCall(id) + if !ok { if err != nil { return err } - if frame.ID == id { - break - } + return errors.New("guest connection closed") } if frame.Error != nil { return frame.Error @@ -268,13 +639,20 @@ func (s *Sandbox) UserTurnCtx(ctx context.Context, text string, opts TurnOptions } func (s *Sandbox) userTurnLocked(ctx context.Context, text string, opts TurnOptions, onEvent func(protocol.AgentEvent), v2API bool) (*TurnOutcome, error) { - s.mu.Lock() - defer s.mu.Unlock() if v2API && opts.needsV2() && s.GuestProtocol < 2 { return nil, fmt.Errorf("%w: need protocol 2, guest speaks %d", ErrGuestTooOld, s.GuestProtocol) } + s.turnMu.Lock() + defer s.turnMu.Unlock() + s.startReading() + + s.mu.Lock() + if s.lifeCtx.Err() != nil { + s.mu.Unlock() + return nil, errors.New("guest connection closed") + } s.nextID++ - id := fmt.Sprintf("%d", s.nextID) + id := strconv.Itoa(s.nextID) params := protocol.UserTurnParams{Text: text} if s.GuestProtocol >= 2 { params.MaxTurns = opts.MaxTurns @@ -283,44 +661,51 @@ func (s *Sandbox) userTurnLocked(ctx context.Context, text string, opts TurnOpti } raw, err := protocol.EncodeParams(params) if err != nil { + s.mu.Unlock() return nil, err } - supportsCancel := v2API && s.GuestProtocol >= 2 - if supportsCancel { - defer s.conn.SetDeadline(time.Time{}) - } else if deadline, ok := ctx.Deadline(); ok { - _ = s.conn.SetDeadline(deadline) - defer s.conn.SetDeadline(time.Time{}) - } - if err := s.writeFrame(protocol.Frame{ID: id, Method: "user_turn", Params: raw}); err != nil { + turnQ := newFrameQueue(turnQueueFrames, turnQueueBytes) + s.activeTurn = id + s.turnQ = turnQ + s.mu.Unlock() + if err := s.writeFrame(protocol.Frame{V: protocol.Version, ID: id, Method: "user_turn", Params: raw}); err != nil { + s.failConnection(err) return nil, err } + defer func() { + s.mu.Lock() + if s.activeTurn == id { + s.activeTurn = "" + s.turnQ = nil + } + s.mu.Unlock() + turnQ.close(context.Canceled) + }() - var cancelOnce sync.Once + supportsCancel := v2API && s.GuestProtocol >= 2 stopWatch := make(chan struct{}) defer close(stopWatch) + waitCtx := ctx if supportsCancel { - go func() { - select { - case <-ctx.Done(): - cancelOnce.Do(func() { - raw, err := protocol.EncodeParams(protocol.CancelTurnParams{ID: id}) - if err != nil { - return - } - _ = s.writeFrame(protocol.Frame{ID: id + "-cancel", Method: "cancel_turn", Params: raw}) - _ = s.conn.SetReadDeadline(time.Now().Add(cancelResponseTimeout)) - }) - case <-stopWatch: - } - }() + var stopWait context.CancelCauseFunc + waitCtx, stopWait = context.WithCancelCause(context.Background()) + defer stopWait(context.Canceled) + go s.watchCancel(ctx, id, stopWatch, stopWait) } out := &TurnOutcome{} for { - frame, err := protocol.ReadFrame(s.conn) - if err != nil { - return out, err + frame, ok, err := turnQ.pop(waitCtx) + if !ok { + if err != nil { + if supportsCancel { + if cause := context.Cause(waitCtx); cause != nil { + return out, cause + } + } + return out, err + } + return out, errors.New("guest connection closed") } if frame.Method == "agent_event" && frame.ID == id { ev, err := protocol.DecodeParams[protocol.AgentEvent](frame.Params) @@ -336,24 +721,86 @@ func (s *Sandbox) userTurnLocked(ctx context.Context, text string, opts TurnOpti } continue } - if frame.ID == id { - if frame.Error != nil { - if frame.Error.Code == "canceled" { - out.Canceled = true - return out, frame.Error - } - return out, frame.Error + if frame.Error != nil { + if frame.Error.Code == "canceled" { + out.Canceled = true } - return out, nil + return out, frame.Error } + return out, nil } } +// watchCancel forwards cancellation to a protocol-2+ guest and fails the turn +// if the guest never acknowledges within cancelResponseTimeout. +func (s *Sandbox) watchCancel(ctx context.Context, turnID string, stopWatch <-chan struct{}, stopWait context.CancelCauseFunc) { + select { + case <-ctx.Done(): + case <-stopWatch: + return + } + raw, err := protocol.EncodeParams(protocol.CancelTurnParams{ID: turnID}) + if err != nil { + return + } + cancelCtx, cancel := context.WithTimeout(context.Background(), cancelResponseTimeout) + defer cancel() + err = s.writeFrameContext(cancelCtx, protocol.Frame{V: protocol.Version, ID: turnID + "-cancel", Method: "cancel_turn", Params: raw}) + if err != nil { + s.failConnection(err) + stopWait(err) + return + } + select { + case <-stopWatch: + case <-cancelCtx.Done(): + stopWait(&protocol.Error{ + Code: "timeout", + Message: "guest did not acknowledge cancel within " + cancelResponseTimeout.String(), + }) + } +} + +func (s *Sandbox) PushSecrets(ctx context.Context, model config.Model, secrets map[string]string) error { + if len(secrets) == 0 { + return nil + } + if s.GuestProtocol < 2 { + return fmt.Errorf("guest image predates secret push; run make image") + } + if s.GuestProtocol == 2 { + fmt.Fprintf(os.Stderr, "abox: guest speaks protocol 2; pushing legacy secrets (run make image to upgrade)\n") + } + modelKey := model.EnvName() + rest := map[string]string{} + for k, v := range secrets { + if k != modelKey { + rest[k] = v + } + } + var modelSecrets map[string]string + if s.GuestProtocol < 3 { + if v, ok := secrets[modelKey]; ok { + modelSecrets = map[string]string{modelKey: v} + } + } + if err := s.SetModel(ctx, model, modelSecrets); err != nil { + return err + } + if len(rest) > 0 { + return s.SetMCPTokens(ctx, rest) + } + return nil +} + func (s *Sandbox) SetMCPTokens(ctx context.Context, secrets map[string]string) error { return s.Call(ctx, "set_mcp_tokens", protocol.SetMCPTokensParams{Secrets: secrets}, nil) } func (s *Sandbox) SetModel(ctx context.Context, model config.Model, secrets map[string]string) error { + if s.GuestProtocol >= 3 { + secrets = nil + } return s.Call(ctx, "set_model", protocol.SetModelParams{ Model: model.ToGuest(), Secrets: secrets, @@ -382,8 +829,13 @@ func (s *Sandbox) TransferArchive(ctx context.Context, archive []byte) error { func (s *Sandbox) Stop() error { if s.conn != nil { - _ = s.Call(context.Background(), "shutdown", map[string]bool{"ok": true}, nil) - s.conn.Close() + ctx, cancel := context.WithTimeout(context.Background(), shutdownTimeout) + _ = s.Call(ctx, "shutdown", map[string]bool{"ok": true}, nil) + cancel() + if s.lifeCancel != nil { + s.lifeCancel() + } + _ = s.conn.Close() } if s.cmd != nil && s.cmd.Process != nil { _ = s.cmd.Process.Signal(os.Interrupt) diff --git a/internal/runtime/runtime_push_test.go b/internal/runtime/runtime_push_test.go new file mode 100644 index 0000000..7933e72 --- /dev/null +++ b/internal/runtime/runtime_push_test.go @@ -0,0 +1,445 @@ +package runtime + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net" + "strings" + "sync" + "testing" + "time" + + "github.com/AdminTurnedDevOps/ABox/internal/config" + "github.com/AdminTurnedDevOps/ABox/protocol" +) + +type fakeGuest struct { + t *testing.T + conn net.Conn + writeMu sync.Mutex + onRequest func(frame protocol.Frame, reply func(protocol.Frame)) + onResponse func(frame protocol.Frame) +} + +func (g *fakeGuest) write(f protocol.Frame) { + g.writeMu.Lock() + defer g.writeMu.Unlock() + if err := protocol.WriteFrame(g.conn, f); err != nil { + g.t.Errorf("guest write: %v", err) + } +} + +func (g *fakeGuest) serve() { + for { + frame, err := protocol.ReadFrame(g.conn) + if err != nil { + return + } + if frame.Method == "" { + if g.onResponse != nil { + g.onResponse(frame) + } + continue + } + if g.onRequest != nil { + g.onRequest(frame, g.write) + } + } +} + +func newPipeSandbox(t *testing.T, guestProtocol int) (*Sandbox, *fakeGuest) { + t.Helper() + host, guest := net.Pipe() + t.Cleanup(func() { host.Close(); guest.Close() }) + s := &Sandbox{conn: host, GuestProtocol: guestProtocol, calls: map[string]*frameQueue{}} + g := &fakeGuest{t: t, conn: guest} + go g.serve() + return s, g +} + +func TestPushSecretsOrderAndSplit(t *testing.T) { + s, g := newPipeSandbox(t, 2) + var order []string + var mu sync.Mutex + g.onRequest = func(frame protocol.Frame, reply func(protocol.Frame)) { + mu.Lock() + order = append(order, frame.Method) + mu.Unlock() + switch frame.Method { + case "set_model": + p, err := protocol.DecodeParams[protocol.SetModelParams](frame.Params) + if err != nil { + t.Errorf("set_model params: %v", err) + } + if p.Secrets["XAI_API_KEY"] != "mk" { + t.Errorf("set_model missing model credential: %v", p.Secrets) + } + if len(p.Secrets) != 1 { + t.Errorf("set_model carries more than the model credential: %v", p.Secrets) + } + case "set_mcp_tokens": + p, err := protocol.DecodeParams[protocol.SetMCPTokensParams](frame.Params) + if err != nil { + t.Errorf("set_mcp_tokens params: %v", err) + } + if p.Secrets["ABOX_MCP_GH_TOKEN"] != "mt" { + t.Errorf("set_mcp_tokens missing mcp token: %v", p.Secrets) + } + if _, ok := p.Secrets["XAI_API_KEY"]; ok { + t.Errorf("model credential leaked into set_mcp_tokens: %v", p.Secrets) + } + default: + t.Errorf("unexpected method %q", frame.Method) + } + ok, _ := protocol.EncodeParams(map[string]bool{"ok": true}) + reply(protocol.Frame{ID: frame.ID, Result: ok}) + } + + model := config.Model{Name: "grok-default", Provider: "xai", CredentialEnv: "XAI_API_KEY"} + err := s.PushSecrets(context.Background(), model, map[string]string{ + "XAI_API_KEY": "mk", "ABOX_MCP_GH_TOKEN": "mt", + }) + if err != nil { + t.Fatal(err) + } + mu.Lock() + defer mu.Unlock() + if len(order) != 2 || order[0] != "set_model" || order[1] != "set_mcp_tokens" { + t.Fatalf("order %v", order) + } +} + +func TestPushSecretsProto3FiltersModelCredential(t *testing.T) { + s, g := newPipeSandbox(t, 3) + g.onRequest = func(frame protocol.Frame, reply func(protocol.Frame)) { + switch frame.Method { + case "set_model": + p, err := protocol.DecodeParams[protocol.SetModelParams](frame.Params) + if err != nil { + t.Errorf("params: %v", err) + } + if len(p.Secrets) != 0 { + t.Errorf("proto-3 set_model must carry no secrets: %v", p.Secrets) + } + case "set_mcp_tokens": + p, _ := protocol.DecodeParams[protocol.SetMCPTokensParams](frame.Params) + if p.Secrets["ABOX_MCP_GH_TOKEN"] != "mt" { + t.Errorf("mcp token missing: %v", p.Secrets) + } + if _, ok := p.Secrets["XAI_API_KEY"]; ok { + t.Errorf("model credential leaked to proto-3 guest: %v", p.Secrets) + } + default: + t.Errorf("unexpected method %q", frame.Method) + } + ok, _ := protocol.EncodeParams(map[string]bool{"ok": true}) + reply(protocol.Frame{ID: frame.ID, Result: ok}) + } + model := config.Model{Name: "grok-default", Provider: "xai", CredentialEnv: "XAI_API_KEY"} + err := s.PushSecrets(context.Background(), model, map[string]string{ + "XAI_API_KEY": "mk", "ABOX_MCP_GH_TOKEN": "mt", + }) + if err != nil { + t.Fatal(err) + } +} + +func TestPushSecretsProto1Refused(t *testing.T) { + s, _ := newPipeSandbox(t, 1) + model := config.Model{Name: "grok-default", Provider: "xai", CredentialEnv: "XAI_API_KEY"} + err := s.PushSecrets(context.Background(), model, map[string]string{"XAI_API_KEY": "k"}) + if err == nil || !strings.Contains(err.Error(), "make image") { + t.Fatalf("got %v", err) + } +} + +func TestPushSecretsEmptyNoop(t *testing.T) { + s, g := newPipeSandbox(t, 1) + g.onRequest = func(frame protocol.Frame, _ func(protocol.Frame)) { + t.Errorf("unexpected frame %q", frame.Method) + } + model := config.Model{Name: "grok-default", Provider: "xai", CredentialEnv: "XAI_API_KEY"} + if err := s.PushSecrets(context.Background(), model, nil); err != nil { + t.Fatal(err) + } +} + +func TestSetModelDropsSecretsOnProto3(t *testing.T) { + s, g := newPipeSandbox(t, 3) + g.onRequest = func(frame protocol.Frame, reply func(protocol.Frame)) { + p, err := protocol.DecodeParams[protocol.SetModelParams](frame.Params) + if err != nil { + t.Errorf("params: %v", err) + } + if len(p.Secrets) != 0 { + t.Errorf("proto-3 set_model secrets: %v", p.Secrets) + } + ok, _ := protocol.EncodeParams(map[string]bool{"ok": true}) + reply(protocol.Frame{ID: frame.ID, Result: ok}) + } + model := config.Model{Name: "g", Provider: "xai", CredentialEnv: "XAI_API_KEY"} + if err := s.SetModel(context.Background(), model, map[string]string{"XAI_API_KEY": "k"}); err != nil { + t.Fatal(err) + } +} + +func TestGuestCallMidTurn(t *testing.T) { + s, g := newPipeSandbox(t, 3) + turnStarted := make(chan string, 1) + gotOpen := make(chan protocol.ProviderOpenResult, 1) + g.onRequest = func(frame protocol.Frame, reply func(protocol.Frame)) { + if frame.Method != "user_turn" { + t.Errorf("unexpected method %q", frame.Method) + ok, _ := protocol.EncodeParams(map[string]bool{"ok": true}) + reply(protocol.Frame{ID: frame.ID, Result: ok}) + return + } + go func() { + open, _ := protocol.EncodeParams(protocol.ProviderOpenParams{Model: "grok-default"}) + g.write(protocol.Frame{V: protocol.Version, ID: "g-1", Method: "provider_open", Params: open}) + turnStarted <- frame.ID + }() + } + g.onResponse = func(frame protocol.Frame) { + if frame.ID != "g-1" { + return + } + var res protocol.ProviderOpenResult + if err := json.Unmarshal(frame.Result, &res); err != nil { + t.Errorf("open result: %v", err) + return + } + gotOpen <- res + } + s.OnGuestCall = handlerFunc(func(ctx context.Context, method string, params json.RawMessage, notify func(string, any) error) (any, *protocol.Error) { + if method != "provider_open" { + return nil, &protocol.Error{Code: "host", Message: "unexpected " + method} + } + return protocol.ProviderOpenResult{StreamID: "s1"}, nil + }) + + done := make(chan error, 1) + go func() { + done <- s.UserTurn(context.Background(), "hi", nil) + }() + + turnID := <-turnStarted + res := <-gotOpen + if res.StreamID != "s1" { + t.Fatalf("open result %+v", res) + } + ok, _ := protocol.EncodeParams(map[string]bool{"ok": true}) + g.write(protocol.Frame{ID: turnID, Result: ok}) + select { + case err := <-done: + if err != nil { + t.Fatal(err) + } + case <-time.After(2 * time.Second): + t.Fatal("turn did not complete") + } +} + +func TestGuestCallUnknownMethodTypedError(t *testing.T) { + s, g := newPipeSandbox(t, 3) + s.startReading() // the test guest speaks before any host call starts the loop + gotReply := make(chan protocol.Frame, 1) + g.onRequest = func(frame protocol.Frame, reply func(protocol.Frame)) { + ok, _ := protocol.EncodeParams(map[string]bool{"ok": true}) + reply(protocol.Frame{ID: frame.ID, Result: ok}) + } + g.onResponse = func(frame protocol.Frame) { + if frame.ID == "g-9" { + gotReply <- frame + } + } + s.OnGuestCall = handlerFunc(func(ctx context.Context, method string, _ json.RawMessage, _ func(string, any) error) (any, *protocol.Error) { + return nil, &protocol.Error{Code: "host", Message: "unknown guest method " + method} + }) + + g.write(protocol.Frame{V: protocol.Version, ID: "g-9", Method: "bogus", Params: []byte(`{}`)}) + select { + case frame := <-gotReply: + if frame.Error == nil || !strings.Contains(frame.Error.Message, "unknown guest method bogus") { + t.Fatalf("frame %+v", frame) + } + case <-time.After(2 * time.Second): + t.Fatal("no reply to unknown guest method") + } +} + +func TestGuestCallRejectedBeforeProtocol3(t *testing.T) { + s, g := newPipeSandbox(t, 2) + s.startReading() + called := make(chan struct{}, 1) + s.OnGuestCall = handlerFunc(func(ctx context.Context, method string, params json.RawMessage, notify func(string, any) error) (any, *protocol.Error) { + called <- struct{}{} + return protocol.ProviderOpenResult{StreamID: "s1"}, nil + }) + gotReply := make(chan protocol.Frame, 1) + g.onResponse = func(frame protocol.Frame) { + if frame.ID == "g-1" { + gotReply <- frame + } + } + g.write(protocol.Frame{V: 2, ID: "g-1", Method: "provider_open", Params: []byte(`{"model":"x"}`)}) + select { + case frame := <-gotReply: + if frame.Error == nil || !strings.Contains(frame.Error.Message, "protocol 3") { + t.Fatalf("frame %+v", frame) + } + if frame.Result != nil { + t.Fatalf("protocol-2 guest received broker result %s", frame.Result) + } + case <-time.After(2 * time.Second): + t.Fatal("no reply rejecting protocol-2 broker call") + } + select { + case <-called: + t.Fatal("broker handler ran for protocol-2 guest") + default: + } +} + +func TestGuestCallNoHandlerErrors(t *testing.T) { + s, g := newPipeSandbox(t, 3) + s.startReading() + s.OnGuestCall = nil + gotReply := make(chan protocol.Frame, 1) + g.onResponse = func(frame protocol.Frame) { + if frame.ID == "g-2" { + gotReply <- frame + } + } + g.write(protocol.Frame{V: protocol.Version, ID: "g-2", Method: "provider_open", Params: []byte(`{}`)}) + select { + case frame := <-gotReply: + if frame.Error == nil || !strings.Contains(frame.Error.Message, "not supported") { + t.Fatalf("frame %+v", frame) + } + case <-time.After(2 * time.Second): + t.Fatal("no reply without handler") + } +} + +func TestGuestCallConcurrencyReturnsBusy(t *testing.T) { + s, g := newPipeSandbox(t, 3) + regularSlots := protocol.MaxGuestCalls - 1 + started := make(chan struct{}, regularSlots) + release := make(chan struct{}) + s.OnGuestCall = handlerFunc(func(ctx context.Context, method string, params json.RawMessage, notify func(string, any) error) (any, *protocol.Error) { + if method == "provider_cancel" { + return map[string]bool{"ok": true}, nil + } + started <- struct{}{} + <-release + return map[string]bool{"ok": true}, nil + }) + busy := make(chan protocol.Frame, 1) + cancelReply := make(chan protocol.Frame, 1) + g.onResponse = func(frame protocol.Frame) { + if frame.Error != nil && frame.Error.Code == "busy" { + busy <- frame + } + if frame.ID == "g-cancel" { + cancelReply <- frame + } + } + s.startReading() + for i := 0; i < regularSlots; i++ { + g.write(protocol.Frame{V: protocol.Version, ID: fmt.Sprintf("g-%d", i), Method: "hold", Params: []byte(`{}`)}) + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("handler did not start") + } + } + g.write(protocol.Frame{V: protocol.Version, ID: "g-cancel", Method: "provider_cancel", Params: []byte(`{"stream_id":"s1"}`)}) + select { + case frame := <-cancelReply: + if frame.Error != nil { + t.Fatalf("provider cancellation was blocked: %+v", frame.Error) + } + case <-time.After(2 * time.Second): + t.Fatal("provider cancellation was blocked by regular guest calls") + } + g.write(protocol.Frame{V: protocol.Version, ID: "g-busy", Method: "hold", Params: []byte(`{}`)}) + select { + case frame := <-busy: + if frame.ID != "g-busy" || frame.Error.Code != "busy" { + t.Fatalf("reply %+v", frame) + } + case <-time.After(2 * time.Second): + t.Fatal("no typed busy reply") + } + close(release) +} + +func TestGuestCallContextEndsWithConnection(t *testing.T) { + s, g := newPipeSandbox(t, 3) + canceled := make(chan struct{}) + s.OnGuestCall = handlerFunc(func(ctx context.Context, method string, params json.RawMessage, notify func(string, any) error) (any, *protocol.Error) { + <-ctx.Done() + close(canceled) + return nil, &protocol.Error{Code: "canceled", Message: ctx.Err().Error()} + }) + s.startReading() + g.write(protocol.Frame{V: protocol.Version, ID: "g-life", Method: "hold", Params: []byte(`{}`)}) + s.failConnection(errors.New("test disconnect")) + select { + case <-canceled: + case <-time.After(2 * time.Second): + t.Fatal("handler context survived connection") + } +} + +func TestConnectionFailureWakesPendingCall(t *testing.T) { + host, guest := net.Pipe() + t.Cleanup(func() { host.Close(); guest.Close() }) + s := &Sandbox{conn: host, GuestProtocol: 3} + requestRead := make(chan struct{}) + go func() { + _, _ = protocol.ReadFrame(guest) + close(requestRead) + _ = guest.Close() + }() + done := make(chan error, 1) + go func() { + done <- s.Call(context.Background(), "wait", map[string]bool{"ok": true}, nil) + }() + <-requestRead + select { + case err := <-done: + if err == nil { + t.Fatal("pending call succeeded after disconnect") + } + case <-time.After(2 * time.Second): + t.Fatal("pending call was not woken") + } +} + +type handlerFunc func(ctx context.Context, method string, params json.RawMessage, notify func(string, any) error) (any, *protocol.Error) + +func (h handlerFunc) Handle(ctx context.Context, method string, params json.RawMessage, notify func(string, any) error) (any, *protocol.Error) { + return h(ctx, method, params, notify) +} + +func TestPushSecretsOnClosedConn(t *testing.T) { + s, g := newPipeSandbox(t, 2) + g.onRequest = func(frame protocol.Frame, reply func(protocol.Frame)) { + ok, _ := protocol.EncodeParams(map[string]bool{"ok": true}) + reply(protocol.Frame{ID: frame.ID, Result: ok}) + } + g.conn.Close() + s.conn.Close() + model := config.Model{Name: "g", CredentialEnv: "XAI_API_KEY"} + err := s.PushSecrets(context.Background(), model, map[string]string{"XAI_API_KEY": "k"}) + if err == nil { + t.Fatal("expected connection error") + } + if !errors.Is(err, context.DeadlineExceeded) && !strings.Contains(err.Error(), "closed") { + t.Fatalf("got %v", err) + } +} diff --git a/internal/runtime/runtime_test.go b/internal/runtime/runtime_test.go index c9488a9..fc05140 100644 --- a/internal/runtime/runtime_test.go +++ b/internal/runtime/runtime_test.go @@ -42,7 +42,7 @@ func TestPrepareResumeDoesNotClobberRoot(t *testing.T) { if err := os.WriteFile(golden, []byte("GOLDEN"), 0o600); err != nil { t.Fatal(err) } - err = Prepare(s, golden, config.Model{Name: "grok", Provider: "xai", Model: "grok-4"}, nil, nil, true) + err = Prepare(s, golden, config.Model{Name: "grok", Provider: "xai", Model: "grok-4"}, nil, true) if err != nil { t.Fatal(err) } @@ -70,7 +70,7 @@ func TestPrepareResumeRewritesReadOnlyConfig(t *testing.T) { if err := os.WriteFile(s.ConfigDisk(), make([]byte, 1<<20), 0o400); err != nil { t.Fatal(err) } - err = Prepare(s, "", config.Model{Name: "grok", Provider: "xai", Model: "grok-4"}, map[string]string{"XAI_API_KEY": "k"}, nil, true) + err = Prepare(s, "", config.Model{Name: "grok", Provider: "xai", Model: "grok-4"}, nil, true) if err != nil { t.Fatal(err) } diff --git a/internal/runtime/runtime_turn_test.go b/internal/runtime/runtime_turn_test.go index 3fcfcaf..3f1a55e 100644 --- a/internal/runtime/runtime_turn_test.go +++ b/internal/runtime/runtime_turn_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "net" + "strings" "sync" "testing" "time" @@ -60,6 +61,142 @@ func TestUserTurnPlainOnPipe(t *testing.T) { } } +func TestUserTurnDoesNotDropBurstFrames(t *testing.T) { + host, guest := net.Pipe() + t.Cleanup(func() { host.Close(); guest.Close() }) + s := &Sandbox{conn: host, GuestProtocol: 2} + + const eventCount = 200 + written := make(chan error, 1) + go func() { + turn, err := protocol.ReadFrame(guest) + if err != nil { + written <- err + return + } + for i := 0; i < eventCount; i++ { + raw, _ := protocol.EncodeParams(protocol.AgentEvent{Kind: "text", Text: "x"}) + if err := protocol.WriteFrame(guest, protocol.Frame{ID: turn.ID, Method: "agent_event", Params: raw}); err != nil { + written <- err + return + } + } + written <- protocol.WriteFrame(guest, protocol.Frame{ID: turn.ID, Result: []byte(`{"ok":true}`)}) + }() + + first := make(chan struct{}) + release := make(chan struct{}) + var once sync.Once + done := make(chan error, 1) + count := 0 + go func() { + done <- s.UserTurn(context.Background(), "hi", func(protocol.AgentEvent) { + count++ + once.Do(func() { + close(first) + <-release + }) + }) + }() + <-first + if err := <-written; err != nil { + t.Fatal(err) + } + close(release) + if err := <-done; err != nil { + t.Fatal(err) + } + if count != eventCount { + t.Fatalf("got %d events, want %d", count, eventCount) + } +} + +func TestFrameQueueBudgets(t *testing.T) { + q := newFrameQueue(1, 128) + if err := q.push(protocol.Frame{ID: "1"}); err != nil { + t.Fatal(err) + } + if err := q.push(protocol.Frame{ID: "2"}); !errors.Is(err, errQueueFull) { + t.Fatalf("count overflow: %v", err) + } + + q = newFrameQueue(2, 32) + if err := q.push(protocol.Frame{ID: "large", Result: []byte(`{"value":"too large for queue"}`)}); !errors.Is(err, errQueueFull) { + t.Fatalf("byte overflow: %v", err) + } +} + +func TestUserTurnQueueOverflowFailsClearly(t *testing.T) { + host, guest := net.Pipe() + t.Cleanup(func() { host.Close(); guest.Close() }) + s := &Sandbox{conn: host, GuestProtocol: 2} + + guestDone := make(chan error, 1) + go func() { + turn, err := protocol.ReadFrame(guest) + if err != nil { + guestDone <- err + return + } + raw, _ := protocol.EncodeParams(protocol.AgentEvent{Kind: "text", Text: "x"}) + for i := 0; i < turnQueueFrames+16; i++ { + if err := protocol.WriteFrame(guest, protocol.Frame{ID: turn.ID, Method: "agent_event", Params: raw}); err != nil { + guestDone <- err + return + } + } + guestDone <- nil + }() + + first := make(chan struct{}) + release := make(chan struct{}) + var once sync.Once + done := make(chan error, 1) + go func() { + done <- s.UserTurn(context.Background(), "hi", func(protocol.AgentEvent) { + once.Do(func() { + close(first) + <-release + }) + }) + }() + <-first + select { + case <-guestDone: + case <-time.After(2 * time.Second): + t.Fatal("overflow did not close the connection") + } + close(release) + select { + case err := <-done: + if err == nil || !strings.Contains(err.Error(), "turn frame queue overflow") { + t.Fatalf("got %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("turn did not fail after overflow") + } +} + +func TestCallWriteHonorsContextUnderBackpressure(t *testing.T) { + host, guest := net.Pipe() + t.Cleanup(func() { host.Close(); guest.Close() }) + s := &Sandbox{conn: host, GuestProtocol: 2} + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + done := make(chan error, 1) + go func() { + done <- s.Call(ctx, "blocked", map[string]bool{"ok": true}, nil) + }() + select { + case err := <-done: + if err == nil { + t.Fatal("blocked write succeeded") + } + case <-time.After(2 * time.Second): + t.Fatal("blocked write ignored context") + } +} + func TestUserTurnCtxCancelWritesCancelTurn(t *testing.T) { host, guest := net.Pipe() t.Cleanup(func() { host.Close(); guest.Close() }) diff --git a/internal/session/scrub.go b/internal/session/scrub.go new file mode 100644 index 0000000..5f9def5 --- /dev/null +++ b/internal/session/scrub.go @@ -0,0 +1,156 @@ +package session + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "path/filepath" + + "github.com/AdminTurnedDevOps/ABox/internal/config" +) + +// ScrubSecrets removes the "secrets" key from guest-config.json and +// config.raw. It never deletes sessions. +func ScrubSecrets(root string) (int, error) { + entries, err := os.ReadDir(root) + if err != nil { + if os.IsNotExist(err) { + return 0, nil + } + return 0, err + } + scrubbed := 0 + var scrubErrs []error + for _, e := range entries { + if !e.IsDir() { + continue + } + dir := filepath.Join(root, e.Name()) + n, err := scrubSessionDir(dir) + if n > 0 { + scrubbed++ + } + if err != nil { + scrubErrs = append(scrubErrs, fmt.Errorf("session %s: %w", e.Name(), err)) + } + } + return scrubbed, errors.Join(scrubErrs...) +} + +func ScrubSecretsEverywhere() (int, error) { + n, err := ScrubSecrets(config.SessionRoot()) + legacyDir := config.LegacyAppSupportDir() + if legacyDir == "" { + return n, err + } + n2, err2 := ScrubSecrets(filepath.Join(legacyDir, "sessions")) + return n + n2, errors.Join(err, err2) +} + +func scrubSessionDir(dir string) (int, error) { + rewritten := 0 + var scrubErrs []error + guestConfig := filepath.Join(dir, "guest-config.json") + if ok, err := scrubJSONFile(guestConfig); err != nil { + scrubErrs = append(scrubErrs, err) + } else if ok { + rewritten++ + } + configDisk := filepath.Join(dir, "config.raw") + if ok, err := scrubConfigDisk(configDisk); err != nil { + scrubErrs = append(scrubErrs, err) + } else if ok { + rewritten++ + } + return rewritten, errors.Join(scrubErrs...) +} + +func scrubJSONFile(path string) (bool, error) { + data, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + return false, nil + } + return false, fmt.Errorf("read %s: %w", path, err) + } + scrubbed, err := scrubbedJSONObject(data) + if err != nil { + return false, fmt.Errorf("parse %s: %w", path, err) + } + if scrubbed == nil { + return false, nil + } + tmp := path + ".tmp" + if err := os.WriteFile(tmp, scrubbed, 0o600); err != nil { + return false, fmt.Errorf("write %s: %w", path, err) + } + if err := os.Rename(tmp, path); err != nil { + return false, fmt.Errorf("replace %s: %w", path, err) + } + if err := os.Chmod(path, 0o600); err != nil { + return true, fmt.Errorf("chmod %s: %w", path, err) + } + return true, nil +} + +func scrubConfigDisk(path string) (bool, error) { + f, err := os.Open(path) + if err != nil { + if os.IsNotExist(err) { + return false, nil + } + return false, fmt.Errorf("open %s: %w", path, err) + } + buf, err := io.ReadAll(io.LimitReader(f, ConfigDiskSize+1)) + closeErr := f.Close() + if err != nil { + return false, fmt.Errorf("read %s: %w", path, err) + } + if closeErr != nil { + return false, fmt.Errorf("close %s: %w", path, closeErr) + } + if len(buf) > ConfigDiskSize { + return false, fmt.Errorf("read %s: config disk exceeds %d bytes", path, ConfigDiskSize) + } + data := buf + if i := bytes.IndexByte(data, 0); i >= 0 { // guest JSON ends at first NUL + data = data[:i] + } + scrubbed, err := scrubbedJSONObject(data) + if err != nil { + return false, fmt.Errorf("parse %s: %w", path, err) + } + if scrubbed == nil { + return false, nil + } + if err := WritePaddedConfig(path, scrubbed); err != nil { + return false, fmt.Errorf("rewrite %s: %w", path, err) + } + return true, nil +} + +func scrubbedJSONObject(data []byte) ([]byte, error) { + data = bytes.TrimSpace(data) + if len(data) == 0 { + return nil, fmt.Errorf("empty JSON document") + } + var obj map[string]json.RawMessage + if err := json.Unmarshal(data, &obj); err != nil { + return nil, fmt.Errorf("malformed JSON object: %w", err) + } + if obj == nil { + return nil, fmt.Errorf("expected JSON object") + } + if _, ok := obj["secrets"]; !ok { + return nil, nil + } + delete(obj, "secrets") + out, err := json.MarshalIndent(obj, "", " ") + if err != nil { + return nil, fmt.Errorf("scrub: %w", err) + } + return out, nil +} diff --git a/internal/session/scrub_test.go b/internal/session/scrub_test.go new file mode 100644 index 0000000..0f2b2c1 --- /dev/null +++ b/internal/session/scrub_test.go @@ -0,0 +1,234 @@ +package session + +import ( + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/AdminTurnedDevOps/ABox/internal/config" +) + +func writeLegacySession(t *testing.T, root, id string, withSecrets bool) { + t.Helper() + dir := filepath.Join(root, id) + if err := os.MkdirAll(dir, 0o700); err != nil { + t.Fatal(err) + } + guestCfg := map[string]any{ + "session_id": id, + "capability": "cap", + "model": map[string]any{"name": "grok"}, + } + if withSecrets { + guestCfg["secrets"] = map[string]string{"XAI_API_KEY": "legacy-key"} + } + data, err := json.MarshalIndent(guestCfg, "", " ") + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "guest-config.json"), data, 0o600); err != nil { + t.Fatal(err) + } + + // config.raw: JSON + NUL padding to 1 MiB. + cfgData := data + if withSecrets { + var obj map[string]json.RawMessage + _ = json.Unmarshal(data, &obj) + obj["secrets"], _ = json.Marshal(map[string]string{"XAI_API_KEY": "legacy-key"}) + cfgData, _ = json.MarshalIndent(obj, "", " ") + } + if err := WritePaddedConfig(filepath.Join(dir, "config.raw"), cfgData); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "root.raw"), []byte("session disk bytes"), 0o600); err != nil { + t.Fatal(err) + } +} + +func TestScrubSecretsEverywhereIncludesLegacyAppSupport(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + aboxHome := t.TempDir() + t.Setenv("ABOX_HOME", aboxHome) + writeLegacySession(t, config.SessionRoot(), "modern", true) + legacyRoot := filepath.Join(config.LegacyAppSupportDir(), "sessions") + writeLegacySession(t, legacyRoot, "old", true) + n, err := ScrubSecretsEverywhere() + if err != nil { + t.Fatal(err) + } + if n != 2 { + t.Fatalf("scrubbed %d", n) + } + for _, dir := range []string{ + filepath.Join(config.SessionRoot(), "modern"), + filepath.Join(legacyRoot, "old"), + } { + body, err := os.ReadFile(filepath.Join(dir, "guest-config.json")) + if err != nil { + t.Fatal(err) + } + if strings.Contains(string(body), `"secrets"`) { + t.Fatalf("%s still has secrets: %s", dir, body) + } + } +} + +func TestScrubSecretsRemovesLegacySecrets(t *testing.T) { + root := t.TempDir() + writeLegacySession(t, root, "aaa", true) + writeLegacySession(t, root, "bbb", true) + + n, err := ScrubSecrets(root) + if err != nil { + t.Fatal(err) + } + if n != 2 { + t.Fatalf("scrubbed %d sessions, want 2", n) + } + + for _, id := range []string{"aaa", "bbb"} { + dir := filepath.Join(root, id) + data, err := os.ReadFile(filepath.Join(dir, "guest-config.json")) + if err != nil { + t.Fatal(err) + } + if strings.Contains(string(data), "secrets") || strings.Contains(string(data), "legacy-key") { + t.Fatalf("guest-config still has secrets: %s", data) + } + var obj map[string]json.RawMessage + if err := json.Unmarshal(data, &obj); err != nil { + t.Fatalf("scrubbed guest-config is not valid JSON: %v", err) + } + if obj["session_id"] == nil || obj["model"] == nil { + t.Fatalf("scrub dropped unrelated fields: %s", data) + } + st, _ := os.Stat(filepath.Join(dir, "guest-config.json")) + if st.Mode().Perm() != 0o600 { + t.Fatalf("guest-config perm %o", st.Mode().Perm()) + } + + raw, err := os.ReadFile(filepath.Join(dir, "config.raw")) + if err != nil { + t.Fatal(err) + } + if len(raw) != ConfigDiskSize { + t.Fatalf("config.raw size %d", len(raw)) + } + if strings.Contains(string(raw), "legacy-key") || strings.Contains(string(raw), "secrets") { + t.Fatalf("config.raw still has secrets") + } + st, _ = os.Stat(filepath.Join(dir, "config.raw")) + if st.Mode().Perm() != 0o400 { + t.Fatalf("config.raw perm %o", st.Mode().Perm()) + } + // root.raw untouched. + disk, err := os.ReadFile(filepath.Join(dir, "root.raw")) + if err != nil { + t.Fatal(err) + } + if string(disk) != "session disk bytes" { + t.Fatalf("root.raw was modified: %q", disk) + } + } +} + +func TestScrubSecretsIdempotent(t *testing.T) { + root := t.TempDir() + writeLegacySession(t, root, "aaa", true) + if n, err := ScrubSecrets(root); err != nil || n != 1 { + t.Fatalf("first pass: n=%d err=%v", n, err) + } + if n, err := ScrubSecrets(root); err != nil || n != 0 { + t.Fatalf("second pass: n=%d err=%v (want no-op)", n, err) + } +} + +func TestScrubSecretsSkipsCleanSessions(t *testing.T) { + root := t.TempDir() + writeLegacySession(t, root, "clean", false) + if n, err := ScrubSecrets(root); err != nil || n != 0 { + t.Fatalf("n=%d err=%v", n, err) + } +} + +func TestScrubSecretsMissingRoot(t *testing.T) { + n, err := ScrubSecrets(filepath.Join(t.TempDir(), "nope")) + if err != nil || n != 0 { + t.Fatalf("n=%d err=%v", n, err) + } +} + +func TestScrubSecretsContinuesAndCombinesSessionErrors(t *testing.T) { + root := t.TempDir() + writeLegacySession(t, root, "good", true) + + badDir := filepath.Join(root, "bad-json") + if err := os.MkdirAll(badDir, 0o700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(badDir, "guest-config.json"), []byte(`{"secrets":`), 0o600); err != nil { + t.Fatal(err) + } + if err := WritePaddedConfig(filepath.Join(badDir, "config.raw"), []byte(`not-json`)); err != nil { + t.Fatal(err) + } + + unreadableDir := filepath.Join(root, "unreadable") + if err := os.MkdirAll(filepath.Join(unreadableDir, "guest-config.json"), 0o700); err != nil { + t.Fatal(err) + } + + n, err := ScrubSecrets(root) + if n != 1 { + t.Fatalf("scrubbed %d sessions, want 1", n) + } + if err == nil { + t.Fatal("expected combined scrub errors") + } + for _, want := range []string{"bad-json", "guest-config.json", "config.raw", "unreadable"} { + if !strings.Contains(err.Error(), want) { + t.Fatalf("error %q does not contain %q", err, want) + } + } + good, readErr := os.ReadFile(filepath.Join(root, "good", "guest-config.json")) + if readErr != nil { + t.Fatal(readErr) + } + if strings.Contains(string(good), "legacy-key") { + t.Fatalf("good session was not scrubbed: %s", good) + } +} + +func TestScrubbedJSONObjectRejectsMalformedJSON(t *testing.T) { + for _, data := range [][]byte{nil, []byte(`{"clean":`), []byte(`null`), []byte(`[]`)} { + if scrubbed, err := scrubbedJSONObject(data); err == nil || scrubbed != nil { + t.Fatalf("data %q: scrubbed=%q err=%v", data, scrubbed, err) + } + } +} + +func TestScrubbedJSONObjectPreservesUnknownRawValues(t *testing.T) { + input := []byte(`{"unknown":{"large":9007199254740993123456789,"future":[true,{"x":"y"}]},"secrets":{"TOKEN":"secret"},"tail":"kept"}`) + out, err := scrubbedJSONObject(input) + if err != nil { + t.Fatal(err) + } + var obj map[string]json.RawMessage + if err := json.Unmarshal(out, &obj); err != nil { + t.Fatal(err) + } + if _, exists := obj["secrets"]; exists { + t.Fatalf("secrets retained: %s", out) + } + var unknown map[string]json.RawMessage + if err := json.Unmarshal(obj["unknown"], &unknown); err != nil { + t.Fatal(err) + } + if string(unknown["large"]) != "9007199254740993123456789" || string(obj["tail"]) != `"kept"` { + t.Fatalf("unknown values changed: %s", out) + } +} diff --git a/internal/session/session.go b/internal/session/session.go index 0949c6b..6439417 100644 --- a/internal/session/session.go +++ b/internal/session/session.go @@ -172,7 +172,7 @@ func WriteTranscript(path string, lines []string) error { return os.WriteFile(path, data, 0o600) } -func (s *Session) WriteGuestConfig(model config.Model, secrets map[string]string, servers []config.MCPServer) error { +func (s *Session) WriteGuestConfig(model config.Model, servers []config.MCPServer) error { var gs []protocol.GuestMCPServer for _, srv := range servers { gs = append(gs, protocol.GuestMCPServer{ @@ -188,7 +188,6 @@ func (s *Session) WriteGuestConfig(model config.Model, secrets map[string]string VsockPort: protocol.RPCPort, RepoDir: protocol.GuestRepoDir, Model: model.ToGuest(), - Secrets: secrets, MCPServers: gs, } data, err := json.MarshalIndent(cfg, "", " ") @@ -198,6 +197,33 @@ func (s *Session) WriteGuestConfig(model config.Model, secrets map[string]string return os.WriteFile(s.GuestConfigJSON(), data, 0o600) } +// ConfigDiskSize is the fixed size of the sealed read-only config disk. +const ConfigDiskSize = 1 << 20 + +// WritePaddedConfig writes data to path zero-padded to ConfigDiskSize bytes +// and leaves the file read-only (mode 0400). It is the shared writer for the +// config.raw disk: creation, resume rewrite, and secret scrubbing all keep +// the same layout the guest parses (JSON, then NUL padding). +func WritePaddedConfig(path string, data []byte) error { + if len(data) > ConfigDiskSize { + return fmt.Errorf("config disk payload too large: %d", len(data)) + } + // Resume rewrites config.raw; the previous run left it mode 0400. + _ = os.Chmod(path, 0o600) + f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600) + if err != nil { + return err + } + defer f.Close() + if _, err := f.Write(data); err != nil { + return err + } + if _, err := f.Write(make([]byte, ConfigDiskSize-len(data))); err != nil { + return err + } + return os.Chmod(path, 0o400) +} + func randomHex(n int) (string, error) { b := make([]byte, n) if _, err := rand.Read(b); err != nil { diff --git a/internal/session/session_test.go b/internal/session/session_test.go index 1e8a6bc..0b25df9 100644 --- a/internal/session/session_test.go +++ b/internal/session/session_test.go @@ -10,13 +10,13 @@ import ( "github.com/AdminTurnedDevOps/ABox/internal/config" ) -func TestWriteGuestConfigIncludesMCP(t *testing.T) { +func TestWriteGuestConfigIncludesMCPNoSecrets(t *testing.T) { t.Setenv("HOME", t.TempDir()) s, err := Create("/repo", "deadbeef") if err != nil { t.Fatal(err) } - err = s.WriteGuestConfig(config.Model{Name: "grok"}, map[string]string{"ABOX_MCP_GH_TOKEN": "tok"}, []config.MCPServer{ + err = s.WriteGuestConfig(config.Model{Name: "grok"}, []config.MCPServer{ {Name: "gh", URL: "https://api.githubcopilot.com/mcp/", CredentialEnv: "ABOX_MCP_GH_TOKEN"}, }) if err != nil { @@ -30,6 +30,48 @@ func TestWriteGuestConfigIncludesMCP(t *testing.T) { if !strings.Contains(body, "api.githubcopilot.com") || !strings.Contains(body, "ABOX_MCP_GH_TOKEN") { t.Fatalf("guest config missing mcp: %s", body) } + if strings.Contains(body, `"secrets"`) { + t.Fatalf("guest config leaked a secrets key: %s", body) + } + st, err := os.Stat(s.GuestConfigJSON()) + if err != nil { + t.Fatal(err) + } + if st.Mode().Perm() != 0o600 { + t.Fatalf("perm %o", st.Mode().Perm()) + } +} + +func TestWritePaddedConfigLayout(t *testing.T) { + path := filepath.Join(t.TempDir(), "config.raw") + data := []byte(`{"session_id":"x"}`) + if err := WritePaddedConfig(path, data); err != nil { + t.Fatal(err) + } + raw, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + if len(raw) != ConfigDiskSize { + t.Fatalf("size %d", len(raw)) + } + if string(raw[:len(data)]) != string(data) { + t.Fatalf("payload %q", raw[:len(data)]) + } + if raw[len(data)] != 0 { + t.Fatal("missing zero padding") + } + st, err := os.Stat(path) + if err != nil { + t.Fatal(err) + } + if st.Mode().Perm() != 0o400 { + t.Fatalf("perm %o", st.Mode().Perm()) + } + // Rewriting a read-only file must still work (resume path). + if err := WritePaddedConfig(path, []byte(`{"session_id":"y"}`)); err != nil { + t.Fatal(err) + } } func TestLoadRequiresRootRaw(t *testing.T) { diff --git a/internal/tui/commands.go b/internal/tui/commands.go index 5a994c3..403b2b4 100644 --- a/internal/tui/commands.go +++ b/internal/tui/commands.go @@ -1,10 +1,12 @@ package tui import ( + "context" + "fmt" "strings" "github.com/AdminTurnedDevOps/ABox/internal/config" - "github.com/AdminTurnedDevOps/ABox/internal/credentials" + "github.com/AdminTurnedDevOps/ABox/internal/credsource" ) type slashCmd struct { @@ -44,22 +46,64 @@ func mcpServers(cfg config.File) []config.MCPServer { return servers } -func applyMCPKey(server config.MCPServer, key string) (string, error) { +var saveCredential = func(envName, value string) (credsource.SaveResult, error) { + return credsource.SavePreferred(context.Background(), envName, value) +} + +func applyMCPKey(cfg config.File, server config.MCPServer, key string) (config.File, string, string, error) { env := config.TokenEnv(server) - if err := credentials.Save(env, key); err != nil { - return "", err + res, err := saveCredential(env, key) + if err != nil { + return cfg, "", "", err + } + found := false + for i, s := range cfg.MCPServers { + if s.Name == server.Name { + found = true + cfg.MCPServers[i].CredentialEnv = "" + cfg.MCPServers[i].Credential = &config.CredentialRef{Source: res.Source, Name: env} + if err := cfg.Save(); err != nil { + return cfg, env, res.Note, fmt.Errorf("config update: %w", err) + } + break + } + } + if !found { + return cfg, env, res.Note, fmt.Errorf("config update did not retain MCP server %q", server.Name) } - credentials.SetEnv(env, key) - return env, nil + return cfg, env, res.Note, nil } -func applyProviderKey(cfg config.File, choice config.ProviderProfile, key string) (config.Model, error) { - if err := credentials.Save(choice.Env, key); err != nil { - return config.Model{}, err +func applyProviderKey(cfg config.File, choice config.ProviderProfile, key string) (config.File, config.Model, string, error) { + res, err := saveCredential(choice.Env, key) + if err != nil { + return cfg, config.Model{}, "", err + } + found := false + changed := false + for i, m := range cfg.Models { + if m.Name == choice.Name { + found = true + cfg.Models[i].CredentialEnv = "" + cfg.Models[i].Credential = &config.CredentialRef{Source: res.Source, Name: choice.Env} + changed = true + } + } + if !found { + model := choice.ModelConfig() + model.CredentialEnv = "" + model.Credential = &config.CredentialRef{Source: res.Source, Name: choice.Env} + cfg.Models = append(cfg.Models, model) + changed = true + } + if changed { + if err := cfg.Save(); err != nil { + return cfg, config.Model{}, res.Note, fmt.Errorf("config update: %w", err) + } } - credentials.SetEnv(choice.Env, key) - if m, ok := cfg.ModelNamed(choice.Name); ok { - return m, nil + sel, ok := cfg.ModelNamed(choice.Name) + if !ok { + return cfg, config.Model{}, res.Note, fmt.Errorf("config update did not retain model profile %q", choice.Name) } - return choice.ModelConfig(), nil + return cfg, sel, res.Note, nil } diff --git a/internal/tui/commands_test.go b/internal/tui/commands_test.go new file mode 100644 index 0000000..1e595e2 --- /dev/null +++ b/internal/tui/commands_test.go @@ -0,0 +1,128 @@ +package tui + +import ( + "context" + "encoding/json" + "testing" + + "github.com/AdminTurnedDevOps/ABox/internal/config" + "github.com/AdminTurnedDevOps/ABox/internal/credsource" + "github.com/AdminTurnedDevOps/ABox/internal/runtime" + "github.com/AdminTurnedDevOps/ABox/protocol" +) + +func TestApplyProviderKeyAddsMissingProfileWithCurrentReference(t *testing.T) { + t.Setenv("ABOX_HOME", t.TempDir()) + original := saveCredential + t.Cleanup(func() { saveCredential = original }) + saveCredential = func(envName, value string) (credsource.SaveResult, error) { + if envName != "OPENAI_API_KEY" || value != "secret" { + t.Fatalf("save %q %q", envName, value) + } + return credsource.SaveResult{Source: "keychain", Keychain: true, Note: "test"}, nil + } + + cfg := config.Defaults() + cfg.Models = append(cfg.Models[:1:1], cfg.Models[2]) + choice := config.DefaultProviders()[1] + got, sel, _, err := applyProviderKey(cfg, choice, "secret") + if err != nil { + t.Fatal(err) + } + if sel.Name != choice.Name || sel.Credential == nil || sel.Credential.Source != "keychain" || sel.Credential.Name != choice.Env { + t.Fatalf("selected model %+v", sel) + } + persisted, _, err := config.Load() + if err != nil { + t.Fatal(err) + } + model, ok := persisted.ModelNamed(choice.Name) + if !ok || model.Credential == nil || model.Credential.Name != choice.Env { + t.Fatalf("persisted model %+v found=%v; cfg=%+v", model, ok, got.Models) + } +} + +func TestApplyProviderKeyFallbackReplacesExplicitCloudReference(t *testing.T) { + t.Setenv("ABOX_HOME", t.TempDir()) + original := saveCredential + t.Cleanup(func() { saveCredential = original }) + saveCredential = func(envName, value string) (credsource.SaveResult, error) { + return credsource.SaveResult{Source: "env", Note: "fallback"}, nil + } + + cfg := config.Defaults() + choice := config.DefaultProviders()[0] + cfg.Models[0].CredentialEnv = "" + cfg.Models[0].Credential = &config.CredentialRef{Source: "vault", Name: "secret/abox/xai"} + got, sel, _, err := applyProviderKey(cfg, choice, "secret") + if err != nil { + t.Fatal(err) + } + if sel.CredentialEnv != "" || sel.Credential == nil || sel.Credential.Source != "env" || sel.Credential.Name != choice.Env { + t.Fatalf("selected model retained stale reference: %+v", sel) + } + if got.Models[0].Credential == nil || got.Models[0].Credential.Source != "env" { + t.Fatalf("updated config %+v", got.Models[0]) + } + persisted, _, err := config.Load() + if err != nil { + t.Fatal(err) + } + persistedModel, ok := persisted.ModelNamed(choice.Name) + if !ok || persistedModel.Credential == nil || persistedModel.Credential.Source != "env" { + t.Fatalf("persisted model retained stale reference: %+v", persistedModel) + } +} + +func TestApplyMCPKeyFallbackReplacesExplicitNonEnvReference(t *testing.T) { + t.Setenv("ABOX_HOME", t.TempDir()) + original := saveCredential + t.Cleanup(func() { saveCredential = original }) + saveCredential = func(envName, value string) (credsource.SaveResult, error) { + return credsource.SaveResult{Source: "env", Note: "fallback"}, nil + } + + server := config.MCPServer{ + Name: "github", URL: "https://api.githubcopilot.com/mcp/", + Credential: &config.CredentialRef{Source: "aws", Name: "abox/github-token"}, + } + cfg := config.Defaults() + cfg.MCPServers = []config.MCPServer{server} + got, env, _, err := applyMCPKey(cfg, server, "secret") + if err != nil { + t.Fatal(err) + } + ref := got.MCPServers[0].Credential + if got.MCPServers[0].CredentialEnv != "" || ref == nil || ref.Source != "env" || ref.Name != env { + t.Fatalf("MCP server retained stale reference: %+v", got.MCPServers[0]) + } + persisted, _, err := config.Load() + if err != nil { + t.Fatal(err) + } + ref = persisted.MCPServers[0].Credential + if ref == nil || ref.Source != "env" || ref.Name != env { + t.Fatalf("persisted MCP server retained stale reference: %+v", persisted.MCPServers[0]) + } +} + +func TestUpdateHostBrokerUsesCurrentModelConfig(t *testing.T) { + cfg := config.Defaults() + cfg.Models = []config.Model{{ + Name: "updated", Provider: "openai", Model: "gpt-current", CredentialEnv: "CURRENT_API_KEY", + }} + sb := &runtime.Sandbox{} + m := model{sandbox: sb, resolver: credsource.NewResolver()} + t.Cleanup(func() { _ = m.resolver.Close() }) + m.updateHostBroker(cfg) + + raw, err := json.Marshal(protocol.ProviderOpenParams{Model: "updated"}) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + if _, perr := sb.OnGuestCall.Handle(ctx, "provider_open", raw, nil); perr != nil { + t.Fatalf("updated broker rejected current model: %v", perr) + } +} diff --git a/internal/tui/tui.go b/internal/tui/tui.go index 2256893..c20263a 100644 --- a/internal/tui/tui.go +++ b/internal/tui/tui.go @@ -3,8 +3,8 @@ package tui import ( "context" "fmt" - "os" "strings" + "time" "charm.land/bubbles/v2/key" "charm.land/bubbles/v2/textarea" @@ -13,6 +13,8 @@ import ( "charm.land/lipgloss/v2" "github.com/AdminTurnedDevOps/ABox/internal/config" + "github.com/AdminTurnedDevOps/ABox/internal/credsource" + "github.com/AdminTurnedDevOps/ABox/internal/llmbroker" "github.com/AdminTurnedDevOps/ABox/internal/runtime" "github.com/AdminTurnedDevOps/ABox/internal/session" "github.com/AdminTurnedDevOps/ABox/protocol" @@ -49,12 +51,24 @@ type model struct { err string cancel context.CancelFunc events <-chan protocol.AgentEvent + resolver *credsource.Resolver + selKeyStatus string + provKeyStatus map[string]string + mcpKeyStatus map[string]string } type evMsg protocol.AgentEvent type errMsg error type doneMsg struct{} +// Presence is cached: the render path must not shell out to keychain or HTTP. +type credStatusMsg struct { + sel string + prov map[string]string + mcp map[string]string + partial bool +} + func New(cfg config.File, sel config.Model, sb *runtime.Sandbox, vmState string, log []string, transcriptPath string) model { ta := textarea.New() ta.Placeholder = "Ask ABox Anything" @@ -73,7 +87,44 @@ func New(cfg config.File, sel config.Model, sb *runtime.Sandbox, vmState string, return model{cfg: cfg, sel: sel, sandbox: sb, ta: ta, keyIn: ki, vmState: vmState, log: log, transcriptPath: transcriptPath} } -func (m model) Init() tea.Cmd { return textarea.Blink } +func (m model) Init() tea.Cmd { + return tea.Batch(textarea.Blink, checkCredStatus(m.cfg, m.sel, m.resolver, false)) +} + +func checkCredStatus(cfg config.File, sel config.Model, r *credsource.Resolver, mcp bool) tea.Cmd { + if r == nil { + r = credsource.NewResolver() + } + return func() tea.Msg { + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + msg := credStatusMsg{prov: map[string]string{}, mcp: map[string]string{}} + msg.sel = credStatusLabel(ctx, r, sel.CredentialReference()) + for _, p := range providerChoices() { + model, ok := cfg.ModelNamed(p.Name) + if !ok { + model = p.ModelConfig() + } + msg.prov[p.Name] = credStatusLabel(ctx, r, model.CredentialReference()) + } + if mcp { + for _, s := range mcpServers(cfg) { + msg.mcp[s.Name] = credStatusLabel(ctx, r, s.CredentialReference()) + } + msg.partial = false + } else { + msg.partial = true + } + return msg + } +} + +func credStatusLabel(ctx context.Context, r *credsource.Resolver, ref config.CredentialRef) string { + if credsource.Present(ctx, r, credsource.FromConfig(ref)) { + return "key ok" + } + return "no key" +} func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { switch msg := msg.(type) { @@ -183,6 +234,18 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { } return m.submit() } + case credStatusMsg: + m.selKeyStatus = msg.sel + for name, status := range msg.prov { + if m.provKeyStatus == nil { + m.provKeyStatus = map[string]string{} + } + m.provKeyStatus[name] = status + } + if !msg.partial { + m.mcpKeyStatus = msg.mcp + } + return m, nil case evMsg: switch msg.Kind { case "text": @@ -274,12 +337,15 @@ func (m model) runSlash(text string) (tea.Model, tea.Cmd) { m.mode = modeProviderPick m.provSel = 0 m.err = "" + if len(m.provKeyStatus) == 0 { + return m, checkCredStatus(m.cfg, m.sel, m.resolver, false) + } return m, nil case "/mcp": m.mode = modeMCPPick m.mcpSel = 0 m.err = "" - return m, nil + return m, checkCredStatus(m.cfg, m.sel, m.resolver, true) case "/help": m.log = append(m.log, "commands:") for _, c := range slashCommands { @@ -338,7 +404,7 @@ func (m model) saveMCPKey() (tea.Model, tea.Cmd) { m.err = "token is empty" return m, nil } - env, err := applyMCPKey(m.mcpPick, key) + cfg, env, note, err := applyMCPKey(m.cfg, m.mcpPick, key) m.keyIn.SetValue("") m.keyIn.Blur() m.ta.Focus() @@ -347,6 +413,8 @@ func (m model) saveMCPKey() (tea.Model, tea.Cmd) { m.err = err.Error() return m, nil } + m.cfg = cfg + m.mcpKeyStatus = map[string]string{m.mcpPick.Name: "key ok"} if m.sandbox != nil { if err := m.sandbox.SetMCPTokens(context.Background(), map[string]string{env: key}); err != nil { m.err = "saved on host but guest MCP update failed: " + err.Error() @@ -354,7 +422,7 @@ func (m model) saveMCPKey() (tea.Model, tea.Cmd) { } } m.err = "" - m.log = append(m.log, "mcp "+m.mcpPick.Name+" token saved (OAuth: abox mcp login "+m.mcpPick.Name+")") + m.log = append(m.log, "mcp "+m.mcpPick.Name+" token saved ("+note+") (OAuth: abox mcp login "+m.mcpPick.Name+")") m.saveTranscript() return m, nil } @@ -365,7 +433,7 @@ func (m model) saveProviderKey() (tea.Model, tea.Cmd) { m.err = "API key is empty" return m, nil } - sel, err := applyProviderKey(m.cfg, m.provPick, key) + cfg, sel, note, err := applyProviderKey(m.cfg, m.provPick, key) m.keyIn.SetValue("") m.keyIn.Blur() m.ta.Focus() @@ -374,20 +442,34 @@ func (m model) saveProviderKey() (tea.Model, tea.Cmd) { m.err = err.Error() return m, nil } - m.sel = sel + m.cfg = cfg if m.sandbox != nil { + // Refresh the broker even if the guest update fails so it is not stuck on startup config. + m.updateHostBroker(cfg) secrets := map[string]string{m.provPick.Env: key} if err := m.sandbox.SetModel(context.Background(), sel, secrets); err != nil { m.err = "saved on host but guest agent update failed: " + err.Error() return m, nil } } + m.sel = sel + m.selKeyStatus = "key ok" + if m.provKeyStatus == nil { + m.provKeyStatus = map[string]string{} + } + m.provKeyStatus[m.provPick.Name] = "key ok" m.err = "" - m.log = append(m.log, "connected "+m.provPick.Label+" (agent in microVM)") + m.log = append(m.log, "connected "+m.provPick.Label+" ("+note+")") m.saveTranscript() return m, nil } +func (m model) updateHostBroker(cfg config.File) { + if m.sandbox != nil { + m.sandbox.OnGuestCall = llmbroker.New(cfg, m.resolver) + } +} + func waitEvent(ch <-chan protocol.AgentEvent) tea.Cmd { return func() tea.Msg { ev, ok := <-ch @@ -440,9 +522,9 @@ func (m model) View() tea.View { muted := lipgloss.NewStyle().Foreground(lipgloss.Color("#71717A")) bar := lipgloss.NewStyle().Foreground(lipgloss.Color("#F4F4F5")).Background(lipgloss.Color("#141416")).Padding(0, 1) - cred := "no key" - if m.sel.CredentialPresent() { - cred = "key ok" + cred := m.selKeyStatus + if cred == "" { + cred = "…" } header := bar.Render(fmt.Sprintf("ABox %s/%s vm:%s net:%s %s", m.sel.Provider, m.sel.Model, m.vmState, m.cfg.Connectivity.Mode, cred)) @@ -472,9 +554,9 @@ func (m model) View() tea.View { if i == m.provSel { mark = "> " } - status := "no key" - if p.Env != "" && strings.TrimSpace(os.Getenv(p.Env)) != "" { - status = "key ok" + status := "…" + if s, ok := m.provKeyStatus[p.Name]; ok { + status = s } b.WriteString(mark + p.Label + " " + status + "\n") } @@ -493,10 +575,13 @@ func (m model) View() tea.View { if i == m.mcpSel { mark = "> " } - status := "no token" - env := config.TokenEnv(s) - if env != "" && strings.TrimSpace(os.Getenv(env)) != "" { - status = "token ok" + status := "…" + if st, ok := m.mcpKeyStatus[s.Name]; ok { + if st == "key ok" { + status = "token ok" + } else { + status = "no token" + } } b.WriteString(mark + s.Name + " " + s.URL + " " + status + "\n") } @@ -596,8 +681,13 @@ func max(a, b int) int { return b } -func Run(cfg config.File, sel config.Model, sb *runtime.Sandbox, vmState string, log []string, transcriptPath string) error { - p := tea.NewProgram(New(cfg, sel, sb, vmState, log, transcriptPath)) +func Run(cfg config.File, sel config.Model, sb *runtime.Sandbox, vmState string, log []string, resolver *credsource.Resolver, transcriptPath string) error { + if resolver == nil { + resolver = credsource.NewResolver() + } + m := New(cfg, sel, sb, vmState, log, transcriptPath) + m.resolver = resolver + p := tea.NewProgram(m) _, err := p.Run() return err } diff --git a/pkg/abox/abox.go b/pkg/abox/abox.go index c617d53..08c5fb8 100644 --- a/pkg/abox/abox.go +++ b/pkg/abox/abox.go @@ -6,17 +6,19 @@ import ( "fmt" "os" "path/filepath" + "sync" "time" "github.com/AdminTurnedDevOps/ABox/internal/config" - "github.com/AdminTurnedDevOps/ABox/internal/credentials" + "github.com/AdminTurnedDevOps/ABox/internal/credsource" + "github.com/AdminTurnedDevOps/ABox/internal/llmbroker" "github.com/AdminTurnedDevOps/ABox/internal/repository" "github.com/AdminTurnedDevOps/ABox/internal/runtime" "github.com/AdminTurnedDevOps/ABox/internal/session" "github.com/AdminTurnedDevOps/ABox/protocol" ) -// ErrGuestTooOld is returned when a v2-only Turn option is used against a v1 guest. +// ErrGuestTooOld is returned when an operation requires a protocol-2+ guest. var ErrGuestTooOld = runtime.ErrGuestTooOld type Options struct { @@ -56,14 +58,21 @@ func open(ctx context.Context, opts Options, resume bool, resumeID string) (*Ses if err != nil { return nil, fmt.Errorf("load config: %w", err) } - if err := credentials.ApplyToEnv(); err != nil { - return nil, fmt.Errorf("credentials: %w", err) + n, scrubErr := session.ScrubSecretsEverywhere() + if n > 0 { + fmt.Fprintf(os.Stderr, "abox: scrubbed plaintext secrets from %d old session(s)\n", n) } + if scrubErr != nil { + return nil, fmt.Errorf("scrub legacy session secrets: %w", scrubErr) + } + resolver := credsource.NewResolver() sel, ok := cfg.ModelNamed(opts.Model) if !ok { + resolver.Close() return nil, fmt.Errorf("no model profile %q (config %s)", opts.Model, cfgPath) } if err := os.MkdirAll(config.SessionRoot(), 0o700); err != nil { + resolver.Close() return nil, err } @@ -72,23 +81,27 @@ func open(ctx context.Context, opts Options, resume bool, resumeID string) (*Ses if resume { loaded, err := loadResume(opts.RepoPath, resumeID) if err != nil { + resolver.Close() return nil, err } sess = loaded } else { created, err := session.Create(opts.RepoPath, "pending") if err != nil { + resolver.Close() return nil, fmt.Errorf("create session: %w", err) } sess = created opened, err := repository.OpenForSession(opts.RepoPath, filepath.Join(sess.Dir, "host-tree")) if err != nil { + resolver.Close() return nil, fmt.Errorf("snapshot repo: %w", err) } snap = opened sess.RepoRoot = snap.Root sess.HEAD = snap.HEAD if err := sess.WriteMeta(); err != nil { + resolver.Close() return nil, err } } @@ -99,9 +112,11 @@ func open(ctx context.Context, opts Options, resume bool, resumeID string) (*Ses } mcpServers, err := cfg.ResolvedMCPServers() if err != nil { + resolver.Close() return nil, err } - if err := runtime.Prepare(sess, image, sel, cfg.SecretsFromEnv(), mcpServers, resume); err != nil { + if err := runtime.Prepare(sess, image, sel, mcpServers, resume); err != nil { + resolver.Close() return nil, err } vcpu, ram := cfg.Resources.Resolved() @@ -123,20 +138,54 @@ func open(ctx context.Context, opts Options, resume bool, resumeID string) (*Ses } sb, err := runtime.Start(bootCtx, sess, vmm, vcpu, ram) if err != nil { + resolver.Close() return nil, fmt.Errorf("start vm: %w", err) } + if sb.GuestProtocol < 2 { + sb.Stop() + resolver.Close() + if resume { + return nil, fmt.Errorf("%w: cannot resume protocol-1 session %q after secretless config rewrite; rebuild the guest image and start a new session", ErrGuestTooOld, sess.ID) + } + return nil, fmt.Errorf("%w: protocol-1 guest cannot use the secretless config; rebuild the guest image", ErrGuestTooOld) + } + sb.OnGuestCall = llmbroker.New(cfg, resolver) + + // Push whatever resolved before returning a partial-resolution error. + pushCtx, pushCancel := context.WithTimeout(ctx, 30*time.Second) + defer pushCancel() + secrets, resolveErr := credsource.ResolveSelected(pushCtx, resolver, cfg, sel) + pushErr := sb.PushSecrets(pushCtx, sel, secrets) + if err := credentialStartupError(resolveErr, pushErr); err != nil { + sb.Stop() + resolver.Close() + return nil, err + } if !resume { archive, err := repository.ArchiveHEAD(snap.Root) if err != nil { sb.Stop() + resolver.Close() return nil, fmt.Errorf("archive repo: %w", err) } if err := sb.TransferArchive(ctx, archive); err != nil { sb.Stop() + resolver.Close() return nil, fmt.Errorf("transfer repo: %w", err) } } - return &Session{cfg: cfg, sess: sess, sb: sb, sel: sel}, nil + return &Session{cfg: cfg, sess: sess, sb: sb, sel: sel, resolver: resolver}, nil +} + +func credentialStartupError(resolveErr, pushErr error) error { + var errs []error + if resolveErr != nil { + errs = append(errs, fmt.Errorf("resolve credentials: %w", resolveErr)) + } + if pushErr != nil { + errs = append(errs, fmt.Errorf("push resolved credentials: %w", pushErr)) + } + return errors.Join(errs...) } func loadResume(repoPath, id string) (*session.Session, error) { @@ -176,10 +225,12 @@ type TurnResult struct { } type Session struct { - cfg config.File - sess *session.Session - sb *runtime.Sandbox - sel config.Model + mu sync.RWMutex + cfg config.File + sess *session.Session + sb *runtime.Sandbox + sel config.Model + resolver *credsource.Resolver } func (s *Session) ID() string { return s.sess.ID } @@ -207,6 +258,8 @@ func (s *Session) TurnOpts(ctx context.Context, prompt string, opts TurnOpts, on } func (s *Session) turn(ctx context.Context, prompt string, opts TurnOpts, onEvent func(Event)) (*TurnResult, error) { + s.mu.RLock() + defer s.mu.RUnlock() rtOpts := runtime.TurnOptions{ MaxTurns: opts.MaxTurns, RichEvents: opts.RichEvents, @@ -227,9 +280,7 @@ func (s *Session) turn(ctx context.Context, prompt string, opts TurnOpts, onEven out, err := s.sb.UserTurnCtx(ctx, prompt, rtOpts, onEvent) res := &TurnResult{} if out != nil { - res.Usage = out.Usage - res.StopReason = out.StopReason - res.Canceled = out.Canceled + res = turnResult(out) } if err != nil && errors.Is(err, runtime.ErrGuestTooOld) { return res, fmt.Errorf("%w: %v", ErrGuestTooOld, err) @@ -237,13 +288,47 @@ func (s *Session) turn(ctx context.Context, prompt string, opts TurnOpts, onEven return res, err } +func turnResult(out *runtime.TurnOutcome) *TurnResult { + if out == nil { + return &TurnResult{} + } + return &TurnResult{Usage: out.Usage, StopReason: out.StopReason, Canceled: out.Canceled} +} + func (s *Session) SetModel(ctx context.Context, model string) error { - sel, ok := s.cfg.ModelNamed(model) + if !s.mu.TryLock() { + return fmt.Errorf("cannot set model while a turn or model update is in progress") + } + defer s.mu.Unlock() + if s.sb.GuestProtocol < 2 { + return fmt.Errorf("%w: set model requires protocol 2, guest speaks %d", ErrGuestTooOld, s.sb.GuestProtocol) + } + cfg, _, err := config.Load() + if err != nil { + return fmt.Errorf("reload config: %w", err) + } + sel, ok := cfg.ModelNamed(model) if !ok { return fmt.Errorf("no model profile %q", model) } + var secrets map[string]string + if s.sb.GuestProtocol == 2 { + ref := sel.CredentialReference() + val, err := s.resolver.Resolve(ctx, credsource.FromConfig(ref)) + if err != nil { + return fmt.Errorf("credential for model %q (%s %s): %w", sel.Name, ref.Source, ref.Name, err) + } + secrets = map[string]string{sel.EnvName(): string(val.Bytes)} + val.Zero() + } + if err := s.sb.SetModel(ctx, sel, secrets); err != nil { + return err + } + // Broker snapshots config at construction; replace it after the guest accepts the model. + s.sb.OnGuestCall = llmbroker.New(cfg, s.resolver) + s.cfg = cfg s.sel = sel - return s.sb.SetModel(ctx, sel, s.cfg.SecretsFromEnv()) + return nil } func (s *Session) SetMCPTokens(ctx context.Context, secrets map[string]string) error { @@ -278,5 +363,7 @@ func (s *Session) Close() error { if s.sb == nil { return nil } - return s.sb.Stop() + err := s.sb.Stop() + s.resolver.Close() + return err } diff --git a/pkg/abox/abox_test.go b/pkg/abox/abox_test.go index 8e6e21c..d0d4ce7 100644 --- a/pkg/abox/abox_test.go +++ b/pkg/abox/abox_test.go @@ -5,9 +5,11 @@ import ( "errors" "os" "path/filepath" + "strings" "testing" "time" + "github.com/AdminTurnedDevOps/ABox/internal/config" "github.com/AdminTurnedDevOps/ABox/internal/runtime" "github.com/AdminTurnedDevOps/ABox/internal/session" "github.com/AdminTurnedDevOps/ABox/protocol" @@ -61,6 +63,69 @@ func TestTurnOptsRejectsOldGuest(t *testing.T) { } } +func TestTurnResultPreservesUsage(t *testing.T) { + usage := &protocol.UsageInfo{InputTokens: 11, OutputTokens: 5} + got := turnResult(&runtime.TurnOutcome{Usage: usage, StopReason: "end_turn", Canceled: true}) + if got.Usage == nil || got.Usage.InputTokens != 11 || got.Usage.OutputTokens != 5 { + t.Fatalf("usage %+v", got.Usage) + } + if got.StopReason != "end_turn" || !got.Canceled { + t.Fatalf("result %+v", got) + } +} + +func TestSetModelRejectsProtocolOne(t *testing.T) { + s := &Session{sb: &runtime.Sandbox{GuestProtocol: 1}} + if err := s.SetModel(context.Background(), "grok-default"); !errors.Is(err, ErrGuestTooOld) { + t.Fatalf("got %v", err) + } +} + +func TestSetModelRejectsConcurrentTurn(t *testing.T) { + s := &Session{sb: &runtime.Sandbox{GuestProtocol: 3}} + s.mu.RLock() + defer s.mu.RUnlock() + if err := s.SetModel(context.Background(), "grok-default"); err == nil { + t.Fatal("expected in-progress error") + } +} + +func TestCredentialStartupErrorReportsResolveAndPushFailures(t *testing.T) { + resolveErr := errors.New("model credential unavailable") + pushErr := errors.New("mcp token push failed") + err := credentialStartupError(resolveErr, pushErr) + if !errors.Is(err, resolveErr) || !errors.Is(err, pushErr) { + t.Fatalf("combined error %v", err) + } + if err == nil || err.Error() != "resolve credentials: model credential unavailable\npush resolved credentials: mcp token push failed" { + t.Fatalf("error text %q", err) + } +} + +func TestCredentialStartupErrorAllowsSuccessfulPartialPush(t *testing.T) { + resolveErr := errors.New("one credential unavailable") + err := credentialStartupError(resolveErr, nil) + if !errors.Is(err, resolveErr) { + t.Fatalf("error %v", err) + } +} + +func TestOpenReturnsLegacySessionScrubError(t *testing.T) { + t.Setenv("ABOX_HOME", t.TempDir()) + legacy := filepath.Join(config.SessionRoot(), "legacy") + if err := os.MkdirAll(legacy, 0o700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(legacy, "guest-config.json"), []byte("{"), 0o600); err != nil { + t.Fatal(err) + } + + sess, err := Open(context.Background(), Options{RepoPath: t.TempDir()}) + if sess != nil || err == nil || !strings.Contains(err.Error(), "scrub legacy session secrets") { + t.Fatalf("session=%v error=%v", sess, err) + } +} + func TestLoadResumeByID(t *testing.T) { t.Setenv("ABOX_HOME", t.TempDir()) created, err := session.Create(t.TempDir(), "head") diff --git a/pkg/abox/doc.go b/pkg/abox/doc.go index 537c803..dfa758c 100644 --- a/pkg/abox/doc.go +++ b/pkg/abox/doc.go @@ -1,7 +1,7 @@ // Package abox is the public Go SDK for embedding an ABox microVM agent session. // // Runtime requirements: Apple Silicon, libkrun/libkrunfw, and a golden guest -// image (`make image`). Resume of a session created before a guest rebuild -// may speak protocol v1; Capabilities() reports what that guest supports. -// Turn options, rich events, and mid-turn cancel require protocol 2. +// image (`make image`). Protocol-1 guests cannot consume current secretless +// session configuration and are rejected with ErrGuestTooOld. Protocol 2 is +// the legacy secret-push path; protocol 3 uses the host provider broker. package abox diff --git a/protocol/protocol.go b/protocol/protocol.go index 66859a7..8dc1b13 100644 --- a/protocol/protocol.go +++ b/protocol/protocol.go @@ -10,12 +10,23 @@ import ( ) const ( - Version = 2 + Version = 3 // host provider broker; LLM credentials stay on the host + MaxFrameBytes = 1 << 20 MaxArchiveChunk = 256 << 10 MaxHistoryBytes = 256 << 10 RPCPort = 1024 GuestRepoDir = "/work/repo" + + MaxProviderChunk = 256 << 10 // provider_send chunk size + MaxProviderRequest = 4 << 20 // reassembled provider_send budget + MaxProviderToolArgs = 512 << 10 // per tool-args bound; larger -> error event + MaxProviderEvent = 512 << 10 // encoded provider_event params budget + MaxProviderMessages = 4096 // messages in one provider request + MaxProviderTools = 256 // tool schemas in one provider request + MaxProviderEvents = 1 << 20 // events in one provider stream + MaxProviderStreams = 2 // concurrent provider streams per session + MaxGuestCalls = 8 // concurrent guest-initiated host RPCs ) // Frame is a length-prefixed JSON message. @@ -77,6 +88,7 @@ type GetContextResult struct { type HelloResult struct { Accepted bool `json:"accepted"` Message string `json:"message,omitempty"` + Protocol int `json:"protocol,omitempty"` // 0 if omitted; old hosts hang without this } type ListFilesParams struct { @@ -165,7 +177,7 @@ type GuestConfig struct { VsockPort uint32 `json:"vsock_port"` RepoDir string `json:"repo_dir"` Model GuestModel `json:"model"` - Secrets map[string]string `json:"secrets,omitempty"` + Secrets map[string]string `json:"secrets,omitempty"` // deprecated; kept so old images still parse MCPServers []GuestMCPServer `json:"mcp_servers,omitempty"` } @@ -204,6 +216,58 @@ type SetMCPTokensParams struct { Secrets map[string]string `json:"secrets"` } +// Model is a configured alias, never a URL, header, or credential name. +type ProviderOpenParams struct { + Model string `json:"model"` + Rich bool `json:"rich,omitempty"` +} + +type ProviderOpenResult struct { + StreamID string `json:"stream_id"` +} + +type ProviderSendParams struct { + StreamID string `json:"stream_id"` + Data []byte `json:"data"` + Last bool `json:"last,omitempty"` // host starts the provider call after Last +} + +type ProviderRequest struct { + Messages []ProviderMessage `json:"messages"` + Tools []ProviderToolSchema `json:"tools,omitempty"` +} + +type ProviderCancelParams struct { + StreamID string `json:"stream_id"` +} + +type ProviderEventParams struct { + StreamID string `json:"stream_id"` + Type string `json:"type"` + Text string `json:"text,omitempty"` + ToolID string `json:"tool_id,omitempty"` + ToolName string `json:"tool_name,omitempty"` + ToolArgs string `json:"tool_args,omitempty"` + Usage *UsageInfo `json:"usage,omitempty"` + StopReason string `json:"stop_reason,omitempty"` + Err string `json:"err,omitempty"` +} + +type ProviderMessage struct { + Role string `json:"role"` + Content string `json:"content,omitempty"` + ToolID string `json:"tool_id,omitempty"` + ToolName string `json:"tool_name,omitempty"` + ToolArgs string `json:"tool_args,omitempty"` + ToolResult string `json:"tool_result,omitempty"` +} + +type ProviderToolSchema struct { + Name string `json:"name"` + Description string `json:"description,omitempty"` + Parameters map[string]any `json:"parameters,omitempty"` +} + func WriteFrame(w io.Writer, f Frame) error { if f.V == 0 { f.V = Version @@ -217,11 +281,26 @@ func WriteFrame(w io.Writer, f Frame) error { } var hdr [4]byte binary.BigEndian.PutUint32(hdr[:], uint32(len(body))) - if _, err := w.Write(hdr[:]); err != nil { + if err := writeFull(w, hdr[:]); err != nil { return err } - _, err = w.Write(body) - return err + return writeFull(w, body) +} + +func writeFull(w io.Writer, p []byte) error { + for len(p) > 0 { + n, err := w.Write(p) + if n > 0 { + p = p[n:] + } + if err != nil { + return err + } + if n == 0 { + return io.ErrShortWrite + } + } + return nil } func ReadFrame(r io.Reader) (Frame, error) { @@ -250,21 +329,25 @@ func ReadFrameLimit(r io.Reader, limit int) (Frame, error) { // TrimHistory keeps the newest lines whose JSON size fits in maxBytes. func TrimHistory(h []HistoryLine, maxBytes int) []HistoryLine { - if maxBytes <= 0 { + if maxBytes < 2 { return nil } - var out []HistoryLine - var size int + out := make([]HistoryLine, 0) + size := 2 // JSON array brackets. for i := len(h) - 1; i >= 0; i-- { b, err := json.Marshal(h[i]) if err != nil { continue } - if size+len(b) > maxBytes && len(out) > 0 { - break + itemSize := len(b) + if len(out) > 0 { + itemSize++ // Comma separator. + } + if size+itemSize > maxBytes { + continue } out = append(out, h[i]) - size += len(b) + size += itemSize } for i, j := 0, len(out)-1; i < j; i, j = i+1, j-1 { out[i], out[j] = out[j], out[i] diff --git a/protocol/protocol_test.go b/protocol/protocol_test.go index 9f0566e..dd56a9e 100644 --- a/protocol/protocol_test.go +++ b/protocol/protocol_test.go @@ -2,10 +2,29 @@ package protocol import ( "bytes" + "encoding/json" + "errors" + "io" "strings" "testing" ) +type shortWriter struct { + buf bytes.Buffer + n int +} + +func (w *shortWriter) Write(p []byte) (int, error) { + if len(p) > w.n { + p = p[:w.n] + } + return w.buf.Write(p) +} + +type zeroWriter struct{} + +func (zeroWriter) Write([]byte) (int, error) { return 0, nil } + func TestFrameRoundTrip(t *testing.T) { var buf bytes.Buffer orig := Frame{V: Version, ID: "1", Method: "list_files"} @@ -21,6 +40,27 @@ func TestFrameRoundTrip(t *testing.T) { } } +func TestWriteFrameUsesFullWrites(t *testing.T) { + w := &shortWriter{n: 3} + orig := Frame{V: Version, ID: "short", Method: "list_files"} + if err := WriteFrame(w, orig); err != nil { + t.Fatal(err) + } + got, err := ReadFrame(&w.buf) + if err != nil { + t.Fatal(err) + } + if got.ID != orig.ID || got.Method != orig.Method { + t.Fatalf("got %+v", got) + } +} + +func TestWriteFrameRejectsNoProgressWriter(t *testing.T) { + if err := WriteFrame(zeroWriter{}, Frame{ID: "1"}); !errors.Is(err, io.ErrShortWrite) { + t.Fatalf("got %v", err) + } +} + func TestTrimHistoryKeepsNewest(t *testing.T) { h := []HistoryLine{ {Kind: "user", Text: strings.Repeat("a", 200)}, @@ -32,6 +72,22 @@ func TestTrimHistoryKeepsNewest(t *testing.T) { } } +func TestTrimHistoryRejectsSingleOversizedItem(t *testing.T) { + const limit = 64 + h := []HistoryLine{{Kind: "text", Text: strings.Repeat("x", 256)}} + got := TrimHistory(h, limit) + if len(got) != 0 { + t.Fatalf("oversized history retained: %#v", got) + } + raw, err := json.Marshal(got) + if err != nil { + t.Fatal(err) + } + if len(raw) > limit { + t.Fatalf("encoded history is %d bytes, limit %d", len(raw), limit) + } +} + func TestHistoryLineAliasesAgentEvent(t *testing.T) { ev := AgentEvent{Kind: "text", Text: "hi", Tool: "search", Status: "ok", Err: "e"} var line HistoryLine = ev @@ -47,3 +103,74 @@ func TestRejectOversizedFrame(t *testing.T) { t.Fatal("expected error") } } + +func TestProviderBrokerTypesRoundTrip(t *testing.T) { + req := ProviderRequest{ + Messages: []ProviderMessage{{ + Role: "assistant", ToolID: "1", ToolName: "list_files", ToolArgs: `{"path":"."}`, + }}, + Tools: []ProviderToolSchema{{ + Name: "list_files", Description: "list", Parameters: map[string]any{"type": "object"}, + }}, + } + raw, err := json.Marshal(req) + if err != nil { + t.Fatal(err) + } + var got ProviderRequest + if err := json.Unmarshal(raw, &got); err != nil { + t.Fatal(err) + } + if len(got.Messages) != 1 || got.Messages[0].ToolName != "list_files" || got.Messages[0].ToolArgs == "" { + t.Fatalf("%+v", got.Messages) + } + if len(got.Tools) != 1 || got.Tools[0].Name != "list_files" { + t.Fatalf("%+v", got.Tools) + } + + ev := ProviderEventParams{ + StreamID: "s1", Type: "text", Text: "hi", + Usage: &UsageInfo{InputTokens: 1, OutputTokens: 2}, StopReason: "end_turn", + } + rawEv, err := json.Marshal(ev) + if err != nil { + t.Fatal(err) + } + var gotEv ProviderEventParams + if err := json.Unmarshal(rawEv, &gotEv); err != nil { + t.Fatal(err) + } + if gotEv.StreamID != "s1" || gotEv.Usage == nil || gotEv.Usage.OutputTokens != 2 { + t.Fatalf("%+v", gotEv) + } + + open := ProviderOpenParams{Model: "grok-default", Rich: true} + rawOpen, _ := json.Marshal(open) + var gotOpen ProviderOpenParams + _ = json.Unmarshal(rawOpen, &gotOpen) + if gotOpen.Model != "grok-default" || !gotOpen.Rich { + t.Fatalf("%+v", gotOpen) + } +} + +func TestHelloResultProtocolField(t *testing.T) { + raw, err := json.Marshal(HelloResult{Accepted: true, Protocol: 3}) + if err != nil { + t.Fatal(err) + } + var got HelloResult + if err := json.Unmarshal(raw, &got); err != nil { + t.Fatal(err) + } + if !got.Accepted || got.Protocol != 3 { + t.Fatalf("%+v", got) + } + // Old hosts send no protocol field: decodes as 0. + var legacy HelloResult + if err := json.Unmarshal([]byte(`{"accepted":true}`), &legacy); err != nil { + t.Fatal(err) + } + if legacy.Protocol != 0 { + t.Fatalf("legacy ack: %+v", legacy) + } +} From 2a4a8ee3a326695c4f6c75d3c96aeaf992bee230 Mon Sep 17 00:00:00 2001 From: adminturneddevops Date: Sun, 6 Sep 2026 18:38:31 -0400 Subject: [PATCH 2/4] read --- README.md | 40 +++++++++++++++++++++++++++------------- 1 file changed, 27 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 5363a93..d114b3d 100644 --- a/README.md +++ b/README.md @@ -250,18 +250,6 @@ func Stream(ctx context.Context, model config.Model, key string, client *http.Cl ![](img/prov1.png) ![](img/prov2.png) -Config lives at `~/.abox/config.yaml`. Keys are **not** stored in that file. Credential sources: `env`, `keychain` (macOS), `vault`, `azure`, `aws`. `/provider` in the TUI saves to the macOS keychain first (service `abox`), falling back to `~/.abox/credentials.env` (mode 0600). LLM keys stay on the host. MCP tokens still go to the guest because the guest makes those HTTPS calls. - -```yaml -credential: - source: keychain # env | keychain | vault | azure | aws - name: ANTHROPIC_API_KEY # env var, keychain account, vault path, Azure secret URI, or AWS secret id - # field: value # vault/aws only - # version: "4" # vault/azure only -``` - -`credential_env: XAI_API_KEY` is the same as `{source: env, name: XAI_API_KEY}`. Vault needs `VAULT_ADDR` + `VAULT_TOKEN`. Azure needs `AZURE_CLIENT_ID` / `AZURE_TENANT_ID` / `AZURE_CLIENT_SECRET` (or `az login`). AWS needs `AWS_ACCESS_KEY_ID` + `AWS_SECRET_ACCESS_KEY`. `abox creds migrate` moves existing `credentials.env` entries into the keychain. - Default profiles: ```yaml @@ -291,6 +279,32 @@ Guest egress is allowlisted for configured MCP origins on HTTPS `:443` via libkr LLM traffic does **not** take the MCP `connectivity.mode` path. Direct vs agentgateway today applies to MCP servers. The host broker hits the provider `base_url` above. +## Credentials + +Config lives at `~/.abox/config.yaml`. Keys are **not** stored in that file. Each model or MCP server points at a source: + +```yaml +credential: + source: keychain # env | keychain | vault | azure | aws + name: ANTHROPIC_API_KEY # env var, keychain account, vault path, Azure secret URI, or AWS secret id + # field: value # vault/aws only + # version: "4" # vault/azure only +``` + +| Source | `name` is | Auth | +| --- | --- | --- | +| `env` | environment variable (also reads `~/.abox/credentials.env`) | — | +| `keychain` | macOS keychain account (service `abox`) | — | +| `vault` | Vault KV v2 path (`secret/abox/anthropic`) | `VAULT_ADDR` + `VAULT_TOKEN` (or `~/.vault-token`) | +| `azure` | Key Vault secret URI (`https://myvault.vault.azure.net/secrets/name`) | `AZURE_CLIENT_ID` / `AZURE_TENANT_ID` / `AZURE_CLIENT_SECRET`, or `az login` | +| `aws` | Secrets Manager secret id | `AWS_ACCESS_KEY_ID` + `AWS_SECRET_ACCESS_KEY` (`AWS_REGION`) | + +`credential_env: XAI_API_KEY` is the same as `{source: env, name: XAI_API_KEY}`. + +`/provider` and `/mcp` in the TUI save to the macOS keychain first, falling back to `credentials.env` (mode 0600) if the keychain is locked or missing. `abox creds migrate` moves existing `credentials.env` entries into the keychain. + +LLM keys stay on the host. MCP tokens still go to the guest because the guest makes those HTTPS calls. + ## MCP Integration ABox is an MCP **client**. Remote tools are Streamable HTTP. Stdio MCP is not implemented. @@ -311,7 +325,7 @@ type StreamableClientTransport struct { ![](img/mcpsandbox.png.png) -Config lives at `~/.abox/config.yaml` (same pattern as `~/.claude`, `~/.codex`). First `abox` run creates `~/.abox/` (mode 0700) and a default `config.yaml` if they are missing. Credentials are `~/.abox/credentials.env`. +Config lives at `~/.abox/config.yaml` (same pattern as `~/.claude`, `~/.codex`). First `abox` run creates `~/.abox/` (mode 0700) and a default `config.yaml` if they are missing. MCP tokens use the same credential sources as LLM keys (see [Credentials](#credentials)). Add a Streamable HTTP server without editing YAML by hand. `--mode` is required: From 2d8ada9647c9da912397f65a0027b9069ded1a4b Mon Sep 17 00:00:00 2001 From: adminturneddevops Date: Sun, 6 Sep 2026 18:39:37 -0400 Subject: [PATCH 3/4] read --- README.md | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index d114b3d..5a56eea 100644 --- a/README.md +++ b/README.md @@ -281,6 +281,16 @@ LLM traffic does **not** take the MCP `connectivity.mode` path. Direct vs agentg ## Credentials +The following credential providers are supported (where your LLM API key lives): + +| Source | `name` is | Auth | +| --- | --- | --- | +| `env` | environment variable (also reads `~/.abox/credentials.env`) | — | +| `keychain` | macOS keychain account (service `abox`) | — | +| `vault` | Vault KV v2 path (`secret/abox/anthropic`) | `VAULT_ADDR` + `VAULT_TOKEN` (or `~/.vault-token`) | +| `azure` | Key Vault secret URI (`https://myvault.vault.azure.net/secrets/name`) | `AZURE_CLIENT_ID` / `AZURE_TENANT_ID` / `AZURE_CLIENT_SECRET`, or `az login` | +| `aws` | Secrets Manager secret id | `AWS_ACCESS_KEY_ID` + `AWS_SECRET_ACCESS_KEY` (`AWS_REGION`) | + Config lives at `~/.abox/config.yaml`. Keys are **not** stored in that file. Each model or MCP server points at a source: ```yaml @@ -291,14 +301,6 @@ credential: # version: "4" # vault/azure only ``` -| Source | `name` is | Auth | -| --- | --- | --- | -| `env` | environment variable (also reads `~/.abox/credentials.env`) | — | -| `keychain` | macOS keychain account (service `abox`) | — | -| `vault` | Vault KV v2 path (`secret/abox/anthropic`) | `VAULT_ADDR` + `VAULT_TOKEN` (or `~/.vault-token`) | -| `azure` | Key Vault secret URI (`https://myvault.vault.azure.net/secrets/name`) | `AZURE_CLIENT_ID` / `AZURE_TENANT_ID` / `AZURE_CLIENT_SECRET`, or `az login` | -| `aws` | Secrets Manager secret id | `AWS_ACCESS_KEY_ID` + `AWS_SECRET_ACCESS_KEY` (`AWS_REGION`) | - `credential_env: XAI_API_KEY` is the same as `{source: env, name: XAI_API_KEY}`. `/provider` and `/mcp` in the TUI save to the macOS keychain first, falling back to `credentials.env` (mode 0600) if the keychain is locked or missing. `abox creds migrate` moves existing `credentials.env` entries into the keychain. From 4b932e1ac46786e076ad29f484705044ae2367a4 Mon Sep 17 00:00:00 2001 From: adminturneddevops Date: Sun, 6 Sep 2026 19:02:59 -0400 Subject: [PATCH 4/4] aws local integration --- README.md | 5 +- internal/credsource/aws.go | 111 +++++++++++++++++++--- internal/credsource/cloud_test.go | 100 ++++++++++++++++++++ internal/tui/commands.go | 54 +++++++++++ internal/tui/commands_test.go | 75 +++++++++++++++ internal/tui/tui.go | 151 ++++++++++++++++++++++++++++-- 6 files changed, 476 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index 5a56eea..0567c8b 100644 --- a/README.md +++ b/README.md @@ -72,6 +72,7 @@ abox ``` - `/provider` sets Grok, OpenAI, or Anthropic API keys +- `/credential` points a model at Vault, Azure Key Vault, or AWS Secrets Manager - `/mcp` lists configured Streamable HTTP MCP servers and accepts a Bearer token (`abox mcp login` for OAuth) - `abox --resume` reopens the latest session for this repo (same `root.raw`, LLM conversation, and TUI transcript). `abox --resume ` picks a session. Plain `abox` still starts a new session. - `ctrl+c` quits @@ -289,7 +290,7 @@ The following credential providers are supported (where your LLM API key lives): | `keychain` | macOS keychain account (service `abox`) | — | | `vault` | Vault KV v2 path (`secret/abox/anthropic`) | `VAULT_ADDR` + `VAULT_TOKEN` (or `~/.vault-token`) | | `azure` | Key Vault secret URI (`https://myvault.vault.azure.net/secrets/name`) | `AZURE_CLIENT_ID` / `AZURE_TENANT_ID` / `AZURE_CLIENT_SECRET`, or `az login` | -| `aws` | Secrets Manager secret id | `AWS_ACCESS_KEY_ID` + `AWS_SECRET_ACCESS_KEY` (`AWS_REGION`) | +| `aws` | Secrets Manager secret id | `AWS_ACCESS_KEY_ID` + `AWS_SECRET_ACCESS_KEY` (`AWS_REGION`), or `~/.aws/credentials` | Config lives at `~/.abox/config.yaml`. Keys are **not** stored in that file. Each model or MCP server points at a source: @@ -303,7 +304,7 @@ credential: `credential_env: XAI_API_KEY` is the same as `{source: env, name: XAI_API_KEY}`. -`/provider` and `/mcp` in the TUI save to the macOS keychain first, falling back to `credentials.env` (mode 0600) if the keychain is locked or missing. `abox creds migrate` moves existing `credentials.env` entries into the keychain. +`/provider` and `/mcp` in the TUI save to the macOS keychain first, falling back to `credentials.env` (mode 0600) if the keychain is locked or missing. `/credential` writes a Vault / Azure Key Vault / AWS Secrets Manager reference into `config.yaml` (it does not store cloud tokens). When the cloud auth env vars are unset, Azure uses the local `az login` session and AWS uses `~/.aws/credentials` (and region from `~/.aws/config`). `abox creds migrate` moves existing `credentials.env` entries into the keychain. LLM keys stay on the host. MCP tokens still go to the guest because the guest makes those HTTPS calls. diff --git a/internal/credsource/aws.go b/internal/credsource/aws.go index 6328505..e338660 100644 --- a/internal/credsource/aws.go +++ b/internal/credsource/aws.go @@ -13,6 +13,7 @@ import ( "net/http" "net/url" "os" + "path/filepath" "strings" "time" ) @@ -30,18 +31,9 @@ var newAWSClient = func() *http.Client { } func (awsSource) Resolve(ctx context.Context, ref Reference) (Value, error) { - accessKey := strings.TrimSpace(os.Getenv("AWS_ACCESS_KEY_ID")) - secretKey := os.Getenv("AWS_SECRET_ACCESS_KEY") - sessionToken := strings.TrimSpace(os.Getenv("AWS_SESSION_TOKEN")) - region := strings.TrimSpace(os.Getenv("AWS_REGION")) - if region == "" { - region = strings.TrimSpace(os.Getenv("AWS_DEFAULT_REGION")) - } - if accessKey == "" || secretKey == "" { - return Value{}, fmt.Errorf("%w: aws source requires AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY", ErrLocked) - } - if region == "" { - return Value{}, fmt.Errorf("aws source requires AWS_REGION or AWS_DEFAULT_REGION") + accessKey, secretKey, sessionToken, region, err := awsAuth() + if err != nil { + return Value{}, err } host := fmt.Sprintf("%s.%s.amazonaws.com", awsService, region) reqURL := "https://" + host + "/" @@ -94,6 +86,101 @@ func (awsSource) Resolve(ctx context.Context, ref Reference) (Value, error) { func (awsSource) Close() error { return nil } +func awsAuth() (accessKey, secretKey, sessionToken, region string, err error) { + accessKey = strings.TrimSpace(os.Getenv("AWS_ACCESS_KEY_ID")) + secretKey = os.Getenv("AWS_SECRET_ACCESS_KEY") + sessionToken = strings.TrimSpace(os.Getenv("AWS_SESSION_TOKEN")) + region = strings.TrimSpace(os.Getenv("AWS_REGION")) + if region == "" { + region = strings.TrimSpace(os.Getenv("AWS_DEFAULT_REGION")) + } + if accessKey == "" || strings.TrimSpace(secretKey) == "" { + var fileRegion string + accessKey, secretKey, sessionToken, fileRegion = awsSharedFileAuth() + if region == "" { + region = fileRegion + } + } else if region == "" { + _, _, _, region = awsSharedFileAuth() + } + if accessKey == "" || strings.TrimSpace(secretKey) == "" { + return "", "", "", "", fmt.Errorf("%w: aws source needs AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY, or ~/.aws/credentials", ErrLocked) + } + if region == "" { + return "", "", "", "", fmt.Errorf("aws source needs AWS_REGION or region in ~/.aws/config") + } + return accessKey, secretKey, sessionToken, region, nil +} + +func awsSharedFileAuth() (accessKey, secretKey, sessionToken, region string) { + profile := strings.TrimSpace(os.Getenv("AWS_PROFILE")) + if profile == "" { + profile = "default" + } + home, _ := os.UserHomeDir() + credPath := strings.TrimSpace(os.Getenv("AWS_SHARED_CREDENTIALS_FILE")) + if credPath == "" && home != "" { + credPath = filepath.Join(home, ".aws", "credentials") + } + configPath := strings.TrimSpace(os.Getenv("AWS_CONFIG_FILE")) + if configPath == "" && home != "" { + configPath = filepath.Join(home, ".aws", "config") + } + if creds := parseAWSINIFile(credPath)[profile]; creds != nil { + accessKey = strings.TrimSpace(creds["aws_access_key_id"]) + secretKey = creds["aws_secret_access_key"] + sessionToken = strings.TrimSpace(creds["aws_session_token"]) + } + cfgSection := "default" + if profile != "default" { + cfgSection = "profile " + profile + } + if cfg := parseAWSINIFile(configPath)[cfgSection]; cfg != nil { + region = strings.TrimSpace(cfg["region"]) + } + return accessKey, secretKey, sessionToken, region +} + +func parseAWSINIFile(path string) map[string]map[string]string { + if path == "" { + return map[string]map[string]string{} + } + data, err := os.ReadFile(path) + if err != nil { + return map[string]map[string]string{} + } + return parseAWSINI(string(data)) +} + +func parseAWSINI(data string) map[string]map[string]string { + out := map[string]map[string]string{} + section := "" + for _, line := range strings.Split(data, "\n") { + line = strings.TrimSpace(line) + if line == "" || strings.HasPrefix(line, "#") || strings.HasPrefix(line, ";") { + continue + } + if strings.HasPrefix(line, "[") && strings.HasSuffix(line, "]") { + section = strings.TrimSpace(line[1 : len(line)-1]) + if out[section] == nil { + out[section] = map[string]string{} + } + continue + } + if section == "" { + continue + } + k, v, ok := strings.Cut(line, "=") + if !ok { + continue + } + k = strings.ToLower(strings.TrimSpace(k)) + v = strings.Trim(strings.TrimSpace(v), `"'`) + out[section][k] = v + } + return out +} + func awsFieldBytes(name, field, secretString string) (Value, error) { if field == "" { return Value{Bytes: []byte(secretString)}, nil diff --git a/internal/credsource/cloud_test.go b/internal/credsource/cloud_test.go index 84dd2b4..f460aac 100644 --- a/internal/credsource/cloud_test.go +++ b/internal/credsource/cloud_test.go @@ -414,6 +414,9 @@ func TestAWSAccessDenied(t *testing.T) { } func TestAWSMissingCredentials(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + t.Setenv("AWS_SHARED_CREDENTIALS_FILE", filepath.Join(t.TempDir(), "missing")) + t.Setenv("AWS_CONFIG_FILE", filepath.Join(t.TempDir(), "missing")) t.Setenv("AWS_ACCESS_KEY_ID", "") t.Setenv("AWS_SECRET_ACCESS_KEY", "") _, err := testResolver().Resolve(context.Background(), Reference{Source: "aws", Name: "x"}) @@ -423,6 +426,8 @@ func TestAWSMissingCredentials(t *testing.T) { } func TestAWSMissingRegion(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + t.Setenv("AWS_CONFIG_FILE", filepath.Join(t.TempDir(), "missing")) t.Setenv("AWS_ACCESS_KEY_ID", "AKID") t.Setenv("AWS_SECRET_ACCESS_KEY", "k") t.Setenv("AWS_REGION", "") @@ -433,6 +438,101 @@ func TestAWSMissingRegion(t *testing.T) { } } +func TestAWSResolveFromSharedCredentials(t *testing.T) { + dir := t.TempDir() + credFile := filepath.Join(dir, "credentials") + cfgFile := filepath.Join(dir, "config") + if err := os.WriteFile(credFile, []byte("[default]\naws_access_key_id = AKID-FILE\naws_secret_access_key = file-secret\n"), 0o600); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(cfgFile, []byte("[default]\nregion = us-west-2\n"), 0o600); err != nil { + t.Fatal(err) + } + t.Setenv("AWS_ACCESS_KEY_ID", "") + t.Setenv("AWS_SECRET_ACCESS_KEY", "") + t.Setenv("AWS_SESSION_TOKEN", "") + t.Setenv("AWS_REGION", "") + t.Setenv("AWS_DEFAULT_REGION", "") + t.Setenv("AWS_PROFILE", "") + t.Setenv("AWS_SHARED_CREDENTIALS_FILE", credFile) + t.Setenv("AWS_CONFIG_FILE", cfgFile) + + var sawAuth string + newAWSServer(t, 200, `{"SecretString":"from-file"}`, func(r *http.Request, auth string) { + sawAuth = auth + }) + v, err := testResolver().Resolve(context.Background(), Reference{Source: "aws", Name: "prod/anthropic"}) + if err != nil { + t.Fatal(err) + } + if string(v.Bytes) != "from-file" { + t.Fatalf("got %q", v.Bytes) + } + if !strings.Contains(sawAuth, "Credential=AKID-FILE/") || !strings.Contains(sawAuth, "us-west-2") { + t.Fatalf("auth %q", sawAuth) + } + if strings.Contains(sawAuth, "file-secret") { + t.Fatal("secret key leaked into Authorization") + } +} + +func TestAWSEnvCredentialsBeatSharedFile(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "credentials"), []byte("[default]\naws_access_key_id = AKID-FILE\naws_secret_access_key = file-secret\n"), 0o600); err != nil { + t.Fatal(err) + } + t.Setenv("AWS_SHARED_CREDENTIALS_FILE", filepath.Join(dir, "credentials")) + t.Setenv("AWS_ACCESS_KEY_ID", "AKID-ENV") + t.Setenv("AWS_SECRET_ACCESS_KEY", "env-secret") + t.Setenv("AWS_SESSION_TOKEN", "") + t.Setenv("AWS_REGION", "us-east-1") + + var sawAuth string + newAWSServer(t, 200, `{"SecretString":"ok"}`, func(r *http.Request, auth string) { + sawAuth = auth + }) + if _, err := testResolver().Resolve(context.Background(), Reference{Source: "aws", Name: "x"}); err != nil { + t.Fatal(err) + } + if !strings.Contains(sawAuth, "Credential=AKID-ENV/") { + t.Fatalf("env credentials not used: %q", sawAuth) + } + if strings.Contains(sawAuth, "AKID-FILE") { + t.Fatalf("shared file leaked into env path: %q", sawAuth) + } +} + +func TestAWSProfileFromSharedCredentials(t *testing.T) { + dir := t.TempDir() + credFile := filepath.Join(dir, "credentials") + cfgFile := filepath.Join(dir, "config") + if err := os.WriteFile(credFile, []byte("[default]\naws_access_key_id = AKID-DEFAULT\naws_secret_access_key = default-secret\n\n[work]\naws_access_key_id = AKID-WORK\naws_secret_access_key = work-secret\n"), 0o600); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(cfgFile, []byte("[default]\nregion = us-east-1\n\n[profile work]\nregion = eu-central-1\n"), 0o600); err != nil { + t.Fatal(err) + } + t.Setenv("AWS_ACCESS_KEY_ID", "") + t.Setenv("AWS_SECRET_ACCESS_KEY", "") + t.Setenv("AWS_SESSION_TOKEN", "") + t.Setenv("AWS_REGION", "") + t.Setenv("AWS_DEFAULT_REGION", "") + t.Setenv("AWS_PROFILE", "work") + t.Setenv("AWS_SHARED_CREDENTIALS_FILE", credFile) + t.Setenv("AWS_CONFIG_FILE", cfgFile) + + var sawAuth string + newAWSServer(t, 200, `{"SecretString":"ok"}`, func(r *http.Request, auth string) { + sawAuth = auth + }) + if _, err := testResolver().Resolve(context.Background(), Reference{Source: "aws", Name: "x"}); err != nil { + t.Fatal(err) + } + if !strings.Contains(sawAuth, "Credential=AKID-WORK/") || !strings.Contains(sawAuth, "eu-central-1") { + t.Fatalf("auth %q", sawAuth) + } +} + // --- ResolveSelected --- type mapSource struct { diff --git a/internal/tui/commands.go b/internal/tui/commands.go index 403b2b4..23b12f3 100644 --- a/internal/tui/commands.go +++ b/internal/tui/commands.go @@ -16,10 +16,27 @@ type slashCmd struct { var slashCommands = []slashCmd{ {Name: "/provider", Help: "Connect Grok, OpenAI, or Anthropic and set an API key"}, + {Name: "/credential", Help: "Point a model at Vault, Azure Key Vault, or AWS Secrets Manager"}, {Name: "/mcp", Help: "List MCP servers and paste a Bearer token (OAuth: abox mcp login)"}, {Name: "/help", Help: "List slash commands"}, } +type cloudCredSource struct { + Source string + Label string + Placeholder string + Prompt string + Note string +} + +func cloudCredentialChoices() []cloudCredSource { + return []cloudCredSource{ + {Source: "vault", Label: "HashiCorp Vault", Placeholder: "secret/abox/anthropic", Prompt: "path> ", Note: "host needs VAULT_ADDR and VAULT_TOKEN (or ~/.vault-token)"}, + {Source: "azure", Label: "Azure Key Vault", Placeholder: "https://myvault.vault.azure.net/secrets/name", Prompt: "uri> ", Note: "host needs AZURE_CLIENT_ID/TENANT_ID/SECRET or az login"}, + {Source: "aws", Label: "AWS Secrets Manager", Placeholder: "abox/anthropic", Prompt: "id> ", Note: "host uses AWS_* env or ~/.aws/credentials"}, + } +} + func providerChoices() []config.ProviderProfile { return config.DefaultProviders() } @@ -107,3 +124,40 @@ func applyProviderKey(cfg config.File, choice config.ProviderProfile, key string } return cfg, sel, res.Note, nil } + +func applyCloudCredential(cfg config.File, choice config.ProviderProfile, ref config.CredentialRef) (config.File, config.Model, string, error) { + ref.Name = strings.TrimSpace(ref.Name) + next := cfg + next.Models = append([]config.Model(nil), cfg.Models...) + found := false + for i, m := range next.Models { + if m.Name == choice.Name { + next.Models[i].CredentialEnv = "" + cred := ref + next.Models[i].Credential = &cred + found = true + } + } + if !found { + model := choice.ModelConfig() + model.CredentialEnv = "" + cred := ref + model.Credential = &cred + next.Models = append(next.Models, model) + } + if err := next.Save(); err != nil { + return cfg, config.Model{}, "", err + } + sel, ok := next.ModelNamed(choice.Name) + if !ok { + return cfg, config.Model{}, "", fmt.Errorf("config update did not retain model profile %q", choice.Name) + } + note := ref.Source + " " + ref.Name + for _, src := range cloudCredentialChoices() { + if src.Source == ref.Source { + note = src.Note + break + } + } + return next, sel, note, nil +} diff --git a/internal/tui/commands_test.go b/internal/tui/commands_test.go index 1e595e2..d8a045a 100644 --- a/internal/tui/commands_test.go +++ b/internal/tui/commands_test.go @@ -3,6 +3,8 @@ package tui import ( "context" "encoding/json" + "os" + "strings" "testing" "github.com/AdminTurnedDevOps/ABox/internal/config" @@ -106,6 +108,79 @@ func TestApplyMCPKeyFallbackReplacesExplicitNonEnvReference(t *testing.T) { } } +func TestApplyCloudCredentialWritesAzureRef(t *testing.T) { + t.Setenv("ABOX_HOME", t.TempDir()) + cfg := config.Defaults() + choice := config.DefaultProviders()[2] + ref := config.CredentialRef{Source: "azure", Name: "https://testkv.vault.azure.net/secrets/anthropic"} + got, sel, note, err := applyCloudCredential(cfg, choice, ref) + if err != nil { + t.Fatal(err) + } + if sel.CredentialEnv != "" || sel.Credential == nil || sel.Credential.Source != "azure" || sel.Credential.Name != ref.Name { + t.Fatalf("selected model %+v", sel) + } + if !strings.Contains(note, "az login") { + t.Fatalf("note %q", note) + } + persisted, _, err := config.Load() + if err != nil { + t.Fatal(err) + } + model, ok := persisted.ModelNamed(choice.Name) + if !ok || model.Credential == nil || model.Credential.Source != "azure" || model.Credential.Name != ref.Name { + t.Fatalf("persisted %+v found=%v cfg=%+v", model, ok, got.Models) + } + body, err := os.ReadFile(config.Path()) + if err != nil { + t.Fatal(err) + } + if strings.Contains(string(body), "sk-") || strings.Contains(string(body), "secret=") { + t.Fatalf("config contained a secret: %s", body) + } +} + +func TestApplyCloudCredentialRejectsBadAzureURI(t *testing.T) { + t.Setenv("ABOX_HOME", t.TempDir()) + cfg := config.Defaults() + if err := cfg.Save(); err != nil { + t.Fatal(err) + } + choice := config.DefaultProviders()[2] + _, _, _, err := applyCloudCredential(cfg, choice, config.CredentialRef{ + Source: "azure", Name: "https://evil.example/secrets/x", + }) + if err == nil { + t.Fatal("expected invalid azure URI to fail") + } + body, err := os.ReadFile(config.Path()) + if err != nil { + t.Fatal(err) + } + if strings.Contains(string(body), "evil.example") { + t.Fatalf("invalid ref was persisted: %s", body) + } +} + +func TestApplyCloudCredentialAddsMissingProfile(t *testing.T) { + t.Setenv("ABOX_HOME", t.TempDir()) + cfg := config.Defaults() + cfg.Models = cfg.Models[:1] + choice := config.DefaultProviders()[1] + got, sel, _, err := applyCloudCredential(cfg, choice, config.CredentialRef{ + Source: "vault", Name: "secret/abox/openai", + }) + if err != nil { + t.Fatal(err) + } + if sel.Name != choice.Name || sel.Credential == nil || sel.Credential.Source != "vault" { + t.Fatalf("selected %+v", sel) + } + if _, ok := got.ModelNamed(choice.Name); !ok { + t.Fatalf("config missing profile: %+v", got.Models) + } +} + func TestUpdateHostBrokerUsesCurrentModelConfig(t *testing.T) { cfg := config.Defaults() cfg.Models = []config.Model{{ diff --git a/internal/tui/tui.go b/internal/tui/tui.go index c20263a..67f4636 100644 --- a/internal/tui/tui.go +++ b/internal/tui/tui.go @@ -28,6 +28,9 @@ const ( modeProviderKey modeMCPPick modeMCPKey + modeCredSourcePick + modeCredModelPick + modeCredName ) type model struct { @@ -40,6 +43,8 @@ type model struct { slashSel int provSel int provPick config.ProviderProfile + credSourceSel int + credSource cloudCredSource mcpSel int mcpPick config.MCPServer log []string @@ -136,6 +141,7 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { switch msg.String() { case "ctrl+c": if m.mode != modeChat { + m.cancelCredInput() m.mode = modeChat m.keyIn.Blur() m.ta.Focus() @@ -150,6 +156,7 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { return m, tea.Quit case "esc": if m.mode != modeChat { + m.cancelCredInput() m.mode = modeChat m.keyIn.Blur() m.ta.Focus() @@ -163,7 +170,13 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { } return m, nil } - if m.mode == modeProviderPick { + if m.mode == modeCredSourcePick { + if m.credSourceSel > 0 { + m.credSourceSel-- + } + return m, nil + } + if m.mode == modeProviderPick || m.mode == modeCredModelPick { if m.provSel > 0 { m.provSel-- } @@ -180,7 +193,13 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { } return m, nil } - if m.mode == modeProviderPick { + if m.mode == modeCredSourcePick { + if m.credSourceSel < len(cloudCredentialChoices())-1 { + m.credSourceSel++ + } + return m, nil + } + if m.mode == modeProviderPick || m.mode == modeCredModelPick { if m.provSel < len(providerChoices())-1 { m.provSel++ } @@ -197,7 +216,13 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { } return m, nil } - if m.mode == modeProviderPick { + if m.mode == modeCredSourcePick { + if m.credSourceSel > 0 { + m.credSourceSel-- + } + return m, nil + } + if m.mode == modeProviderPick || m.mode == modeCredModelPick { if m.provSel > 0 { m.provSel-- } @@ -210,7 +235,13 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { } return m, nil } - if m.mode == modeProviderPick { + if m.mode == modeCredSourcePick { + if m.credSourceSel < len(cloudCredentialChoices())-1 { + m.credSourceSel++ + } + return m, nil + } + if m.mode == modeProviderPick || m.mode == modeCredModelPick { if m.provSel < len(providerChoices())-1 { m.provSel++ } @@ -226,6 +257,15 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { if m.mode == modeProviderKey { return m.saveProviderKey() } + if m.mode == modeCredSourcePick { + return m.acceptCredSource() + } + if m.mode == modeCredModelPick { + return m.acceptCredModel() + } + if m.mode == modeCredName { + return m.saveCloudCredential() + } if m.mode == modeMCPPick { return m.acceptMCP() } @@ -277,7 +317,7 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.saveTranscript() return m, nil } - if m.mode == modeProviderKey || m.mode == modeMCPKey { + if m.mode == modeProviderKey || m.mode == modeMCPKey || m.mode == modeCredName { var cmd tea.Cmd m.keyIn, cmd = m.keyIn.Update(msg) return m, cmd @@ -341,6 +381,12 @@ func (m model) runSlash(text string) (tea.Model, tea.Cmd) { return m, checkCredStatus(m.cfg, m.sel, m.resolver, false) } return m, nil + case "/credential": + m.mode = modeCredSourcePick + m.credSourceSel = 0 + m.provSel = 0 + m.err = "" + return m, nil case "/mcp": m.mode = modeMCPPick m.mcpSel = 0 @@ -368,6 +414,75 @@ func slashExact(name string) bool { return false } +func (m *model) cancelCredInput() { + m.keyIn.SetValue("") + m.keyIn.EchoMode = textinput.EchoPassword + m.keyIn.EchoCharacter = '•' + m.keyIn.Placeholder = "paste API key" + m.keyIn.Prompt = "key> " +} + +func (m model) acceptCredSource() (tea.Model, tea.Cmd) { + choices := cloudCredentialChoices() + if m.credSourceSel < 0 || m.credSourceSel >= len(choices) { + return m, nil + } + m.credSource = choices[m.credSourceSel] + m.mode = modeCredModelPick + m.provSel = 0 + m.err = "" + return m, nil +} + +func (m model) acceptCredModel() (tea.Model, tea.Cmd) { + choices := providerChoices() + if m.provSel < 0 || m.provSel >= len(choices) { + return m, nil + } + m.provPick = choices[m.provSel] + m.mode = modeCredName + m.keyIn.EchoMode = textinput.EchoNormal + m.keyIn.Placeholder = m.credSource.Placeholder + m.keyIn.Prompt = m.credSource.Prompt + m.keyIn.SetValue("") + m.keyIn.Focus() + m.ta.Blur() + m.err = "" + return m, textinput.Blink +} + +func (m model) saveCloudCredential() (tea.Model, tea.Cmd) { + name := strings.TrimSpace(m.keyIn.Value()) + m.cancelCredInput() + m.keyIn.Blur() + m.ta.Focus() + m.mode = modeChat + if name == "" { + m.err = "credential name is empty" + return m, nil + } + cfg, sel, note, err := applyCloudCredential(m.cfg, m.provPick, config.CredentialRef{ + Source: m.credSource.Source, Name: name, + }) + if err != nil { + m.err = err.Error() + return m, nil + } + m.cfg = cfg + if m.sandbox != nil { + m.updateHostBroker(cfg) + if err := m.sandbox.SetModel(context.Background(), sel, nil); err != nil { + m.err = "saved on host but guest agent update failed: " + err.Error() + return m, nil + } + } + m.sel = sel + m.err = "" + m.log = append(m.log, "credential "+m.provPick.Label+" -> "+m.credSource.Label+" "+name+" ("+note+")") + m.saveTranscript() + return m, checkCredStatus(m.cfg, m.sel, m.resolver, false) +} + func (m model) acceptProvider() (tea.Model, tea.Cmd) { choices := providerChoices() if m.provSel < 0 || m.provSel >= len(choices) { @@ -539,7 +654,7 @@ func (m model) View() tea.View { wrapped := wrapLog(m.log, wrapW) body := strings.Join(tail(wrapped, bodyH), "\n") if body == "" { - body = muted.Render("Type / for commands. /provider sets Grok, OpenAI, or Anthropic keys.") + body = muted.Render("Type / for commands.") } composer := m.ta.View() @@ -563,6 +678,30 @@ func (m model) View() tea.View { composer = b.String() case modeProviderKey: composer = "API key for " + m.provPick.Label + "\n" + m.keyIn.View() + case modeCredSourcePick: + var b strings.Builder + b.WriteString("Credential store (reference only; tokens stay in the environment)\n") + for i, s := range cloudCredentialChoices() { + mark := " " + if i == m.credSourceSel { + mark = "> " + } + b.WriteString(mark + s.Label + "\n") + } + composer = b.String() + case modeCredModelPick: + var b strings.Builder + b.WriteString("Model for " + m.credSource.Label + "\n") + for i, p := range providerChoices() { + mark := " " + if i == m.provSel { + mark = "> " + } + b.WriteString(mark + p.Label + "\n") + } + composer = b.String() + case modeCredName: + composer = m.credSource.Label + " for " + m.provPick.Label + "\n" + m.keyIn.View() case modeMCPPick: servers := mcpServers(m.cfg) var b strings.Builder