diff --git a/README.md b/README.md index 24928cc..f499366 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ [![Actions Status](https://github.com/epithet-ssh/epithet/workflows/build/badge.svg)](https://github.com/epithet-ssh/epithet/actions) [![Go Reportcard](https://goreportcard.com/badge/github.com/epithet-ssh/epithet)](https://goreportcard.com/report/github.com/epithet-ssh/epithet) -Epithet is an SSH certificate authority that replaces static authorized_keys with short-lived certificates (2-10 minutes). It creates on-demand SSH agents for each outbound connection, enabling real-time policy enforcement without touching your target hosts. +Epithet is an SSH agent and certificate authority that replaces static authorized_keys with short-lived certificates (2-10 minutes). It creates on-demand SSH agents for each outbound connection, enabling real-time policy enforcement without touching your target hosts. ## Quick start diff --git a/cmd/epithet/agent.go b/cmd/epithet/agent.go index 7a55b00..d8cba39 100644 --- a/cmd/epithet/agent.go +++ b/cmd/epithet/agent.go @@ -4,9 +4,11 @@ import ( "context" "crypto/sha256" "encoding/hex" + "errors" "fmt" "log/slog" "os" + "os/exec" "os/signal" "path/filepath" "strconv" @@ -14,6 +16,7 @@ import ( "syscall" "time" + "github.com/epithet-ssh/epithet/pkg/agent" "github.com/epithet-ssh/epithet/pkg/broker" "github.com/epithet-ssh/epithet/pkg/caclient" "github.com/epithet-ssh/epithet/pkg/tlsconfig" @@ -32,74 +35,273 @@ type AgentCLI struct { } // AgentStartCLI is the default subcommand that starts the agent/broker. -type AgentStartCLI struct{} +// When Shell is provided, it runs in wrapper mode: wrapping a shell with an +// epithet-aware agent proxy (mirroring ssh-agent's "ssh-agent bash" pattern). +type AgentStartCLI struct { + Shell string `arg:"" optional:"" help:"Shell to wrap with epithet agent (enables wrapper mode)"` +} func (s *AgentStartCLI) Run(parent *AgentCLI, logger *slog.Logger, tlsCfg tlsconfig.Config) error { - // Validate required fields for start + if s.Shell != "" { + return s.runWrapper(parent, logger, tlsCfg) + } + return s.runDaemon(parent, logger, tlsCfg) +} + +// runDaemon runs the broker as a long-lived daemon (original behavior). +func (s *AgentStartCLI) runDaemon(parent *AgentCLI, logger *slog.Logger, tlsCfg tlsconfig.Config) error { + b, tempDir, brokerSock, agentDir, homeDir, err := setupBrokerWithOptions(parent, logger, tlsCfg, false) + if err != nil { + return err + } + + // Clean up temp directory on exit. + defer func() { + if err := os.RemoveAll(tempDir); err != nil { + logger.Warn("failed to remove temp directory", "error", err, "path", tempDir) + } else { + logger.Debug("removed temp directory", "path", tempDir) + } + }() + + // Set up context with cancellation on signals. + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + sigChan := make(chan os.Signal, 1) + signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM) + go func() { + <-sigChan + logger.Info("received shutdown signal") + cancel() + }() + + // Generate SSH config. + sshConfigPath := filepath.Join(tempDir, "ssh-config.conf") + if err := parent.generateSSHConfig(sshConfigPath, agentDir, brokerSock, homeDir); err != nil { + logger.Warn("failed to generate SSH config", "error", err, "path", sshConfigPath) + } else { + includePattern := filepath.Join(homeDir, ".epithet", "run", "*", "ssh-config.conf") + if err := checkSSHConfigInclude(homeDir, includePattern, logger); err != nil { + logger.Warn(fmt.Sprintf("Add 'Include %s' to ~/.ssh/config", includePattern)) + } + logger.Debug("generated SSH config", "path", sshConfigPath) + } + + // Start broker. + logger.Info("starting broker", "socket", brokerSock) + err = b.Serve(ctx) + if err != nil && err != context.Canceled { + return fmt.Errorf("broker serve error: %w", err) + } + + logger.Info("broker shutdown complete") + return nil +} + +// runWrapper wraps a shell with an epithet-aware agent proxy. It probes the +// current SSH_AUTH_SOCK for an upstream epithet agent, starts the broker, and +// creates a proxy listener that handles epithet protocol extensions while +// delegating standard SSH agent operations to the upstream agent. +func (s *AgentStartCLI) runWrapper(parent *AgentCLI, logger *slog.Logger, tlsCfg tlsconfig.Config) error { + upstreamSocket := os.Getenv("SSH_AUTH_SOCK") + + // Probe for upstream epithet agent. + var depth int + var brokerOpts []broker.Option + if upstreamSocket != "" { + hello, err := agent.ProbeUpstream(upstreamSocket) + if err != nil { + logger.Warn("failed to probe upstream agent", "error", err) + upstreamSocket = "" + } else if hello != nil { + depth = hello.ChainDepth + 1 + brokerOpts = append(brokerOpts, broker.WithUpstream(upstreamSocket)) + logger.Info("detected upstream epithet agent", "chain_depth", depth) + } else { + logger.Debug("upstream is vanilla ssh-agent") + } + } else { + logger.Debug("no SSH_AUTH_SOCK set") + } + + b, tempDir, brokerSock, agentDir, homeDir, err := setupBrokerWithOptions(parent, logger, tlsCfg, true, brokerOpts...) + if err != nil { + return err + } + + // Clean up temp directory on exit. + defer func() { + if err := os.RemoveAll(tempDir); err != nil { + logger.Warn("failed to remove temp directory", "error", err, "path", tempDir) + } + }() + + // Set up context with cancellation. + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // Start broker in background. + brokerErr := make(chan error, 1) + go func() { + brokerErr <- b.Serve(ctx) + }() + if err := waitForBrokerStartup(b, brokerErr); err != nil { + return err + } + + // Generate SSH config. + sshConfigPath := filepath.Join(tempDir, "ssh-config.conf") + if err := parent.generateSSHConfig(sshConfigPath, agentDir, brokerSock, homeDir); err != nil { + logger.Warn("failed to generate SSH config", "error", err, "path", sshConfigPath) + } else { + includePattern := filepath.Join(homeDir, ".epithet", "run", "*", "ssh-config.conf") + if err := checkSSHConfigInclude(homeDir, includePattern, logger); err != nil { + logger.Warn(fmt.Sprintf("Add 'Include %s' to ~/.ssh/config", includePattern)) + } + } + + proxySock := filepath.Join(tempDir, "proxy.sock") + setup := func(p *agent.ProxyAgent) { + p.RegisterExtension(agent.ExtensionHello, agent.HelloHandler(depth)) + p.RegisterExtension(agent.ExtensionAuth, agent.AuthHandler(func() (string, error) { + return b.Authenticate(nil) + })) + } + proxyListener := agent.NewProxyListener(logger, proxySock, upstreamSocket, setup) + go func() { + if err := proxyListener.Serve(ctx); err != nil && err != context.Canceled { + logger.Error("proxy listener error", "error", err) + } + }() + <-proxyListener.Ready() + upstreamSocket = proxySock + logger.Info("proxy agent listening", "socket", proxySock) + + // Build child environment with updated SSH_AUTH_SOCK. + childEnv := os.Environ() + if upstreamSocket != "" { + childEnv = replaceEnv(childEnv, "SSH_AUTH_SOCK", upstreamSocket) + } + + // Run the shell as a child process. + cmd := exec.Command(s.Shell) + cmd.Env = childEnv + cmd.Stdin = os.Stdin + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + + // Forward signals to child. + sigChan := make(chan os.Signal, 1) + signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM) + go func() { + for sig := range sigChan { + if cmd.Process != nil { + cmd.Process.Signal(sig) + } + } + }() + + logger.Info("starting shell", "shell", s.Shell) + shellErr := cmd.Run() + + // Shell exited — shut down broker and drain its error. + cancel() + signal.Stop(sigChan) + if err := <-brokerErr; err != nil && err != context.Canceled { + logger.Error("broker error", "error", err) + } + + if shellErr != nil { + var exitErr *exec.ExitError + if errors.As(shellErr, &exitErr) { + os.Exit(exitErr.ExitCode()) + } + return fmt.Errorf("shell failed: %w", shellErr) + } + + return nil +} + +func waitForBrokerStartup(b *broker.Broker, brokerErr <-chan error) error { + select { + case <-b.Ready(): + return nil + case err := <-brokerErr: + if err == nil || err == context.Canceled { + return fmt.Errorf("broker exited before becoming ready") + } + return fmt.Errorf("broker failed to start: %w", err) + } +} + +// setupBrokerWithOptions creates the broker, temp directories, and resolves auth. +// When wrapperMode is true, uses a per-process temp directory to avoid collisions +// between multiple wrapped shells. When false (daemon mode), uses a deterministic +// path derived from the CA URLs so SSH config includes can find it. +// Returns (broker, tempDir, brokerSock, agentDir, homeDir, error). +func setupBrokerWithOptions(parent *AgentCLI, logger *slog.Logger, tlsCfg tlsconfig.Config, wrapperMode bool, opts ...broker.Option) (*broker.Broker, string, string, string, string, error) { if len(parent.CaURL) == 0 { - return fmt.Errorf("--ca-url is required (at least one)") + return nil, "", "", "", "", fmt.Errorf("--ca-url is required (at least one)") } - // Parse CA URLs into endpoints caEndpoints, err := caclient.ParseCAURLs(parent.CaURL) if err != nil { - return fmt.Errorf("invalid CA URL: %w", err) + return nil, "", "", "", "", fmt.Errorf("invalid CA URL: %w", err) } - logger.Debug("agent start command received", "ca_urls", parent.CaURL, "ca_timeout", parent.CaTimeout, "ca_cooldown", parent.CaCooldown) + logger.Debug("agent command received", "ca_urls", parent.CaURL, "ca_timeout", parent.CaTimeout, "ca_cooldown", parent.CaCooldown) - // Validate all CA URLs require TLS (unless --insecure) for _, ep := range caEndpoints { if err := tlsCfg.ValidateURL(ep.URL); err != nil { - return fmt.Errorf("CA URL %q: %w", ep.URL, err) + return nil, "", "", "", "", fmt.Errorf("CA URL %q: %w", ep.URL, err) } } - // Get home directory homeDir, err := os.UserHomeDir() if err != nil { - return fmt.Errorf("failed to get home directory: %w", err) + return nil, "", "", "", "", fmt.Errorf("failed to get home directory: %w", err) } - // Create a unique temporary directory for this broker instance - // Use a hash of the CA URLs to make it deterministic - instanceID := hashString(fmt.Sprintf("%v", parent.CaURL)) runDir := filepath.Join(homeDir, ".epithet", "run") - tempDir := filepath.Join(runDir, instanceID) - - // Clean up stale run directories from dead processes cleanupStaleRunDirs(runDir, logger) - // Clean up temp directory on exit - defer func() { - if err := os.RemoveAll(tempDir); err != nil { - logger.Warn("failed to remove temp directory", "error", err, "path", tempDir) - } else { - logger.Debug("removed temp directory", "path", tempDir) + var tempDir string + if wrapperMode { + // Wrapper mode: each invocation gets its own directory to avoid + // collisions between multiple wrapped shells against the same CA. + tempDir, err = os.MkdirTemp(runDir, "wrap-") + if err != nil { + // Ensure runDir exists and retry. + if mkErr := os.MkdirAll(runDir, 0700); mkErr != nil { + return nil, "", "", "", "", fmt.Errorf("failed to create run directory: %w", mkErr) + } + tempDir, err = os.MkdirTemp(runDir, "wrap-") + if err != nil { + return nil, "", "", "", "", fmt.Errorf("failed to create temp directory: %w", err) + } + } + } else { + // Daemon mode: deterministic path so SSH config includes can find it. + instanceID := hashString(fmt.Sprintf("%v", parent.CaURL)) + tempDir = filepath.Join(runDir, instanceID) + if err := os.MkdirAll(tempDir, 0700); err != nil { + return nil, "", "", "", "", fmt.Errorf("failed to create temp directory: %w", err) } - }() - - // Create temp directory - if err := os.MkdirAll(tempDir, 0700); err != nil { - return fmt.Errorf("failed to create temp directory: %w", err) } - // Write PID file for stale detection pidFile := filepath.Join(tempDir, "broker.pid") if err := os.WriteFile(pidFile, []byte(strconv.Itoa(os.Getpid())), 0600); err != nil { - return fmt.Errorf("failed to write PID file: %w", err) + return nil, "", "", "", "", fmt.Errorf("failed to write PID file: %w", err) } - // Define paths within temp directory brokerSock := filepath.Join(tempDir, "broker.sock") agentDir := filepath.Join(tempDir, "agent") - // Create agent directory if err := os.MkdirAll(agentDir, 0700); err != nil { - return fmt.Errorf("failed to create agent directory: %w", err) + return nil, "", "", "", "", fmt.Errorf("failed to create agent directory: %w", err) } - // Create CA client caClientOpts := []caclient.Option{ caclient.WithLogger(logger), caclient.WithTimeout(parent.CaTimeout), @@ -108,81 +310,60 @@ func (s *AgentStartCLI) Run(parent *AgentCLI, logger *slog.Logger, tlsCfg tlscon } caClient, err := caclient.New(caEndpoints, caClientOpts...) if err != nil { - return fmt.Errorf("failed to create CA client: %w", err) + return nil, "", "", "", "", fmt.Errorf("failed to create CA client: %w", err) } - // Resolve auth command: use --auth if provided, otherwise discover from CA. authCommand := parent.Auth if authCommand == "" { logger.Debug("no --auth provided, discovering from CA") - // Fetch public key (also caches discovery URL). _, err := caClient.GetPublicKey(context.Background()) if err != nil { - return fmt.Errorf("failed to get CA public key: %w", err) + return nil, "", "", "", "", fmt.Errorf("failed to get CA public key: %w", err) } - // Fetch unauthenticated discovery (auth config only, no token needed). discovery, err := caClient.GetDiscovery(context.Background(), "") if err != nil { - return fmt.Errorf("failed to get discovery config (try --auth to specify manually): %w", err) + return nil, "", "", "", "", fmt.Errorf("failed to get discovery config (try --auth to specify manually): %w", err) } if discovery == nil || discovery.Auth == nil { - return fmt.Errorf("CA did not return auth config in discovery (try --auth to specify manually)") + return nil, "", "", "", "", fmt.Errorf("CA did not return auth config in discovery (try --auth to specify manually)") } - // Convert auth config to command. authCommand, err = broker.AuthConfigToCommand(*discovery.Auth) if err != nil { - return fmt.Errorf("failed to convert discovery auth config: %w", err) + return nil, "", "", "", "", fmt.Errorf("failed to convert discovery auth config: %w", err) } logger.Info("discovered auth config from CA", "type", discovery.Auth.Type, "issuer", discovery.Auth.Issuer) } - // Create broker - b, err := broker.New(*logger, brokerSock, authCommand, caClient, agentDir) + b, err := broker.New(*logger, brokerSock, authCommand, caClient, agentDir, opts...) if err != nil { - return fmt.Errorf("failed to create broker: %w", err) + return nil, "", "", "", "", fmt.Errorf("failed to create broker: %w", err) } - // Set up context with cancellation on signals - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - sigChan := make(chan os.Signal, 1) - signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM) - go func() { - <-sigChan - logger.Info("received shutdown signal") - cancel() - }() - - // Generate SSH config file in the temp directory - sshConfigPath := filepath.Join(tempDir, "ssh-config.conf") - - if err := parent.generateSSHConfig(sshConfigPath, agentDir, brokerSock, homeDir); err != nil { - logger.Warn("failed to generate SSH config", "error", err, "path", sshConfigPath) - // Don't fail startup, just warn - } else { - // Check if ~/.ssh/config has the Include directive - includePattern := filepath.Join(homeDir, ".epithet", "run", "*", "ssh-config.conf") - if err := checkSSHConfigInclude(homeDir, includePattern, logger); err != nil { + return b, tempDir, brokerSock, agentDir, homeDir, nil +} - logger.Warn(fmt.Sprintf("Add 'Include %s' to ~/.ssh/config", includePattern)) +// replaceEnv returns a copy of environ with key set to value. +// If key already exists, it is replaced; otherwise it is appended. +func replaceEnv(environ []string, key, value string) []string { + prefix := key + "=" + result := make([]string, 0, len(environ)+1) + found := false + for _, e := range environ { + if strings.HasPrefix(e, prefix) { + result = append(result, prefix+value) + found = true + } else { + result = append(result, e) } - logger.Debug("generated SSH config", "path", sshConfigPath) } - - // Start broker - logger.Info("starting broker", "socket", brokerSock) - err = b.Serve(ctx) - if err != nil && err != context.Canceled { - return fmt.Errorf("broker serve error: %w", err) + if !found { + result = append(result, prefix+value) } - - logger.Info("broker shutdown complete") - return nil + return result } // generateSSHConfig writes an SSH config file for epithet diff --git a/cmd/epithet/agent_test.go b/cmd/epithet/agent_test.go new file mode 100644 index 0000000..af1c863 --- /dev/null +++ b/cmd/epithet/agent_test.go @@ -0,0 +1,99 @@ +package main + +import ( + "context" + "io" + "log/slog" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/epithet-ssh/epithet/pkg/broker" + "github.com/epithet-ssh/epithet/pkg/caclient" +) + +func TestReplaceEnv_ExistingKey(t *testing.T) { + env := []string{"HOME=/home/user", "SSH_AUTH_SOCK=/old/path", "SHELL=/bin/fish"} + result := replaceEnv(env, "SSH_AUTH_SOCK", "/new/path") + + found := false + for _, e := range result { + if e == "SSH_AUTH_SOCK=/new/path" { + found = true + } + if e == "SSH_AUTH_SOCK=/old/path" { + t.Error("old value still present") + } + } + if !found { + t.Error("new value not found") + } + if len(result) != 3 { + t.Errorf("expected 3 entries, got %d", len(result)) + } +} + +func TestReplaceEnv_NewKey(t *testing.T) { + env := []string{"HOME=/home/user", "SHELL=/bin/fish"} + result := replaceEnv(env, "SSH_AUTH_SOCK", "/new/path") + + found := false + for _, e := range result { + if e == "SSH_AUTH_SOCK=/new/path" { + found = true + } + } + if !found { + t.Error("new key not appended") + } + if len(result) != 3 { + t.Errorf("expected 3 entries, got %d", len(result)) + } +} + +func TestReplaceEnv_EmptyEnv(t *testing.T) { + result := replaceEnv(nil, "KEY", "value") + if len(result) != 1 || result[0] != "KEY=value" { + t.Errorf("expected [KEY=value], got %v", result) + } +} + +func TestWaitForBrokerStartup_ReturnsServeError(t *testing.T) { + t.Parallel() + + tempDir := t.TempDir() + agentDir := filepath.Join(tempDir, "agent") + if err := os.MkdirAll(agentDir, 0700); err != nil { + t.Fatalf("mkdir agent dir: %v", err) + } + + endpoints := []caclient.CAEndpoint{{URL: "http://127.0.0.1", Priority: caclient.DefaultPriority}} + caClient, err := caclient.New(endpoints) + if err != nil { + t.Fatalf("create CA client: %v", err) + } + + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + socketPath := filepath.Join(tempDir, strings.Repeat("sock", 40)) + b, err := broker.New(*logger, socketPath, "true", caClient, agentDir) + if err != nil { + t.Fatalf("create broker: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + brokerErr := make(chan error, 1) + go func() { + brokerErr <- b.Serve(ctx) + }() + + err = waitForBrokerStartup(b, brokerErr) + if err == nil { + t.Fatal("expected startup error, got nil") + } + if !strings.Contains(err.Error(), "broker failed to start") { + t.Fatalf("expected startup failure wrapper, got %v", err) + } +} diff --git a/pkg/agent/extensions.go b/pkg/agent/extensions.go new file mode 100644 index 0000000..5d08bc4 --- /dev/null +++ b/pkg/agent/extensions.go @@ -0,0 +1,53 @@ +package agent + +import ( + "encoding/json" + "fmt" +) + +// Extension type constants follow the OpenSSH vendor extension naming convention. +const ( + ExtensionHello = "epithet-hello@epithet.dev" + ExtensionAuth = "epithet-auth@epithet.dev" +) + +// HelloResponse is returned by the epithet-hello extension. +type HelloResponse struct { + ProtocolVersion int `json:"protocol_version"` + ChainDepth int `json:"chain_depth"` +} + +// AuthResponse is returned by the epithet-auth extension. +type AuthResponse struct { + Token string `json:"token"` +} + +// HelloHandler returns an ExtensionHandler for epithet-hello@epithet.dev. +func HelloHandler(depth int) ExtensionHandler { + return func(_ []byte) ([]byte, error) { + return json.Marshal(HelloResponse{ + ProtocolVersion: 1, + ChainDepth: depth, + }) + } +} + +// AuthHandler returns an ExtensionHandler for epithet-auth@epithet.dev. +// The authenticate function is called to obtain a token — typically wired +// to the broker's auth flow. +// +// Security: the returned token is a bearer credential that transits the SSH +// agent forwarding channel. Any host in the forwarding chain (i.e., the +// remote sshd process) can observe the token. This is the same trust model +// as standard SSH agent forwarding — the remote host can use the forwarded +// agent to sign challenges on the user's behalf. Users should only forward +// agents to hosts they trust, same as with ssh-agent -A. +func AuthHandler(authenticate func() (string, error)) ExtensionHandler { + return func(_ []byte) ([]byte, error) { + token, err := authenticate() + if err != nil { + return nil, fmt.Errorf("auth failed: %w", err) + } + return json.Marshal(AuthResponse{Token: token}) + } +} diff --git a/pkg/agent/proxy.go b/pkg/agent/proxy.go new file mode 100644 index 0000000..dbf953c --- /dev/null +++ b/pkg/agent/proxy.go @@ -0,0 +1,85 @@ +package agent + +import ( + "golang.org/x/crypto/ssh" + "golang.org/x/crypto/ssh/agent" +) + +const agentSuccess = 6 + +// ExtensionHandler handles a single epithet protocol extension. +// contents is the raw request payload; returns raw response payload (without the type byte). +type ExtensionHandler func(contents []byte) ([]byte, error) + +// ProxyAgent wraps an upstream ssh agent, proxying standard operations and +// dispatching epithet protocol extensions to registered handlers. +// Unregistered extensions are forwarded to the upstream agent. +type ProxyAgent struct { + upstream agent.ExtendedAgent + extensions map[string]ExtensionHandler +} + +// NewProxyAgent creates a ProxyAgent that proxies to the given upstream agent. +func NewProxyAgent(upstream agent.ExtendedAgent) *ProxyAgent { + return &ProxyAgent{ + upstream: upstream, + extensions: make(map[string]ExtensionHandler), + } +} + +// RegisterExtension registers a handler for the given extension type. +func (p *ProxyAgent) RegisterExtension(name string, handler ExtensionHandler) { + p.extensions[name] = handler +} + +// Extension dispatches to a registered handler or forwards to upstream. +func (p *ProxyAgent) Extension(extensionType string, contents []byte) ([]byte, error) { + if handler, ok := p.extensions[extensionType]; ok { + payload, err := handler(contents) + if err != nil { + return nil, err + } + // Prepend SSH_AGENT_SUCCESS byte — ServeAgent sends this as-is, + // and the client strips the type byte before returning. + return append([]byte{agentSuccess}, payload...), nil + } + return p.upstream.Extension(extensionType, contents) +} + +// Standard agent methods — all delegate to upstream. + +func (p *ProxyAgent) List() ([]*agent.Key, error) { + return p.upstream.List() +} + +func (p *ProxyAgent) Sign(key ssh.PublicKey, data []byte) (*ssh.Signature, error) { + return p.upstream.Sign(key, data) +} + +func (p *ProxyAgent) Add(key agent.AddedKey) error { + return p.upstream.Add(key) +} + +func (p *ProxyAgent) Remove(key ssh.PublicKey) error { + return p.upstream.Remove(key) +} + +func (p *ProxyAgent) RemoveAll() error { + return p.upstream.RemoveAll() +} + +func (p *ProxyAgent) Lock(passphrase []byte) error { + return p.upstream.Lock(passphrase) +} + +func (p *ProxyAgent) Unlock(passphrase []byte) error { + return p.upstream.Unlock(passphrase) +} + +func (p *ProxyAgent) Signers() ([]ssh.Signer, error) { + return p.upstream.Signers() +} + +func (p *ProxyAgent) SignWithFlags(key ssh.PublicKey, data []byte, flags agent.SignatureFlags) (*ssh.Signature, error) { + return p.upstream.SignWithFlags(key, data, flags) +} diff --git a/pkg/agent/proxy_listener.go b/pkg/agent/proxy_listener.go new file mode 100644 index 0000000..ef17303 --- /dev/null +++ b/pkg/agent/proxy_listener.go @@ -0,0 +1,135 @@ +package agent + +import ( + "context" + "errors" + "io" + "log/slog" + "net" + "os" + "sync" + + "golang.org/x/crypto/ssh/agent" +) + +// ProxyListener accepts connections on a Unix socket and serves each with a +// ProxyAgent that proxies to an upstream agent. Each connection gets a fresh +// upstream dial (the SSH agent protocol is connection-scoped). +// +// The setup callback is called for each new ProxyAgent, allowing the caller +// to register extension handlers. +type ProxyListener struct { + socketPath string + upstreamSocketPath string + setup func(*ProxyAgent) + log *slog.Logger + + listener net.Listener + ready chan struct{} // Closed when listener is accepting connections. + done chan struct{} + closeOnce sync.Once +} + +// NewProxyListener creates a ProxyListener. Call Serve() to begin accepting connections. +func NewProxyListener(log *slog.Logger, socketPath, upstreamSocketPath string, setup func(*ProxyAgent)) *ProxyListener { + return &ProxyListener{ + socketPath: socketPath, + upstreamSocketPath: upstreamSocketPath, + setup: setup, + log: log, + ready: make(chan struct{}), + done: make(chan struct{}), + } +} + +// Serve starts the listener and blocks until the context is cancelled. +// The ready channel is always closed before Serve returns, whether +// successfully or on error, so callers blocking on Ready() won't hang. +func (p *ProxyListener) Serve(ctx context.Context) error { + os.Remove(p.socketPath) + listener, err := net.Listen("unix", p.socketPath) + if err != nil { + close(p.ready) + return err + } + if err := os.Chmod(p.socketPath, 0600); err != nil { + listener.Close() + close(p.ready) + return err + } + p.listener = listener + close(p.ready) + + go p.acceptLoop() + + <-ctx.Done() + p.Close() + return ctx.Err() +} + +func (p *ProxyListener) acceptLoop() { + for { + conn, err := p.listener.Accept() + if err != nil { + if errors.Is(err, net.ErrClosed) { + return + } + select { + case <-p.done: + return + default: + p.log.Warn("proxy listener accept error", "error", err) + continue + } + } + go p.serveConn(conn) + } +} + +func (p *ProxyListener) serveConn(conn net.Conn) { + defer conn.Close() + + var upstream agent.ExtendedAgent + if p.upstreamSocketPath == "" { + // Preserve standard ssh-agent behavior even without an upstream socket. + upstream = agent.NewKeyring().(agent.ExtendedAgent) + } else { + // Dial a fresh connection to the upstream agent for this client. + upstreamConn, err := net.Dial("unix", p.upstreamSocketPath) + if err != nil { + p.log.Warn("failed to connect to upstream agent", "error", err) + return + } + defer upstreamConn.Close() + + upstream = agent.NewClient(upstreamConn) + } + + proxy := NewProxyAgent(upstream) + p.setup(proxy) + + err := agent.ServeAgent(proxy, conn) + if err != nil && err != io.EOF { + p.log.Debug("proxy agent connection ended", "error", err) + } +} + +// Ready returns a channel that is closed when the listener is accepting connections. +func (p *ProxyListener) Ready() <-chan struct{} { + return p.ready +} + +// SocketPath returns the path to the proxy listener's Unix socket. +func (p *ProxyListener) SocketPath() string { + return p.socketPath +} + +// Close stops the listener and cleans up. +func (p *ProxyListener) Close() { + p.closeOnce.Do(func() { + if p.listener != nil { + p.listener.Close() + } + close(p.done) + }) +} diff --git a/pkg/agent/proxy_listener_test.go b/pkg/agent/proxy_listener_test.go new file mode 100644 index 0000000..2d18fd7 --- /dev/null +++ b/pkg/agent/proxy_listener_test.go @@ -0,0 +1,345 @@ +package agent + +import ( + "context" + "encoding/json" + "log/slog" + "net" + "os" + "sync/atomic" + "testing" + + "golang.org/x/crypto/ssh/agent" +) + +// startVanillaAgent starts a plain ssh-agent keyring on a Unix socket. +func startVanillaAgent(t *testing.T) (string, func()) { + t.Helper() + + f, err := os.CreateTemp("", "vanilla-agent.*") + if err != nil { + t.Fatal(err) + } + sockPath := f.Name() + f.Close() + os.Remove(sockPath) + + listener, err := net.Listen("unix", sockPath) + if err != nil { + t.Fatal(err) + } + + keyring := agent.NewKeyring() + go func() { + for { + conn, err := listener.Accept() + if err != nil { + return + } + go func(c net.Conn) { + defer c.Close() + agent.ServeAgent(keyring, c) + }(conn) + } + }() + + return sockPath, func() { + listener.Close() + os.Remove(sockPath) + } +} + +func TestProxyListenerEndToEnd(t *testing.T) { + // Set up a vanilla upstream agent. + upstreamSock, cleanupUpstream := startVanillaAgent(t) + defer cleanupUpstream() + + // Create proxy listener that adds epithet extensions. + proxySock := tempSocketPath(t) + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) + + var authCalled atomic.Bool + setup := func(p *ProxyAgent) { + p.RegisterExtension(ExtensionHello, HelloHandler(0)) + p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) { + authCalled.Store(true) + return "proxy-listener-token", nil + })) + } + proxy := NewProxyListener(logger, proxySock, upstreamSock, setup) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + go proxy.Serve(ctx) + <-proxy.Ready() + + // Connect to the proxy and test extensions. + conn, err := net.Dial("unix", proxySock) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + + client := agent.NewClient(conn).(agent.ExtendedAgent) + + // Test hello. + resp, err := client.Extension(ExtensionHello, nil) + if err != nil { + t.Fatalf("hello failed: %v", err) + } + var hello HelloResponse + if err := json.Unmarshal(stripSuccessByte(t, resp), &hello); err != nil { + t.Fatalf("unmarshal hello: %v", err) + } + if hello.ProtocolVersion != 1 { + t.Errorf("expected protocol_version=1, got %d", hello.ProtocolVersion) + } + + // Test auth. + resp, err = client.Extension(ExtensionAuth, nil) + if err != nil { + t.Fatalf("auth failed: %v", err) + } + var authResp AuthResponse + if err := json.Unmarshal(stripSuccessByte(t, resp), &authResp); err != nil { + t.Fatalf("unmarshal auth: %v", err) + } + if authResp.Token != "proxy-listener-token" { + t.Errorf("expected 'proxy-listener-token', got %q", authResp.Token) + } + if !authCalled.Load() { + t.Error("auth handler was not called") + } + + // Test standard agent ops pass through to upstream. + keys, err := client.List() + if err != nil { + t.Fatalf("List failed: %v", err) + } + if len(keys) != 0 { + t.Errorf("expected 0 keys, got %d", len(keys)) + } +} + +func TestProxyListenerMultipleConnections(t *testing.T) { + // Each connection should get its own upstream dial. + upstreamSock, cleanupUpstream := startVanillaAgent(t) + defer cleanupUpstream() + + proxySock := tempSocketPath(t) + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) + + var callCount atomic.Int32 + setup := func(p *ProxyAgent) { + p.RegisterExtension(ExtensionHello, HelloHandler(0)) + p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) { + callCount.Add(1) + return "token", nil + })) + } + proxy := NewProxyListener(logger, proxySock, upstreamSock, setup) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + go proxy.Serve(ctx) + <-proxy.Ready() + + // Make two independent connections. + for i := range 3 { + conn, err := net.Dial("unix", proxySock) + if err != nil { + t.Fatalf("connection %d: %v", i, err) + } + client := agent.NewClient(conn).(agent.ExtendedAgent) + _, err = client.Extension(ExtensionAuth, nil) + if err != nil { + t.Fatalf("connection %d auth: %v", i, err) + } + conn.Close() + } + + if callCount.Load() != 3 { + t.Errorf("expected 3 auth calls, got %d", callCount.Load()) + } +} + +func TestProxyListenerWithoutUpstream(t *testing.T) { + proxySock := tempSocketPath(t) + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) + + var authCalled atomic.Bool + setup := func(p *ProxyAgent) { + p.RegisterExtension(ExtensionHello, HelloHandler(0)) + p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) { + authCalled.Store(true) + return "local-token", nil + })) + } + proxy := NewProxyListener(logger, proxySock, "", setup) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + go proxy.Serve(ctx) + <-proxy.Ready() + + conn, err := net.Dial("unix", proxySock) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + + client := agent.NewClient(conn).(agent.ExtendedAgent) + + resp, err := client.Extension(ExtensionHello, nil) + if err != nil { + t.Fatalf("hello failed: %v", err) + } + var hello HelloResponse + if err := json.Unmarshal(stripSuccessByte(t, resp), &hello); err != nil { + t.Fatalf("unmarshal hello: %v", err) + } + if hello.ChainDepth != 0 { + t.Errorf("expected chain depth 0, got %d", hello.ChainDepth) + } + + resp, err = client.Extension(ExtensionAuth, nil) + if err != nil { + t.Fatalf("auth failed: %v", err) + } + var authResp AuthResponse + if err := json.Unmarshal(stripSuccessByte(t, resp), &authResp); err != nil { + t.Fatalf("unmarshal auth: %v", err) + } + if authResp.Token != "local-token" { + t.Errorf("expected local token, got %q", authResp.Token) + } + if !authCalled.Load() { + t.Fatal("expected auth handler to be called") + } + + keys, err := client.List() + if err != nil { + t.Fatalf("List failed: %v", err) + } + if len(keys) != 0 { + t.Errorf("expected empty local keyring, got %d keys", len(keys)) + } +} + +func TestMultiHopChain(t *testing.T) { + // Simulate: laptop → shell1 → shell2 + // Each hop has its own proxy listener wrapping the previous. + // An auth request on shell2 should traverse back to the laptop. + + // Laptop: vanilla agent + proxy with auth handler (the origin). + laptopUpstream, cleanupLaptop := startVanillaAgent(t) + defer cleanupLaptop() + + laptopProxy := tempSocketPath(t) + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) + + var laptopAuthCalled atomic.Bool + laptopSetup := func(p *ProxyAgent) { + p.RegisterExtension(ExtensionHello, HelloHandler(0)) + p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) { + laptopAuthCalled.Store(true) + return "laptop-origin-token", nil + })) + } + laptopListener := NewProxyListener(logger, laptopProxy, laptopUpstream, laptopSetup) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + go laptopListener.Serve(ctx) + <-laptopListener.Ready() + + // Verify laptop probe shows depth 0. + hello, err := ProbeUpstream(laptopProxy) + if err != nil { + t.Fatalf("probe laptop: %v", err) + } + if hello == nil { + t.Fatal("expected epithet agent at laptop, got nil") + } + if hello.ChainDepth != 0 { + t.Errorf("laptop depth: expected 0, got %d", hello.ChainDepth) + } + + // Shell1: proxy wrapping laptop, with depth 1. + // Auth requests are forwarded to laptop via RequestAuth. + shell1Proxy := tempSocketPath(t) + shell1Setup := func(p *ProxyAgent) { + p.RegisterExtension(ExtensionHello, HelloHandler(1)) + p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) { + // Shell1 proxies auth to laptop. + return RequestAuth(laptopProxy) + })) + } + shell1Listener := NewProxyListener(logger, shell1Proxy, laptopProxy, shell1Setup) + + go shell1Listener.Serve(ctx) + <-shell1Listener.Ready() + + // Verify shell1 probe shows depth 1. + hello, err = ProbeUpstream(shell1Proxy) + if err != nil { + t.Fatalf("probe shell1: %v", err) + } + if hello.ChainDepth != 1 { + t.Errorf("shell1 depth: expected 1, got %d", hello.ChainDepth) + } + + // Shell2: proxy wrapping shell1, with depth 2. + shell2Proxy := tempSocketPath(t) + shell2Setup := func(p *ProxyAgent) { + p.RegisterExtension(ExtensionHello, HelloHandler(2)) + p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) { + // Shell2 proxies auth to shell1. + return RequestAuth(shell1Proxy) + })) + } + shell2Listener := NewProxyListener(logger, shell2Proxy, shell1Proxy, shell2Setup) + + go shell2Listener.Serve(ctx) + <-shell2Listener.Ready() + + // Verify shell2 probe shows depth 2. + hello, err = ProbeUpstream(shell2Proxy) + if err != nil { + t.Fatalf("probe shell2: %v", err) + } + if hello.ChainDepth != 2 { + t.Errorf("shell2 depth: expected 2, got %d", hello.ChainDepth) + } + + // Request auth from shell2 — should chain all the way back to laptop. + token, err := RequestAuth(shell2Proxy) + if err != nil { + t.Fatalf("multi-hop auth: %v", err) + } + if token != "laptop-origin-token" { + t.Errorf("expected 'laptop-origin-token', got %q", token) + } + if !laptopAuthCalled.Load() { + t.Error("laptop auth handler was never called — chain is broken") + } +} + +// tempSocketPath returns a unique temporary socket path. +func tempSocketPath(t *testing.T) string { + t.Helper() + // Use /tmp for short paths (macOS 104-byte socket path limit). + f, err := os.CreateTemp("/tmp", "ep-test.*") + if err != nil { + t.Fatal(err) + } + path := f.Name() + f.Close() + os.Remove(path) + t.Cleanup(func() { os.Remove(path) }) + return path +} diff --git a/pkg/agent/proxy_test.go b/pkg/agent/proxy_test.go new file mode 100644 index 0000000..aadc7b1 --- /dev/null +++ b/pkg/agent/proxy_test.go @@ -0,0 +1,246 @@ +package agent + +import ( + "encoding/json" + "errors" + "net" + "os" + "testing" + + "golang.org/x/crypto/ssh/agent" +) + +// startProxyServer creates a ProxyAgent over a real Unix socket, returning the +// client-side ExtendedAgent and a cleanup function. +func startProxyServer(t *testing.T, setup func(*ProxyAgent)) (agent.ExtendedAgent, func()) { + t.Helper() + + upstream := agent.NewKeyring() + + f, err := os.CreateTemp("", "proxy-test.*") + if err != nil { + t.Fatal(err) + } + sockPath := f.Name() + f.Close() + os.Remove(sockPath) + + listener, err := net.Listen("unix", sockPath) + if err != nil { + t.Fatal(err) + } + + go func() { + for { + conn, err := listener.Accept() + if err != nil { + return + } + go func(c net.Conn) { + defer c.Close() + proxy := NewProxyAgent(upstream.(agent.ExtendedAgent)) + setup(proxy) + agent.ServeAgent(proxy, c) + }(conn) + } + }() + + conn, err := net.Dial("unix", sockPath) + if err != nil { + listener.Close() + os.Remove(sockPath) + t.Fatal(err) + } + + client := agent.NewClient(conn) + + cleanup := func() { + conn.Close() + listener.Close() + os.Remove(sockPath) + } + + return client.(agent.ExtendedAgent), cleanup +} + +// stripSuccessByte removes the leading SSH_AGENT_SUCCESS byte from an +// extension response. The client returns the raw wire buffer which starts +// with the type byte. +func stripSuccessByte(t *testing.T, resp []byte) []byte { + t.Helper() + if len(resp) == 0 { + t.Fatal("empty response") + } + if resp[0] != agentSuccess { + t.Fatalf("expected first byte %d (SSH_AGENT_SUCCESS), got %d", agentSuccess, resp[0]) + } + return resp[1:] +} + +func TestHelloHandler(t *testing.T) { + client, cleanup := startProxyServer(t, func(p *ProxyAgent) { + p.RegisterExtension(ExtensionHello, HelloHandler(0)) + }) + defer cleanup() + + resp, err := client.Extension(ExtensionHello, nil) + if err != nil { + t.Fatal(err) + } + + var hello HelloResponse + if err := json.Unmarshal(stripSuccessByte(t, resp), &hello); err != nil { + t.Fatalf("failed to unmarshal hello response: %v", err) + } + + if hello.ProtocolVersion != 1 { + t.Errorf("expected protocol_version=1, got %d", hello.ProtocolVersion) + } + if hello.ChainDepth != 0 { + t.Errorf("expected chain_depth=0, got %d", hello.ChainDepth) + } +} + +func TestHelloHandlerChainDepth(t *testing.T) { + client, cleanup := startProxyServer(t, func(p *ProxyAgent) { + p.RegisterExtension(ExtensionHello, HelloHandler(3)) + }) + defer cleanup() + + resp, err := client.Extension(ExtensionHello, nil) + if err != nil { + t.Fatal(err) + } + + var hello HelloResponse + if err := json.Unmarshal(stripSuccessByte(t, resp), &hello); err != nil { + t.Fatalf("failed to unmarshal: %v", err) + } + + if hello.ChainDepth != 3 { + t.Errorf("expected chain_depth=3, got %d", hello.ChainDepth) + } +} + +func TestAuthHandler(t *testing.T) { + called := false + client, cleanup := startProxyServer(t, func(p *ProxyAgent) { + p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) { + called = true + return "test-token-123", nil + })) + }) + defer cleanup() + + resp, err := client.Extension(ExtensionAuth, nil) + if err != nil { + t.Fatal(err) + } + + if !called { + t.Error("authenticate function was not called") + } + + var authResp AuthResponse + if err := json.Unmarshal(stripSuccessByte(t, resp), &authResp); err != nil { + t.Fatalf("failed to unmarshal: %v", err) + } + + if authResp.Token != "test-token-123" { + t.Errorf("expected token 'test-token-123', got %q", authResp.Token) + } +} + +func TestAuthHandlerError(t *testing.T) { + client, cleanup := startProxyServer(t, func(p *ProxyAgent) { + p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) { + return "", errors.New("user cancelled") + })) + }) + defer cleanup() + + _, err := client.Extension(ExtensionAuth, nil) + if err == nil { + t.Fatal("expected error from failed auth, got nil") + } +} + +func TestStandardOpsPassthrough(t *testing.T) { + client, cleanup := startProxyServer(t, func(p *ProxyAgent) { + p.RegisterExtension(ExtensionHello, HelloHandler(0)) + }) + defer cleanup() + + // List should work (empty keyring). + keys, err := client.List() + if err != nil { + t.Fatalf("List failed: %v", err) + } + if len(keys) != 0 { + t.Errorf("expected 0 keys, got %d", len(keys)) + } +} + +func TestUnregisteredExtensionForwardedToUpstream(t *testing.T) { + // The upstream keyring's Extension returns ErrExtensionUnsupported, + // which ServeAgent translates to SSH_AGENT_FAILURE. + client, cleanup := startProxyServer(t, func(p *ProxyAgent) { + // Register nothing — all extensions forward to upstream. + }) + defer cleanup() + + _, err := client.Extension("unknown@example.com", nil) + if err == nil { + t.Fatal("expected error for unknown extension, got nil") + } +} + +func TestCustomExtensionHandler(t *testing.T) { + client, cleanup := startProxyServer(t, func(p *ProxyAgent) { + p.RegisterExtension("custom@example.com", func(contents []byte) ([]byte, error) { + return []byte("custom-response"), nil + }) + }) + defer cleanup() + + resp, err := client.Extension("custom@example.com", nil) + if err != nil { + t.Fatal(err) + } + + payload := stripSuccessByte(t, resp) + if string(payload) != "custom-response" { + t.Errorf("expected 'custom-response', got %q", string(payload)) + } +} + +func TestDirectProxyAgentExtension(t *testing.T) { + // Test the ProxyAgent directly without going through the socket, + // verifying the raw wire format (includes success byte). + upstream := agent.NewKeyring() + proxy := NewProxyAgent(upstream.(agent.ExtendedAgent)) + proxy.RegisterExtension(ExtensionHello, HelloHandler(0)) + + resp, err := proxy.Extension(ExtensionHello, nil) + if err != nil { + t.Fatal(err) + } + + // First byte should be SSH_AGENT_SUCCESS (6). + if len(resp) == 0 { + t.Fatal("empty response") + } + if resp[0] != agentSuccess { + t.Errorf("expected first byte %d (SSH_AGENT_SUCCESS), got %d", agentSuccess, resp[0]) + } + + // Rest should be valid JSON. + var hello HelloResponse + if err := json.Unmarshal(resp[1:], &hello); err != nil { + t.Fatalf("failed to unmarshal payload: %v", err) + } + if hello.ProtocolVersion != 1 { + t.Errorf("expected protocol_version=1, got %d", hello.ProtocolVersion) + } +} + diff --git a/pkg/agent/upstream.go b/pkg/agent/upstream.go new file mode 100644 index 0000000..b6f28a5 --- /dev/null +++ b/pkg/agent/upstream.go @@ -0,0 +1,110 @@ +package agent + +import ( + "encoding/json" + "errors" + "fmt" + "net" + + "golang.org/x/crypto/ssh/agent" +) + +// ProbeUpstream connects to the agent socket and sends epithet-hello. +// Returns the hello response if the agent supports epithet extensions, +// or nil if it's a vanilla ssh-agent. +func ProbeUpstream(socketPath string) (*HelloResponse, error) { + conn, err := net.Dial("unix", socketPath) + if err != nil { + return nil, fmt.Errorf("failed to connect to upstream agent: %w", err) + } + defer conn.Close() + + client := agent.NewClient(conn) + extClient, ok := client.(agent.ExtendedAgent) + if !ok { + return nil, nil + } + + resp, err := extClient.Extension(ExtensionHello, nil) + if err != nil { + // ErrExtensionUnsupported means it's a vanilla ssh-agent. + if errors.Is(err, agent.ErrExtensionUnsupported) { + return nil, nil + } + return nil, fmt.Errorf("hello extension failed: %w", err) + } + + // Strip the SSH_AGENT_SUCCESS byte. + if len(resp) > 0 && resp[0] == agentSuccess { + resp = resp[1:] + } + + var hello HelloResponse + if err := json.Unmarshal(resp, &hello); err != nil { + return nil, fmt.Errorf("failed to parse hello response: %w", err) + } + + return &hello, nil +} + +// UpstreamUnavailableError indicates that the upstream agent could not be +// reached or does not support epithet extensions. This is the only class +// of error where falling back to local auth is appropriate. +type UpstreamUnavailableError struct { + Err error +} + +func (e *UpstreamUnavailableError) Error() string { + return fmt.Sprintf("upstream unavailable: %v", e.Err) +} + +func (e *UpstreamUnavailableError) Unwrap() error { + return e.Err +} + +// RequestAuth sends an epithet-auth extension request to the upstream agent. +// The upstream runs its own auth plugin with its own state — only the token +// is returned. +// +// Returns UpstreamUnavailableError if the upstream can't be reached or doesn't +// support the extension (safe to fall back to local auth). Other errors indicate +// the upstream attempted auth and failed (should not fall back). +func RequestAuth(socketPath string) (string, error) { + conn, err := net.Dial("unix", socketPath) + if err != nil { + return "", &UpstreamUnavailableError{Err: err} + } + defer conn.Close() + + client := agent.NewClient(conn) + extClient, ok := client.(agent.ExtendedAgent) + if !ok { + return "", &UpstreamUnavailableError{Err: fmt.Errorf("agent does not support extensions")} + } + + resp, err := extClient.Extension(ExtensionAuth, nil) + if err != nil { + if errors.Is(err, agent.ErrExtensionUnsupported) { + return "", &UpstreamUnavailableError{Err: err} + } + // Extension was dispatched but handler returned an error — this is a + // genuine auth failure (user cancelled, plugin error, etc.). + return "", fmt.Errorf("upstream auth failed: %w", err) + } + + // Strip the SSH_AGENT_SUCCESS byte. + if len(resp) > 0 && resp[0] == agentSuccess { + resp = resp[1:] + } + + var authResp AuthResponse + if err := json.Unmarshal(resp, &authResp); err != nil { + return "", fmt.Errorf("failed to parse auth response: %w", err) + } + + if authResp.Token == "" { + return "", fmt.Errorf("upstream returned empty token") + } + + return authResp.Token, nil +} diff --git a/pkg/agent/upstream_test.go b/pkg/agent/upstream_test.go new file mode 100644 index 0000000..fd54ae9 --- /dev/null +++ b/pkg/agent/upstream_test.go @@ -0,0 +1,152 @@ +package agent + +import ( + "net" + "os" + "testing" + + "golang.org/x/crypto/ssh/agent" +) + +// startProxyServerSocket creates a ProxyAgent over a real Unix socket, returning +// the socket path and a cleanup function. Used to test upstream client functions. +func startProxyServerSocket(t *testing.T, setup func(*ProxyAgent)) (string, func()) { + t.Helper() + + upstream := agent.NewKeyring() + + f, err := os.CreateTemp("", "upstream-test.*") + if err != nil { + t.Fatal(err) + } + sockPath := f.Name() + f.Close() + os.Remove(sockPath) + + listener, err := net.Listen("unix", sockPath) + if err != nil { + t.Fatal(err) + } + + go func() { + for { + conn, err := listener.Accept() + if err != nil { + return + } + go func(c net.Conn) { + defer c.Close() + proxy := NewProxyAgent(upstream.(agent.ExtendedAgent)) + setup(proxy) + agent.ServeAgent(proxy, c) + }(conn) + } + }() + + cleanup := func() { + listener.Close() + os.Remove(sockPath) + } + + return sockPath, cleanup +} + +func TestProbeUpstreamEpithet(t *testing.T) { + sockPath, cleanup := startProxyServerSocket(t, func(p *ProxyAgent) { + p.RegisterExtension(ExtensionHello, HelloHandler(2)) + }) + defer cleanup() + + hello, err := ProbeUpstream(sockPath) + if err != nil { + t.Fatal(err) + } + if hello == nil { + t.Fatal("expected hello response, got nil") + } + if hello.ProtocolVersion != 1 { + t.Errorf("expected protocol_version=1, got %d", hello.ProtocolVersion) + } + if hello.ChainDepth != 2 { + t.Errorf("expected chain_depth=2, got %d", hello.ChainDepth) + } +} + +func TestProbeUpstreamVanillaAgent(t *testing.T) { + // Serve a plain keyring (no extensions registered) — behaves like a vanilla ssh-agent. + f, err := os.CreateTemp("", "vanilla-test.*") + if err != nil { + t.Fatal(err) + } + sockPath := f.Name() + f.Close() + os.Remove(sockPath) + + listener, err := net.Listen("unix", sockPath) + if err != nil { + t.Fatal(err) + } + defer func() { + listener.Close() + os.Remove(sockPath) + }() + + keyring := agent.NewKeyring() + go func() { + for { + conn, err := listener.Accept() + if err != nil { + return + } + go func(c net.Conn) { + defer c.Close() + agent.ServeAgent(keyring, c) + }(conn) + } + }() + + hello, err := ProbeUpstream(sockPath) + if err != nil { + t.Fatal(err) + } + if hello != nil { + t.Errorf("expected nil for vanilla agent, got %+v", hello) + } +} + +func TestProbeUpstreamBadSocket(t *testing.T) { + _, err := ProbeUpstream("/nonexistent/socket.sock") + if err == nil { + t.Fatal("expected error for bad socket path") + } +} + +func TestRequestAuth(t *testing.T) { + sockPath, cleanup := startProxyServerSocket(t, func(p *ProxyAgent) { + p.RegisterExtension(ExtensionAuth, AuthHandler(func() (string, error) { + return "upstream-token-456", nil + })) + }) + defer cleanup() + + token, err := RequestAuth(sockPath) + if err != nil { + t.Fatal(err) + } + if token != "upstream-token-456" { + t.Errorf("expected 'upstream-token-456', got %q", token) + } +} + +func TestRequestAuthNoExtensionSupport(t *testing.T) { + // Serve a proxy with no auth handler registered — upstream returns ErrExtensionUnsupported. + sockPath, cleanup := startProxyServerSocket(t, func(p *ProxyAgent) { + // No extensions registered. + }) + defer cleanup() + + _, err := RequestAuth(sockPath) + if err == nil { + t.Fatal("expected error when auth extension not supported") + } +} diff --git a/pkg/broker/auth.go b/pkg/broker/auth.go index cd69366..e0a8901 100644 --- a/pkg/broker/auth.go +++ b/pkg/broker/auth.go @@ -50,6 +50,15 @@ func (h *Auth) Token() string { return h.token } +// SetToken injects a token obtained from an upstream agent. +// Used when the broker receives a token via the agent extension protocol +// instead of running a local auth command. +func (h *Auth) SetToken(token string) { + h.lock.Lock() + defer h.lock.Unlock() + h.token = token +} + // ClearToken clears the stored authentication token. // Keeps the state intact (refresh token may still be valid). // Used when receiving HTTP 401 from CA to force re-authentication. diff --git a/pkg/broker/broker.go b/pkg/broker/broker.go index 9d63a79..a381a73 100644 --- a/pkg/broker/broker.go +++ b/pkg/broker/broker.go @@ -69,6 +69,7 @@ type Broker struct { caClient *caclient.Client // Immutable after New() agentSocketDir string // Immutable after New() + upstreamSocket string // Immutable after New(); empty if no upstream // For graceful shutdown: track in-flight RPC connections activeRPC sync.WaitGroup @@ -86,6 +87,15 @@ func (f optionFunc) apply(b *Broker) error { return f(b) } +// WithUpstream configures the broker to proxy auth requests to an upstream +// epithet agent at the given socket path before falling back to local auth. +func WithUpstream(socketPath string) Option { + return optionFunc(func(b *Broker) error { + b.upstreamSocket = socketPath + return nil + }) +} + // New creates a new Broker instance. This does not start listening - call Serve() to begin accepting connections. func New(log slog.Logger, socketPath string, authCommand string, caClient *caclient.Client, agentSocketDir string, options ...Option) (*Broker, error) { if caClient == nil { @@ -276,7 +286,7 @@ func (b *Broker) MatchWithUserOutput(conn policy.Connection, userOutput io.Write token := b.auth.Token() if token == "" { b.log.Debug("no auth token, authenticating") - token, err = b.auth.Run(nil, userOutput) + token, err = b.Authenticate(userOutput) if err != nil { // Auth command failed - don't retry automatically. // User should fix the issue and retry the SSH connection. @@ -454,6 +464,34 @@ func (b *Broker) ensureAgent(connectionHash policy.ConnectionHash, credential ag return nil } +// Authenticate obtains an auth token, trying upstream first if configured. +// This is also used by the agent proxy's extension handler to run the local +// auth flow when responding to downstream epithet-auth requests. +// +// When an upstream is configured, only falls back to local auth if the upstream +// is unreachable or doesn't support extensions (UpstreamUnavailableError). +// Genuine upstream auth failures (user cancelled, plugin error) are not retried +// locally — the remote host typically cannot complete browser/FIDO flows. +func (b *Broker) Authenticate(userOutput io.Writer) (string, error) { + if b.upstreamSocket != "" { + token, err := agent.RequestAuth(b.upstreamSocket) + if err == nil { + b.auth.SetToken(token) + return token, nil + } + + var unavailable *agent.UpstreamUnavailableError + if errors.As(err, &unavailable) { + b.log.Warn("upstream unavailable, falling back to local auth", "error", err) + } else { + // Upstream was reached and auth was attempted but failed. + // Don't fall back — the failure is authoritative. + return "", err + } + } + return b.auth.Run(nil, userOutput) +} + // shouldHandle checks if the given hostname matches discovery patterns. // Always fetches discovery patterns, authenticating if needed. // Returns true if epithet should handle this connection, false otherwise. @@ -505,7 +543,7 @@ func (b *Broker) getDiscoveryPatterns(userOutput io.Writer) (*caclient.Discovery if token == "" { b.log.Debug("no auth token, authenticating to get discovery") var err error - token, err = b.auth.Run(nil, userOutput) + token, err = b.Authenticate(userOutput) if err != nil { return nil, fmt.Errorf("authentication failed: %w", err) } @@ -520,7 +558,7 @@ func (b *Broker) getDiscoveryPatterns(userOutput io.Writer) (*caclient.Discovery if errors.As(err, &invalidToken) { b.log.Debug("token expired during Hello, re-authenticating") b.auth.ClearToken() - token, err = b.auth.Run(nil, userOutput) + token, err = b.Authenticate(userOutput) if err != nil { return nil, fmt.Errorf("authentication failed: %w", err) } diff --git a/pkg/broker/broker_test.go b/pkg/broker/broker_test.go index 5719cb2..cff3b3b 100644 --- a/pkg/broker/broker_test.go +++ b/pkg/broker/broker_test.go @@ -7,6 +7,7 @@ import ( "fmt" "io" "log/slog" + "net" "net/http" "net/http/httptest" "os" @@ -14,11 +15,13 @@ import ( "testing" "time" + "github.com/epithet-ssh/epithet/pkg/agent" pb "github.com/epithet-ssh/epithet/pkg/brokerv1" "github.com/epithet-ssh/epithet/pkg/caclient" "github.com/epithet-ssh/epithet/pkg/policy" "github.com/lmittmann/tint" "github.com/stretchr/testify/require" + sshagent "golang.org/x/crypto/ssh/agent" "google.golang.org/grpc" "google.golang.org/grpc/credentials/insecure" ) @@ -326,6 +329,211 @@ printf '%s' "test-token" require.Empty(t, result2.Error) } +func TestAuthenticate_WithUpstream(t *testing.T) { + t.Parallel() + + // Start an upstream agent that serves epithet-auth. + upstreamSock := startUpstreamAuthAgent(t, "upstream-token-from-test") + + tmpDir := shortTempDir(t) + socketPath := tmpDir + "/b.sock" + agentDir := tmpDir + "/a" + + // Local auth should NOT be called — upstream provides the token. + localAuthScript := writeTestScript(t, `#!/bin/sh +echo "ERROR: local auth should not have been called" >&2 +exit 1 +`) + b, err := New(*testLogger(t), socketPath, localAuthScript, testCAClient(t, "http://localhost:9999"), agentDir, WithUpstream(upstreamSock)) + require.NoError(t, err) + + token, err := b.Authenticate(nil) + require.NoError(t, err) + require.Equal(t, "upstream-token-from-test", token) + + // Token should be cached in auth. + require.Equal(t, "upstream-token-from-test", b.auth.Token()) +} + +func TestAuthenticate_FallsBackToLocal(t *testing.T) { + t.Parallel() + + // Start an upstream agent that does NOT serve epithet-auth (no extensions). + upstreamSock := startPlainAgent(t) + + tmpDir := shortTempDir(t) + socketPath := tmpDir + "/b.sock" + agentDir := tmpDir + "/a" + + // Local auth will be called as fallback. + localAuthScript := writeTestScript(t, `#!/bin/sh +cat > /dev/null +printf '%s' "local-fallback-token" +`) + b, err := New(*testLogger(t), socketPath, localAuthScript, testCAClient(t, "http://localhost:9999"), agentDir, WithUpstream(upstreamSock)) + require.NoError(t, err) + + token, err := b.Authenticate(nil) + require.NoError(t, err) + // Local auth base64url-encodes the token. + require.NotEmpty(t, token) +} + +func TestAuthenticate_NoUpstream(t *testing.T) { + t.Parallel() + + tmpDir := shortTempDir(t) + socketPath := tmpDir + "/b.sock" + agentDir := tmpDir + "/a" + + localAuthScript := writeTestScript(t, `#!/bin/sh +cat > /dev/null +printf '%s' "local-only-token" +`) + // No WithUpstream — should go straight to local auth. + b, err := New(*testLogger(t), socketPath, localAuthScript, testCAClient(t, "http://localhost:9999"), agentDir) + require.NoError(t, err) + + token, err := b.Authenticate(nil) + require.NoError(t, err) + require.NotEmpty(t, token) +} + +func TestAuthenticate_UpstreamAuthFailure_NoFallback(t *testing.T) { + t.Parallel() + + // Start an upstream agent whose auth handler returns an error. + upstreamSock := startUpstreamFailingAuthAgent(t) + + tmpDir := shortTempDir(t) + socketPath := tmpDir + "/b.sock" + agentDir := tmpDir + "/a" + + // Local auth should NOT be called — upstream failure should propagate. + localAuthScript := writeTestScript(t, `#!/bin/sh +echo "ERROR: local auth should not have been called" >&2 +exit 1 +`) + b, err := New(*testLogger(t), socketPath, localAuthScript, testCAClient(t, "http://localhost:9999"), agentDir, WithUpstream(upstreamSock)) + require.NoError(t, err) + + _, err = b.Authenticate(nil) + require.Error(t, err) + require.Contains(t, err.Error(), "upstream auth failed") +} + +// startUpstreamFailingAuthAgent starts a proxy agent whose auth handler always fails. +func startUpstreamFailingAuthAgent(t *testing.T) string { + t.Helper() + + tmpDir := shortTempDir(t) + upstreamSock := tmpDir + "/upstream.sock" + vanillaSock := tmpDir + "/vanilla.sock" + + vanillaListener, err := net.Listen("unix", vanillaSock) + require.NoError(t, err) + t.Cleanup(func() { vanillaListener.Close() }) + + keyring := sshagent.NewKeyring() + go func() { + for { + conn, err := vanillaListener.Accept() + if err != nil { + return + } + go func(c net.Conn) { + defer c.Close() + sshagent.ServeAgent(keyring, c) + }(conn) + } + }() + + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) + setup := func(p *agent.ProxyAgent) { + p.RegisterExtension(agent.ExtensionAuth, agent.AuthHandler(func() (string, error) { + return "", fmt.Errorf("user cancelled authentication") + })) + } + proxy := agent.NewProxyListener(logger, upstreamSock, vanillaSock, setup) + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + go proxy.Serve(ctx) + <-proxy.Ready() + + return upstreamSock +} + +// startUpstreamAuthAgent starts a proxy agent that handles epithet-auth and returns the given token. +func startUpstreamAuthAgent(t *testing.T, token string) string { + t.Helper() + + tmpDir := shortTempDir(t) + upstreamSock := tmpDir + "/upstream.sock" + vanillaSock := tmpDir + "/vanilla.sock" + + // Start a vanilla agent as the upstream's upstream. + vanillaListener, err := net.Listen("unix", vanillaSock) + require.NoError(t, err) + t.Cleanup(func() { vanillaListener.Close() }) + + keyring := sshagent.NewKeyring() + go func() { + for { + conn, err := vanillaListener.Accept() + if err != nil { + return + } + go func(c net.Conn) { + defer c.Close() + sshagent.ServeAgent(keyring, c) + }(conn) + } + }() + + // Start the proxy listener with auth handler. + logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})) + setup := func(p *agent.ProxyAgent) { + p.RegisterExtension(agent.ExtensionAuth, agent.AuthHandler(func() (string, error) { + return token, nil + })) + } + proxy := agent.NewProxyListener(logger, upstreamSock, vanillaSock, setup) + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + go proxy.Serve(ctx) + <-proxy.Ready() + + return upstreamSock +} + +// startPlainAgent starts a vanilla ssh-agent (no epithet extensions). +func startPlainAgent(t *testing.T) string { + t.Helper() + + tmpDir := shortTempDir(t) + sockPath := tmpDir + "/plain.sock" + + listener, err := net.Listen("unix", sockPath) + require.NoError(t, err) + t.Cleanup(func() { listener.Close() }) + + keyring := sshagent.NewKeyring() + go func() { + for { + conn, err := listener.Accept() + if err != nil { + return + } + go func(c net.Conn) { + defer c.Close() + sshagent.ServeAgent(keyring, c) + }(conn) + } + }() + + return sockPath +} + func testLogger(t *testing.T) *slog.Logger { logger := slog.New(tint.NewHandler(t.Output(), &tint.Options{ Level: slog.LevelDebug,