mirror of
https://github.com/Twigpine/zero.git
synced 2026-11-08 09:49:34 -10:00
Compare commits
4
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
8b7b8296ac | ||
|
|
f857c4afb4 | ||
|
|
bb7de29c80 | ||
|
|
de5d402e1f |
No files matched your search
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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) }
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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 ""
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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{Title: item.Title, URL: item.URL, Snippet: item.Snippet})
|
||||
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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
|
||||
Reference in new issue
Block a user