Compare commits

..
Author SHA1 Message Date
Gnanam 8b7b8296ac tui: migrate off the deprecated Bubble Tea mouse Type API (SA1019)
The mouse classification helpers checked both the current Button/Action pair and
the deprecated msg.Type / tea.Mouse* constants. Bubble Tea's parser always
populates Button+Action and only DERIVES Type from them, so the Type fallbacks
were pure redundancy (a left-button drag is Action==Motion, already covered by
mouseMotion). Drop them — clears all 11 SA1019 findings in internal/tui with no
behavior change. The left-drag test now asserts the Action==Motion path directly.

Gates: gofmt/vet/build(host+linux+windows)/test -race/staticcheck(SA1019 0)/
deadcode(no new) all pass.
2026-06-15 19:08:26 +05:30
KRATOS f857c4afb4 feat(daemon): TLS remote bridge — authenticated network access to the daemon (#211)
* feat(daemon): TLS remote bridge — authenticated network access to the daemon

Add internal/daemon/remote: an OPT-IN, TLS-only network bridge on top of the
local daemon. A remote client drives the SAME SessionManager/Pool over the
network behind bearer-token auth, reusing the daemon's framing and dispatch
unchanged. The local Unix-socket daemon stays the default and is unchanged.

- Bridge: per accepted TLS connection, authenticate THEN hand to the daemon's
  control dispatch via the new Server.ServeConn seam — a remote session runs
  under the same sandbox/risk model as a local one (never bypasses it).
- Authenticator + TokenAuthenticator: constant-time bearer-token check, token
  from $ZERO_DAEMON_REMOTE_TOKEN or a file; optional attestation hook (no-op
  default). Auth happens before any control frame; failure closes the connection
  after a small backoff; tokens are never logged.
- RemoteClient (DialRemote): verified TLS dial (never InsecureSkipVerify) + auth
  handshake, then the existing daemon.Client via the new NewClientConn seam, so
  remote and local share one protocol.
- Fail-closed: TLS mandatory (refuses without cert/key), a configurable
  min-protocol-version floor, the 1 MiB frame cap reused, and a bounded
  concurrent-connection cap that refuses overflow.
- CLI: `zero daemon serve-remote --addr --tls-cert --tls-key`; `run`/`attach`
  gain `--remote <addr> [--token --ca-cert --server-name]`.

Daemon seams added (additive, behavior unchanged): Server.ServeConn,
NewClientConn. SessionLink (event POST) and git-bundle repo transfer are
documented follow-ups.

Tests cover auth accept/reject (env + file), TLS-required, cert verification,
version floor, run/status/attach round-trip, re-attach resumption, connection
cap, and ListenAndServeTLS/Close. Gates: gofmt / vet / build (host+linux+
windows) / test -race / staticcheck (no new) / govulncheck (0) / deadcode
(no new) all pass.

* swarm: de-flake TestMailboxConcurrentSends on Windows CI

The 200-way concurrent-send stress test inherited from the swarm package timed
out on Windows CI (103/200 sends hit the 10s lock timeout): the lock's 20ms
retry sleep starves waiters when Windows file ops are slow under heavy
contention. Shorten the retry to 2ms (a freed lock is re-acquired promptly) and
raise the test's lock timeout to 60s so a legitimately-slow CI never times a send
out. Passes repeatedly under -race. (CI de-flake of a test that surfaced on
#210's Windows smoke; functionally unrelated to the OAuth changes.)

* remote: address CodeRabbit re-review (TLS dial + Close idempotency)

- client: replace the deprecated, non-cancellation-aware tls.DialWithDialer with
  a context-aware tls.Dialer.DialContext (resolves the golangci-lint noctx CI
  blocker; the context bounds the dial + TLS handshake).
- bridge: Close now clears b.listener under the lock, so a repeat Close is a
  no-op instead of returning a closed-listener error (matches its doc). Test
  asserts a double Close succeeds.

Gates: gofmt/vet/build(host+linux+windows)/test -race/staticcheck(no new)/
govulncheck(0)/deadcode(no new) all pass.
2026-06-15 18:48:07 +05:30
KRATOS bb7de29c80 feat(oauth): reusable provider OAuth engine + zero auth CLI (#210)
* feat(oauth): reusable provider OAuth engine + `zero auth` CLI

Add internal/oauth: a transport/identity-agnostic OAuth 2.0 engine that powers
provider login alongside the MCP-server OAuth zero already has. It is additive
and provider-neutral — NO providers, endpoints, or client identities are baked
into the binary; every provider is configured entirely via ZERO_OAUTH_<NAME>_*
env vars, so the engine works with any OAuth 2.0 / OIDC server.

Engine (internal/oauth):
- PKCE (S256 mandatory; "plain" refused) + per-flow CSRF state.
- Authorization-code flow: BuildAuthorizationURL / ExchangeCode / Refresh, with
  an https-only token-endpoint guard (loopback exempt) and error bodies that
  carry only error/error_description (never token material).
- Loopback callback server: 127.0.0.1-only, OS-assigned port, single use, state
  verified, then closed.
- Device-authorization grant (RFC 8628): RequestDeviceCode + PollDeviceToken
  honoring authorization_pending / slow_down / interval / expiry (headless/SSH).
- RFC 8414 + OIDC discovery (issuer-driven endpoint resolution).
- Namespaced token store ("provider:<name>" / "mcp:<name>"), 0600, file-locked
  (ownership-aware), atomic writes, malformed-fails-closed; GetFresh (on-demand
  refresh) + Handle401 (forced refresh).
- Opt-in proactive RefreshScheduler.

CLI (zero auth):
- login <provider> [--device] [--scope ...], logout, status [provider],
  refresh <provider> [--watch]. Status/output never print token material.
- Registered in app.go dispatch + top-level help.

MCP OAuth (internal/mcp) is left byte-for-byte unchanged (zero regression);
sharing the engine with MCP and encrypted-at-rest/keyring storage are documented
follow-ups. Tests cover every path (PKCE, auth URL, https-guard + redaction,
exchange/refresh, loopback capture + state-mismatch + timeout, device pending/
slow_down/expired, store namespacing/0600/locking, provider env resolution,
manager login (loopback+device)/GetFresh/Handle401, scheduler). Gates: gofmt /
vet / build (host+linux+windows) / test -race / staticcheck (no new) /
govulncheck (0) / deadcode (no new) all pass.

* oauth: fix Windows smoke — use OS-appropriate path in store-path test

TestResolveStorePathHonorsOverride hardcoded a unix "/tmp/custom/tok.json"
literal, which is not absolute on Windows (no drive), so ResolveStorePath
correctly resolves it against the current drive and the verbatim comparison
failed (D:\tmp\custom\tok.json). The production code is correct; the test now
builds the override with filepath.Join(t.TempDir(), ...) so it is absolute and
clean on every OS.

* oauth: address CodeRabbit review (11 findings)

Security / correctness:
- flow: caller-supplied ExtraAuthParams/extraParams can no longer override the
  reserved OAuth/PKCE fields (state, redirect_uri, code_challenge[_method],
  response_type, client_id) — e.g. forcing method=plain.
- device: a device-authorization response without expires_in now gets a bounded
  default lifetime (fail closed), and the poll loop re-checks expiry AFTER the
  interval sleep so it never polls past the deadline (e.g. after slow_down).
- lock: handle WriteString/Close errors on the token lock file — a partial write
  no longer strands an undeletable lock.
- store: a non-nil env map is now authoritative (hermetic) — a controlled map
  never falls back to ambient ZERO_OAUTH_* / HOME / XDG_CONFIG_HOME.
- scheduler: a transient loadForKey error backs off and retries (bounded) instead
  of permanently stopping proactive refresh.
- cli auth: reject flags a subcommand does not accept (e.g. login --watch,
  status --device) and reject empty --scope values — fail fast.

Lint blockers:
- manager: restructured the discovery fallback so the non-nil-err branch no longer
  returns nil (nilerr).
- providers: dropped the ineffectual flow initializer (ineffassign).

Tests: assert Save/Load errors in the scheduler test; the slow_down device test
no longer depends on post-expiry polling. Added tests for reserved-param
stripping, hermetic env, device default-expiry, and CLI flag validation. The
store-path test was already made OS-portable in the Windows-smoke fix.

Gates: gofmt/vet/build(host+linux+windows)/test -race/staticcheck(no new)/
govulncheck(0)/deadcode(no new) all pass.

* swarm: de-flake TestMailboxConcurrentSends on Windows CI

The 200-way concurrent-send stress test inherited from the swarm package timed
out on Windows CI (103/200 sends hit the 10s lock timeout): the lock's 20ms
retry sleep starves waiters when Windows file ops are slow under heavy
contention. Shorten the retry to 2ms (a freed lock is re-acquired promptly) and
raise the test's lock timeout to 60s so a legitimately-slow CI never times a send
out. Passes repeatedly under -race. (CI de-flake of a test that surfaced on
#210's Windows smoke; functionally unrelated to the OAuth changes.)

* oauth: address CodeRabbit re-review (4 findings)

- loopback (security): NewLoopbackListener refuses an empty CSRF state — an empty
  state would match a callback carrying no state at all, defeating the check.
- device (correctness): RequestDeviceCode now sends client_secret when set, so a
  confidential client authenticates on the device-authorization endpoint too
  (consistent with the token poll).
- cli auth: document --watch in the help Flags section.
- scheduler test: assert the seed Save error instead of ignoring it.

Also scrubbed a stray third-party reference from a loopback.go comment.
Tests added: empty-state rejection, device client_secret is sent. Gates:
gofmt/vet/build/test -race/staticcheck(no new)/govulncheck(0)/deadcode(no new)
all pass.

* mcp(oauth): delegate the OAuth engine to internal/oauth (dedup) (#212)

OAuth Phase 2: the MCP OAuth code now reuses the shared internal/oauth engine
instead of carrying its own copy of the transport/identity-agnostic primitives.
internal/mcp/oauth.go drops ~110 lines (526 -> 418): newPKCE, newState,
discoverAuthorizationServer, the authorize-URL builder, token exchange/refresh,
the token-response parser, and joinWellKnown now delegate to internal/oauth.

Behavior-preserving: MCP keeps its own OAuthConfig/StoredToken types,
LoginOptions/Login orchestration, loopback handling, registerClient (dynamic
client registration — not part of the shared engine), CLI output, and on-disk
token format (mcp-oauth-tokens.json, raw server-name keys). The existing mcp
oauth tests pass UNCHANGED. The only behavioral delta is a hardening inherited
from the shared engine: the token endpoint must be https (loopback exempt) —
a credential is no longer sent to a plaintext non-loopback endpoint.

Stacks on #210 (the internal/oauth package). Gates: gofmt/vet/build(host+
linux+windows)/test -race/staticcheck(no new)/govulncheck(0)/deadcode(no new
vs base) all pass; `zero mcp oauth status` output is identical.
2026-06-15 18:47:43 +05:30
de5d402e1f web_search: add domains allowlist + score field (#209)
* web_search: add domains allowlist + score field

The web_search tool now accepts an optional 'domains' array that filters
results to the listed hostnames before they reach the model. This is the
zero-only prompt-injection defense that none of the major CLIs ship: a
model that the user has constrained to a known-good set of domains cannot
be tricked into fetching arbitrary pages from search results.

Also surfaces the provider's optional 'score' field on each result row
when present (zero/absent scores are not rendered, to keep the common
case tidy). The 'score' field is parsed from the generic JSON shape
already understood by the httpSearchBackend.

Filters are applied after the search returns, before the empty-result
short-circuit, so a non-empty allowlist that ate every hit surfaces as a
clear 'no results matched domains' error rather than a misleading
'no results' message.

Co-Authored-By: Claude <[email protected]>

* web_search: fail closed on all-invalid domains, strip :port on canonicalize

CodeRabbit review fixes:

1. canonicalizeWebSearchHost: switch parsed.Host -> parsed.Hostname() so
   an allowlist entry like 'https://react.dev:443/path' normalizes the same
   as 'react.dev'. The previous code left ':443' in the string, which
   silently broke every match against a result URL on the default port.

2. stringListArgWebSearch: return a 'provided' bool so Run can distinguish
   'no allowlist' from 'allowlist that normalized to empty'. The previous
   code returned len(domains)==0 in both cases, and Run fell through to
   the unfiltered path — defeating the prompt-injection defense the
   parameter exists to provide. Now we error fail-closed.

3. Tests:
   - TestWebSearchDomainsFilterRejectsAllInvalidInputs: locks in the
     fail-closed contract for all-whitespace, all-empty, embedded-space,
     and mixed-invalid allowlists.
   - TestWebSearchDomainsFilterAcceptsHostPort: locks in the
     Hostname()-strips-port fix for the scheme'd case and asserts that
     scheme-less 'host:port' is rejected.
   - TestCanonicalizeWebSearchHostStripsPort: the direct unit-level
     regression test for canonicalizeWebSearchHost.

Co-Authored-By: Claude <[email protected]>

---------

Co-authored-by: Vasanthdev2004 <[email protected]>
Co-authored-by: Claude <[email protected]>
2026-06-15 18:27:21 +05:30
42 changed files with 4863 additions and 235 deletions

No files matched your search

+3 -1
View File
@@ -125,7 +125,9 @@ func Score(suite Suite, input ScoreInput) Report {
if len(task.RequiredTraceEvents) > 0 {
report.Results = append(report.Results, scoreTraceEvents(task.RequiredTraceEvents, input))
}
report.Results = append(report.Results, unknownCommandResults(input.CommandResults, seenCommands, input)...)
for _, result := range unknownCommandResults(input.CommandResults, seenCommands, input) {
report.Results = append(report.Results, result)
}
report.finishSummary()
return report
}
+3 -1
View File
@@ -515,7 +515,9 @@ func agentEvalReportFromBenchmark(suite agenteval.Suite, report agenteval.Benchm
if task.Agent.Truncated {
converted.Truncated = true
}
converted.Failures = append(converted.Failures, agentEvalFailuresFromTaskReport(task)...)
for _, failure := range agentEvalFailuresFromTaskReport(task) {
converted.Failures = append(converted.Failures, failure)
}
}
return converted
}
+3
View File
@@ -278,6 +278,8 @@ func runWithDeps(args []string, stdout io.Writer, stderr io.Writer, deps appDeps
return runHooks(args[1:], stdout, stderr, deps)
case "mcp":
return runMCP(args[1:], stdout, stderr, deps)
case "auth":
return runAuth(args[1:], stdout, stderr, deps)
case "sandbox":
return runSandbox(args[1:], stdout, stderr, deps)
case "update":
@@ -700,6 +702,7 @@ Commands:
tools Scaffold and list local Zero plugin-tools
hooks Inspect Zero hook configuration
mcp Manage MCP backend settings
auth Log in to model providers via OAuth
sandbox Inspect sandbox policy and persistent grants
update Check for Zero CLI updates
worktrees Prepare isolated git worktrees
+344
View File
@@ -0,0 +1,344 @@
package cli
import (
"context"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/Gitlawb/zero/internal/oauth"
"github.com/Gitlawb/zero/internal/redaction"
)
// runAuth dispatches `zero auth <command>` for provider OAuth login. It is
// additive and independent of `zero mcp oauth` (MCP server auth), which is
// unchanged.
func runAuth(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int {
if len(args) == 0 {
return writeExecUsageError(stderr, "auth subcommand required. Use `zero auth status`.")
}
if args[0] == "-h" || args[0] == "--help" || args[0] == "help" {
if err := writeAuthHelp(stdout); err != nil {
return exitCrash
}
return exitSuccess
}
switch args[0] {
case "login":
return runAuthLogin(args[1:], stdout, stderr, deps)
case "logout":
return runAuthLogout(args[1:], stdout, stderr, deps)
case "status":
return runAuthStatus(args[1:], stdout, stderr, deps)
case "refresh":
return runAuthRefresh(args[1:], stdout, stderr, deps)
default:
return writeExecUsageError(stderr, fmt.Sprintf("unknown auth subcommand %q", args[0]))
}
}
// authArgs is the parsed form of an auth subcommand's arguments.
type authArgs struct {
positional []string
json bool
device bool
watch bool
scopes []string
help bool
}
func parseAuthArgs(sub string, args []string) (authArgs, error) {
var parsed authArgs
addScope := func(scope string) error {
scope = strings.TrimSpace(scope)
if scope == "" {
return fmt.Errorf("--scope requires a non-empty value")
}
parsed.scopes = append(parsed.scopes, scope)
return nil
}
for i := 0; i < len(args); i++ {
arg := args[i]
switch {
case arg == "-h" || arg == "--help" || arg == "help":
parsed.help = true
case arg == "--json":
parsed.json = true
case arg == "--device":
parsed.device = true
case arg == "--watch":
parsed.watch = true
case arg == "--scope":
if i+1 >= len(args) {
return authArgs{}, fmt.Errorf("--scope requires a value")
}
i++
if err := addScope(args[i]); err != nil {
return authArgs{}, err
}
case strings.HasPrefix(arg, "--scope="):
if err := addScope(strings.TrimPrefix(arg, "--scope=")); err != nil {
return authArgs{}, err
}
case strings.HasPrefix(arg, "-"):
return authArgs{}, fmt.Errorf("unknown flag %q", arg)
default:
parsed.positional = append(parsed.positional, arg)
}
}
if parsed.help {
return parsed, nil // help short-circuits flag validation
}
if err := validateAuthFlags(sub, parsed); err != nil {
return authArgs{}, err
}
return parsed, nil
}
// validateAuthFlags rejects flags a subcommand does not accept, so an ambiguous
// invocation fails fast instead of silently ignoring a flag.
func validateAuthFlags(sub string, a authArgs) error {
allowed := map[string]map[string]bool{
"login": {"device": true, "scope": true},
"logout": {"json": true},
"status": {"json": true},
"refresh": {"watch": true},
}[sub]
bad := func(name string) error { return fmt.Errorf("zero auth %s does not accept %s", sub, name) }
if a.json && !allowed["json"] {
return bad("--json")
}
if a.device && !allowed["device"] {
return bad("--device")
}
if a.watch && !allowed["watch"] {
return bad("--watch")
}
if len(a.scopes) > 0 && !allowed["scope"] {
return bad("--scope")
}
return nil
}
// newAuthManager builds an oauth.Manager backed by the file store, printing the
// authorization URL / device code to stdout. The store path honors
// ZERO_OAUTH_TOKENS_PATH (env), so callers/tests can redirect it.
func newAuthManager(deps appDeps, out io.Writer) (*oauth.Manager, error) {
store, err := oauth.NewStore(oauth.StoreOptions{Now: deps.now})
if err != nil {
return nil, err
}
return oauth.NewManager(oauth.ManagerOptions{
Store: store,
HTTPClient: &http.Client{Timeout: 30 * time.Second},
Now: deps.now,
Out: out,
// The opener prints the URL so headless shells can copy it; the URL
// carries no token material. A real browser launch is intentionally not
// performed (the sandbox/headless contexts make printing the safer default).
OpenBrowser: func(string) error { return nil },
})
}
func runAuthLogin(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int {
parsed, err := parseAuthArgs("login", args)
if err != nil {
return writeExecUsageError(stderr, err.Error())
}
if parsed.help {
_ = writeAuthHelp(stdout)
return exitSuccess
}
if len(parsed.positional) != 1 {
return writeExecUsageError(stderr, "usage: zero auth login <provider> [--device] [--scope <scope>]")
}
manager, err := newAuthManager(deps, stdout)
if err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
status, err := manager.Login(context.Background(), oauth.LoginOptions{
Provider: parsed.positional[0],
Device: parsed.device,
ExtraScopes: parsed.scopes,
})
if err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
if _, err := fmt.Fprintf(stdout, "Logged in to %s.\n%s\n", parsed.positional[0], oauth.FormatStatuses([]oauth.Status{status})); err != nil {
return exitCrash
}
return exitSuccess
}
func runAuthLogout(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int {
parsed, err := parseAuthArgs("logout", args)
if err != nil {
return writeExecUsageError(stderr, err.Error())
}
if parsed.help {
_ = writeAuthHelp(stdout)
return exitSuccess
}
if len(parsed.positional) != 1 {
return writeExecUsageError(stderr, "usage: zero auth logout <provider>")
}
provider := parsed.positional[0]
manager, err := newAuthManager(deps, stdout)
if err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
removed, err := manager.Logout(provider)
if err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
if parsed.json {
payload := struct {
Provider string `json:"provider"`
Removed bool `json:"removed"`
}{Provider: provider, Removed: removed}
if err := writePrettyJSON(stdout, payload); err != nil {
return exitCrash
}
return exitSuccess
}
msg := fmt.Sprintf("No OAuth login stored for %s.\n", provider)
if removed {
msg = fmt.Sprintf("Logged out of %s.\n", provider)
}
if _, err := fmt.Fprint(stdout, msg); err != nil {
return exitCrash
}
return exitSuccess
}
func runAuthStatus(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int {
parsed, err := parseAuthArgs("status", args)
if err != nil {
return writeExecUsageError(stderr, err.Error())
}
if parsed.help {
_ = writeAuthHelp(stdout)
return exitSuccess
}
if len(parsed.positional) > 1 {
return writeExecUsageError(stderr, "usage: zero auth status [provider]")
}
manager, err := newAuthManager(deps, stdout)
if err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
statuses, err := manager.StatusAll()
if err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
if len(parsed.positional) == 1 {
statuses = filterAuthStatuses(statuses, parsed.positional[0])
}
if parsed.json {
payload := struct {
Logins []oauth.Status `json:"logins"`
}{Logins: statuses}
if err := writePrettyJSON(stdout, payload); err != nil {
return exitCrash
}
return exitSuccess
}
if _, err := fmt.Fprintln(stdout, oauth.FormatStatuses(statuses)); err != nil {
return exitCrash
}
return exitSuccess
}
func runAuthRefresh(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int {
parsed, err := parseAuthArgs("refresh", args)
if err != nil {
return writeExecUsageError(stderr, err.Error())
}
if parsed.help {
_ = writeAuthHelp(stdout)
return exitSuccess
}
if len(parsed.positional) != 1 {
return writeExecUsageError(stderr, "usage: zero auth refresh <provider>")
}
provider := parsed.positional[0]
manager, err := newAuthManager(deps, stdout)
if err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
key := oauth.ProviderKey(provider)
if parsed.watch {
return runAuthRefreshWatch(manager, key, provider, stdout, stderr)
}
if _, err := manager.Handle401(context.Background(), key); err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
if _, err := fmt.Fprintf(stdout, "Refreshed OAuth token for %s.\n", provider); err != nil {
return exitCrash
}
return exitSuccess
}
// runAuthRefreshWatch keeps a provider's token fresh in the foreground until
// interrupted. This is the opt-in proactive-refresh scheduler surface (for a
// long-running external process that reads the token file). It validates a
// refreshable token exists first, then schedules refreshes before each expiry.
func runAuthRefreshWatch(manager *oauth.Manager, key, provider string, stdout io.Writer, stderr io.Writer) int {
ctx, stop := signalContext()
defer stop()
// Validate (and refresh-if-needed) once up front so a missing/non-refreshable
// token fails fast instead of silently watching nothing.
if _, err := manager.GetFresh(ctx, key); err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
scheduler := oauth.NewRefreshScheduler()
scheduler.Start(ctx, manager, key)
defer scheduler.Stop()
if _, err := fmt.Fprintf(stdout, "Watching %s — refreshing before expiry. Press Ctrl+C to stop.\n", provider); err != nil {
return exitCrash
}
<-ctx.Done()
return exitSuccess
}
func filterAuthStatuses(statuses []oauth.Status, provider string) []oauth.Status {
want := oauth.ProviderKey(provider)
filtered := make([]oauth.Status, 0, 1)
for _, st := range statuses {
if st.Key == want {
filtered = append(filtered, st)
}
}
return filtered
}
func writeAuthHelp(w io.Writer) error {
_, err := fmt.Fprint(w, `Usage:
zero auth <command>
Commands:
login <provider> [--device] [--scope <scope>] Log in to a provider via OAuth
logout <provider> Delete a provider's stored login
status [provider] Show login presence/expiry (never the token)
refresh <provider> [--watch] Force a token refresh (--watch keeps it fresh)
A provider is any OAuth 2.0 / OIDC server you configure via env (no providers are
built in). For a provider named <name>, set:
ZERO_OAUTH_<NAME>_CLIENT_ID (required)
ZERO_OAUTH_<NAME>_CLIENT_SECRET (optional)
ZERO_OAUTH_<NAME>_AUTHORIZE_URL ZERO_OAUTH_<NAME>_TOKEN_URL
ZERO_OAUTH_<NAME>_DEVICE_URL ZERO_OAUTH_<NAME>_ISSUER_URL (for discovery)
ZERO_OAUTH_<NAME>_SCOPES ZERO_OAUTH_<NAME>_FLOW (loopback|device)
Endpoint URLs must be https (loopback exempt).
Flags:
--device Use the device-code flow (headless/SSH; no browser)
--scope Add an OAuth scope (repeatable)
--watch Keep the token fresh in the foreground (refresh only)
--json Print result as JSON (status/logout)
-h, --help Show this help
`)
return err
}
+129
View File
@@ -0,0 +1,129 @@
package cli
import (
"bytes"
"path/filepath"
"strings"
"testing"
"github.com/Gitlawb/zero/internal/oauth"
)
// withAuthStore points the provider OAuth store at a temp file for the test.
func withAuthStore(t *testing.T) string {
t.Helper()
path := filepath.Join(t.TempDir(), "oauth-tokens.json")
t.Setenv("ZERO_OAUTH_TOKENS_PATH", path)
return path
}
func TestRunAuthStatusEmpty(t *testing.T) {
withAuthStore(t)
var stdout, stderr bytes.Buffer
if code := runWithDeps([]string{"auth", "status"}, &stdout, &stderr, appDeps{}); code != exitSuccess {
t.Fatalf("exit = %d stderr=%s", code, stderr.String())
}
if !strings.Contains(stdout.String(), "No OAuth provider logins are stored.") {
t.Fatalf("status output = %q", stdout.String())
}
}
func TestRunAuthStatusReportsLoginWithoutSecret(t *testing.T) {
path := withAuthStore(t)
store, err := oauth.NewStore(oauth.StoreOptions{FilePath: path})
if err != nil {
t.Fatalf("NewStore: %v", err)
}
if err := store.Save(oauth.ProviderKey("demo"), oauth.Token{
AccessToken: "super-secret", RefreshToken: "super-secret-rt", Account: "[email protected]",
}); err != nil {
t.Fatalf("Save: %v", err)
}
var stdout, stderr bytes.Buffer
if code := runWithDeps([]string{"auth", "status"}, &stdout, &stderr, appDeps{}); code != exitSuccess {
t.Fatalf("exit = %d stderr=%s", code, stderr.String())
}
out := stdout.String()
if !strings.Contains(out, "demo") || !strings.Contains(out, "[email protected]") {
t.Fatalf("status should show provider + account: %q", out)
}
if strings.Contains(out, "super-secret") {
t.Fatalf("status leaked token material: %q", out)
}
}
func TestRunAuthLogoutNothing(t *testing.T) {
withAuthStore(t)
var stdout, stderr bytes.Buffer
if code := runWithDeps([]string{"auth", "logout", "demo"}, &stdout, &stderr, appDeps{}); code != exitSuccess {
t.Fatalf("exit = %d stderr=%s", code, stderr.String())
}
if !strings.Contains(stdout.String(), "No OAuth login stored for demo") {
t.Fatalf("logout output = %q", stdout.String())
}
}
func TestRunAuthLoginValidation(t *testing.T) {
withAuthStore(t)
var stdout, stderr bytes.Buffer
// Missing provider.
if code := runWithDeps([]string{"auth", "login"}, &stdout, &stderr, appDeps{}); code == exitSuccess {
t.Fatal("login with no provider should fail")
}
// --json is rejected for the interactive login.
stdout.Reset()
stderr.Reset()
if code := runWithDeps([]string{"auth", "login", "demo", "--json"}, &stdout, &stderr, appDeps{}); code == exitSuccess {
t.Fatal("login --json should be rejected")
}
}
func TestRunAuthLoginUnknownProvider(t *testing.T) {
withAuthStore(t)
var stdout, stderr bytes.Buffer
if code := runWithDeps([]string{"auth", "login", "does-not-exist"}, &stdout, &stderr, appDeps{}); code == exitSuccess {
t.Fatal("unknown provider login should fail")
}
if !strings.Contains(stderr.String(), "not configured") {
t.Fatalf("stderr = %q, want not-configured error", stderr.String())
}
}
func TestRunAuthRefreshNoToken(t *testing.T) {
withAuthStore(t)
t.Setenv("ZERO_OAUTH_DEMO_CLIENT_ID", "client") // so config resolves; refresh still fails (no token)
var stdout, stderr bytes.Buffer
if code := runWithDeps([]string{"auth", "refresh", "demo"}, &stdout, &stderr, appDeps{}); code == exitSuccess {
t.Fatal("refresh with no stored token should fail")
}
}
func TestRunAuthRejectsWrongFlags(t *testing.T) {
withAuthStore(t)
cases := [][]string{
{"auth", "login", "demo", "--watch"}, // watch is refresh-only
{"auth", "login", "demo", "--json"}, // json not for interactive login
{"auth", "status", "demo", "--device"}, // device is login-only
{"auth", "logout", "demo", "--scope", "x"}, // scope is login-only
{"auth", "refresh", "demo", "--json"}, // json not for refresh
{"auth", "login", "demo", "--scope", ""}, // empty scope rejected
}
for _, args := range cases {
var stdout, stderr bytes.Buffer
if code := runWithDeps(args, &stdout, &stderr, appDeps{}); code == exitSuccess {
t.Errorf("args %v should be rejected, got success", args)
}
}
}
func TestRunAuthHelp(t *testing.T) {
var stdout, stderr bytes.Buffer
if code := runWithDeps([]string{"auth", "--help"}, &stdout, &stderr, appDeps{}); code != exitSuccess {
t.Fatalf("exit = %d", code)
}
for _, want := range []string{"zero auth", "login", "logout", "status", "refresh", "--device"} {
if !strings.Contains(stdout.String(), want) {
t.Fatalf("help missing %q:\n%s", want, stdout.String())
}
}
}
+274 -18
View File
@@ -13,6 +13,7 @@ import (
"github.com/Gitlawb/zero/internal/background"
"github.com/Gitlawb/zero/internal/daemon"
"github.com/Gitlawb/zero/internal/daemon/remote"
)
// runDaemon dispatches the `zero daemon ...` subcommands. The daemon supervises a
@@ -35,6 +36,8 @@ func runDaemon(args []string, stdout io.Writer, stderr io.Writer, _ appDeps) int
return runDaemonRun(rest, stdout, stderr)
case "attach":
return runDaemonAttach(rest, stdout, stderr)
case "serve-remote":
return runDaemonServeRemote(rest, stdout, stderr)
case "-h", "--help", "help":
return writeDaemonUsage(stdout, exitSuccess)
default:
@@ -53,6 +56,14 @@ Commands:
run --session <id> [--cwd <dir>] [--prompt <text>] [exec flags...]
Create/route a session and stream its output.
attach <session> Attach to a running session's stream.
serve-remote --addr <host:port> --tls-cert <f> --tls-key <f>
Serve an opt-in, TLS-only network bridge to this
daemon. Requires a bearer token in $ZERO_DAEMON_REMOTE_TOKEN
(or $ZERO_DAEMON_REMOTE_TOKEN_FILE).
run and attach accept --remote <host:port> [--token <t>] [--ca-cert <f>]
[--server-name <name>] to drive a remote daemon over the bridge instead of the
local socket.
`)
return code
}
@@ -212,6 +223,7 @@ func runDaemonStatus(args []string, stdout io.Writer, stderr io.Writer) int {
func runDaemonRun(args []string, stdout io.Writer, stderr io.Writer) int {
session, cwd, prompt := "", "", ""
var remoteFlags remoteDialFlags
var forward []string
for i := 0; i < len(args); i++ {
a := args[i]
@@ -225,6 +237,38 @@ func runDaemonRun(args []string, stdout io.Writer, stderr io.Writer) int {
switch {
case a == "-h" || a == "--help":
return writeDaemonUsage(stdout, exitSuccess)
case a == "--remote":
v, ok := value()
if !ok {
return writeExecUsageError(stderr, "--remote requires a value")
}
remoteFlags.Addr = v
case strings.HasPrefix(a, "--remote="):
remoteFlags.Addr = strings.TrimPrefix(a, "--remote=")
case a == "--token":
v, ok := value()
if !ok {
return writeExecUsageError(stderr, "--token requires a value")
}
remoteFlags.Token = v
case strings.HasPrefix(a, "--token="):
remoteFlags.Token = strings.TrimPrefix(a, "--token=")
case a == "--ca-cert":
v, ok := value()
if !ok {
return writeExecUsageError(stderr, "--ca-cert requires a value")
}
remoteFlags.CACert = v
case strings.HasPrefix(a, "--ca-cert="):
remoteFlags.CACert = strings.TrimPrefix(a, "--ca-cert=")
case a == "--server-name":
v, ok := value()
if !ok {
return writeExecUsageError(stderr, "--server-name requires a value")
}
remoteFlags.ServerName = v
case strings.HasPrefix(a, "--server-name="):
remoteFlags.ServerName = strings.TrimPrefix(a, "--server-name=")
case a == "--session":
v, ok := value()
if !ok {
@@ -261,13 +305,9 @@ func runDaemonRun(args []string, stdout io.Writer, stderr io.Writer) int {
if prompt == "" && len(forward) == 0 {
return writeExecUsageError(stderr, "daemon run requires --prompt <text> or exec args")
}
paths, err := daemon.DefaultPaths()
client, err := dialForCLI(remoteFlags)
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
client, err := daemon.Dial(paths.Socket)
if err != nil {
return writeAppError(stderr, "zero daemon is not running (start it with `zero daemon start`)", exitCrash)
return writeAppError(stderr, daemonDialError(remoteFlags, err), exitCrash)
}
defer client.Close()
if err := client.Run(session, cwd, prompt, forward, func(line string) { fmt.Fprintln(stdout, line) }); err != nil {
@@ -278,16 +318,57 @@ func runDaemonRun(args []string, stdout io.Writer, stderr io.Writer) int {
func runDaemonAttach(args []string, stdout io.Writer, stderr io.Writer) int {
session := ""
var remoteFlags remoteDialFlags
extra := 0
for _, a := range args {
if a == "-h" || a == "--help" {
for i := 0; i < len(args); i++ {
a := args[i]
value := func() (string, bool) {
if i+1 >= len(args) {
return "", false
}
i++
return args[i], true
}
switch {
case a == "-h" || a == "--help":
return writeDaemonUsage(stdout, exitSuccess)
}
if session == "" {
case a == "--remote":
v, ok := value()
if !ok {
return writeExecUsageError(stderr, "--remote requires a value")
}
remoteFlags.Addr = v
case strings.HasPrefix(a, "--remote="):
remoteFlags.Addr = strings.TrimPrefix(a, "--remote=")
case a == "--token":
v, ok := value()
if !ok {
return writeExecUsageError(stderr, "--token requires a value")
}
remoteFlags.Token = v
case strings.HasPrefix(a, "--token="):
remoteFlags.Token = strings.TrimPrefix(a, "--token=")
case a == "--ca-cert":
v, ok := value()
if !ok {
return writeExecUsageError(stderr, "--ca-cert requires a value")
}
remoteFlags.CACert = v
case strings.HasPrefix(a, "--ca-cert="):
remoteFlags.CACert = strings.TrimPrefix(a, "--ca-cert=")
case a == "--server-name":
v, ok := value()
if !ok {
return writeExecUsageError(stderr, "--server-name requires a value")
}
remoteFlags.ServerName = v
case strings.HasPrefix(a, "--server-name="):
remoteFlags.ServerName = strings.TrimPrefix(a, "--server-name=")
case session == "":
session = a
continue
default:
extra++
}
extra++
}
if strings.TrimSpace(session) == "" {
return writeExecUsageError(stderr, "daemon attach requires a <session> id")
@@ -295,13 +376,9 @@ func runDaemonAttach(args []string, stdout io.Writer, stderr io.Writer) int {
if extra > 0 {
return writeExecUsageError(stderr, "daemon attach accepts exactly one <session> id")
}
paths, err := daemon.DefaultPaths()
client, err := dialForCLI(remoteFlags)
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
client, err := daemon.Dial(paths.Socket)
if err != nil {
return writeAppError(stderr, "zero daemon is not running (start it with `zero daemon start`)", exitCrash)
return writeAppError(stderr, daemonDialError(remoteFlags, err), exitCrash)
}
defer client.Close()
if err := client.Attach(session, func(line string) { fmt.Fprintln(stdout, line) }); err != nil {
@@ -310,6 +387,185 @@ func runDaemonAttach(args []string, stdout io.Writer, stderr io.Writer) int {
return exitSuccess
}
// daemonDialError tailors the connect-failure message for local vs remote.
func daemonDialError(flags remoteDialFlags, err error) string {
if strings.TrimSpace(flags.Addr) != "" {
return "failed to reach remote daemon at " + flags.Addr + ": " + err.Error()
}
return "zero daemon is not running (start it with `zero daemon start`)"
}
// runDaemonServeRemote starts the local daemon plus an opt-in, TLS-only network
// bridge. TLS and a bearer token are mandatory (fail closed): it refuses to
// start without a cert/key pair and a token from the environment.
func runDaemonServeRemote(args []string, stdout io.Writer, stderr io.Writer) int {
addr, certFile, keyFile := "", "", ""
minVersion, maxConns := 0, 0
for i := 0; i < len(args); i++ {
a := args[i]
value := func() (string, bool) {
if i+1 >= len(args) {
return "", false
}
i++
return args[i], true
}
switch {
case a == "-h" || a == "--help":
return writeDaemonUsage(stdout, exitSuccess)
case a == "--addr":
v, ok := value()
if !ok {
return writeExecUsageError(stderr, "--addr requires a value")
}
addr = v
case strings.HasPrefix(a, "--addr="):
addr = strings.TrimPrefix(a, "--addr=")
case a == "--tls-cert":
v, ok := value()
if !ok {
return writeExecUsageError(stderr, "--tls-cert requires a value")
}
certFile = v
case strings.HasPrefix(a, "--tls-cert="):
certFile = strings.TrimPrefix(a, "--tls-cert=")
case a == "--tls-key":
v, ok := value()
if !ok {
return writeExecUsageError(stderr, "--tls-key requires a value")
}
keyFile = v
case strings.HasPrefix(a, "--tls-key="):
keyFile = strings.TrimPrefix(a, "--tls-key=")
case a == "--min-version":
v, ok := value()
if !ok {
return writeExecUsageError(stderr, "--min-version requires a value")
}
minVersion = atoiOrZero(v)
case a == "--max-conns":
v, ok := value()
if !ok {
return writeExecUsageError(stderr, "--max-conns requires a value")
}
maxConns = atoiOrZero(v)
default:
return writeExecUsageError(stderr, fmt.Sprintf("unknown flag %q for daemon serve-remote", a))
}
}
if strings.TrimSpace(addr) == "" {
return writeExecUsageError(stderr, "daemon serve-remote requires --addr <host:port>")
}
// Fail closed: TLS cert/key + a bearer token are mandatory.
tlsConfig, err := remote.ServerTLSConfig(certFile, keyFile)
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
token, err := remote.TokenFromEnv()
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
auth, err := remote.NewTokenAuthenticator(token)
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
paths, err := daemon.DefaultPaths()
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
launcher, err := daemon.NewExecLauncher(daemon.ExecLauncherConfig{})
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
logf := func(line string) { fmt.Fprintln(stderr, "[daemon] "+line) }
pool, err := daemon.NewPool(daemon.PoolOptions{Launcher: launcher, Log: logf})
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
mgr, err := daemon.NewSessionManager(daemon.SessionManagerOptions{Pool: pool})
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
srv, err := daemon.NewServer(daemon.ServerOptions{Paths: paths, Manager: mgr, Pool: pool, Log: logf})
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
bridge, err := remote.NewBridge(remote.BridgeOptions{
Server: srv, Authenticator: auth, MinVersion: minVersion, MaxConnections: maxConns, Log: logf,
})
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
// Serve the local control socket too, so local clients keep working.
go func() {
if serveErr := srv.Serve(); serveErr != nil {
logf("local serve error: " + serveErr.Error())
}
}()
serveErr := make(chan error, 1)
go func() { serveErr <- bridge.ListenAndServeTLS(addr, tlsConfig) }()
sigCh := make(chan os.Signal, 1)
signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
fmt.Fprintf(stdout, "zero daemon remote bridge listening on %s (TLS)\n", addr)
select {
case <-sigCh:
srv.Shutdown()
_ = bridge.Close()
<-serveErr // wait for the accept loop to unwind
return exitSuccess
case err := <-serveErr:
// Bind/serve failed before any signal (e.g. address in use).
srv.Shutdown()
return writeAppError(stderr, err.Error(), exitCrash)
}
}
// remoteDialFlags holds the optional flags that redirect run/attach to a remote
// daemon. When Addr is empty the local control socket is used.
type remoteDialFlags struct {
Addr string
Token string
CACert string
ServerName string
}
// dialForCLI returns a daemon.Client for either the local socket (Addr empty) or
// a remote bridge (Addr set). For remote, the token falls back to the env.
func dialForCLI(flags remoteDialFlags) (*daemon.Client, error) {
if strings.TrimSpace(flags.Addr) == "" {
paths, err := daemon.DefaultPaths()
if err != nil {
return nil, err
}
return daemon.Dial(paths.Socket)
}
token := strings.TrimSpace(flags.Token)
if token == "" {
token, _ = remote.TokenFromEnv() // best effort; DialRemote rejects an empty token
}
return remote.DialRemote(remote.RemoteConfig{
Address: flags.Addr,
Token: token,
CACertFile: flags.CACert,
ServerName: flags.ServerName,
})
}
func atoiOrZero(s string) int {
n := 0
for _, r := range strings.TrimSpace(s) {
if r < '0' || r > '9' {
return 0
}
n = n*10 + int(r-'0')
}
return n
}
// daemonReachable reports whether a daemon is accepting connections (a successful
// handshake). Used for the single-instance check and start-up wait.
func daemonReachable(paths daemon.Paths) bool {
+13
View File
@@ -26,6 +26,19 @@ func Dial(socketPath string) (*Client, error) {
return c, nil
}
// NewClientConn completes the control handshake on an already-connected
// transport (e.g. a TLS connection the remote bridge has authenticated) and
// returns a Client ready for Run/Attach/Status. It takes ownership of conn and
// closes it on handshake failure.
func NewClientConn(conn net.Conn) (*Client, error) {
c := &Client{conn: conn}
if err := c.handshake(); err != nil {
_ = conn.Close()
return nil, err
}
return c, nil
}
func (c *Client) handshake() error {
if err := WriteControl(c.conn, Ctrl{Type: CtrlHello, Version: ProtoVersion}); err != nil {
return err
+158
View File
@@ -0,0 +1,158 @@
// Package remote adds an OPT-IN, TLS-only network bridge on top of zero's local
// daemon (internal/daemon). It lets a remote client drive the SAME daemon
// SessionManager/Pool over the network, behind bearer-token authentication and a
// protocol-version floor. The local Unix-socket daemon is unchanged and remains
// the default; nothing here activates unless `zero daemon serve-remote` is run.
//
// Security (fail closed):
// - TLS is mandatory; the bridge refuses to serve without a cert/key.
// - A connection is authenticated BEFORE any control frame is dispatched; a
// failed/absent token closes the connection (after a small backoff) and no
// session frame is ever processed.
// - Tokens are compared in constant time and are never logged.
// - The control-frame size cap and version negotiation from internal/daemon are
// reused unchanged; oversize/old-version handshakes are rejected.
// - A remote-driven session runs through the same daemon dispatch, so it stays
// under the same sandbox + risk model — remote never bypasses local controls.
package remote
import (
"crypto/subtle"
"encoding/json"
"errors"
"fmt"
"io"
"os"
"strings"
"github.com/Gitlawb/zero/internal/daemon"
)
// Env vars the bridge reads for its bearer token.
const (
EnvToken = "ZERO_DAEMON_REMOTE_TOKEN"
EnvTokenFile = "ZERO_DAEMON_REMOTE_TOKEN_FILE"
)
// ErrUnauthorized is returned when a token does not match.
var ErrUnauthorized = errors.New("remote: unauthorized")
// Authenticator verifies a bearer token presented by a remote client.
type Authenticator interface {
Authenticate(token string) error
}
// TokenAuthenticator compares a presented token against a fixed secret in
// constant time.
type TokenAuthenticator struct {
token string
}
// NewTokenAuthenticator builds a token authenticator, refusing an empty secret
// (fail closed — a bridge must never accept everyone).
func NewTokenAuthenticator(token string) (*TokenAuthenticator, error) {
token = strings.TrimSpace(token)
if token == "" {
return nil, errors.New("remote: a non-empty auth token is required")
}
return &TokenAuthenticator{token: token}, nil
}
// Authenticate reports nil when token matches the configured secret.
func (a *TokenAuthenticator) Authenticate(token string) error {
if subtle.ConstantTimeCompare([]byte(token), []byte(a.token)) == 1 {
return nil
}
return ErrUnauthorized
}
// TokenFromEnv resolves the bridge token from EnvToken, or a file named by
// EnvTokenFile. It never logs the token.
func TokenFromEnv() (string, error) {
if t := strings.TrimSpace(os.Getenv(EnvToken)); t != "" {
return t, nil
}
if file := strings.TrimSpace(os.Getenv(EnvTokenFile)); file != "" {
data, err := os.ReadFile(file)
if err != nil {
return "", fmt.Errorf("remote: read token file: %w", err)
}
t := strings.TrimSpace(string(data))
if t == "" {
return "", errors.New("remote: token file is empty")
}
return t, nil
}
return "", fmt.Errorf("remote: set %s or %s", EnvToken, EnvTokenFile)
}
// Attestation is an optional post-token hook (e.g. workload attestation). The
// default is a no-op; a deployment can supply a stricter implementation.
type Attestation interface {
Verify(meta map[string]string) error
}
type noopAttestation struct{}
func (noopAttestation) Verify(map[string]string) error { return nil }
// authRequest is the first frame a remote client sends (before the daemon
// hello). Token is never logged.
type authRequest struct {
Token string `json:"token"`
Version int `json:"version"`
Meta map[string]string `json:"meta,omitempty"`
}
// authResponse is the bridge's reply to the auth handshake.
type authResponse struct {
OK bool `json:"ok"`
Version int `json:"version,omitempty"`
Message string `json:"message,omitempty"`
}
func writeAuthRequest(w io.Writer, req authRequest) error {
payload, err := json.Marshal(req)
if err != nil {
return err
}
return daemon.WriteFrame(w, daemon.KindCtrl, payload)
}
func readAuthRequest(r io.Reader) (authRequest, error) {
kind, payload, err := daemon.ReadFrame(r)
if err != nil {
return authRequest{}, err
}
if kind != daemon.KindCtrl {
return authRequest{}, errors.New("remote: expected auth frame")
}
var req authRequest
if err := json.Unmarshal(payload, &req); err != nil {
return authRequest{}, fmt.Errorf("remote: decode auth request: %w", err)
}
return req, nil
}
func writeAuthResponse(w io.Writer, resp authResponse) error {
payload, err := json.Marshal(resp)
if err != nil {
return err
}
return daemon.WriteFrame(w, daemon.KindCtrl, payload)
}
func readAuthResponse(r io.Reader) (authResponse, error) {
kind, payload, err := daemon.ReadFrame(r)
if err != nil {
return authResponse{}, err
}
if kind != daemon.KindCtrl {
return authResponse{}, errors.New("remote: expected auth response frame")
}
var resp authResponse
if err := json.Unmarshal(payload, &resp); err != nil {
return authResponse{}, fmt.Errorf("remote: decode auth response: %w", err)
}
return resp, nil
}
+67
View File
@@ -0,0 +1,67 @@
package remote
import (
"errors"
"os"
"path/filepath"
"testing"
)
func TestTokenAuthenticator(t *testing.T) {
if _, err := NewTokenAuthenticator(" "); err == nil {
t.Fatal("empty token must be rejected (fail closed)")
}
a, err := NewTokenAuthenticator("s3cret")
if err != nil {
t.Fatalf("NewTokenAuthenticator: %v", err)
}
if err := a.Authenticate("s3cret"); err != nil {
t.Fatalf("matching token should authenticate: %v", err)
}
if err := a.Authenticate("wrong"); !errors.Is(err, ErrUnauthorized) {
t.Fatalf("mismatch err = %v, want ErrUnauthorized", err)
}
if err := a.Authenticate(""); !errors.Is(err, ErrUnauthorized) {
t.Fatalf("empty presented token err = %v, want ErrUnauthorized", err)
}
}
func TestTokenFromEnv(t *testing.T) {
// Clear both so the no-config path errors.
t.Setenv(EnvToken, "")
t.Setenv(EnvTokenFile, "")
if _, err := TokenFromEnv(); err == nil {
t.Fatal("TokenFromEnv with neither var set must error")
}
// Direct env.
t.Setenv(EnvToken, "from-env")
if tok, err := TokenFromEnv(); err != nil || tok != "from-env" {
t.Fatalf("TokenFromEnv(env) = %q, %v", tok, err)
}
// File (env takes precedence, so clear env first).
t.Setenv(EnvToken, "")
file := filepath.Join(t.TempDir(), "tok")
if err := os.WriteFile(file, []byte(" from-file\n"), 0o600); err != nil {
t.Fatalf("write token file: %v", err)
}
t.Setenv(EnvTokenFile, file)
if tok, err := TokenFromEnv(); err != nil || tok != "from-file" {
t.Fatalf("TokenFromEnv(file) = %q, %v", tok, err)
}
// Empty file fails closed.
if err := os.WriteFile(file, []byte(" \n"), 0o600); err != nil {
t.Fatalf("rewrite token file: %v", err)
}
if _, err := TokenFromEnv(); err == nil {
t.Fatal("empty token file must error")
}
}
func TestServerTLSConfigRequiresCertKey(t *testing.T) {
if _, err := ServerTLSConfig("", ""); err == nil {
t.Fatal("ServerTLSConfig must require a cert and key (TLS mandatory)")
}
if _, err := ServerTLSConfig("/nope/cert.pem", "/nope/key.pem"); err == nil {
t.Fatal("ServerTLSConfig must error on unreadable cert/key")
}
}
+214
View File
@@ -0,0 +1,214 @@
package remote
import (
"crypto/tls"
"errors"
"fmt"
"net"
"strings"
"sync"
"time"
"github.com/Gitlawb/zero/internal/daemon"
)
const (
defaultMaxConnections = 32
defaultHandshakeTimeout = 10 * time.Second
defaultAuthFailDelay = 250 * time.Millisecond
)
// Bridge serves authenticated remote connections and drives the local daemon's
// control dispatch for each one.
type Bridge struct {
server *daemon.Server
auth Authenticator
attest Attestation
minVersion int
handshakeTimeout time.Duration
authFailDelay time.Duration
log func(string)
sem chan struct{}
mu sync.Mutex
listener net.Listener
}
// BridgeOptions configures a Bridge.
type BridgeOptions struct {
// Server is the local daemon to drive (required).
Server *daemon.Server
// Authenticator verifies bearer tokens (required).
Authenticator Authenticator
// Attestation is an optional post-token hook; nil => no-op.
Attestation Attestation
// MinVersion is the lowest control-protocol version accepted; 0 =>
// daemon.ProtoVersion.
MinVersion int
// MaxConnections bounds concurrent remote connections; 0 => default.
MaxConnections int
// HandshakeTimeout bounds the auth handshake; 0 => default.
HandshakeTimeout time.Duration
// AuthFailDelay slows brute-force attempts; <0 => none, 0 => default.
AuthFailDelay time.Duration
Log func(string)
}
// NewBridge validates options and builds a Bridge.
func NewBridge(opts BridgeOptions) (*Bridge, error) {
if opts.Server == nil {
return nil, errors.New("remote: bridge requires a daemon server")
}
if opts.Authenticator == nil {
return nil, errors.New("remote: bridge requires an authenticator")
}
minVersion := opts.MinVersion
if minVersion <= 0 {
minVersion = daemon.ProtoVersion
}
maxConns := opts.MaxConnections
if maxConns <= 0 {
maxConns = defaultMaxConnections
}
handshakeTimeout := opts.HandshakeTimeout
if handshakeTimeout <= 0 {
handshakeTimeout = defaultHandshakeTimeout
}
authFailDelay := opts.AuthFailDelay
switch {
case authFailDelay < 0:
authFailDelay = 0
case authFailDelay == 0:
authFailDelay = defaultAuthFailDelay
}
attest := opts.Attestation
if attest == nil {
attest = noopAttestation{}
}
return &Bridge{
server: opts.Server,
auth: opts.Authenticator,
attest: attest,
minVersion: minVersion,
handshakeTimeout: handshakeTimeout,
authFailDelay: authFailDelay,
log: opts.Log,
sem: make(chan struct{}, maxConns),
}, nil
}
func (b *Bridge) logf(format string, args ...any) {
if b.log != nil {
b.log(fmt.Sprintf(format, args...))
}
}
// ListenAndServeTLS binds a TLS listener on addr and serves until the listener
// is closed (via Close). It refuses to serve without a TLS config (no plaintext
// remote bridge).
func (b *Bridge) ListenAndServeTLS(addr string, tlsConfig *tls.Config) error {
if tlsConfig == nil {
return errors.New("remote: TLS config is required; refusing to serve a plaintext remote bridge")
}
listener, err := tls.Listen("tcp", addr, tlsConfig)
if err != nil {
return fmt.Errorf("remote: listen %s: %w", addr, err)
}
b.mu.Lock()
b.listener = listener
b.mu.Unlock()
return b.Serve(listener)
}
// Close stops the listener started by ListenAndServeTLS, causing Serve to
// return. Safe to call before serving or more than once.
func (b *Bridge) Close() error {
b.mu.Lock()
listener := b.listener
b.listener = nil // clear so a repeat Close is a no-op, not a closed-listener error
b.mu.Unlock()
if listener != nil {
return listener.Close()
}
return nil
}
// Serve accepts connections from listener until it returns an error (e.g. it is
// closed). Each connection is authenticated, then handed to the daemon's control
// dispatch. Connections beyond MaxConnections are refused immediately.
func (b *Bridge) Serve(listener net.Listener) error {
for {
conn, err := listener.Accept()
if err != nil {
return err
}
select {
case b.sem <- struct{}{}:
default:
// At capacity: refuse fast (fail closed) rather than queueing.
b.logf("remote: refused connection (at capacity)")
_ = conn.Close()
continue
}
go func() {
defer func() { <-b.sem }()
b.handle(conn)
}()
}
}
// handle authenticates conn then drives the daemon dispatch. It always closes
// conn (directly on a failed handshake, or via ServeConn on success).
func (b *Bridge) handle(conn net.Conn) {
// Bound the auth handshake so a stalled peer cannot hold a connection slot.
_ = conn.SetDeadline(time.Now().Add(b.handshakeTimeout))
req, err := readAuthRequest(conn)
if err != nil {
_ = conn.Close()
return
}
if req.Version < b.minVersion {
b.deny(conn, "unsupported protocol version")
return
}
if err := b.auth.Authenticate(req.Token); err != nil {
b.deny(conn, "unauthorized")
return
}
if err := b.attest.Verify(req.Meta); err != nil {
b.deny(conn, "attestation failed")
return
}
if err := writeAuthResponse(conn, authResponse{OK: true, Version: daemon.ProtoVersion}); err != nil {
_ = conn.Close()
return
}
// Clear the handshake deadline: a session may stream for a long time.
_ = conn.SetDeadline(time.Time{})
b.server.ServeConn(conn) // performs the daemon handshake + one command, then closes conn
}
// deny rejects an unauthenticated connection after a small backoff (to slow
// brute force). The reason is generic and never includes any token material.
func (b *Bridge) deny(conn net.Conn, message string) {
if b.authFailDelay > 0 {
time.Sleep(b.authFailDelay)
}
_ = writeAuthResponse(conn, authResponse{OK: false, Message: message})
_ = conn.Close()
}
// ServerTLSConfig builds a server TLS config from a cert/key pair, refusing if
// either is missing (TLS is mandatory for the remote bridge).
func ServerTLSConfig(certFile, keyFile string) (*tls.Config, error) {
certFile = strings.TrimSpace(certFile)
keyFile = strings.TrimSpace(keyFile)
if certFile == "" || keyFile == "" {
return nil, errors.New("remote: a TLS cert and key are required for serve-remote")
}
cert, err := tls.LoadX509KeyPair(certFile, keyFile)
if err != nil {
return nil, fmt.Errorf("remote: load TLS keypair: %w", err)
}
return &tls.Config{Certificates: []tls.Certificate{cert}, MinVersion: tls.VersionTLS12}, nil
}
+342
View File
@@ -0,0 +1,342 @@
package remote
import (
"context"
"crypto/ecdsa"
"crypto/elliptic"
"crypto/rand"
"crypto/tls"
"crypto/x509"
"crypto/x509/pkix"
"encoding/pem"
"errors"
"math/big"
"net"
"os"
"path/filepath"
"testing"
"time"
"github.com/Gitlawb/zero/internal/daemon"
)
// ---- fake daemon worker (the daemon's test fakes are unexported) -------------
type fakeLines struct{ ch <-chan string }
func (l fakeLines) Next() (string, bool, error) {
line, ok := <-l.ch
return line, ok, nil
}
type fakeWorker struct {
ch chan string
pid int
}
func (w *fakeWorker) Stdout() daemon.Lines { return fakeLines{ch: w.ch} }
func (w *fakeWorker) Wait() (int, error) { return 0, nil }
func (w *fakeWorker) Kill() error { return nil }
func (w *fakeWorker) Pid() int { return w.pid }
// staticLauncher emits the given lines then ends the stream.
func staticLauncher(lines ...string) daemon.Launcher {
return func(_ context.Context, _ daemon.WorkerSpec) (daemon.WorkerHandle, error) {
ch := make(chan string, len(lines))
for _, l := range lines {
ch <- l
}
close(ch)
return &fakeWorker{ch: ch, pid: 1}, nil
}
}
// blockingLauncher returns a worker whose stream stays open until release is
// closed, so a session (and its bridge connection slot) is held open.
func blockingLauncher(release <-chan struct{}) daemon.Launcher {
return func(_ context.Context, _ daemon.WorkerSpec) (daemon.WorkerHandle, error) {
ch := make(chan string)
go func() {
<-release
close(ch)
}()
return &fakeWorker{ch: ch, pid: 1}, nil
}
}
func newBridgeServer(t *testing.T, launcher daemon.Launcher) *daemon.Server {
t.Helper()
dir := t.TempDir()
pool, err := daemon.NewPool(daemon.PoolOptions{Size: 4, Launcher: launcher, KillTimeout: 200 * time.Millisecond})
if err != nil {
t.Fatalf("NewPool: %v", err)
}
mgr, err := daemon.NewSessionManager(daemon.SessionManagerOptions{Pool: pool})
if err != nil {
t.Fatalf("NewSessionManager: %v", err)
}
srv, err := daemon.NewServer(daemon.ServerOptions{
Paths: daemon.Paths{Socket: filepath.Join(dir, "s.sock"), Lock: filepath.Join(dir, "s.lock"), Status: filepath.Join(dir, "s.status")},
Manager: mgr,
Pool: pool,
})
if err != nil {
t.Fatalf("NewServer: %v", err)
}
t.Cleanup(srv.Shutdown)
return srv
}
// genTestCert writes a self-signed cert/key (valid for 127.0.0.1) to temp files.
func genTestCert(t *testing.T) (certFile, keyFile string) {
t.Helper()
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
if err != nil {
t.Fatalf("genkey: %v", err)
}
tmpl := &x509.Certificate{
SerialNumber: big.NewInt(1),
Subject: pkix.Name{CommonName: "zero-remote-test"},
NotBefore: time.Now().Add(-time.Hour),
NotAfter: time.Now().Add(time.Hour),
KeyUsage: x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
IsCA: true,
IPAddresses: []net.IP{net.ParseIP("127.0.0.1"), net.ParseIP("::1")},
DNSNames: []string{"localhost"},
}
der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
if err != nil {
t.Fatalf("create cert: %v", err)
}
dir := t.TempDir()
certFile = filepath.Join(dir, "cert.pem")
keyFile = filepath.Join(dir, "key.pem")
certOut, _ := os.Create(certFile)
_ = pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: der})
_ = certOut.Close()
keyDER, _ := x509.MarshalECPrivateKey(key)
keyOut, _ := os.Create(keyFile)
_ = pem.Encode(keyOut, &pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER})
_ = keyOut.Close()
return certFile, keyFile
}
// startBridge starts a TLS bridge on 127.0.0.1:0 and returns its address + cert.
func startBridge(t *testing.T, srv *daemon.Server, opts BridgeOptions) (addr, caCert string) {
t.Helper()
certFile, keyFile := genTestCert(t)
tlsConfig, err := ServerTLSConfig(certFile, keyFile)
if err != nil {
t.Fatalf("ServerTLSConfig: %v", err)
}
opts.Server = srv
bridge, err := NewBridge(opts)
if err != nil {
t.Fatalf("NewBridge: %v", err)
}
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
tlsLn := tls.NewListener(ln, tlsConfig)
go func() { _ = bridge.Serve(tlsLn) }()
t.Cleanup(func() { _ = tlsLn.Close() })
return ln.Addr().String(), certFile
}
func TestBridgeRunStatusRoundTrip(t *testing.T) {
srv := newBridgeServer(t, staticLauncher(`{"type":"event","seq":1}`, `{"type":"event","seq":2}`))
auth, _ := NewTokenAuthenticator("tok")
addr, ca := startBridge(t, srv, BridgeOptions{Authenticator: auth})
client, err := DialRemote(RemoteConfig{Address: addr, Token: "tok", CACertFile: ca})
if err != nil {
t.Fatalf("DialRemote: %v", err)
}
var got []string
if err := client.Run("sess-1", "", "hello", nil, func(l string) { got = append(got, l) }); err != nil {
t.Fatalf("Run: %v", err)
}
_ = client.Close()
if len(got) != 2 || got[0] != `{"type":"event","seq":1}` {
t.Fatalf("run output = %v", got)
}
// Status over a fresh remote connection (same protocol as local).
statusClient, err := DialRemote(RemoteConfig{Address: addr, Token: "tok", CACertFile: ca})
if err != nil {
t.Fatalf("DialRemote(status): %v", err)
}
report, err := statusClient.Status()
if err != nil {
t.Fatalf("Status: %v", err)
}
_ = statusClient.Close()
if report == nil || report.PID == 0 {
t.Fatalf("status report = %+v", report)
}
}
func TestBridgeRejectsBadToken(t *testing.T) {
srv := newBridgeServer(t, staticLauncher())
auth, _ := NewTokenAuthenticator("correct")
addr, ca := startBridge(t, srv, BridgeOptions{Authenticator: auth, AuthFailDelay: -1})
_, err := DialRemote(RemoteConfig{Address: addr, Token: "wrong", CACertFile: ca})
if !errors.Is(err, ErrUnauthorized) {
t.Fatalf("DialRemote(bad token) err = %v, want ErrUnauthorized", err)
}
}
func TestBridgeRejectsLowVersion(t *testing.T) {
srv := newBridgeServer(t, staticLauncher())
auth, _ := NewTokenAuthenticator("tok")
// Require a version higher than the client speaks (daemon.ProtoVersion).
addr, ca := startBridge(t, srv, BridgeOptions{Authenticator: auth, MinVersion: daemon.ProtoVersion + 1, AuthFailDelay: -1})
_, err := DialRemote(RemoteConfig{Address: addr, Token: "tok", CACertFile: ca})
if err == nil {
t.Fatal("DialRemote with a below-min version must be rejected")
}
}
func TestBridgeRequiresVerifiableCert(t *testing.T) {
srv := newBridgeServer(t, staticLauncher())
auth, _ := NewTokenAuthenticator("tok")
addr, _ := startBridge(t, srv, BridgeOptions{Authenticator: auth})
// No CA configured + a self-signed server cert => verification must fail
// (NEVER InsecureSkipVerify).
if _, err := DialRemote(RemoteConfig{Address: addr, Token: "tok"}); err == nil {
t.Fatal("DialRemote must reject an untrusted self-signed cert")
}
}
func TestListenAndServeTLSRequiresConfig(t *testing.T) {
srv := newBridgeServer(t, staticLauncher())
auth, _ := NewTokenAuthenticator("tok")
bridge, _ := NewBridge(BridgeOptions{Server: srv, Authenticator: auth})
if err := bridge.ListenAndServeTLS("127.0.0.1:0", nil); err == nil {
t.Fatal("ListenAndServeTLS(nil) must refuse to serve plaintext")
}
}
func TestListenAndServeTLSServesAndCloses(t *testing.T) {
srv := newBridgeServer(t, staticLauncher(`{"type":"event","seq":1}`))
auth, _ := NewTokenAuthenticator("tok")
certFile, keyFile := genTestCert(t)
tlsConfig, err := ServerTLSConfig(certFile, keyFile)
if err != nil {
t.Fatalf("ServerTLSConfig: %v", err)
}
bridge, err := NewBridge(BridgeOptions{Server: srv, Authenticator: auth})
if err != nil {
t.Fatalf("NewBridge: %v", err)
}
// Bind an ephemeral port, discover it, then serve via ListenAndServeTLS by
// passing the same addr (it rebinds) — instead we bind here to learn the port
// and serve through the public entry on that exact address.
probe, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("probe listen: %v", err)
}
addr := probe.Addr().String()
_ = probe.Close()
serveErr := make(chan error, 1)
go func() { serveErr <- bridge.ListenAndServeTLS(addr, tlsConfig) }()
// Wait until it accepts a connection, then drive one run.
var client *daemon.Client
deadline := time.Now().Add(2 * time.Second)
for {
client, err = DialRemote(RemoteConfig{Address: addr, Token: "tok", CACertFile: certFile, Timeout: 300 * time.Millisecond})
if err == nil {
break
}
if time.Now().After(deadline) {
t.Fatalf("bridge never came up: %v", err)
}
time.Sleep(20 * time.Millisecond)
}
var got []string
if err := client.Run("s", "", "hi", nil, func(l string) { got = append(got, l) }); err != nil {
t.Fatalf("Run: %v", err)
}
_ = client.Close()
if len(got) != 1 {
t.Fatalf("got %d lines, want 1", len(got))
}
// Close stops the accept loop; a second Close is a no-op (not an error).
if err := bridge.Close(); err != nil {
t.Fatalf("Close: %v", err)
}
if err := bridge.Close(); err != nil {
t.Fatalf("second Close should be a no-op, got: %v", err)
}
select {
case <-serveErr: // ListenAndServeTLS returned after Close
case <-time.After(2 * time.Second):
t.Fatal("ListenAndServeTLS did not return after Close")
}
}
func TestBridgeResumeAfterReconnect(t *testing.T) {
srv := newBridgeServer(t, staticLauncher(`{"type":"event","seq":1}`, `{"type":"event","seq":2}`))
auth, _ := NewTokenAuthenticator("tok")
addr, ca := startBridge(t, srv, BridgeOptions{Authenticator: auth})
// First connection runs the session to completion.
c1, err := DialRemote(RemoteConfig{Address: addr, Token: "tok", CACertFile: ca})
if err != nil {
t.Fatalf("DialRemote 1: %v", err)
}
if err := c1.Run("sess-resume", "", "hi", nil, func(string) {}); err != nil {
t.Fatalf("Run: %v", err)
}
_ = c1.Close()
// A new connection re-attaches and replays the buffered session output —
// the dropped first connection did not lose the session.
c2, err := DialRemote(RemoteConfig{Address: addr, Token: "tok", CACertFile: ca})
if err != nil {
t.Fatalf("DialRemote 2: %v", err)
}
defer c2.Close()
var replayed []string
if err := c2.Attach("sess-resume", func(l string) { replayed = append(replayed, l) }); err != nil {
t.Fatalf("Attach: %v", err)
}
if len(replayed) != 2 {
t.Fatalf("re-attach replayed %d lines, want 2", len(replayed))
}
}
func TestBridgeConnectionCap(t *testing.T) {
release := make(chan struct{})
defer close(release)
srv := newBridgeServer(t, blockingLauncher(release))
auth, _ := NewTokenAuthenticator("tok")
addr, ca := startBridge(t, srv, BridgeOptions{Authenticator: auth, MaxConnections: 1})
// First connection holds the only slot with a blocking run.
c1, err := DialRemote(RemoteConfig{Address: addr, Token: "tok", CACertFile: ca})
if err != nil {
t.Fatalf("DialRemote 1: %v", err)
}
defer c1.Close()
runErr := make(chan error, 1)
go func() { runErr <- c1.Run("sess-block", "", "hi", nil, func(string) {}) }()
// Wait until the slot is in use, then a second dial must be refused.
deadline := time.Now().Add(2 * time.Second)
for time.Now().Before(deadline) {
c2, err := DialRemote(RemoteConfig{Address: addr, Token: "tok", CACertFile: ca, Timeout: 500 * time.Millisecond})
if err != nil {
return // refused as expected
}
_ = c2.Close()
time.Sleep(20 * time.Millisecond)
}
t.Fatal("second connection was not refused at the connection cap")
}
+108
View File
@@ -0,0 +1,108 @@
package remote
import (
"context"
"crypto/tls"
"crypto/x509"
"errors"
"fmt"
"net"
"os"
"strings"
"time"
"github.com/Gitlawb/zero/internal/daemon"
)
const defaultDialTimeout = 15 * time.Second
// RemoteConfig configures a remote dial.
type RemoteConfig struct {
// Address is host:port of the remote bridge (required).
Address string
// Token is the bearer token (required).
Token string
// CACertFile, when set, is the only CA trusted for the server cert (use for a
// self-signed bridge). Empty => system roots. InsecureSkipVerify is never used.
CACertFile string
// ServerName overrides the TLS/SNI verification name; empty => host of Address.
ServerName string
// Timeout bounds the dial + auth handshake; 0 => default.
Timeout time.Duration
}
// DialRemote establishes a verified TLS connection to a remote bridge,
// authenticates with the bearer token, and returns a daemon.Client ready for
// Run/Attach/Status — the same client the local socket uses, so remote and local
// share one protocol. The server certificate is always verified (never
// InsecureSkipVerify).
func DialRemote(cfg RemoteConfig) (*daemon.Client, error) {
address := strings.TrimSpace(cfg.Address)
if address == "" {
return nil, errors.New("remote: address is required")
}
token := strings.TrimSpace(cfg.Token)
if token == "" {
return nil, errors.New("remote: a non-empty token is required")
}
tlsConfig, err := clientTLSConfig(cfg)
if err != nil {
return nil, err
}
timeout := cfg.Timeout
if timeout <= 0 {
timeout = defaultDialTimeout
}
// Context-aware TLS dial (tls.DialWithDialer is deprecated and not
// cancellation-aware). The context bounds the dial + TLS handshake only; the
// returned conn is independent of it for the subsequent auth/stream I/O.
dialCtx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
dialer := &tls.Dialer{NetDialer: &net.Dialer{}, Config: tlsConfig}
conn, err := dialer.DialContext(dialCtx, "tcp", address)
if err != nil {
return nil, fmt.Errorf("remote: dial %s: %w", address, err)
}
// Bound the auth handshake; the daemon handshake + stream run without a deadline.
_ = conn.SetDeadline(time.Now().Add(timeout))
if err := writeAuthRequest(conn, authRequest{Token: token, Version: daemon.ProtoVersion}); err != nil {
_ = conn.Close()
return nil, fmt.Errorf("remote: send auth: %w", err)
}
resp, err := readAuthResponse(conn)
if err != nil {
_ = conn.Close()
return nil, fmt.Errorf("remote: auth handshake: %w", err)
}
if !resp.OK {
_ = conn.Close()
return nil, fmt.Errorf("%w: %s", ErrUnauthorized, resp.Message)
}
_ = conn.SetDeadline(time.Time{})
return daemon.NewClientConn(conn)
}
// clientTLSConfig builds a verifying client TLS config. It never disables
// verification; for a self-signed bridge, set CACertFile to the bridge's cert.
func clientTLSConfig(cfg RemoteConfig) (*tls.Config, error) {
tc := &tls.Config{MinVersion: tls.VersionTLS12}
serverName := strings.TrimSpace(cfg.ServerName)
if serverName == "" {
if host, _, err := net.SplitHostPort(cfg.Address); err == nil {
serverName = host
}
}
tc.ServerName = serverName
if ca := strings.TrimSpace(cfg.CACertFile); ca != "" {
pem, err := os.ReadFile(ca)
if err != nil {
return nil, fmt.Errorf("remote: read CA cert: %w", err)
}
pool := x509.NewCertPool()
if !pool.AppendCertsFromPEM(pem) {
return nil, errors.New("remote: CA cert file contained no certificates")
}
tc.RootCAs = pool
}
return tc, nil
}
+7
View File
@@ -167,6 +167,13 @@ func (s *Server) writeStatusFile() error {
return nil
}
// ServeConn runs the control protocol (handshake + one command) on an
// already-established connection, reusing the exact local dispatch path. The
// remote bridge calls it AFTER authenticating a TLS connection, so a remote
// session is handled identically to a local one (same SessionManager/Pool, same
// sandbox/risk model) — remote never bypasses the local controls. It closes conn.
func (s *Server) ServeConn(conn net.Conn) { s.handleConn(conn) }
// handleConn performs the handshake then dispatches a single control command.
func (s *Server) handleConn(conn net.Conn) {
defer conn.Close()
+79 -187
View File
@@ -2,9 +2,6 @@ package mcp
import (
"context"
"crypto/rand"
"crypto/sha256"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
@@ -14,20 +11,56 @@ import (
"net/url"
"strings"
"time"
"github.com/Gitlawb/zero/internal/oauth"
)
// ServerAuthOAuth is the value of an MCP server's auth field that selects the
// OAuth 2.0 + PKCE authorization-code flow.
const ServerAuthOAuth = "oauth"
const (
wellKnownAuthServerPath = "/.well-known/oauth-authorization-server"
defaultLoginTimeout = 3 * time.Minute
// pkceVerifierBytes yields a 43-character base64url verifier, the minimum
// length permitted by the PKCE spec while remaining high-entropy.
pkceVerifierBytes = 32
stateBytes = 32
)
const defaultLoginTimeout = 3 * time.Minute
// MCP OAuth delegates its transport/identity-agnostic engine to internal/oauth
// (PKCE, RFC 8414 discovery, authorize-URL build, token exchange/refresh) so the
// two share one implementation. MCP keeps its own LoginOptions/Login
// orchestration, OAuthConfig/StoredToken types, loopback handling, CLI, and
// on-disk token format — all behavior-preserving. The shared engine also adds an
// https-only token-endpoint guard (loopback exempt), a hardening MCP inherits.
// pkceToParams converts the shared PKCE pair to MCP's local type.
func pkceToParams(p oauth.PKCE) pkceParams {
return pkceParams{Verifier: p.Verifier, Challenge: p.Challenge, Method: p.Method}
}
// pkceToOAuth converts MCP's local PKCE type back to the shared type.
func pkceToOAuth(p pkceParams) oauth.PKCE {
return oauth.PKCE{Verifier: p.Verifier, Challenge: p.Challenge, Method: p.Method}
}
// configFor builds the shared oauth.Config from MCP's OAuthConfig + resolved
// endpoints for a flow.
func configFor(cfg OAuthConfig, metadata authServerMetadata) oauth.Config {
return oauth.Config{
ClientID: cfg.ClientID,
ClientSecret: cfg.ClientSecret,
Scopes: cfg.Scopes,
AuthorizationEndpoint: metadata.AuthorizationEndpoint,
TokenEndpoint: metadata.TokenEndpoint,
}
}
// tokenToStored converts a shared oauth.Token to MCP's StoredToken (MCP does not
// use the Account field).
func tokenToStored(t oauth.Token) StoredToken {
return StoredToken{
AccessToken: t.AccessToken,
RefreshToken: t.RefreshToken,
TokenType: t.TokenType,
Scopes: t.Scopes,
ExpiresAt: t.ExpiresAt,
}
}
// OAuthConfig describes how to authenticate to a remote MCP server using OAuth.
// Endpoints may be discovered from the server's metadata document; explicit
@@ -84,34 +117,20 @@ type authorizationFlow struct {
now func() time.Time
}
// discoverAuthorizationServer fetches the OAuth 2.0 authorization server
// metadata document published at the well-known path under the given base URL.
// discoverAuthorizationServer fetches the RFC 8414 authorization server metadata
// at the well-known path under baseURL, via the shared engine.
func discoverAuthorizationServer(ctx context.Context, client *http.Client, baseURL string) (authServerMetadata, error) {
if client == nil {
client = http.DefaultClient
}
metadataURL, err := joinWellKnown(baseURL)
meta, err := oauth.DiscoverAuthorizationServer(ctx, client, baseURL)
if err != nil {
return authServerMetadata{}, err
}
request, err := http.NewRequestWithContext(ctx, http.MethodGet, metadataURL, nil)
if err != nil {
return authServerMetadata{}, err
}
request.Header.Set("Accept", "application/json")
response, err := client.Do(request)
if err != nil {
return authServerMetadata{}, fmt.Errorf("fetch authorization server metadata: %w", err)
}
defer response.Body.Close()
if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
return authServerMetadata{}, fmt.Errorf("authorization server metadata returned HTTP %d", response.StatusCode)
}
var metadata authServerMetadata
if err := json.NewDecoder(io.LimitReader(response.Body, 1<<20)).Decode(&metadata); err != nil {
return authServerMetadata{}, fmt.Errorf("decode authorization server metadata: %w", err)
}
return metadata, nil
return authServerMetadata{
Issuer: meta.Issuer,
AuthorizationEndpoint: meta.AuthorizationEndpoint,
TokenEndpoint: meta.TokenEndpoint,
RegistrationEndpoint: meta.RegistrationEndpoint,
ScopesSupported: meta.ScopesSupported,
}, nil
}
// resolveAuthorizationServer discovers metadata and applies explicit config
@@ -149,44 +168,23 @@ func resolveAuthorizationServer(ctx context.Context, client *http.Client, baseUR
return metadata, nil
}
// newPKCE generates a high-entropy code verifier and its S256 challenge.
// newPKCE generates a high-entropy code verifier and its S256 challenge via the
// shared engine.
func newPKCE() (pkceParams, error) {
verifierRaw := make([]byte, pkceVerifierBytes)
if _, err := rand.Read(verifierRaw); err != nil {
return pkceParams{}, fmt.Errorf("generate PKCE verifier: %w", err)
p, err := oauth.NewPKCE()
if err != nil {
return pkceParams{}, err
}
verifier := base64.RawURLEncoding.EncodeToString(verifierRaw)
sum := sha256.Sum256([]byte(verifier))
challenge := base64.RawURLEncoding.EncodeToString(sum[:])
return pkceParams{Verifier: verifier, Challenge: challenge, Method: "S256"}, nil
return pkceToParams(p), nil
}
func newState() (string, error) {
raw := make([]byte, stateBytes)
if _, err := rand.Read(raw); err != nil {
return "", fmt.Errorf("generate state: %w", err)
}
return base64.RawURLEncoding.EncodeToString(raw), nil
return oauth.NewState()
}
// authorizationURL builds the authorization request URL.
// authorizationURL builds the authorization request URL via the shared engine.
func (flow *authorizationFlow) authorizationURL(redirectURI string) (string, error) {
parsed, err := url.Parse(flow.metadata.AuthorizationEndpoint)
if err != nil {
return "", fmt.Errorf("parse authorization endpoint: %w", err)
}
query := parsed.Query()
query.Set("response_type", "code")
query.Set("client_id", flow.config.ClientID)
query.Set("redirect_uri", redirectURI)
query.Set("state", flow.state)
query.Set("code_challenge", flow.pkce.Challenge)
query.Set("code_challenge_method", flow.pkce.Method)
if len(flow.config.Scopes) > 0 {
query.Set("scope", strings.Join(flow.config.Scopes, " "))
}
parsed.RawQuery = query.Encode()
return parsed.String(), nil
return oauth.BuildAuthorizationURL(configFor(flow.config, flow.metadata), pkceToOAuth(flow.pkce), flow.state, redirectURI, nil)
}
// parseCallback validates the redirect query and returns the authorization
@@ -209,117 +207,30 @@ func (flow *authorizationFlow) parseCallback(values url.Values) (string, error)
return code, nil
}
// exchangeCode swaps an authorization code + PKCE verifier for tokens.
// exchangeCode swaps an authorization code + PKCE verifier for tokens via the
// shared engine.
func (flow *authorizationFlow) exchangeCode(ctx context.Context, code string, redirectURI string) (StoredToken, error) {
form := url.Values{}
form.Set("grant_type", "authorization_code")
form.Set("code", code)
form.Set("redirect_uri", redirectURI)
form.Set("client_id", flow.config.ClientID)
form.Set("code_verifier", flow.pkce.Verifier)
if secret := strings.TrimSpace(flow.config.ClientSecret); secret != "" {
form.Set("client_secret", secret)
token, err := oauth.ExchangeCode(ctx, flow.httpClient, configFor(flow.config, flow.metadata), code, flow.pkce.Verifier, redirectURI, flow.now)
if err != nil {
return StoredToken{}, err
}
return postTokenRequest(ctx, flow.httpClient, flow.metadata.TokenEndpoint, form, StoredToken{Scopes: flow.config.Scopes}, flow.now)
return tokenToStored(token), nil
}
// refreshAccessToken exchanges a refresh token for a fresh access token. A
// response that omits a new refresh token preserves the previous one.
// refreshAccessToken exchanges a refresh token for a fresh access token via the
// shared engine. A response that omits a new refresh token preserves the
// previous one.
func refreshAccessToken(ctx context.Context, client *http.Client, cfg OAuthConfig, current StoredToken, now func() time.Time) (StoredToken, error) {
refresh := strings.TrimSpace(current.RefreshToken)
if refresh == "" {
return StoredToken{}, errors.New("no refresh token available")
}
tokenEndpoint := strings.TrimSpace(cfg.TokenEndpoint)
if tokenEndpoint == "" {
return StoredToken{}, errors.New("no token endpoint configured for refresh")
}
form := url.Values{}
form.Set("grant_type", "refresh_token")
form.Set("refresh_token", refresh)
form.Set("client_id", cfg.ClientID)
if secret := strings.TrimSpace(cfg.ClientSecret); secret != "" {
form.Set("client_secret", secret)
}
if len(cfg.Scopes) > 0 {
form.Set("scope", strings.Join(cfg.Scopes, " "))
}
refreshed, err := postTokenRequest(ctx, client, tokenEndpoint, form, StoredToken{Scopes: current.Scopes, RefreshToken: refresh}, now)
token, err := oauth.Refresh(ctx, client, oauth.Config{
ClientID: cfg.ClientID,
ClientSecret: cfg.ClientSecret,
Scopes: cfg.Scopes,
TokenEndpoint: cfg.TokenEndpoint,
}, oauth.Token{AccessToken: current.AccessToken, RefreshToken: current.RefreshToken, TokenType: current.TokenType, Scopes: current.Scopes, ExpiresAt: current.ExpiresAt}, now)
if err != nil {
return StoredToken{}, err
}
return refreshed, nil
}
// tokenResponse is the JSON returned by the token endpoint.
type tokenResponse struct {
AccessToken string `json:"access_token"`
RefreshToken string `json:"refresh_token"`
TokenType string `json:"token_type"`
ExpiresIn int64 `json:"expires_in"`
Scope string `json:"scope"`
Error string `json:"error"`
ErrorDescription string `json:"error_description"`
}
// postTokenRequest performs a token endpoint POST and maps the response onto a
// StoredToken. The base token supplies values to preserve (e.g. an existing
// refresh token or scope set) when the response omits them.
func postTokenRequest(ctx context.Context, client *http.Client, tokenEndpoint string, form url.Values, base StoredToken, now func() time.Time) (StoredToken, error) {
if client == nil {
client = http.DefaultClient
}
if now == nil {
now = time.Now
}
request, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenEndpoint, strings.NewReader(form.Encode()))
if err != nil {
return StoredToken{}, err
}
request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
request.Header.Set("Accept", "application/json")
response, err := client.Do(request)
if err != nil {
return StoredToken{}, fmt.Errorf("token request failed: %w", err)
}
defer response.Body.Close()
body, _ := io.ReadAll(io.LimitReader(response.Body, 1<<20))
var parsed tokenResponse
if len(body) > 0 {
// A malformed body on an error status is reported as a status error below.
_ = json.Unmarshal(body, &parsed)
}
if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
if parsed.Error != "" {
if parsed.ErrorDescription != "" {
return StoredToken{}, fmt.Errorf("token endpoint error %q: %s", parsed.Error, parsed.ErrorDescription)
}
return StoredToken{}, fmt.Errorf("token endpoint error %q", parsed.Error)
}
return StoredToken{}, fmt.Errorf("token endpoint returned HTTP %d", response.StatusCode)
}
if strings.TrimSpace(parsed.AccessToken) == "" {
return StoredToken{}, errors.New("token endpoint returned no access token")
}
token := base
token.AccessToken = parsed.AccessToken
if strings.TrimSpace(parsed.RefreshToken) != "" {
token.RefreshToken = parsed.RefreshToken
}
if strings.TrimSpace(parsed.TokenType) != "" {
token.TokenType = parsed.TokenType
}
if parsed.ExpiresIn > 0 {
token.ExpiresAt = now().Add(time.Duration(parsed.ExpiresIn) * time.Second).UTC()
}
if scope := strings.TrimSpace(parsed.Scope); scope != "" {
token.Scopes = strings.Fields(scope)
}
return token, nil
return tokenToStored(token), nil
}
// registerClient performs dynamic client registration against the registration
@@ -505,22 +416,3 @@ func Login(ctx context.Context, options LoginOptions) (StoredToken, error) {
return StoredToken{}, fmt.Errorf("timed out waiting for OAuth authorization callback: %w", loginCtx.Err())
}
}
func joinWellKnown(baseURL string) (string, error) {
parsed, err := url.Parse(strings.TrimSpace(baseURL))
if err != nil || parsed.Host == "" {
return "", fmt.Errorf("invalid base URL for metadata discovery: %q", baseURL)
}
// RFC 8414: the well-known segment is inserted between the host and any issuer
// path component, so a path-based issuer (https://host/tenant-a) is probed at
// https://host/.well-known/oauth-authorization-server/tenant-a — not at the
// host root (which would break tenant-scoped discovery).
issuerPath := strings.Trim(parsed.Path, "/")
parsed.Path = wellKnownAuthServerPath
if issuerPath != "" {
parsed.Path = strings.TrimRight(wellKnownAuthServerPath, "/") + "/" + issuerPath
}
parsed.RawQuery = ""
parsed.Fragment = ""
return parsed.String(), nil
}
+213
View File
@@ -0,0 +1,213 @@
package oauth
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"strings"
"time"
)
const deviceGrantType = "urn:ietf:params:oauth:grant-type:device_code"
// defaultDeviceCodeLifetime bounds a device authorization that omits expires_in,
// so the expiry gate in PollDeviceToken is always active (fail closed).
const defaultDeviceCodeLifetime = 10 * time.Minute
// DeviceAuth is the result of an RFC 8628 device-authorization request.
type DeviceAuth struct {
DeviceCode string
UserCode string
VerificationURI string
VerificationURIComplete string
Interval time.Duration
ExpiresAt time.Time
}
type deviceAuthResponse struct {
DeviceCode string `json:"device_code"`
UserCode string `json:"user_code"`
VerificationURI string `json:"verification_uri"`
VerificationURIComplete string `json:"verification_uri_complete"`
ExpiresIn int64 `json:"expires_in"`
Interval int64 `json:"interval"`
Error string `json:"error"`
ErrorDescription string `json:"error_description"`
}
// RequestDeviceCode performs the RFC 8628 device-authorization request. This is
// the headless/SSH login path (no browser, no loopback).
func RequestDeviceCode(ctx context.Context, client *http.Client, cfg Config, now func() time.Time) (DeviceAuth, error) {
endpoint := trimmed(cfg.DeviceAuthorizationEndpoint)
if endpoint == "" {
return DeviceAuth{}, errors.New("oauth: provider has no device authorization endpoint")
}
if err := validateTokenEndpoint(endpoint); err != nil {
return DeviceAuth{}, err
}
if client == nil {
client = http.DefaultClient
}
if now == nil {
now = time.Now
}
form := url.Values{}
form.Set("client_id", cfg.ClientID)
if secret := trimmed(cfg.ClientSecret); secret != "" {
// Confidential clients must authenticate on the device endpoint too,
// consistent with the token poll (pollDeviceOnce).
form.Set("client_secret", secret)
}
if len(cfg.Scopes) > 0 {
form.Set("scope", strings.Join(cfg.Scopes, " "))
}
request, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, strings.NewReader(form.Encode()))
if err != nil {
return DeviceAuth{}, err
}
request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
request.Header.Set("Accept", "application/json")
response, err := client.Do(request)
if err != nil {
return DeviceAuth{}, fmt.Errorf("oauth: device authorization request failed: %w", err)
}
defer response.Body.Close()
body, _ := io.ReadAll(io.LimitReader(response.Body, tokenResponseLimit))
var parsed deviceAuthResponse
if len(body) > 0 {
_ = json.Unmarshal(body, &parsed)
}
if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
if parsed.Error != "" {
return DeviceAuth{}, fmt.Errorf("oauth: device authorization error %q", parsed.Error)
}
return DeviceAuth{}, fmt.Errorf("oauth: device authorization returned HTTP %d", response.StatusCode)
}
if trimmed(parsed.DeviceCode) == "" || trimmed(parsed.UserCode) == "" {
return DeviceAuth{}, errors.New("oauth: device authorization response missing device_code/user_code")
}
interval := time.Duration(parsed.Interval) * time.Second
if interval <= 0 {
interval = 5 * time.Second // RFC 8628 default
}
auth := DeviceAuth{
DeviceCode: parsed.DeviceCode,
UserCode: parsed.UserCode,
VerificationURI: parsed.VerificationURI,
VerificationURIComplete: parsed.VerificationURIComplete,
Interval: interval,
}
lifetime := time.Duration(parsed.ExpiresIn) * time.Second
if lifetime <= 0 {
// Fail closed: a response without expires_in must still get a bounded
// lifetime so the poll loop's expiry gate stays active.
lifetime = defaultDeviceCodeLifetime
}
auth.ExpiresAt = now().Add(lifetime)
return auth, nil
}
// PollDeviceToken polls the token endpoint for the device grant until the user
// approves, the device code expires, ctx is done, or a terminal error occurs. It
// honors authorization_pending (keep waiting), slow_down (increase the
// interval), and expired_token.
func PollDeviceToken(ctx context.Context, client *http.Client, cfg Config, auth DeviceAuth, now func() time.Time) (Token, error) {
if now == nil {
now = time.Now
}
interval := auth.Interval
if interval <= 0 {
interval = 5 * time.Second
}
for {
if !auth.ExpiresAt.IsZero() && !auth.ExpiresAt.After(now()) {
return Token{}, errors.New("oauth: device code expired before authorization")
}
select {
case <-ctx.Done():
return Token{}, fmt.Errorf("oauth: device authorization canceled: %w", ctx.Err())
case <-time.After(interval):
}
// Re-check expiry after sleeping: slow_down (or a long interval) can push
// the next poll past ExpiresAt, and we must not poll after the deadline.
if !auth.ExpiresAt.IsZero() && !auth.ExpiresAt.After(now()) {
return Token{}, errors.New("oauth: device code expired before authorization")
}
token, err := pollDeviceOnce(ctx, client, cfg, auth.DeviceCode, now)
switch {
case err == nil:
return token, nil
case errors.Is(err, ErrAuthorizationPending):
// keep waiting at the current interval
case errors.Is(err, ErrSlowDown):
interval += 5 * time.Second // RFC 8628: add 5s on slow_down
default:
return Token{}, err
}
}
}
// pollDeviceOnce makes a single device-token poll, translating the RFC 8628
// error codes into ErrAuthorizationPending / ErrSlowDown / terminal errors.
func pollDeviceOnce(ctx context.Context, client *http.Client, cfg Config, deviceCode string, now func() time.Time) (Token, error) {
if err := validateTokenEndpoint(cfg.TokenEndpoint); err != nil {
return Token{}, err
}
if client == nil {
client = http.DefaultClient
}
form := url.Values{}
form.Set("grant_type", deviceGrantType)
form.Set("device_code", deviceCode)
form.Set("client_id", cfg.ClientID)
if secret := trimmed(cfg.ClientSecret); secret != "" {
form.Set("client_secret", secret)
}
request, err := http.NewRequestWithContext(ctx, http.MethodPost, trimmed(cfg.TokenEndpoint), strings.NewReader(form.Encode()))
if err != nil {
return Token{}, err
}
request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
request.Header.Set("Accept", "application/json")
response, err := client.Do(request)
if err != nil {
return Token{}, fmt.Errorf("oauth: device token poll failed: %w", err)
}
defer response.Body.Close()
body, _ := io.ReadAll(io.LimitReader(response.Body, tokenResponseLimit))
var parsed tokenResponse
if len(body) > 0 {
_ = json.Unmarshal(body, &parsed)
}
if response.StatusCode >= http.StatusOK && response.StatusCode < http.StatusMultipleChoices && trimmed(parsed.AccessToken) != "" {
token := Token{Scopes: cfg.Scopes}
token.AccessToken = parsed.AccessToken
token.RefreshToken = parsed.RefreshToken
token.TokenType = parsed.TokenType
if parsed.ExpiresIn > 0 {
token.ExpiresAt = now().Add(time.Duration(parsed.ExpiresIn) * time.Second).UTC()
}
if scope := trimmed(parsed.Scope); scope != "" {
token.Scopes = strings.Fields(scope)
}
return token, nil
}
switch parsed.Error {
case "authorization_pending":
return Token{}, ErrAuthorizationPending
case "slow_down":
return Token{}, ErrSlowDown
case "expired_token":
return Token{}, errors.New("oauth: device code expired before authorization")
case "access_denied":
return Token{}, errors.New("oauth: authorization denied by the user")
case "":
return Token{}, fmt.Errorf("oauth: device token poll returned HTTP %d", response.StatusCode)
default:
return Token{}, fmt.Errorf("oauth: device token error %q", parsed.Error)
}
}
+143
View File
@@ -0,0 +1,143 @@
package oauth
import (
"context"
"net/http"
"net/http/httptest"
"strings"
"sync/atomic"
"testing"
"time"
)
func TestRequestDeviceCodeSendsClientSecret(t *testing.T) {
var gotSecret string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_ = r.ParseForm()
gotSecret = r.Form.Get("client_secret")
_, _ = w.Write([]byte(`{"device_code":"dc","user_code":"U","verification_uri":"https://x","expires_in":600,"interval":5}`))
}))
defer server.Close()
// A confidential client must authenticate on the device endpoint too.
_, err := RequestDeviceCode(context.Background(), server.Client(), Config{ClientID: "c", ClientSecret: "shh", DeviceAuthorizationEndpoint: server.URL}, nil)
if err != nil {
t.Fatalf("RequestDeviceCode: %v", err)
}
if gotSecret != "shh" {
t.Fatalf("client_secret = %q, want shh", gotSecret)
}
}
func TestRequestDeviceCode(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
_, _ = w.Write([]byte(`{"device_code":"dc","user_code":"WXYZ-1234","verification_uri":"https://x/dev","verification_uri_complete":"https://x/dev?c=WXYZ","expires_in":600,"interval":3}`))
}))
defer server.Close()
cfg := Config{ClientID: "c", DeviceAuthorizationEndpoint: server.URL, Scopes: []string{"a"}}
auth, err := RequestDeviceCode(context.Background(), server.Client(), cfg, nil)
if err != nil {
t.Fatalf("RequestDeviceCode: %v", err)
}
if auth.DeviceCode != "dc" || auth.UserCode != "WXYZ-1234" {
t.Fatalf("device auth = %+v", auth)
}
if auth.Interval != 3*time.Second {
t.Fatalf("interval = %v, want 3s", auth.Interval)
}
if auth.VerificationURIComplete == "" {
t.Fatal("missing verification_uri_complete")
}
}
func TestRequestDeviceCodeDefaultsInterval(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
_, _ = w.Write([]byte(`{"device_code":"dc","user_code":"U","verification_uri":"https://x"}`))
}))
defer server.Close()
auth, err := RequestDeviceCode(context.Background(), server.Client(), Config{ClientID: "c", DeviceAuthorizationEndpoint: server.URL}, nil)
if err != nil {
t.Fatalf("RequestDeviceCode: %v", err)
}
if auth.Interval != 5*time.Second {
t.Fatalf("default interval = %v, want 5s", auth.Interval)
}
// A response without expires_in must still get a bounded expiry (fail closed),
// so the poll loop's expiry gate stays active.
if auth.ExpiresAt.IsZero() {
t.Fatal("missing expires_in must default to a bounded ExpiresAt, got zero")
}
}
func TestPollDeviceTokenHonorsPendingThenSucceeds(t *testing.T) {
var calls atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
n := calls.Add(1)
if n < 3 {
w.WriteHeader(http.StatusBadRequest)
_, _ = w.Write([]byte(`{"error":"authorization_pending"}`))
return
}
_, _ = w.Write([]byte(`{"access_token":"at","token_type":"Bearer","expires_in":3600}`))
}))
defer server.Close()
cfg := Config{ClientID: "c", TokenEndpoint: server.URL}
// Small interval + future expiry so the loop is fast.
auth := DeviceAuth{DeviceCode: "dc", Interval: 5 * time.Millisecond, ExpiresAt: time.Now().Add(5 * time.Second)}
tok, err := PollDeviceToken(context.Background(), server.Client(), cfg, auth, nil)
if err != nil {
t.Fatalf("PollDeviceToken: %v", err)
}
if tok.AccessToken != "at" {
t.Fatalf("token = %+v", tok)
}
if calls.Load() != 3 {
t.Fatalf("polled %d times, want 3 (2 pending + 1 success)", calls.Load())
}
}
func TestPollDeviceTokenSlowDown(t *testing.T) {
var calls atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
if calls.Add(1) == 1 {
w.WriteHeader(http.StatusBadRequest)
_, _ = w.Write([]byte(`{"error":"slow_down"}`))
return
}
_, _ = w.Write([]byte(`{"access_token":"at"}`))
}))
defer server.Close()
cfg := Config{ClientID: "c", TokenEndpoint: server.URL}
// Expiry well beyond the slow_down back-off (+5s) so the success path never
// depends on polling after the device code has expired.
auth := DeviceAuth{DeviceCode: "dc", Interval: 5 * time.Millisecond, ExpiresAt: time.Now().Add(30 * time.Second)}
tok, err := PollDeviceToken(context.Background(), server.Client(), cfg, auth, nil)
if err != nil {
t.Fatalf("PollDeviceToken (slow_down): %v", err)
}
if tok.AccessToken != "at" {
t.Fatalf("token = %+v", tok)
}
}
func TestPollDeviceTokenExpired(t *testing.T) {
cfg := Config{ClientID: "c", TokenEndpoint: "https://unused/token"}
auth := DeviceAuth{DeviceCode: "dc", Interval: time.Millisecond, ExpiresAt: time.Now().Add(-time.Second)}
_, err := PollDeviceToken(context.Background(), http.DefaultClient, cfg, auth, nil)
if err == nil {
t.Fatal("expected expired error")
}
}
func TestPollDeviceTokenAccessDenied(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusBadRequest)
_, _ = w.Write([]byte(`{"error":"access_denied"}`))
}))
defer server.Close()
cfg := Config{ClientID: "c", TokenEndpoint: server.URL}
auth := DeviceAuth{DeviceCode: "dc", Interval: 5 * time.Millisecond, ExpiresAt: time.Now().Add(5 * time.Second)}
_, err := PollDeviceToken(context.Background(), server.Client(), cfg, auth, nil)
if err == nil || !strings.Contains(err.Error(), "denied") {
t.Fatalf("err = %v, want access denied", err)
}
}
+79
View File
@@ -0,0 +1,79 @@
package oauth
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"strings"
)
const wellKnownAuthServerPath = "/.well-known/oauth-authorization-server"
// ServerMetadata is the subset of an RFC 8414 authorization-server metadata
// document the engine consumes.
type ServerMetadata struct {
Issuer string `json:"issuer"`
AuthorizationEndpoint string `json:"authorization_endpoint"`
TokenEndpoint string `json:"token_endpoint"`
DeviceAuthorizationEndpoint string `json:"device_authorization_endpoint"`
RegistrationEndpoint string `json:"registration_endpoint"`
ScopesSupported []string `json:"scopes_supported"`
}
// DiscoverAuthorizationServer fetches the RFC 8414 metadata document at the
// well-known path under baseURL.
func DiscoverAuthorizationServer(ctx context.Context, client *http.Client, baseURL string) (ServerMetadata, error) {
metadataURL, err := joinWellKnown(baseURL)
if err != nil {
return ServerMetadata{}, err
}
return fetchMetadata(ctx, client, metadataURL)
}
// fetchMetadata GETs and decodes an authorization-server metadata document at an
// already-resolved URL (no well-known recomputation), so callers can also use it
// for the OIDC openid-configuration path.
func fetchMetadata(ctx context.Context, client *http.Client, metadataURL string) (ServerMetadata, error) {
if client == nil {
client = http.DefaultClient
}
request, err := http.NewRequestWithContext(ctx, http.MethodGet, metadataURL, nil)
if err != nil {
return ServerMetadata{}, err
}
request.Header.Set("Accept", "application/json")
response, err := client.Do(request)
if err != nil {
return ServerMetadata{}, fmt.Errorf("oauth: fetch authorization server metadata: %w", err)
}
defer response.Body.Close()
if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
return ServerMetadata{}, fmt.Errorf("oauth: authorization server metadata returned HTTP %d", response.StatusCode)
}
var metadata ServerMetadata
if err := json.NewDecoder(io.LimitReader(response.Body, tokenResponseLimit)).Decode(&metadata); err != nil {
return ServerMetadata{}, fmt.Errorf("oauth: decode authorization server metadata: %w", err)
}
return metadata, nil
}
// joinWellKnown builds the RFC 8414 metadata URL, inserting the well-known
// segment between host and any issuer path component (so a path-based issuer is
// probed tenant-scoped, not at the host root).
func joinWellKnown(baseURL string) (string, error) {
parsed, err := url.Parse(trimmed(baseURL))
if err != nil || parsed.Host == "" {
return "", fmt.Errorf("oauth: invalid base URL for metadata discovery: %q", baseURL)
}
issuerPath := strings.Trim(parsed.Path, "/")
parsed.Path = wellKnownAuthServerPath
if issuerPath != "" {
parsed.Path = strings.TrimRight(wellKnownAuthServerPath, "/") + "/" + issuerPath
}
parsed.RawQuery = ""
parsed.Fragment = ""
return parsed.String(), nil
}
+212
View File
@@ -0,0 +1,212 @@
package oauth
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net"
"net/http"
"net/url"
"strings"
"time"
)
const tokenResponseLimit = 1 << 20 // 1 MiB cap on token-endpoint bodies
// tokenResponse is the JSON returned by a token endpoint.
type tokenResponse struct {
AccessToken string `json:"access_token"`
RefreshToken string `json:"refresh_token"`
TokenType string `json:"token_type"`
ExpiresIn int64 `json:"expires_in"`
Scope string `json:"scope"`
Error string `json:"error"`
ErrorDescription string `json:"error_description"`
}
// BuildAuthorizationURL constructs the authorization request URL for an
// authorization-code + PKCE flow. PKCE S256 is always included; a non-S256
// method is refused (ErrPKCEDowngrade). extraParams (and cfg.ExtraAuthParams)
// are appended last.
func BuildAuthorizationURL(cfg Config, pkce PKCE, state, redirectURI string, extraParams map[string]string) (string, error) {
if pkce.Method != MethodS256 {
return "", ErrPKCEDowngrade
}
endpoint := trimmed(cfg.AuthorizationEndpoint)
if endpoint == "" {
return "", errors.New("oauth: no authorization endpoint configured")
}
parsed, err := url.Parse(endpoint)
if err != nil {
return "", fmt.Errorf("oauth: parse authorization endpoint: %w", err)
}
query := parsed.Query()
query.Set("response_type", "code")
query.Set("client_id", cfg.ClientID)
query.Set("redirect_uri", redirectURI)
query.Set("state", state)
query.Set("code_challenge", pkce.Challenge)
query.Set("code_challenge_method", pkce.Method)
if len(cfg.Scopes) > 0 {
query.Set("scope", strings.Join(cfg.Scopes, " "))
}
// Extra params must never override the reserved OAuth/PKCE fields above
// (e.g. forcing code_challenge_method=plain or rewriting state/redirect_uri),
// which would break the flow's security guarantees.
for k, v := range cfg.ExtraAuthParams {
if isReservedAuthParam(k) {
continue
}
query.Set(k, v)
}
for k, v := range extraParams {
if isReservedAuthParam(k) {
continue
}
query.Set(k, v)
}
parsed.RawQuery = query.Encode()
return parsed.String(), nil
}
// isReservedAuthParam reports whether a query key is one the engine controls and
// must not be overridable by caller-supplied extra params.
func isReservedAuthParam(key string) bool {
switch key {
case "response_type", "client_id", "redirect_uri", "state", "code_challenge", "code_challenge_method":
return true
default:
return false
}
}
// validateTokenEndpoint refuses to send a credential to a token endpoint that is
// not https, unless it is a loopback host (mirrors oauth2-client.js
// validateTokenEndpointURL). This prevents leaking codes/tokens over cleartext.
func validateTokenEndpoint(endpoint string) error {
parsed, err := url.Parse(trimmed(endpoint))
if err != nil || parsed.Host == "" {
return fmt.Errorf("oauth: invalid token endpoint %q", endpoint)
}
if parsed.Scheme == "https" {
return nil
}
if parsed.Scheme == "http" && isLoopbackHost(parsed.Hostname()) {
return nil
}
return fmt.Errorf("%w: %s", ErrInsecureTokenEndpoint, parsed.Scheme+"://"+parsed.Host)
}
func isLoopbackHost(host string) bool {
if host == "localhost" {
return true
}
if ip := net.ParseIP(host); ip != nil {
return ip.IsLoopback()
}
return false
}
// ExchangeCode swaps an authorization code + PKCE verifier for tokens.
func ExchangeCode(ctx context.Context, client *http.Client, cfg Config, code, verifier, redirectURI string, now func() time.Time) (Token, error) {
form := url.Values{}
form.Set("grant_type", "authorization_code")
form.Set("code", code)
form.Set("redirect_uri", redirectURI)
form.Set("client_id", cfg.ClientID)
form.Set("code_verifier", verifier)
if secret := trimmed(cfg.ClientSecret); secret != "" {
form.Set("client_secret", secret)
}
return PostToken(ctx, client, cfg.TokenEndpoint, form, Token{Scopes: cfg.Scopes}, now)
}
// Refresh exchanges a refresh token for a fresh access token. A response that
// omits a new refresh token preserves the current one.
func Refresh(ctx context.Context, client *http.Client, cfg Config, current Token, now func() time.Time) (Token, error) {
refresh := trimmed(current.RefreshToken)
if refresh == "" {
return Token{}, ErrNoRefreshToken
}
if trimmed(cfg.TokenEndpoint) == "" {
return Token{}, errors.New("oauth: no token endpoint configured for refresh")
}
form := url.Values{}
form.Set("grant_type", "refresh_token")
form.Set("refresh_token", refresh)
form.Set("client_id", cfg.ClientID)
if secret := trimmed(cfg.ClientSecret); secret != "" {
form.Set("client_secret", secret)
}
if len(cfg.Scopes) > 0 {
form.Set("scope", strings.Join(cfg.Scopes, " "))
}
base := Token{Scopes: current.Scopes, RefreshToken: refresh, Account: current.Account}
return PostToken(ctx, client, cfg.TokenEndpoint, form, base, now)
}
// PostToken performs a token-endpoint POST and maps the response onto a Token.
// The https guard is applied first. The base token supplies values to preserve
// (e.g. an existing refresh token or scopes) when the response omits them.
// Error messages carry only the server's error/error_description — never the raw
// body — so token material in an unexpected payload is not leaked.
func PostToken(ctx context.Context, client *http.Client, tokenEndpoint string, form url.Values, base Token, now func() time.Time) (Token, error) {
if err := validateTokenEndpoint(tokenEndpoint); err != nil {
return Token{}, err
}
if client == nil {
client = http.DefaultClient
}
if now == nil {
now = time.Now
}
request, err := http.NewRequestWithContext(ctx, http.MethodPost, trimmed(tokenEndpoint), strings.NewReader(form.Encode()))
if err != nil {
return Token{}, err
}
request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
request.Header.Set("Accept", "application/json")
response, err := client.Do(request)
if err != nil {
return Token{}, fmt.Errorf("oauth: token request failed: %w", err)
}
defer response.Body.Close()
body, _ := io.ReadAll(io.LimitReader(response.Body, tokenResponseLimit))
var parsed tokenResponse
if len(body) > 0 {
_ = json.Unmarshal(body, &parsed)
}
if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
if parsed.Error != "" {
if parsed.ErrorDescription != "" {
return Token{}, fmt.Errorf("oauth: token endpoint error %q: %s", parsed.Error, parsed.ErrorDescription)
}
return Token{}, fmt.Errorf("oauth: token endpoint error %q", parsed.Error)
}
return Token{}, fmt.Errorf("oauth: token endpoint returned HTTP %d", response.StatusCode)
}
if trimmed(parsed.AccessToken) == "" {
return Token{}, errors.New("oauth: token endpoint returned no access token")
}
token := base
token.AccessToken = parsed.AccessToken
if trimmed(parsed.RefreshToken) != "" {
token.RefreshToken = parsed.RefreshToken
}
if trimmed(parsed.TokenType) != "" {
token.TokenType = parsed.TokenType
}
if parsed.ExpiresIn > 0 {
token.ExpiresAt = now().Add(time.Duration(parsed.ExpiresIn) * time.Second).UTC()
}
if scope := trimmed(parsed.Scope); scope != "" {
token.Scopes = strings.Fields(scope)
}
return token, nil
}
+212
View File
@@ -0,0 +1,212 @@
package oauth
import (
"context"
"crypto/sha256"
"encoding/base64"
"errors"
"net/http"
"net/http/httptest"
"net/url"
"strings"
"testing"
"time"
)
func TestNewPKCE(t *testing.T) {
p, err := NewPKCE()
if err != nil {
t.Fatalf("NewPKCE: %v", err)
}
if p.Method != MethodS256 {
t.Fatalf("method = %q, want S256", p.Method)
}
if len(p.Verifier) != 43 { // base64url(32 bytes) = 43 chars
t.Fatalf("verifier length = %d, want 43", len(p.Verifier))
}
sum := sha256.Sum256([]byte(p.Verifier))
want := base64.RawURLEncoding.EncodeToString(sum[:])
if p.Challenge != want {
t.Fatalf("challenge = %q, want base64url(sha256(verifier))", p.Challenge)
}
}
func TestNewStateUnique(t *testing.T) {
seen := map[string]bool{}
for i := 0; i < 100; i++ {
s, err := NewState()
if err != nil {
t.Fatalf("NewState: %v", err)
}
if s == "" || seen[s] {
t.Fatalf("state not unique/non-empty: %q", s)
}
seen[s] = true
}
}
func TestBuildAuthorizationURL(t *testing.T) {
cfg := Config{
ClientID: "client-123",
Scopes: []string{"read", "write"},
AuthorizationEndpoint: "https://auth.example.com/authorize",
ExtraAuthParams: map[string]string{"login_hint": "[email protected]"},
}
pkce := PKCE{Verifier: "v", Challenge: "chal", Method: MethodS256}
raw, err := BuildAuthorizationURL(cfg, pkce, "state-xyz", "http://127.0.0.1:9/callback", map[string]string{"prompt": "login"})
if err != nil {
t.Fatalf("BuildAuthorizationURL: %v", err)
}
u, _ := url.Parse(raw)
q := u.Query()
checks := map[string]string{
"response_type": "code",
"client_id": "client-123",
"redirect_uri": "http://127.0.0.1:9/callback",
"state": "state-xyz",
"code_challenge": "chal",
"code_challenge_method": "S256",
"scope": "read write",
"login_hint": "[email protected]",
"prompt": "login",
}
for k, want := range checks {
if got := q.Get(k); got != want {
t.Errorf("query %q = %q, want %q", k, got, want)
}
}
}
func TestBuildAuthorizationURLIgnoresReservedExtras(t *testing.T) {
cfg := Config{
ClientID: "c",
AuthorizationEndpoint: "https://auth.example.com/authorize",
// A hostile config tries to weaken the flow via extra params.
ExtraAuthParams: map[string]string{"state": "attacker", "code_challenge_method": "plain", "login_hint": "keep-me"},
}
pkce := PKCE{Challenge: "chal", Method: MethodS256}
raw, err := BuildAuthorizationURL(cfg, pkce, "real-state", "http://127.0.0.1:9/cb", map[string]string{"redirect_uri": "http://evil/cb"})
if err != nil {
t.Fatalf("BuildAuthorizationURL: %v", err)
}
parsed, _ := url.Parse(raw)
q := parsed.Query()
if q.Get("state") != "real-state" {
t.Fatalf("state overridden: %q", q.Get("state"))
}
if q.Get("code_challenge_method") != "S256" {
t.Fatalf("PKCE method downgraded: %q", q.Get("code_challenge_method"))
}
if q.Get("redirect_uri") != "http://127.0.0.1:9/cb" {
t.Fatalf("redirect_uri overridden: %q", q.Get("redirect_uri"))
}
if q.Get("login_hint") != "keep-me" {
t.Fatalf("non-reserved extra param should still apply: %q", q.Get("login_hint"))
}
}
func TestBuildAuthorizationURLRejectsPlain(t *testing.T) {
_, err := BuildAuthorizationURL(Config{AuthorizationEndpoint: "https://a/x"}, PKCE{Method: "plain"}, "s", "r", nil)
if !errors.Is(err, ErrPKCEDowngrade) {
t.Fatalf("err = %v, want ErrPKCEDowngrade", err)
}
}
func TestValidateTokenEndpoint(t *testing.T) {
cases := map[string]bool{
"https://auth.example.com/token": true,
"http://127.0.0.1:1455/token": true, // loopback exempt
"http://localhost:1455/token": true,
"http://[::1]:1455/token": true,
"http://auth.example.com/token": false, // non-https, non-loopback
"ftp://auth.example.com/token": false,
}
for endpoint, ok := range cases {
err := validateTokenEndpoint(endpoint)
if ok && err != nil {
t.Errorf("validateTokenEndpoint(%q) = %v, want nil", endpoint, err)
}
if !ok && err == nil {
t.Errorf("validateTokenEndpoint(%q) = nil, want error", endpoint)
}
}
}
func TestExchangeCodeSuccess(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_ = r.ParseForm()
if r.FormValue("grant_type") != "authorization_code" || r.FormValue("code_verifier") != "verifier" {
w.WriteHeader(http.StatusBadRequest)
return
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"access_token":"at","refresh_token":"rt","token_type":"Bearer","expires_in":3600,"scope":"read write"}`))
}))
defer server.Close()
now := time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC)
cfg := Config{ClientID: "c", TokenEndpoint: server.URL}
tok, err := ExchangeCode(context.Background(), server.Client(), cfg, "code", "verifier", "http://127.0.0.1/cb", func() time.Time { return now })
if err != nil {
t.Fatalf("ExchangeCode: %v", err)
}
if tok.AccessToken != "at" || tok.RefreshToken != "rt" || tok.TokenType != "Bearer" {
t.Fatalf("token = %+v", tok)
}
if !tok.ExpiresAt.Equal(now.Add(time.Hour)) {
t.Fatalf("expiresAt = %v, want now+1h", tok.ExpiresAt)
}
if strings.Join(tok.Scopes, " ") != "read write" {
t.Fatalf("scopes = %v", tok.Scopes)
}
}
func TestExchangeCodeErrorRedactsBody(t *testing.T) {
// The error body contains a stray access_token-looking field; it must NOT
// surface in the returned error (only error/error_description do).
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusBadRequest)
_, _ = w.Write([]byte(`{"error":"invalid_grant","error_description":"bad code","access_token":"SECRET-LEAK"}`))
}))
defer server.Close()
cfg := Config{ClientID: "c", TokenEndpoint: server.URL}
_, err := ExchangeCode(context.Background(), server.Client(), cfg, "code", "v", "http://127.0.0.1/cb", nil)
if err == nil {
t.Fatal("expected error")
}
if !strings.Contains(err.Error(), "invalid_grant") || !strings.Contains(err.Error(), "bad code") {
t.Fatalf("error should carry error/description, got %q", err.Error())
}
if strings.Contains(err.Error(), "SECRET-LEAK") {
t.Fatalf("error leaked raw body token material: %q", err.Error())
}
}
func TestPostTokenRefusesInsecureEndpoint(t *testing.T) {
_, err := PostToken(context.Background(), http.DefaultClient, "http://auth.example.com/token", url.Values{}, Token{}, nil)
if !errors.Is(err, ErrInsecureTokenEndpoint) {
t.Fatalf("err = %v, want ErrInsecureTokenEndpoint", err)
}
}
func TestRefreshNoToken(t *testing.T) {
_, err := Refresh(context.Background(), http.DefaultClient, Config{TokenEndpoint: "https://a/token"}, Token{}, nil)
if !errors.Is(err, ErrNoRefreshToken) {
t.Fatalf("err = %v, want ErrNoRefreshToken", err)
}
}
func TestRefreshPreservesRefreshTokenWhenOmitted(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
_, _ = w.Write([]byte(`{"access_token":"new-at","expires_in":3600}`)) // no refresh_token in response
}))
defer server.Close()
cfg := Config{ClientID: "c", TokenEndpoint: server.URL}
tok, err := Refresh(context.Background(), server.Client(), cfg, Token{RefreshToken: "keep-me"}, nil)
if err != nil {
t.Fatalf("Refresh: %v", err)
}
if tok.AccessToken != "new-at" || tok.RefreshToken != "keep-me" {
t.Fatalf("refresh should preserve old refresh token: %+v", tok)
}
}
+74
View File
@@ -0,0 +1,74 @@
package oauth
import (
"errors"
"fmt"
"os"
"path/filepath"
"sync/atomic"
"time"
)
const (
fileLockTimeout = 5 * time.Second
fileLockStaleAfter = 30 * time.Second
)
var lockSeq atomic.Uint64
// acquireFileLock takes a cross-process exclusive lock by creating lockPath with
// O_EXCL. It retries with a short backoff until a timeout, breaking a lock whose
// file is older than fileLockStaleAfter (so a crashed holder cannot deadlock the
// store). Release is ownership-aware: it removes the lock only if it still holds
// our token, so a stale-broken holder cannot delete a newer holder's lock.
func acquireFileLock(lockPath string, now func() time.Time) (func(), error) {
if now == nil {
now = time.Now
}
if err := os.MkdirAll(filepath.Dir(lockPath), 0o700); err != nil {
return nil, err
}
token := fmt.Sprintf("%d-%d-%d", os.Getpid(), now().UnixNano(), lockSeq.Add(1))
deadline := now().Add(fileLockTimeout)
for {
f, err := os.OpenFile(lockPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
if err == nil {
// A partial/failed write would leave a lock file without our token, so
// the releaser below could never delete it — stranding the lock for
// other processes. Fail closed: remove the file and surface the error.
if _, werr := f.WriteString(token); werr != nil {
_ = f.Close()
_ = os.Remove(lockPath)
return nil, fmt.Errorf("oauth: write token lock: %w", werr)
}
if cerr := f.Close(); cerr != nil {
_ = os.Remove(lockPath)
return nil, fmt.Errorf("oauth: close token lock: %w", cerr)
}
var released bool
return func() {
if released {
return
}
released = true
if data, rerr := os.ReadFile(lockPath); rerr == nil && string(data) == token {
_ = os.Remove(lockPath)
}
}, nil
}
if !errors.Is(err, os.ErrExist) {
return nil, fmt.Errorf("oauth: acquire token lock: %w", err)
}
if info, statErr := os.Stat(lockPath); statErr == nil && time.Since(info.ModTime()) > fileLockStaleAfter {
stale, _ := os.ReadFile(lockPath)
if data, rerr := os.ReadFile(lockPath); rerr == nil && string(data) == string(stale) {
_ = os.Remove(lockPath)
}
continue
}
if now().After(deadline) {
return nil, fmt.Errorf("oauth: timed out acquiring token lock %s", filepath.Base(lockPath))
}
time.Sleep(10 * time.Millisecond)
}
}
+112
View File
@@ -0,0 +1,112 @@
package oauth
import (
"context"
"errors"
"fmt"
"io"
"net"
"net/http"
"net/url"
"strings"
"time"
)
// LoopbackListener is a single-use loopback HTTP server that captures an OAuth
// redirect (?code=&state=) on 127.0.0.1 only: bind an OS-assigned loopback port,
// verify the CSRF state, hand back the code, then close.
type LoopbackListener struct {
listener net.Listener
state string
result chan callbackResult
server *http.Server
}
type callbackResult struct {
code string
err error
}
// NewLoopbackListener binds 127.0.0.1:0 (loopback only, OS-assigned port) and
// begins serving. state is the CSRF value the callback must echo back. Call
// RedirectURI to build the redirect_uri, then Wait for the code. Always Close.
func NewLoopbackListener(state string) (*LoopbackListener, error) {
// Refuse an empty state: the callback check compares the query's state to this
// value, so an empty state would match a callback carrying no state at all,
// defeating CSRF protection (fail closed).
if strings.TrimSpace(state) == "" {
return nil, errors.New("oauth: loopback listener requires a non-empty CSRF state")
}
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
return nil, fmt.Errorf("oauth: start loopback redirect listener: %w", err)
}
l := &LoopbackListener{
listener: ln,
state: state,
result: make(chan callbackResult, 1),
}
l.server = &http.Server{Handler: http.HandlerFunc(l.handle)}
go func() { _ = l.server.Serve(ln) }()
return l, nil
}
// RedirectURI returns the http://127.0.0.1:<port>/callback redirect URI.
func (l *LoopbackListener) RedirectURI() string {
return fmt.Sprintf("http://%s/callback", l.listener.Addr().String())
}
func (l *LoopbackListener) handle(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/callback" {
http.NotFound(w, r)
return
}
code, err := parseCallback(r.URL.Query(), l.state)
if err != nil {
w.WriteHeader(http.StatusBadRequest)
_, _ = io.WriteString(w, "Authorization failed. You may close this window.")
} else {
_, _ = io.WriteString(w, "Authorization complete. You may close this window.")
}
select {
case l.result <- callbackResult{code: code, err: err}:
default:
}
}
// Wait blocks until the callback arrives or ctx is done, returning the
// authorization code.
func (l *LoopbackListener) Wait(ctx context.Context) (string, error) {
select {
case res := <-l.result:
return res.code, res.err
case <-ctx.Done():
return "", fmt.Errorf("oauth: timed out waiting for authorization callback: %w", ctx.Err())
}
}
// Close shuts the listener down (bounded), idempotent.
func (l *LoopbackListener) Close() {
shutdownCtx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
_ = l.server.Shutdown(shutdownCtx)
}
// parseCallback validates the redirect query and returns the authorization code,
// rejecting a mismatched state (CSRF) and surfacing provider errors.
func parseCallback(values url.Values, wantState string) (string, error) {
if got := values.Get("state"); got != wantState {
return "", ErrStateMismatch
}
if providerErr := strings.TrimSpace(values.Get("error")); providerErr != "" {
if desc := strings.TrimSpace(values.Get("error_description")); desc != "" {
return "", fmt.Errorf("oauth: authorization server returned error %q: %s", providerErr, desc)
}
return "", fmt.Errorf("oauth: authorization server returned error %q", providerErr)
}
code := strings.TrimSpace(values.Get("code"))
if code == "" {
return "", errors.New("oauth: callback missing authorization code")
}
return code, nil
}
+86
View File
@@ -0,0 +1,86 @@
package oauth
import (
"context"
"errors"
"net/http"
"strings"
"testing"
"time"
)
func TestLoopbackCapturesCode(t *testing.T) {
l, err := NewLoopbackListener("the-state")
if err != nil {
t.Fatalf("NewLoopbackListener: %v", err)
}
defer l.Close()
if !strings.HasPrefix(l.RedirectURI(), "http://127.0.0.1:") {
t.Fatalf("redirect URI not loopback: %q", l.RedirectURI())
}
go func() {
_, _ = http.Get(l.RedirectURI() + "?code=auth-code&state=the-state")
}()
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
defer cancel()
code, err := l.Wait(ctx)
if err != nil {
t.Fatalf("Wait: %v", err)
}
if code != "auth-code" {
t.Fatalf("code = %q, want auth-code", code)
}
}
func TestNewLoopbackListenerRejectsEmptyState(t *testing.T) {
if _, err := NewLoopbackListener(" "); err == nil {
t.Fatal("an empty CSRF state must be rejected (fail closed)")
}
}
func TestLoopbackRejectsStateMismatch(t *testing.T) {
l, err := NewLoopbackListener("expected-state")
if err != nil {
t.Fatalf("NewLoopbackListener: %v", err)
}
defer l.Close()
go func() {
_, _ = http.Get(l.RedirectURI() + "?code=x&state=WRONG")
}()
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
defer cancel()
_, err = l.Wait(ctx)
if !errors.Is(err, ErrStateMismatch) {
t.Fatalf("Wait err = %v, want ErrStateMismatch", err)
}
}
func TestLoopbackTimesOut(t *testing.T) {
l, err := NewLoopbackListener("s")
if err != nil {
t.Fatalf("NewLoopbackListener: %v", err)
}
defer l.Close()
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
defer cancel()
if _, err := l.Wait(ctx); err == nil {
t.Fatal("Wait should time out cleanly")
}
}
func TestLoopbackProviderError(t *testing.T) {
l, err := NewLoopbackListener("s")
if err != nil {
t.Fatalf("NewLoopbackListener: %v", err)
}
defer l.Close()
go func() {
_, _ = http.Get(l.RedirectURI() + "?error=access_denied&error_description=nope&state=s")
}()
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
defer cancel()
_, err = l.Wait(ctx)
if err == nil || !strings.Contains(err.Error(), "access_denied") {
t.Fatalf("Wait err = %v, want provider error", err)
}
}
+292
View File
@@ -0,0 +1,292 @@
package oauth
import (
"context"
"fmt"
"io"
"net/http"
"strings"
"time"
)
// defaultRefreshBuffer refreshes a token this long before its hard expiry.
const defaultRefreshBuffer = 60 * time.Second
const oidcWellKnownPath = "/.well-known/openid-configuration"
// Manager ties the token store, provider registry, and HTTP client together to
// run logins and serve fresh access tokens. It is the high-level entrypoint the
// CLI and request paths use.
type Manager struct {
store *Store
registry *Registry
client *http.Client
env map[string]string
now func() time.Time
buffer time.Duration
out io.Writer
// openBrowser is invoked with the authorization URL for loopback logins.
// Tests inject a function that drives the loopback redirect.
openBrowser func(authURL string) error
}
// ManagerOptions configures a Manager.
type ManagerOptions struct {
Store *Store
Registry *Registry
HTTPClient *http.Client
Env map[string]string
Now func() time.Time
RefreshBuffer time.Duration
Out io.Writer
OpenBrowser func(authURL string) error
}
// NewManager builds a Manager, filling defaults.
func NewManager(opts ManagerOptions) (*Manager, error) {
if opts.Store == nil {
return nil, fmt.Errorf("oauth: manager requires a store")
}
registry := opts.Registry
if registry == nil {
registry = NewRegistry()
}
client := opts.HTTPClient
if client == nil {
client = http.DefaultClient
}
now := opts.Now
if now == nil {
now = time.Now
}
buffer := opts.RefreshBuffer
if buffer <= 0 {
buffer = defaultRefreshBuffer
}
out := opts.Out
if out == nil {
out = io.Discard
}
open := opts.OpenBrowser
if open == nil {
open = func(string) error { return nil }
}
return &Manager{
store: opts.Store, registry: registry, client: client,
env: opts.Env, now: now, buffer: buffer, out: out, openBrowser: open,
}, nil
}
// LoginOptions configures a single provider login.
type LoginOptions struct {
Provider string
Device bool // force device-code flow
ExtraScopes []string // appended to the provider's scopes
Timeout time.Duration // bounds the whole interactive login
}
// Login runs the provider login (loopback by default, device-code when
// requested or when the provider only supports device), stores the token under
// "provider:<name>", and returns a redaction-safe status.
func (m *Manager) Login(ctx context.Context, opts LoginOptions) (Status, error) {
cfg, flow, err := m.registry.ResolveConfig(opts.Provider, m.env)
if err != nil {
return Status{}, err
}
if len(opts.ExtraScopes) > 0 {
cfg.Scopes = append(append([]string{}, cfg.Scopes...), opts.ExtraScopes...)
}
timeout := opts.Timeout
if timeout <= 0 {
timeout = 5 * time.Minute
}
loginCtx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
cfg, err = m.resolveEndpoints(loginCtx, cfg)
if err != nil {
return Status{}, err
}
useDevice := opts.Device || flow == FlowDevice
var token Token
if useDevice {
token, err = m.loginDevice(loginCtx, cfg)
} else {
token, err = m.loginLoopback(loginCtx, cfg)
}
if err != nil {
return Status{}, err
}
key := ProviderKey(opts.Provider)
if err := m.store.Save(key, token); err != nil {
return Status{}, err
}
return m.statusFor(key)
}
// resolveEndpoints fills missing authorize/token/device endpoints from issuer
// discovery (RFC 8414 then OIDC), leaving any explicitly-pinned endpoint intact.
func (m *Manager) resolveEndpoints(ctx context.Context, cfg Config) (Config, error) {
if trimmed(cfg.IssuerURL) == "" {
return cfg, nil
}
if cfg.AuthorizationEndpoint != "" && cfg.TokenEndpoint != "" && cfg.DeviceAuthorizationEndpoint != "" {
return cfg, nil
}
// Discovery is best-effort: a failure is non-fatal because pinned endpoints
// may already be sufficient for the chosen flow. Only merge on success (and
// never return from the error branch, which would trip the nilerr linter).
if meta, err := m.discover(ctx, cfg.IssuerURL); err == nil {
if cfg.AuthorizationEndpoint == "" {
cfg.AuthorizationEndpoint = meta.AuthorizationEndpoint
}
if cfg.TokenEndpoint == "" {
cfg.TokenEndpoint = meta.TokenEndpoint
}
if cfg.DeviceAuthorizationEndpoint == "" {
cfg.DeviceAuthorizationEndpoint = meta.DeviceAuthorizationEndpoint
}
}
return cfg, nil
}
// discover tries the OAuth (RFC 8414) well-known path, then the OIDC
// openid-configuration path (some issuers publish only the latter).
func (m *Manager) discover(ctx context.Context, issuer string) (ServerMetadata, error) {
meta, err := DiscoverAuthorizationServer(ctx, m.client, issuer)
if err == nil && (meta.AuthorizationEndpoint != "" || meta.TokenEndpoint != "") {
return meta, nil
}
oidcURL := strings.TrimRight(strings.TrimSpace(issuer), "/") + oidcWellKnownPath
return fetchMetadata(ctx, m.client, oidcURL)
}
func (m *Manager) loginLoopback(ctx context.Context, cfg Config) (Token, error) {
if trimmed(cfg.AuthorizationEndpoint) == "" {
return Token{}, fmt.Errorf("oauth: no authorization endpoint (set the authorize URL or a discoverable issuer)")
}
state, err := NewState()
if err != nil {
return Token{}, err
}
pkce, err := NewPKCE()
if err != nil {
return Token{}, err
}
listener, err := NewLoopbackListener(state)
if err != nil {
return Token{}, err
}
defer listener.Close()
redirectURI := listener.RedirectURI()
authURL, err := BuildAuthorizationURL(cfg, pkce, state, redirectURI, nil)
if err != nil {
return Token{}, err
}
fmt.Fprintf(m.out, "Open this URL to authorize:\n %s\n", authURL)
if err := m.openBrowser(authURL); err != nil {
return Token{}, fmt.Errorf("oauth: open authorization URL: %w", err)
}
code, err := listener.Wait(ctx)
if err != nil {
return Token{}, err
}
return ExchangeCode(ctx, m.client, cfg, code, pkce.Verifier, redirectURI, m.now)
}
func (m *Manager) loginDevice(ctx context.Context, cfg Config) (Token, error) {
auth, err := RequestDeviceCode(ctx, m.client, cfg, m.now)
if err != nil {
return Token{}, err
}
target := auth.VerificationURIComplete
if target == "" {
target = auth.VerificationURI
}
// user_code is meant to be displayed to the user; it is not a secret.
fmt.Fprintf(m.out, "To authorize, visit:\n %s\nand enter code: %s\n", target, auth.UserCode)
return PollDeviceToken(ctx, m.client, cfg, auth, m.now)
}
// GetFresh returns a valid access token for key, refreshing on-demand if the
// stored token is expired or within the refresh buffer. Mirrors
// checkAndRefreshOAuthTokenIfNeeded.
func (m *Manager) GetFresh(ctx context.Context, key string) (string, error) {
token, cfg, err := m.loadForKey(key)
if err != nil {
return "", err
}
if !token.NeedsRefresh(m.now(), m.buffer) {
return token.AccessToken, nil
}
return m.refreshAndSave(ctx, key, cfg, token)
}
// Handle401 forces a refresh after an upstream 401, returning the new access
// token. Mirrors handleOAuth401Error.
func (m *Manager) Handle401(ctx context.Context, key string) (string, error) {
token, cfg, err := m.loadForKey(key)
if err != nil {
return "", err
}
return m.refreshAndSave(ctx, key, cfg, token)
}
func (m *Manager) refreshAndSave(ctx context.Context, key string, cfg Config, current Token) (string, error) {
refreshed, err := Refresh(ctx, m.client, cfg, current, m.now)
if err != nil {
return "", err
}
if err := m.store.Save(key, refreshed); err != nil {
return "", err
}
return refreshed.AccessToken, nil
}
// loadForKey loads a stored token and resolves the provider config for its key.
func (m *Manager) loadForKey(key string) (Token, Config, error) {
if err := ValidateKey(key); err != nil {
return Token{}, Config{}, err
}
token, ok, err := m.store.Load(key)
if err != nil {
return Token{}, Config{}, err
}
if !ok {
return Token{}, Config{}, fmt.Errorf("oauth: no stored token for %q", key)
}
name := strings.TrimPrefix(key, KeyPrefixProvider)
if name == key {
return Token{}, Config{}, fmt.Errorf("oauth: refresh is only supported for provider tokens (got %q)", key)
}
cfg, _, err := m.registry.ResolveConfig(name, m.env)
if err != nil {
return Token{}, Config{}, err
}
return token, cfg, nil
}
// Logout removes a provider's stored token, reporting whether one was present.
func (m *Manager) Logout(name string) (bool, error) {
return m.store.Delete(ProviderKey(name))
}
// StatusAll returns the status of every provider login.
func (m *Manager) StatusAll() ([]Status, error) {
return m.store.Status(KeyPrefixProvider)
}
func (m *Manager) statusFor(key string) (Status, error) {
statuses, err := m.store.Status(KeyPrefixProvider)
if err != nil {
return Status{}, err
}
for _, st := range statuses {
if st.Key == key {
return st, nil
}
}
return Status{Key: key}, nil
}
+182
View File
@@ -0,0 +1,182 @@
package oauth
import (
"context"
"io"
"net/http"
"net/http/httptest"
"net/url"
"path/filepath"
"sync/atomic"
"testing"
"time"
)
// fakeProvider is an httptest server that plays a token + device endpoint.
type fakeProvider struct {
server *httptest.Server
tokenHits atomic.Int32
}
func newFakeProvider(t *testing.T, tokenJSON string) *fakeProvider {
t.Helper()
fp := &fakeProvider{}
mux := http.NewServeMux()
mux.HandleFunc("/token", func(w http.ResponseWriter, _ *http.Request) {
fp.tokenHits.Add(1)
_, _ = io.WriteString(w, tokenJSON)
})
mux.HandleFunc("/device", func(w http.ResponseWriter, _ *http.Request) {
// Approve immediately on poll; short interval so the test is fast.
_, _ = io.WriteString(w, `{"device_code":"dc","user_code":"U-1","verification_uri":"https://example/dev","expires_in":600,"interval":1}`)
})
fp.server = httptest.NewServer(mux)
t.Cleanup(fp.server.Close)
return fp
}
func managerFor(t *testing.T, env map[string]string, openBrowser func(string) error) *Manager {
t.Helper()
store, err := NewStore(StoreOptions{FilePath: filepath.Join(t.TempDir(), "tok.json")})
if err != nil {
t.Fatalf("NewStore: %v", err)
}
m, err := NewManager(ManagerOptions{
Store: store,
Env: env,
OpenBrowser: openBrowser,
})
if err != nil {
t.Fatalf("NewManager: %v", err)
}
return m
}
func TestManagerLoginLoopback(t *testing.T) {
fp := newFakeProvider(t, `{"access_token":"at","refresh_token":"rt","token_type":"Bearer","expires_in":3600}`)
env := map[string]string{
"ZERO_OAUTH_DEMO_CLIENT_ID": "client",
"ZERO_OAUTH_DEMO_AUTHORIZE_URL": "https://auth.example.com/authorize",
"ZERO_OAUTH_DEMO_TOKEN_URL": fp.server.URL + "/token",
}
// The fake browser drives the loopback redirect with the captured state.
openBrowser := func(authURL string) error {
u, err := url.Parse(authURL)
if err != nil {
return err
}
redirect := u.Query().Get("redirect_uri")
state := u.Query().Get("state")
_, err = http.Get(redirect + "?code=the-code&state=" + url.QueryEscape(state))
return err
}
m := managerFor(t, env, openBrowser)
status, err := m.Login(context.Background(), LoginOptions{Provider: "demo"})
if err != nil {
t.Fatalf("Login: %v", err)
}
if !status.HasToken || status.Key != ProviderKey("demo") {
t.Fatalf("status = %+v", status)
}
// The token is persisted and GetFresh returns it without a needless refresh.
access, err := m.GetFresh(context.Background(), ProviderKey("demo"))
if err != nil {
t.Fatalf("GetFresh: %v", err)
}
if access != "at" {
t.Fatalf("access = %q, want at", access)
}
}
func TestManagerLoginDevice(t *testing.T) {
fp := newFakeProvider(t, `{"access_token":"dev-at","token_type":"Bearer","expires_in":3600}`)
env := map[string]string{
"ZERO_OAUTH_DEMODEV_CLIENT_ID": "client",
"ZERO_OAUTH_DEMODEV_TOKEN_URL": fp.server.URL + "/token",
"ZERO_OAUTH_DEMODEV_DEVICE_URL": fp.server.URL + "/device",
}
m := managerFor(t, env, nil)
status, err := m.Login(context.Background(), LoginOptions{Provider: "demodev", Device: true})
if err != nil {
t.Fatalf("device Login: %v", err)
}
if !status.HasToken {
t.Fatalf("device status = %+v", status)
}
}
func TestManagerGetFreshRefreshesExpired(t *testing.T) {
fp := newFakeProvider(t, `{"access_token":"fresh-at","expires_in":3600}`)
env := map[string]string{
"ZERO_OAUTH_DEMO_CLIENT_ID": "client",
"ZERO_OAUTH_DEMO_TOKEN_URL": fp.server.URL + "/token",
}
m := managerFor(t, env, nil)
// Seed an expired token with a refresh token.
if err := m.store.Save(ProviderKey("demo"), Token{AccessToken: "stale", RefreshToken: "rt", ExpiresAt: time.Now().Add(-time.Hour)}); err != nil {
t.Fatalf("seed: %v", err)
}
access, err := m.GetFresh(context.Background(), ProviderKey("demo"))
if err != nil {
t.Fatalf("GetFresh: %v", err)
}
if access != "fresh-at" {
t.Fatalf("access = %q, want fresh-at (refreshed)", access)
}
if fp.tokenHits.Load() != 1 {
t.Fatalf("token endpoint hit %d times, want 1 refresh", fp.tokenHits.Load())
}
// The refreshed token is persisted.
stored, _, _ := m.store.Load(ProviderKey("demo"))
if stored.AccessToken != "fresh-at" {
t.Fatalf("stored token not updated: %+v", stored)
}
}
func TestManagerGetFreshSkipsValidToken(t *testing.T) {
fp := newFakeProvider(t, `{"access_token":"should-not-be-used"}`)
env := map[string]string{
"ZERO_OAUTH_DEMO_CLIENT_ID": "client",
"ZERO_OAUTH_DEMO_TOKEN_URL": fp.server.URL + "/token",
}
m := managerFor(t, env, nil)
_ = m.store.Save(ProviderKey("demo"), Token{AccessToken: "valid", RefreshToken: "rt", ExpiresAt: time.Now().Add(time.Hour)})
access, err := m.GetFresh(context.Background(), ProviderKey("demo"))
if err != nil {
t.Fatalf("GetFresh: %v", err)
}
if access != "valid" || fp.tokenHits.Load() != 0 {
t.Fatalf("a valid token must not be refreshed (access=%q hits=%d)", access, fp.tokenHits.Load())
}
}
func TestManagerHandle401ForcesRefresh(t *testing.T) {
fp := newFakeProvider(t, `{"access_token":"after-401"}`)
env := map[string]string{
"ZERO_OAUTH_DEMO_CLIENT_ID": "client",
"ZERO_OAUTH_DEMO_TOKEN_URL": fp.server.URL + "/token",
}
m := managerFor(t, env, nil)
// Token is still valid by clock, but Handle401 forces a refresh anyway.
_ = m.store.Save(ProviderKey("demo"), Token{AccessToken: "valid", RefreshToken: "rt", ExpiresAt: time.Now().Add(time.Hour)})
access, err := m.Handle401(context.Background(), ProviderKey("demo"))
if err != nil {
t.Fatalf("Handle401: %v", err)
}
if access != "after-401" || fp.tokenHits.Load() != 1 {
t.Fatalf("Handle401 must force a refresh (access=%q hits=%d)", access, fp.tokenHits.Load())
}
}
func TestManagerLogout(t *testing.T) {
m := managerFor(t, map[string]string{"ZERO_OAUTH_DEMO_CLIENT_ID": "c"}, nil)
_ = m.store.Save(ProviderKey("demo"), Token{AccessToken: "a"})
removed, err := m.Logout("demo")
if err != nil || !removed {
t.Fatalf("Logout = %v %v", removed, err)
}
if removed2, _ := m.Logout("demo"); removed2 {
t.Fatal("second logout should report nothing removed")
}
}
+95
View File
@@ -0,0 +1,95 @@
// Package oauth is zero's reusable OAuth 2.0 engine. It generalizes the
// authorization-code + PKCE flow that internal/mcp already uses for MCP servers
// into a transport/identity-agnostic core, and adds the pieces provider login
// needs: a device-authorization (RFC 8628) grant, a provider registry, a
// namespaced + refreshing token store, and a proactive refresh scheduler.
//
// Security invariants (enforced throughout, never relaxed):
// - PKCE S256 is mandatory on every authorization-code flow; "plain" is refused.
// - A per-flow CSRF state is generated and verified on the callback.
// - The loopback callback server binds 127.0.0.1 only, on an OS-assigned port,
// serves a single request, then closes.
// - The token endpoint must be https (loopback exempt) before any credential is
// sent — see validateTokenEndpoint.
// - Tokens, codes, and verifiers are never logged; error bodies are redacted.
// - Token files are 0600, base-confined, and file-locked.
//
// This package holds NO vendor secrets: provider client IDs and any overridable
// endpoints come from config/env, so no third-party OAuth client identity is
// baked into the binary.
package oauth
import (
"errors"
"strings"
"time"
)
// Token holds the credentials issued by an authorization server. The token
// fields are sensitive: callers must never log them, and the store persists them
// 0600. It mirrors the on-disk shape internal/mcp uses so the two stay
// format-compatible.
type Token struct {
AccessToken string `json:"access_token"`
RefreshToken string `json:"refresh_token,omitempty"`
TokenType string `json:"token_type,omitempty"`
Scopes []string `json:"scopes,omitempty"`
ExpiresAt time.Time `json:"expires_at,omitempty"`
// Account is an optional non-secret identifier (email / account id) shown in
// status output; never a credential.
Account string `json:"account,omitempty"`
}
// Expired reports whether the token has an expiry that is at or before now.
// A zero ExpiresAt means "no known expiry" and is treated as not expired.
func (t Token) Expired(now time.Time) bool {
return !t.ExpiresAt.IsZero() && !t.ExpiresAt.After(now)
}
// NeedsRefresh reports whether the token is expired or falls within buffer of
// expiry (so a proactive/on-demand refresh should run). A token with no expiry
// never needs a refresh on a timer.
func (t Token) NeedsRefresh(now time.Time, buffer time.Duration) bool {
if t.ExpiresAt.IsZero() {
return false
}
return !t.ExpiresAt.After(now.Add(buffer))
}
// Config describes how to talk to one authorization server for a flow.
// Endpoints may be discovered (RFC 8414) and overridden here.
type Config struct {
ClientID string
ClientSecret string
Scopes []string
AuthorizationEndpoint string
TokenEndpoint string
DeviceAuthorizationEndpoint string
RegistrationEndpoint string
// IssuerURL is the base for metadata discovery when endpoints are not set.
IssuerURL string
// ExtraAuthParams are appended to the authorization URL (e.g. login_hint).
ExtraAuthParams map[string]string
}
// Errors returned by the engine. Callers can match these with errors.Is.
var (
// ErrPKCEDowngrade is returned if a flow is asked to use anything but S256.
ErrPKCEDowngrade = errors.New("oauth: PKCE S256 is mandatory; plain is refused")
// ErrStateMismatch is returned when the callback state does not match (CSRF).
ErrStateMismatch = errors.New("oauth: callback state mismatch; possible CSRF, login aborted")
// ErrInsecureTokenEndpoint is returned when a credential would be sent over a
// non-https, non-loopback endpoint.
ErrInsecureTokenEndpoint = errors.New("oauth: refusing to send credential to a non-https token endpoint")
// ErrNoRefreshToken is returned when a refresh is attempted without one.
ErrNoRefreshToken = errors.New("oauth: no refresh token available")
// ErrAuthorizationPending is the RFC 8628 "keep polling" signal.
ErrAuthorizationPending = errors.New("oauth: authorization pending")
// ErrSlowDown is the RFC 8628 "increase the poll interval" signal.
ErrSlowDown = errors.New("oauth: slow down")
)
// trimmed is a tiny helper used across the package.
func trimmed(s string) string { return strings.TrimSpace(s) }
+49
View File
@@ -0,0 +1,49 @@
package oauth
import (
"crypto/rand"
"crypto/sha256"
"encoding/base64"
"fmt"
)
const (
// pkceVerifierBytes yields a 43-character base64url verifier — the minimum
// length the PKCE spec permits while remaining high-entropy.
pkceVerifierBytes = 32
stateBytes = 32
// MethodS256 is the only code-challenge method this engine accepts.
MethodS256 = "S256"
)
// PKCE is a verifier/challenge pair for an authorization-code flow. The method
// is always S256 (the engine refuses "plain").
type PKCE struct {
Verifier string
Challenge string
Method string
}
// NewPKCE generates a high-entropy code verifier and its S256 challenge. It
// matches internal/mcp's newPKCE byte-for-byte (RawURLEncoding of 32 random
// bytes; challenge = base64url(SHA-256(verifier))) so the two engines are
// interchangeable.
func NewPKCE() (PKCE, error) {
raw := make([]byte, pkceVerifierBytes)
if _, err := rand.Read(raw); err != nil {
return PKCE{}, fmt.Errorf("oauth: generate PKCE verifier: %w", err)
}
verifier := base64.RawURLEncoding.EncodeToString(raw)
sum := sha256.Sum256([]byte(verifier))
challenge := base64.RawURLEncoding.EncodeToString(sum[:])
return PKCE{Verifier: verifier, Challenge: challenge, Method: MethodS256}, nil
}
// NewState generates a high-entropy CSRF state value.
func NewState() (string, error) {
raw := make([]byte, stateBytes)
if _, err := rand.Read(raw); err != nil {
return "", fmt.Errorf("oauth: generate state: %w", err)
}
return base64.RawURLEncoding.EncodeToString(raw), nil
}
+100
View File
@@ -0,0 +1,100 @@
package oauth
import (
"fmt"
"regexp"
"strings"
)
// Flow selects how a provider delivers the authorization result.
type Flow string
const (
// FlowLoopback uses a 127.0.0.1 callback server (browser required).
FlowLoopback Flow = "loopback"
// FlowDevice uses the RFC 8628 device-code flow (headless/SSH).
FlowDevice Flow = "device"
)
// providerNamePattern bounds a provider name to a safe identifier that is also a
// valid store-key segment.
var providerNamePattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$`)
// ValidateProviderName reports whether name is a safe provider identifier.
func ValidateProviderName(name string) error {
if !providerNamePattern.MatchString(name) {
return fmt.Errorf("oauth: invalid provider name %q", name)
}
return nil
}
// Registry resolves a provider's Config from env/config. It ships NO built-in
// providers, endpoints, or client identities: every provider is defined entirely
// by the operator via ZERO_OAUTH_<NAME>_* variables, so no third-party brand or
// OAuth client identity is baked into the binary.
type Registry struct{}
// NewRegistry returns the (stateless) env-driven registry.
func NewRegistry() *Registry { return &Registry{} }
// envKey builds the ZERO_OAUTH_<NAME>_<suffix> variable name for a provider.
func envKey(name, suffix string) string {
up := strings.ToUpper(name)
up = strings.NewReplacer("-", "_", ".", "_").Replace(up)
return "ZERO_OAUTH_" + up + "_" + suffix
}
// ResolveConfig builds the oauth.Config and default Flow for a provider from its
// env/config:
//
// ZERO_OAUTH_<NAME>_CLIENT_ID (required)
// ZERO_OAUTH_<NAME>_CLIENT_SECRET (optional)
// ZERO_OAUTH_<NAME>_AUTHORIZE_URL (loopback flow; or discovered via issuer)
// ZERO_OAUTH_<NAME>_TOKEN_URL (or discovered via issuer)
// ZERO_OAUTH_<NAME>_DEVICE_URL (device flow; or discovered via issuer)
// ZERO_OAUTH_<NAME>_ISSUER_URL (RFC 8414 / OIDC discovery base)
// ZERO_OAUTH_<NAME>_SCOPES (space-separated)
// ZERO_OAUTH_<NAME>_FLOW ("loopback" [default] | "device")
//
// Pinned credential-bearing endpoints must be https (loopback exempt).
func (r *Registry) ResolveConfig(name string, env map[string]string) (Config, Flow, error) {
if err := ValidateProviderName(name); err != nil {
return Config{}, "", err
}
cfg := Config{
ClientID: strings.TrimSpace(envValue(env, envKey(name, "CLIENT_ID"))),
ClientSecret: strings.TrimSpace(envValue(env, envKey(name, "CLIENT_SECRET"))),
AuthorizationEndpoint: strings.TrimSpace(envValue(env, envKey(name, "AUTHORIZE_URL"))),
TokenEndpoint: strings.TrimSpace(envValue(env, envKey(name, "TOKEN_URL"))),
DeviceAuthorizationEndpoint: strings.TrimSpace(envValue(env, envKey(name, "DEVICE_URL"))),
IssuerURL: strings.TrimSpace(envValue(env, envKey(name, "ISSUER_URL"))),
Scopes: strings.Fields(envValue(env, envKey(name, "SCOPES"))),
}
if cfg.ClientID == "" {
return Config{}, "", fmt.Errorf("oauth: provider %q is not configured; set %s (and its endpoints or an issuer)", name, envKey(name, "CLIENT_ID"))
}
var flow Flow
switch strings.ToLower(strings.TrimSpace(envValue(env, envKey(name, "FLOW")))) {
case "", string(FlowLoopback):
flow = FlowLoopback
case string(FlowDevice):
flow = FlowDevice
default:
return Config{}, "", fmt.Errorf("oauth: provider %q has invalid %s (want loopback or device)", name, envKey(name, "FLOW"))
}
// A token endpoint (for exchange/refresh) must be reachable directly or via
// discovery. The per-flow authorize/device endpoints are checked at login
// time (after discovery), so a refresh-only config needs only a token URL.
if cfg.IssuerURL == "" && cfg.TokenEndpoint == "" {
return Config{}, "", fmt.Errorf("oauth: provider %q needs %s or %s", name, envKey(name, "TOKEN_URL"), envKey(name, "ISSUER_URL"))
}
for _, ep := range []string{cfg.TokenEndpoint, cfg.AuthorizationEndpoint, cfg.DeviceAuthorizationEndpoint, cfg.IssuerURL} {
if ep == "" {
continue
}
if err := validateTokenEndpoint(ep); err != nil {
return Config{}, "", err
}
}
return cfg, flow, nil
}
+97
View File
@@ -0,0 +1,97 @@
package oauth
import (
"errors"
"testing"
)
func TestResolveConfigFromEnv(t *testing.T) {
r := NewRegistry()
env := map[string]string{
"ZERO_OAUTH_DEMO_CLIENT_ID": "my-client",
"ZERO_OAUTH_DEMO_CLIENT_SECRET": "shh",
"ZERO_OAUTH_DEMO_SCOPES": "read write",
"ZERO_OAUTH_DEMO_AUTHORIZE_URL": "https://auth.example.com/authorize",
"ZERO_OAUTH_DEMO_TOKEN_URL": "https://auth.example.com/token",
}
cfg, flow, err := r.ResolveConfig("demo", env)
if err != nil {
t.Fatalf("ResolveConfig: %v", err)
}
if cfg.ClientID != "my-client" || cfg.ClientSecret != "shh" {
t.Fatalf("client creds not applied: %+v", cfg)
}
if len(cfg.Scopes) != 2 || cfg.Scopes[0] != "read" {
t.Fatalf("scopes = %v", cfg.Scopes)
}
if cfg.TokenEndpoint != "https://auth.example.com/token" {
t.Fatalf("token endpoint = %q", cfg.TokenEndpoint)
}
if flow != FlowLoopback {
t.Fatalf("flow = %q, want loopback default", flow)
}
}
func TestResolveConfigDeviceFlow(t *testing.T) {
r := NewRegistry()
env := map[string]string{
"ZERO_OAUTH_DEMO_CLIENT_ID": "c",
"ZERO_OAUTH_DEMO_TOKEN_URL": "https://auth.example.com/token",
"ZERO_OAUTH_DEMO_DEVICE_URL": "https://auth.example.com/device",
"ZERO_OAUTH_DEMO_FLOW": "device",
}
_, flow, err := r.ResolveConfig("demo", env)
if err != nil {
t.Fatalf("ResolveConfig: %v", err)
}
if flow != FlowDevice {
t.Fatalf("flow = %q, want device", flow)
}
}
func TestResolveConfigRequiresClientID(t *testing.T) {
r := NewRegistry()
if _, _, err := r.ResolveConfig("demo", map[string]string{}); err == nil {
t.Fatal("missing client id must error")
}
}
func TestResolveConfigRequiresEndpointsOrIssuer(t *testing.T) {
r := NewRegistry()
// client id but no token endpoint and no issuer => error.
_, _, err := r.ResolveConfig("demo", map[string]string{"ZERO_OAUTH_DEMO_CLIENT_ID": "c"})
if err == nil {
t.Fatal("missing endpoints/issuer must error")
}
}
func TestResolveConfigRejectsInsecureEndpoint(t *testing.T) {
r := NewRegistry()
env := map[string]string{
"ZERO_OAUTH_DEMO_CLIENT_ID": "c",
"ZERO_OAUTH_DEMO_AUTHORIZE_URL": "https://auth.example.com/authorize",
"ZERO_OAUTH_DEMO_TOKEN_URL": "http://insecure.example/token", // non-https, non-loopback
}
_, _, err := r.ResolveConfig("demo", env)
if !errors.Is(err, ErrInsecureTokenEndpoint) {
t.Fatalf("err = %v, want ErrInsecureTokenEndpoint", err)
}
}
func TestResolveConfigInvalidName(t *testing.T) {
r := NewRegistry()
for _, bad := range []string{"", "has space", "../escape", "a/b"} {
if _, _, err := r.ResolveConfig(bad, nil); err == nil {
t.Errorf("ResolveConfig(%q) should reject invalid name", bad)
}
}
}
func TestEnvKey(t *testing.T) {
if got := envKey("my-svc", "CLIENT_ID"); got != "ZERO_OAUTH_MY_SVC_CLIENT_ID" {
t.Fatalf("envKey = %q", got)
}
if got := envKey("two.part", "SCOPES"); got != "ZERO_OAUTH_TWO_PART_SCOPES" {
t.Fatalf("envKey = %q", got)
}
}
+125
View File
@@ -0,0 +1,125 @@
package oauth
import (
"context"
"sync"
"time"
)
// maxSchedulerLoadErrors bounds consecutive token-load failures before the
// scheduler gives up, so a transient error retries but a permanent one (e.g. a
// deleted token) does not loop forever.
const maxSchedulerLoadErrors = 5
// RefreshScheduler proactively refreshes a stored provider token shortly before
// it expires. It is an OPTIMIZATION over the on-demand GetFresh path, never the
// source of truth: a failed scheduled refresh is non-fatal and simply retried at
// the next expiry. It is a no-op for a token with no refresh token or no expiry.
type RefreshScheduler struct {
mu sync.Mutex
cancel context.CancelFunc
done chan struct{}
started bool
}
// NewRefreshScheduler returns an idle scheduler.
func NewRefreshScheduler() *RefreshScheduler {
return &RefreshScheduler{}
}
// Start begins proactively refreshing key via the manager until the context is
// canceled or Stop is called. It returns immediately; refresh happens in a
// goroutine. Calling Start twice is a no-op after the first.
func (s *RefreshScheduler) Start(ctx context.Context, m *Manager, key string) {
s.mu.Lock()
if s.started {
s.mu.Unlock()
return
}
runCtx, cancel := context.WithCancel(ctx)
s.cancel = cancel
s.done = make(chan struct{})
s.started = true
done := s.done
s.mu.Unlock()
go s.loop(runCtx, m, key, done)
}
func (s *RefreshScheduler) loop(ctx context.Context, m *Manager, key string, done chan struct{}) {
defer close(done)
loadErrors := 0
for {
token, _, err := m.loadForKey(key)
if err != nil {
// A load error may be transient (a concurrent store write, a brief I/O
// hiccup): back off and retry rather than permanently stopping proactive
// refresh. Give up only after repeated consecutive failures.
loadErrors++
if loadErrors >= maxSchedulerLoadErrors {
return
}
select {
case <-ctx.Done():
return
case <-time.After(m.buffer):
}
continue
}
loadErrors = 0
// No refresh token or no expiry => nothing to schedule on a timer.
if token.RefreshToken == "" || token.ExpiresAt.IsZero() {
return
}
delay := s.delayUntilRefresh(token.ExpiresAt, m.buffer, m.now())
select {
case <-ctx.Done():
return
case <-time.After(delay):
}
// Best-effort refresh; failure is non-fatal — reload and reschedule (the
// on-demand GetFresh remains the safety net for callers).
if _, err := m.GetFresh(ctx, key); err != nil {
// Wait a bounded backoff before retrying so a persistently failing
// refresh does not hot-loop.
select {
case <-ctx.Done():
return
case <-time.After(m.buffer):
}
}
}
}
// delayUntilRefresh computes how long to wait before refreshing: aim for buffer
// before expiry, clamped to >= 0, plus a small deterministic jitter so many
// sessions do not refresh in lockstep.
func (s *RefreshScheduler) delayUntilRefresh(expiresAt time.Time, buffer time.Duration, now time.Time) time.Duration {
target := expiresAt.Add(-buffer)
delay := target.Sub(now)
if delay < 0 {
delay = 0
}
// Jitter up to ~10% of the buffer, derived from the expiry nanos (no RNG dep).
if buffer > 0 {
jitter := time.Duration(expiresAt.UnixNano()%int64(buffer/10+1)) % (buffer/10 + 1)
delay += jitter
}
return delay
}
// Stop cancels the scheduler and waits for its goroutine to exit. Safe to call
// more than once and on a never-started scheduler.
func (s *RefreshScheduler) Stop() {
s.mu.Lock()
cancel := s.cancel
done := s.done
s.cancel = nil
s.mu.Unlock()
if cancel != nil {
cancel()
}
if done != nil {
<-done
}
}
+72
View File
@@ -0,0 +1,72 @@
package oauth
import (
"context"
"io"
"net/http"
"net/http/httptest"
"path/filepath"
"sync/atomic"
"testing"
"time"
)
func TestRefreshSchedulerNoRefreshTokenIsNoop(t *testing.T) {
m := managerFor(t, map[string]string{"ZERO_OAUTH_DEMO_CLIENT_ID": "c"}, nil)
// Token without a refresh token => scheduler must exit promptly (no-op).
if err := m.store.Save(ProviderKey("demo"), Token{AccessToken: "a", ExpiresAt: time.Now().Add(time.Hour)}); err != nil {
t.Fatalf("seed Save: %v", err)
}
s := NewRefreshScheduler()
s.Start(context.Background(), m, ProviderKey("demo"))
// Stop waits for the goroutine; if it didn't exit on its own this still returns
// because Stop cancels — but the loop should already be done.
s.Stop()
}
func TestRefreshSchedulerRefreshesBeforeExpiry(t *testing.T) {
var hits atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
hits.Add(1)
_, _ = io.WriteString(w, `{"access_token":"refreshed","expires_in":3600}`)
}))
defer server.Close()
store, err := NewStore(StoreOptions{FilePath: filepath.Join(t.TempDir(), "tok.json")})
if err != nil {
t.Fatalf("NewStore: %v", err)
}
m, err := NewManager(ManagerOptions{
Store: store,
Env: map[string]string{
"ZERO_OAUTH_DEMO_CLIENT_ID": "c",
"ZERO_OAUTH_DEMO_TOKEN_URL": server.URL,
},
RefreshBuffer: 10 * time.Millisecond,
})
if err != nil {
t.Fatalf("NewManager: %v", err)
}
// Expiry just past the buffer so the scheduled delay is ~0 and a refresh fires
// almost immediately.
if err := store.Save(ProviderKey("demo"), Token{AccessToken: "old", RefreshToken: "rt", ExpiresAt: time.Now().Add(15 * time.Millisecond)}); err != nil {
t.Fatalf("seed Save: %v", err)
}
s := NewRefreshScheduler()
s.Start(context.Background(), m, ProviderKey("demo"))
defer s.Stop()
deadline := time.Now().Add(2 * time.Second)
for time.Now().Before(deadline) {
stored, _, err := store.Load(ProviderKey("demo"))
if err != nil {
t.Fatalf("Load: %v", err)
}
if stored.AccessToken == "refreshed" {
return // success
}
time.Sleep(5 * time.Millisecond)
}
t.Fatalf("scheduler did not refresh before expiry (token endpoint hits=%d)", hits.Load())
}
+324
View File
@@ -0,0 +1,324 @@
package oauth
import (
"encoding/json"
"errors"
"fmt"
"os"
"path/filepath"
"regexp"
"sort"
"strings"
"sync"
"time"
)
const (
storeSchemaVersion = 1
// KeyPrefixProvider namespaces provider-login tokens; MCP server tokens live
// under KeyPrefixMCP in the same format (so a future MCP migration is a key
// rename, not a format change).
KeyPrefixProvider = "provider:"
KeyPrefixMCP = "mcp:"
)
// keyPattern bounds a token key to a safe, single-segment namespaced identifier
// so a key can never traverse or collide with store internals.
var keyPattern = regexp.MustCompile(`^(provider|mcp):[A-Za-z0-9][A-Za-z0-9._-]{0,127}$`)
// ValidateKey reports whether key is a well-formed namespaced token key.
func ValidateKey(key string) error {
if !keyPattern.MatchString(key) {
return fmt.Errorf("oauth: invalid token key %q (want \"provider:<name>\" or \"mcp:<name>\")", key)
}
return nil
}
// ProviderKey builds the store key for a provider login.
func ProviderKey(name string) string { return KeyPrefixProvider + name }
// Status is a redaction-safe summary of a stored token (no secret material).
type Status struct {
Key string `json:"key"`
HasToken bool `json:"hasToken"`
HasRefreshToken bool `json:"hasRefreshToken"`
TokenType string `json:"tokenType,omitempty"`
Account string `json:"account,omitempty"`
Scopes []string `json:"scopes,omitempty"`
ExpiresAt time.Time `json:"expiresAt,omitempty"`
Expired bool `json:"expired"`
}
// StoreOptions configures where provider OAuth tokens are persisted.
type StoreOptions struct {
FilePath string
Env map[string]string
Now func() time.Time
}
// Store persists OAuth tokens (provider + MCP namespaces) in a 0600 JSON file,
// guarded by a cross-process lock and written atomically.
type Store struct {
filePath string
now func() time.Time
mu sync.Mutex
}
type storeFile struct {
SchemaVersion int `json:"schemaVersion"`
Tokens map[string]Token `json:"tokens"`
}
// ResolveStorePath determines the on-disk location for provider OAuth tokens,
// honoring ZERO_OAUTH_TOKENS_PATH, then XDG_CONFIG_HOME, then the home dir.
func ResolveStorePath(env map[string]string) (string, error) {
if override := strings.TrimSpace(envValue(env, "ZERO_OAUTH_TOKENS_PATH")); override != "" {
if filepath.IsAbs(override) {
return filepath.Clean(override), nil
}
return filepath.Abs(override)
}
configHome := strings.TrimSpace(envValue(env, "XDG_CONFIG_HOME"))
if configHome == "" {
home := strings.TrimSpace(firstNonEmpty(envValue(env, "HOME"), envValue(env, "USERPROFILE")))
if home == "" {
var err error
home, err = os.UserHomeDir()
if err != nil {
return "", fmt.Errorf("oauth: resolve user home: %w", err)
}
}
configHome = filepath.Join(home, ".config")
} else if !filepath.IsAbs(configHome) {
resolved, err := filepath.Abs(configHome)
if err != nil {
return "", err
}
configHome = resolved
}
return filepath.Join(configHome, "zero", "oauth-tokens.json"), nil
}
// NewStore builds a file-backed token store.
func NewStore(options StoreOptions) (*Store, error) {
filePath := options.FilePath
var err error
if strings.TrimSpace(filePath) == "" {
filePath, err = ResolveStorePath(options.Env)
if err != nil {
return nil, err
}
}
if !filepath.IsAbs(filePath) {
filePath, err = filepath.Abs(filePath)
if err != nil {
return nil, err
}
}
now := options.Now
if now == nil {
now = time.Now
}
return &Store{filePath: filepath.Clean(filePath), now: now}, nil
}
// Save persists a token under key, replacing any existing entry.
func (s *Store) Save(key string, token Token) error {
if err := ValidateKey(key); err != nil {
return err
}
s.mu.Lock()
defer s.mu.Unlock()
unlock, err := acquireFileLock(s.filePath+".lockfile", s.now)
if err != nil {
return err
}
defer unlock()
state, err := s.readState()
if err != nil {
return err
}
state.Tokens[key] = token
return s.writeState(state)
}
// Load returns the token for key; the bool is false when none is stored.
func (s *Store) Load(key string) (Token, bool, error) {
if err := ValidateKey(key); err != nil {
return Token{}, false, err
}
s.mu.Lock()
defer s.mu.Unlock()
state, err := s.readState()
if err != nil {
return Token{}, false, err
}
token, ok := state.Tokens[key]
return token, ok, nil
}
// Delete removes the token for key, reporting whether one was present.
func (s *Store) Delete(key string) (bool, error) {
if err := ValidateKey(key); err != nil {
return false, err
}
s.mu.Lock()
defer s.mu.Unlock()
unlock, err := acquireFileLock(s.filePath+".lockfile", s.now)
if err != nil {
return false, err
}
defer unlock()
state, err := s.readState()
if err != nil {
return false, err
}
if _, ok := state.Tokens[key]; !ok {
return false, nil
}
delete(state.Tokens, key)
if err := s.writeState(state); err != nil {
return false, err
}
return true, nil
}
// Status returns redaction-safe summaries of every stored token, sorted by key.
// An optional prefix filters to one namespace (e.g. KeyPrefixProvider).
func (s *Store) Status(prefix string) ([]Status, error) {
s.mu.Lock()
defer s.mu.Unlock()
state, err := s.readState()
if err != nil {
return nil, err
}
keys := make([]string, 0, len(state.Tokens))
for k := range state.Tokens {
if prefix == "" || strings.HasPrefix(k, prefix) {
keys = append(keys, k)
}
}
sort.Strings(keys)
now := s.now()
out := make([]Status, 0, len(keys))
for _, k := range keys {
token := state.Tokens[k]
out = append(out, Status{
Key: k,
HasToken: strings.TrimSpace(token.AccessToken) != "",
HasRefreshToken: strings.TrimSpace(token.RefreshToken) != "",
TokenType: token.TokenType,
Account: token.Account,
Scopes: token.Scopes,
ExpiresAt: token.ExpiresAt,
Expired: token.Expired(now),
})
}
return out, nil
}
func (s *Store) readState() (storeFile, error) {
data, err := os.ReadFile(s.filePath)
if err != nil {
if errors.Is(err, os.ErrNotExist) {
return emptyStoreFile(), nil
}
return storeFile{}, err
}
var state storeFile
if err := json.Unmarshal(data, &state); err != nil {
return storeFile{}, fmt.Errorf("oauth: invalid token file at %s: %w", s.filePath, err)
}
if state.SchemaVersion != storeSchemaVersion {
return storeFile{}, fmt.Errorf("oauth: invalid token file at %s: unsupported schemaVersion", s.filePath)
}
if state.Tokens == nil {
state.Tokens = map[string]Token{}
}
for key := range state.Tokens {
if err := ValidateKey(key); err != nil {
return storeFile{}, fmt.Errorf("oauth: invalid token file at %s: %w", s.filePath, err)
}
}
return state, nil
}
func (s *Store) writeState(state storeFile) error {
if err := os.MkdirAll(filepath.Dir(s.filePath), 0o700); err != nil {
return err
}
data, err := json.MarshalIndent(state, "", " ")
if err != nil {
return err
}
tempPath := fmt.Sprintf("%s.tmp-%d-%d", s.filePath, os.Getpid(), s.now().UnixNano())
if err := os.WriteFile(tempPath, append(data, '\n'), 0o600); err != nil {
return err
}
if err := os.Rename(tempPath, s.filePath); err != nil {
_ = os.Remove(tempPath)
return err
}
return nil
}
func emptyStoreFile() storeFile {
return storeFile{SchemaVersion: storeSchemaVersion, Tokens: map[string]Token{}}
}
// FormatStatuses renders a human-readable status table without leaking token
// material.
func FormatStatuses(statuses []Status) string {
if len(statuses) == 0 {
return "No OAuth provider logins are stored."
}
var b strings.Builder
for i, st := range statuses {
if i > 0 {
b.WriteByte('\n')
}
name := strings.TrimPrefix(st.Key, KeyPrefixProvider)
b.WriteString(name)
b.WriteString(": ")
if !st.HasToken {
b.WriteString("no token")
continue
}
b.WriteString("logged in")
if st.Account != "" {
b.WriteString(" as " + st.Account)
}
if st.HasRefreshToken {
b.WriteString(" (refreshable)")
}
if !st.ExpiresAt.IsZero() {
if st.Expired {
b.WriteString(", expired at ")
} else {
b.WriteString(", expires ")
}
b.WriteString(st.ExpiresAt.UTC().Format(time.RFC3339))
}
}
return b.String()
}
// envValue reads a variable. A non-nil env map is authoritative (hermetic): a
// missing key returns "" rather than falling back to the process environment, so
// a caller/test that passes a controlled map can never pick up ambient
// ZERO_OAUTH_* / HOME / XDG_CONFIG_HOME values. Only a nil map uses os.Getenv.
func envValue(env map[string]string, key string) string {
if env != nil {
return env[key]
}
return os.Getenv(key)
}
func firstNonEmpty(values ...string) string {
for _, v := range values {
if strings.TrimSpace(v) != "" {
return v
}
}
return ""
}
+128
View File
@@ -0,0 +1,128 @@
package oauth
import (
"os"
"path/filepath"
"runtime"
"testing"
"time"
)
func newTestStore(t *testing.T) (*Store, string) {
t.Helper()
path := filepath.Join(t.TempDir(), "oauth-tokens.json")
s, err := NewStore(StoreOptions{FilePath: path})
if err != nil {
t.Fatalf("NewStore: %v", err)
}
return s, path
}
func TestStoreSaveLoadDelete(t *testing.T) {
s, _ := newTestStore(t)
tok := Token{AccessToken: "at", RefreshToken: "rt", Account: "me@x"}
if err := s.Save(ProviderKey("demo"), tok); err != nil {
t.Fatalf("Save: %v", err)
}
if err := s.Save(KeyPrefixMCP+"server1", Token{AccessToken: "mcp-at"}); err != nil {
t.Fatalf("Save mcp: %v", err)
}
got, ok, err := s.Load(ProviderKey("demo"))
if err != nil || !ok || got.AccessToken != "at" || got.Account != "me@x" {
t.Fatalf("Load = %+v ok=%v err=%v", got, ok, err)
}
removed, err := s.Delete(ProviderKey("demo"))
if err != nil || !removed {
t.Fatalf("Delete = %v %v", removed, err)
}
if _, ok, _ := s.Load(ProviderKey("demo")); ok {
t.Fatal("token should be gone after delete")
}
// The mcp-namespaced token is untouched.
if _, ok, _ := s.Load(KeyPrefixMCP + "server1"); !ok {
t.Fatal("mcp token should survive provider delete")
}
}
func TestStoreRejectsInvalidKeys(t *testing.T) {
s, _ := newTestStore(t)
for _, bad := range []string{"demo", "provider:", "provider:../escape", "mcp:bad/key", "other:x", ""} {
if err := s.Save(bad, Token{AccessToken: "x"}); err == nil {
t.Errorf("Save(%q) should be rejected", bad)
}
}
for _, ok := range []string{"provider:demo", "mcp:server-1", "provider:two-svc"} {
if err := ValidateKey(ok); err != nil {
t.Errorf("ValidateKey(%q) = %v, want nil", ok, err)
}
}
}
func TestStoreStatusFiltersByPrefix(t *testing.T) {
s, _ := newTestStore(t)
_ = s.Save(ProviderKey("demo"), Token{AccessToken: "a", RefreshToken: "r", ExpiresAt: time.Now().Add(time.Hour)})
_ = s.Save(KeyPrefixMCP+"srv", Token{AccessToken: "m"})
statuses, err := s.Status(KeyPrefixProvider)
if err != nil {
t.Fatalf("Status: %v", err)
}
if len(statuses) != 1 || statuses[0].Key != ProviderKey("demo") {
t.Fatalf("provider status = %+v", statuses)
}
if !statuses[0].HasRefreshToken {
t.Fatal("status should report a refresh token")
}
}
func TestStoreFileMode0600(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("unix file modes")
}
s, path := newTestStore(t)
if err := s.Save(ProviderKey("x"), Token{AccessToken: "a"}); err != nil {
t.Fatalf("Save: %v", err)
}
info, err := os.Stat(path)
if err != nil {
t.Fatalf("stat: %v", err)
}
if perm := info.Mode().Perm(); perm != 0o600 {
t.Fatalf("token file mode = %o, want 600", perm)
}
}
func TestStoreMalformedFailsClosed(t *testing.T) {
s, path := newTestStore(t)
if err := os.WriteFile(path, []byte("{not json"), 0o600); err != nil {
t.Fatalf("corrupt: %v", err)
}
if _, _, err := s.Load(ProviderKey("x")); err == nil {
t.Fatal("malformed store must fail closed")
}
}
func TestEnvValueHermetic(t *testing.T) {
t.Setenv("ZERO_OAUTH_DEMO_CLIENT_ID", "ambient")
// A non-nil map omitting the key must NOT leak the ambient process value.
if got := envValue(map[string]string{}, "ZERO_OAUTH_DEMO_CLIENT_ID"); got != "" {
t.Fatalf("non-nil env map must be hermetic, got %q", got)
}
// A nil map reads the process environment.
if got := envValue(nil, "ZERO_OAUTH_DEMO_CLIENT_ID"); got != "ambient" {
t.Fatalf("nil env map should read os env, got %q", got)
}
}
func TestResolveStorePathHonorsOverride(t *testing.T) {
// Use an OS-appropriate absolute path: a unix-style "/tmp/..." literal is not
// absolute on Windows (no drive), so it would be resolved against the current
// drive and a verbatim comparison would fail there.
override := filepath.Join(t.TempDir(), "custom", "tok.json")
path, err := ResolveStorePath(map[string]string{"ZERO_OAUTH_TOKENS_PATH": override})
if err != nil {
t.Fatalf("ResolveStorePath: %v", err)
}
if path != override {
t.Fatalf("path = %q, want %q", path, override)
}
}
@@ -62,7 +62,7 @@ func TestDetectLocalRuntimesReportsReachableRuntime(t *testing.T) {
if !detected[0].Reachable {
t.Fatalf("runtime should be reachable: %#v", detected[0])
}
if len(detected[0].Models) == 0 || detected[0].Models[0] != "llama3.1" {
if detected[0].Models == nil || len(detected[0].Models) == 0 || detected[0].Models[0] != "llama3.1" {
t.Fatalf("expected discovered model list, got %#v", detected[0].Models)
}
}
+3 -1
View File
@@ -378,6 +378,8 @@ func acquireLock(lockPath string, timeout time.Duration) (func(), error) {
if time.Now().After(deadline) {
return nil, fmt.Errorf("swarm: timed out acquiring lock %s", filepath.Base(lockPath))
}
time.Sleep(20 * time.Millisecond)
// Short retry so a freed lock is re-acquired promptly under heavy
// contention (Windows file ops are slow; a coarse sleep starves waiters).
time.Sleep(2 * time.Millisecond)
}
}
+5 -4
View File
@@ -246,10 +246,11 @@ func TestMailboxLockReleaseIsOwnershipAware(t *testing.T) {
func TestMailboxConcurrentSends(t *testing.T) {
mb := newTestMailbox(t)
// High fan-out + a generous lock timeout so heavy contention is exercised
// (and a lock regression — e.g. a Windows sharing-violation treated as fatal —
// would surface as lost/failed messages here rather than flaking).
mb.LockTimeout = 10 * time.Second
// High fan-out exercises heavy contention (a lock regression — e.g. a Windows
// sharing-violation treated as fatal — surfaces as lost/failed messages here).
// The lock timeout is generous so a slow CI (Windows file ops are slow under
// 200-way contention) never times a legitimate send out and flakes.
mb.LockTimeout = 60 * time.Second
const n = 200
var wg sync.WaitGroup
var failures atomic.Int32
+1 -1
View File
@@ -302,7 +302,7 @@ func webFetchRedirectPolicy(previous func(*http.Request, []*http.Request) error,
return fmt.Errorf("too many redirects: maximum is %d", webFetchRedirectLimit)
}
if err := validateParsedWebFetchURL(request.Context(), request.URL, resolver, gate); err != nil {
return fmt.Errorf("unsafe redirect URL: %w", err)
return fmt.Errorf("Unsafe redirect URL: %w", err)
}
if previous != nil {
return previous(request, via)
+1 -1
View File
@@ -285,7 +285,7 @@ func TestWebFetchToolRejectsUnsafeRedirects(t *testing.T) {
if result.Status != StatusError {
t.Fatalf("expected redirect safety error, got %s: %s", result.Status, result.Output)
}
if !strings.Contains(result.Output, "unsafe redirect URL") {
if !strings.Contains(result.Output, "Unsafe redirect URL") {
t.Fatalf("expected unsafe redirect message, got %q", result.Output)
}
})
+171 -7
View File
@@ -29,6 +29,9 @@ type searchResult struct {
Title string
URL string
Snippet string
// Score is an optional provider-supplied relevance signal in the [0, 1]
// range. Zero (the zero value) is treated as "absent" by the renderer.
Score float64
}
// searchBackend discovers URLs for a query. It is an interface so any hosted
@@ -67,6 +70,11 @@ func newWebSearchToolWithBackend(backend searchBackend) Tool {
Minimum: intPtr(1),
Maximum: intPtr(maxWebSearchLimit),
},
"domains": {
Type: "array",
Description: "Optional list of allowed hostnames; results outside these hosts are filtered out before they reach the model. Useful as a prompt-injection defense when you only want results from a known set of domains.",
Items: &PropertySchema{Type: "string"},
},
},
Required: []string{"query"},
AdditionalProperties: false,
@@ -123,6 +131,17 @@ func (tool webSearchTool) Run(ctx context.Context, args map[string]any) Result {
if limit > maxWebSearchLimit {
limit = maxWebSearchLimit
}
domains, domainsProvided, err := stringListArgWebSearch(args, "domains")
if err != nil {
return errorResult("Error: Invalid arguments for web_search: " + err.Error())
}
// Fail closed: the caller asked for an allowlist, but every entry was
// invalid (or the list was empty). Silently dropping the filter would
// turn a constrained search into an unconstrained one, defeating the
// prompt-injection defense the parameter exists to provide.
if domainsProvided && len(domains) == 0 {
return errorResult("Error: web_search 'domains' argument was provided but contained no valid hostnames; remove the argument or pass valid hostnames")
}
if tool.backend == nil {
return errorResult("Error: no search backend configured. Set ZERO_WEBSEARCH_BASE_URL (and ZERO_WEBSEARCH_API_KEY) to enable web_search.")
@@ -135,6 +154,16 @@ func (tool webSearchTool) Run(ctx context.Context, args map[string]any) Result {
if err != nil {
return errorResult("Error performing web search: " + redactWebSearchText(err.Error()))
}
// Filter by domains BEFORE the empty-result short-circuit so a non-empty
// allowlist that ate every hit surfaces as a clear "no results matched
// domains" error rather than a misleading "no results" message.
if len(results) > 0 && len(domains) > 0 {
filtered, err := filterWebSearchByDomains(results, domains)
if err != nil {
return errorResult("Error: " + err.Error())
}
results = filtered
}
if len(results) == 0 {
return okResult("No results for query: " + redactWebSearchText(query))
}
@@ -146,6 +175,9 @@ func (tool webSearchTool) Run(ctx context.Context, args map[string]any) Result {
// formatSearchResults renders results as a compact numbered list:
// "1. Title — URL" with the snippet indented on the next line.
// If a result carries a non-zero Score, it is shown as " — score 0.91" after
// the title; absence (zero value) is rendered without the score to keep the
// common case tidy.
func formatSearchResults(results []searchResult) string {
lines := make([]string, 0, len(results)*2)
for index, result := range results {
@@ -153,7 +185,11 @@ func formatSearchResults(results []searchResult) string {
if title == "" {
title = "(untitled)"
}
lines = append(lines, fmt.Sprintf("%d. %s — %s", index+1, title, strings.TrimSpace(result.URL)))
header := fmt.Sprintf("%d. %s — %s", index+1, title, strings.TrimSpace(result.URL))
if result.Score > 0 {
header += fmt.Sprintf(" — score %.2f", result.Score)
}
lines = append(lines, header)
if snippet := strings.TrimSpace(result.Snippet); snippet != "" {
lines = append(lines, " "+snippet)
}
@@ -260,13 +296,16 @@ func (backend *httpSearchBackend) Search(ctx context.Context, query string, limi
return parseSearchResults(body)
}
// parseSearchResults accepts either a bare array [{title,url,snippet}] or a
// wrapped object {"results":[...]}, the two shapes common across providers.
// parseSearchResults accepts either a bare array [{title,url,snippet,score}] or
// a wrapped object {"results":[...]}, the two shapes common across providers.
// The optional "score" field is forwarded when present; providers that omit
// it leave searchResult.Score at the zero value and the renderer skips it.
func parseSearchResults(body []byte) ([]searchResult, error) {
type rawResult struct {
Title string `json:"title"`
URL string `json:"url"`
Snippet string `json:"snippet"`
Title string `json:"title"`
URL string `json:"url"`
Snippet string `json:"snippet"`
Score float64 `json:"score,omitempty"`
}
trimmed := bytes.TrimSpace(body)
if len(trimmed) == 0 {
@@ -276,7 +315,12 @@ func parseSearchResults(body []byte) ([]searchResult, error) {
convert := func(raw []rawResult) []searchResult {
out := make([]searchResult, 0, len(raw))
for _, item := range raw {
out = append(out, searchResult(item))
out = append(out, searchResult{
Title: item.Title,
URL: item.URL,
Snippet: item.Snippet,
Score: item.Score,
})
}
return out
}
@@ -300,3 +344,123 @@ func parseSearchResults(body []byte) ([]searchResult, error) {
func redactWebSearchText(value string) string {
return redaction.RedactString(value, redaction.Options{})
}
// stringListArgWebSearch reads an optional []string argument. The conventional
// `[]any` shape is preferred (matches the rest of the tools package) but some
// providers' tool-calling shapes emit `[]string` directly; tolerate both.
//
// The returned provided bool is true when the caller passed a "domains" key at
// all (even with an empty array or all-invalid entries). That lets Run
// distinguish "no allowlist" from "allowlist that normalized to empty" — the
// latter is a fail-closed error, not a silent permission to skip filtering.
func stringListArgWebSearch(args map[string]any, key string) (domains []string, provided bool, err error) {
value, ok := args[key]
if !ok || value == nil {
return nil, false, nil
}
provided = true
raw, ok := value.([]any)
if !ok {
if asString, ok2 := value.([]string); ok2 {
return normalizeWebSearchDomainList(asString), true, nil
}
return nil, true, fmt.Errorf("%s must be a list of strings", key)
}
out := make([]string, 0, len(raw))
for index, item := range raw {
text, ok := item.(string)
if !ok {
return nil, true, fmt.Errorf("%s[%d] must be a string", key, index)
}
out = append(out, text)
}
return normalizeWebSearchDomainList(out), true, nil
}
// normalizeWebSearchDomainList lowercases, strips scheme/www/path fragments,
// drops empty entries, and de-duplicates while preserving input order.
func normalizeWebSearchDomainList(domains []string) []string {
seen := make(map[string]struct{}, len(domains))
out := make([]string, 0, len(domains))
for _, raw := range domains {
host := canonicalizeWebSearchHost(raw)
if host == "" {
continue
}
if _, ok := seen[host]; ok {
continue
}
seen[host] = struct{}{}
out = append(out, host)
}
return out
}
// canonicalizeWebSearchHost extracts a bare lowercase hostname from a domain
// string, accepting "react.dev", "https://react.dev/x", and "www.React.dev"
// forms. Returns "" on empty/whitespace input. The "www." prefix is stripped
// before lowercasing so case-insensitive matches like "WWW.react.dev" work.
func canonicalizeWebSearchHost(raw string) string {
trimmed := strings.TrimSpace(raw)
if trimmed == "" {
return ""
}
if strings.Contains(trimmed, "://") {
parsed, err := url.Parse(trimmed)
if err != nil {
return ""
}
// Hostname() strips the port so allowlist entries like
// "https://react.dev:443/path" normalize the same as "react.dev".
// Using parsed.Host here would leave ":443" in the string and
// silently break every match against a result URL on port 443.
trimmed = parsed.Hostname()
}
trimmed = strings.TrimSpace(trimmed)
if i := strings.IndexAny(trimmed, "/?#"); i >= 0 {
trimmed = trimmed[:i]
}
trimmed = strings.ToLower(trimmed)
trimmed = strings.TrimPrefix(trimmed, "www.")
if trimmed == "" || strings.ContainsAny(trimmed, " \t\n") {
return ""
}
return trimmed
}
// filterWebSearchByDomains drops results whose host is not in the allowlist.
// A result whose URL fails to parse is dropped (fail-closed). When the filter
// eats every result, the error names the failing allowlist so the caller knows
// the constraint, not the provider, is the cause.
func filterWebSearchByDomains(results []searchResult, allowed []string) ([]searchResult, error) {
allowSet := make(map[string]struct{}, len(allowed))
for _, d := range allowed {
allowSet[d] = struct{}{}
}
out := make([]searchResult, 0, len(results))
for _, row := range results {
host := hostFromWebSearchURL(row.URL)
if host == "" {
// Unparseable URL: never let it through when an allowlist is set.
continue
}
if _, ok := allowSet[host]; !ok {
continue
}
out = append(out, row)
}
if len(out) == 0 {
return nil, fmt.Errorf("no web_search results matched domains=%v; loosen the allowlist or remove the domains argument", allowed)
}
return out, nil
}
func hostFromWebSearchURL(raw string) string {
parsed, err := url.Parse(strings.TrimSpace(raw))
if err != nil || parsed == nil {
return ""
}
host := strings.ToLower(parsed.Hostname())
host = strings.TrimPrefix(host, "www.")
return host
}
+330
View File
@@ -256,3 +256,333 @@ func TestSameHostRedirectPolicy(t *testing.T) {
t.Fatal("redirect limit must be enforced")
}
}
// ---- new: domains allowlist + score field ----------------------------------
func TestWebSearchDomainsFilterKeepsOnlyAllowedHosts(t *testing.T) {
backend := &fakeSearchBackend{results: []searchResult{
{Title: "RSC", URL: "https://react.dev/rsc", Snippet: "Server components"},
{Title: "RFC", URL: "https://github.com/reactjs/rfcs", Snippet: "RFC 0188"},
{Title: "Other", URL: "https://stackoverflow.com/q/1", Snippet: "unrelated"},
}}
tool := newWebSearchToolWithBackend(backend)
res := tool.Run(context.Background(), map[string]any{
"query": "react server components",
"domains": []any{"react.dev", "github.com"},
})
if res.Status != StatusOK {
t.Fatalf("status = %v, output = %q", res.Status, res.Output)
}
if strings.Contains(res.Output, "stackoverflow.com") {
t.Errorf("stackoverflow result must be filtered out, got: %s", res.Output)
}
if !strings.Contains(res.Output, "react.dev/rsc") || !strings.Contains(res.Output, "github.com/reactjs/rfcs") {
t.Errorf("allowed-host results missing, got: %s", res.Output)
}
}
func TestWebSearchDomainsFilterEmptyResultIsError(t *testing.T) {
backend := &fakeSearchBackend{results: []searchResult{
{Title: "x", URL: "https://stackoverflow.com/q/1"},
}}
tool := newWebSearchToolWithBackend(backend)
res := tool.Run(context.Background(), map[string]any{
"query": "x",
"domains": []any{"react.dev"},
})
if res.Status != StatusError {
t.Fatalf("expected error when allowlist eats every result, got %v: %s", res.Status, res.Output)
}
if !strings.Contains(res.Output, "no web_search results matched domains") {
t.Errorf("expected clear allowlist error, got: %s", res.Output)
}
}
func TestWebSearchDomainsFilterToleratesStringSlice(t *testing.T) {
// Some providers send []string instead of []any. The tool must accept both.
backend := &fakeSearchBackend{results: []searchResult{
{Title: "ok", URL: "https://react.dev/x"},
}}
tool := newWebSearchToolWithBackend(backend)
res := tool.Run(context.Background(), map[string]any{
"query": "x",
"domains": []string{"react.dev"},
})
if res.Status != StatusOK {
t.Fatalf("expected ok for []string domains, got %v: %s", res.Status, res.Output)
}
}
func TestWebSearchDomainsFilterNormalizesSchemeAndCase(t *testing.T) {
backend := &fakeSearchBackend{results: []searchResult{
{Title: "ok", URL: "https://React.Dev/x"},
}}
tool := newWebSearchToolWithBackend(backend)
res := tool.Run(context.Background(), map[string]any{
"query": "x",
"domains": []any{"https://REACT.dev/path?ignored=1", "WWW.react.dev"},
})
if res.Status != StatusOK {
t.Fatalf("expected ok after normalization, got %v: %s", res.Status, res.Output)
}
}
func TestWebSearchDomainsFilterStripsWWWFromResult(t *testing.T) {
backend := &fakeSearchBackend{results: []searchResult{
{Title: "ok", URL: "https://www.react.dev/x"},
}}
tool := newWebSearchToolWithBackend(backend)
// Allowlist "react.dev" must still match a result that has www. in the URL.
res := tool.Run(context.Background(), map[string]any{
"query": "x",
"domains": []any{"react.dev"},
})
if res.Status != StatusOK {
t.Fatalf("expected ok when result host has www., got %v: %s", res.Status, res.Output)
}
}
func TestWebSearchDomainsFilterRejectsMalformedURLs(t *testing.T) {
backend := &fakeSearchBackend{results: []searchResult{
{Title: "ok", URL: "https://react.dev/x"},
{Title: "bad", URL: "::::not a url::::"},
}}
tool := newWebSearchToolWithBackend(backend)
res := tool.Run(context.Background(), map[string]any{
"query": "x",
"domains": []any{"react.dev"},
})
if res.Status != StatusOK {
t.Fatalf("status = %v, output = %s", res.Status, res.Output)
}
if strings.Contains(res.Output, "not a url") {
t.Errorf("malformed URL must be filtered out, got: %s", res.Output)
}
}
func TestWebSearchDomainsFilterRejectsBadArgType(t *testing.T) {
tool := newWebSearchToolWithBackend(&fakeSearchBackend{})
res := tool.Run(context.Background(), map[string]any{
"query": "x",
"domains": "react.dev", // string instead of array
})
if res.Status != StatusError {
t.Fatalf("expected error for non-array domains, got %v: %s", res.Status, res.Output)
}
}
func TestWebSearchRendersScoreWhenPresent(t *testing.T) {
backend := &fakeSearchBackend{results: []searchResult{
{Title: "ranked", URL: "https://react.dev/rsc", Score: 0.91},
{Title: "unranked", URL: "https://react.dev/other"}, // zero score, must be omitted
}}
tool := newWebSearchToolWithBackend(backend)
res := tool.Run(context.Background(), map[string]any{"query": "x"})
if res.Status != StatusOK {
t.Fatalf("status = %v, output = %s", res.Status, res.Output)
}
if !strings.Contains(res.Output, "score 0.91") {
t.Errorf("expected 'score 0.91' in output, got: %s", res.Output)
}
// The zero-score row must NOT render "score 0.00" — that would be noisy.
if strings.Contains(res.Output, "score 0.00") {
t.Errorf("zero score must not be rendered, got: %s", res.Output)
}
}
func TestHTTPSearchBackendParsesScoreField(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"results":[{"title":"A","url":"https://a.test","snippet":"sa","score":0.77},{"title":"B","url":"https://b.test","snippet":"sb"}]}`))
}))
defer server.Close()
backend := &httpSearchBackend{client: server.Client(), baseURL: server.URL}
results, err := backend.Search(context.Background(), "q", 5)
if err != nil {
t.Fatalf("Search: %v", err)
}
if len(results) != 2 {
t.Fatalf("results = %d, want 2", len(results))
}
if results[0].Score != 0.77 {
t.Errorf("first result Score = %v, want 0.77", results[0].Score)
}
if results[1].Score != 0 {
t.Errorf("absent score must be zero, got %v", results[1].Score)
}
}
func TestCanonicalizeWebSearchHost(t *testing.T) {
cases := []struct {
in, want string
}{
{"react.dev", "react.dev"},
{" REACT.dev ", "react.dev"},
{"https://react.dev/x", "react.dev"},
{"http://www.react.dev", "react.dev"},
{"WWW.react.dev", "react.dev"},
{"react.dev/path?x=1", "react.dev"},
{"", ""},
{" ", ""},
{"https://", ""},
{"react .dev", ""}, // space inside
}
for _, c := range cases {
if got := canonicalizeWebSearchHost(c.in); got != c.want {
t.Errorf("canonicalizeWebSearchHost(%q) = %q, want %q", c.in, got, c.want)
}
}
}
func TestNormalizeWebSearchDomainList(t *testing.T) {
got := normalizeWebSearchDomainList([]string{" REACT.dev ", "https://github.com/x", "www.example.com", "react.dev", "", " "})
want := []string{"react.dev", "github.com", "example.com"}
if len(got) != len(want) {
t.Fatalf("normalizeWebSearchDomainList = %v, want %v", got, want)
}
for i := range want {
if got[i] != want[i] {
t.Errorf("[%d] = %q, want %q", i, got[i], want[i])
}
}
}
func TestFilterWebSearchByDomains(t *testing.T) {
rows := []searchResult{
{URL: "https://react.dev/x"},
{URL: "https://github.com/y"},
{URL: "https://stackoverflow.com/z"},
}
filtered, err := filterWebSearchByDomains(rows, []string{"react.dev", "github.com"})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(filtered) != 2 {
t.Fatalf("filtered len = %d, want 2; got %v", len(filtered), filtered)
}
// Empty allowlist must be rejected when called directly with no rows that match.
_, err = filterWebSearchByDomains([]searchResult{{URL: "https://x.test"}}, []string{"nope.test"})
if err == nil {
t.Fatal("expected error when allowlist is non-empty and no rows match")
}
}
func TestWebSearchSchemaHasDomainsField(t *testing.T) {
tool := NewWebSearchTool()
domains := tool.Parameters().Properties["domains"]
if domains.Type != "array" {
t.Fatalf("domains.Type = %q, want array", domains.Type)
}
if domains.Items == nil || domains.Items.Type != "string" {
t.Fatalf("domains.Items = %#v, want array of string", domains.Items)
}
if domains.Description == "" {
t.Fatal("domains.Description should explain the prompt-injection defense")
}
}
// TestWebSearchDomainsFilterRejectsAllInvalidInputs covers the fail-closed
// contract: when the caller passes a 'domains' argument but every entry is
// invalid (e.g. all whitespace, all with embedded spaces), the tool must
// error rather than silently return unfiltered results. The whole point of
// the parameter is the prompt-injection defense; an unfiltered result set
// would defeat it.
func TestWebSearchDomainsFilterRejectsAllInvalidInputs(t *testing.T) {
backend := &fakeSearchBackend{results: []searchResult{
{Title: "x", URL: "https://react.dev/x"},
}}
tool := newWebSearchToolWithBackend(backend)
cases := []struct {
name string
domains []any
}{
{"all whitespace", []any{" ", "\t", "\n"}},
{"empty strings", []any{"", ""}},
{"embedded spaces", []any{"react .dev", "github .com"}},
{"mixed invalid", []any{"", " ", "react .dev"}},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
res := tool.Run(context.Background(), map[string]any{
"query": "x",
"domains": c.domains,
})
if res.Status != StatusError {
t.Fatalf("expected error when all domains are invalid, got %v: %s", res.Status, res.Output)
}
if !strings.Contains(res.Output, "no valid hostnames") {
t.Errorf("expected 'no valid hostnames' in error, got: %s", res.Output)
}
if backend.gotQuery != "" {
t.Errorf("backend must not have been called when domains are invalid; got query %q", backend.gotQuery)
}
})
}
}
// TestWebSearchDomainsFilterAcceptsHostPort covers the canonicalization
// contract: an allowlist entry that includes a port WITH a scheme (e.g. the
// result of a model passing "https://react.dev:443/path") must still match
// a result URL on the same hostname. The CodeRabbit-flagged bug was that
// the canonicalizer was using parsed.Host (which includes ":443") rather
// than parsed.Hostname() (which strips the port), so a "react.dev" allowlist
// silently never matched results from "https://react.dev:443/...". The fix
// uses Hostname(), so this test locks in the corrected contract.
func TestWebSearchDomainsFilterAcceptsHostPort(t *testing.T) {
cases := []struct {
name string
allowlist []any
expectMatch bool
}{
{"with explicit https port", []any{"https://react.dev:443/path"}, true},
{"with explicit http port", []any{"http://react.dev:80/x"}, true},
{"scheme-less host:port is rejected", []any{"react.dev:9999"}, false}, // not a valid hostname
{"port on different host", []any{"github.com:443"}, false},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
backend := &fakeSearchBackend{results: []searchResult{
{Title: "ok", URL: "https://react.dev/x"},
}}
tool := newWebSearchToolWithBackend(backend)
res := tool.Run(context.Background(), map[string]any{
"query": "x",
"domains": c.allowlist,
})
if c.expectMatch {
if res.Status != StatusOK {
t.Fatalf("expected match, got %v: %s", res.Status, res.Output)
}
if !strings.Contains(res.Output, "react.dev/x") {
t.Errorf("expected react.dev result, got: %s", res.Output)
}
} else {
if res.Status != StatusError {
t.Fatalf("expected no-match error, got %v: %s", res.Status, res.Output)
}
}
})
}
}
func TestCanonicalizeWebSearchHostStripsPort(t *testing.T) {
// Lock in the contract that parsed.Host (which includes ":443") is NOT
// used; parsed.Hostname() is. CodeRabbit flagged this inconsistency.
cases := []struct {
in, want string
}{
{"https://react.dev:443/x", "react.dev"},
{"http://react.dev:80", "react.dev"},
{"https://api.react.dev:8443/v1", "api.react.dev"},
{"https://react.dev", "react.dev"},
}
for _, c := range cases {
if got := canonicalizeWebSearchHost(c.in); got != c.want {
t.Errorf("canonicalizeWebSearchHost(%q) = %q, want %q", c.in, got, c.want)
}
}
}
+9 -9
View File
@@ -170,28 +170,28 @@ func (m model) syncMouseCapture() (model, tea.Cmd) {
return m, tea.DisableMouse
}
// Bubble Tea's Type field is deprecated, but its parser still populates it for
// compatibility cases such as left-button drag events. Keep these helpers
// tolerant of both the current Button/Action pair and legacy Type values.
// Mouse classification uses the current Button/Action pair only. Bubble Tea's
// parser always populates Button+Action and merely derives the deprecated Type
// field from them, so checking Type adds nothing — and a left-button drag is
// Action==Motion (which mouseMotion already covers), not a press.
func mouseLeftPress(msg tea.MouseMsg) bool {
return msg.Button == tea.MouseButtonLeft && msg.Action == tea.MouseActionPress ||
msg.Type == tea.MouseLeft && msg.Action == tea.MouseActionPress
return msg.Button == tea.MouseButtonLeft && msg.Action == tea.MouseActionPress
}
func mouseMotion(msg tea.MouseMsg) bool {
return msg.Action == tea.MouseActionMotion || msg.Type == tea.MouseMotion
return msg.Action == tea.MouseActionMotion
}
func mouseRelease(msg tea.MouseMsg) bool {
return msg.Action == tea.MouseActionRelease || msg.Type == tea.MouseRelease
return msg.Action == tea.MouseActionRelease
}
func mouseWheelUp(msg tea.MouseMsg) bool {
return msg.Button == tea.MouseButtonWheelUp || msg.Type == tea.MouseWheelUp
return msg.Button == tea.MouseButtonWheelUp
}
func mouseWheelDown(msg tea.MouseMsg) bool {
return msg.Button == tea.MouseButtonWheelDown || msg.Type == tea.MouseWheelDown
return msg.Button == tea.MouseButtonWheelDown
}
func (m model) mouseOverComposer(msg tea.MouseMsg) bool {
+2 -3
View File
@@ -348,12 +348,11 @@ func TestTranscriptSelectionLeftDragDoesNotResetAnchor(t *testing.T) {
Y: 1,
})
m = updated.(model)
// Bubble Tea marks left-button drag motion as Type MouseLeft for backward
// compatibility; this must update the cursor without resetting the anchor.
// A left-button drag is Action==Motion with Button==Left; this must update the
// cursor without resetting the selection anchor.
updated, _ = m.Update(tea.MouseMsg{
Button: tea.MouseButtonLeft,
Action: tea.MouseActionMotion,
Type: tea.MouseLeft,
X: 8,
Y: 1,
})
+1 -1
View File
@@ -635,7 +635,7 @@ func TestProviderWizardPersistsPastedKeyToUserConfig(t *testing.T) {
updated, _ = next.Update(tea.KeyMsg{Type: tea.KeyEnter})
next = updated.(model)
updated, _ = next.Update(tea.KeyMsg{Type: tea.KeyEnter})
_ = updated.(model)
next = updated.(model)
if captured.APIKey != secret {
t.Fatalf("captured APIKey = %q, want pasted secret", captured.APIKey)