1 files changed, 152 insertions, 0 deletions
diff --git a/internal/ssh/ssh_agent_test.go b/internal/ssh/ssh_agent_test.go
new file mode 100644
index 0000000..b71c078
--- /dev/null
+++ b/internal/ssh/ssh_agent_test.go
@@ -0,0 +1,152 @@
+package ssh
+
+import (
+	"errors"
+	"io"
+	"net"
+	"sync/atomic"
+	"testing"
+
+	"golang.org/x/crypto/ssh/agent"
+)
+
+// countingConn wraps a net.Conn and tracks how many times Close was called.
+type countingConn struct {
+	net.Conn
+	closeCount atomic.Int32
+}
+
+func (c *countingConn) Close() error {
+	c.closeCount.Add(1)
+	return c.Conn.Close()
+}
+
+// pipePair returns a client/server net.Conn pair and wraps the client side
+// in a counting closer so tests can assert Close was invoked.
+func pipePair() (*countingConn, net.Conn) {
+	client, server := net.Pipe()
+	return &countingConn{Conn: client}, server
+}
+
+// newFakeAgent starts a serving ssh-agent on the server side of the pipe.
+// Callers must close the returned server conn when done (or rely on client Close
+// propagating through net.Pipe).
+func newFakeAgent(t *testing.T, server net.Conn, keyring agent.Agent) {
+	t.Helper()
+	go func() {
+		_ = agent.ServeAgent(keyring, server)
+		_ = server.Close()
+	}()
+}
+
+func withDialAgent(t *testing.T, dial func() (net.Conn, error)) {
+	t.Helper()
+	orig := dialAgent
+	dialAgent = func(_ string) (net.Conn, error) {
+		return dial()
+	}
+	t.Cleanup(func() { dialAgent = orig })
+}
+
+// TestAgentSignersWithKeyIndexClosesConnOnDialListError exercises the error
+// path where agent.List fails. The ssh-agent connection must not be leaked.
+func TestAgentSignersWithKeyIndexClosesConnOnListError(t *testing.T) {
+	client, server := pipePair()
+	// Close the server side immediately so agentClient.List() returns an error.
+	_ = server.Close()
+
+	withDialAgent(t, func() (net.Conn, error) { return client, nil })
+
+	signers, closer, err := AgentSignersWithKeyIndex(-1)
+	if err == nil {
+		t.Fatalf("expected error from agent.List when server closed, got nil")
+	}
+	if signers != nil {
+		t.Fatalf("expected nil signers on error, got %d", len(signers))
+	}
+	if closer == nil {
+		t.Fatalf("expected non-nil closer even on error path")
+	}
+	// closer may be a no-op on error (ownership already released internally),
+	// but calling Close must be safe.
+	_ = closer.Close()
+
+	if got := client.closeCount.Load(); got < 1 {
+		t.Fatalf("expected underlying agent conn to be closed on error path, closeCount=%d", got)
+	}
+}
+
+// TestAgentSignersWithKeyIndexClosesConnOnDialError verifies that when the
+// initial dial fails no conn is ever created and the returned closer is safe.
+func TestAgentSignersWithKeyIndexClosesConnOnDialError(t *testing.T) {
+	dialErr := errors.New("dial failed")
+	withDialAgent(t, func() (net.Conn, error) { return nil, dialErr })
+
+	signers, closer, err := AgentSignersWithKeyIndex(-1)
+	if err == nil {
+		t.Fatalf("expected dial error, got nil")
+	}
+	if signers != nil {
+		t.Fatalf("expected nil signers on dial error, got %d", len(signers))
+	}
+	if closer == nil {
+		t.Fatalf("expected non-nil closer even on dial error")
+	}
+	if err := closer.Close(); err != nil {
+		t.Fatalf("closer.Close on dial error should be a no-op, got %v", err)
+	}
+}
+
+// TestAgentSignersWithKeyIndexReturnsOwnerCloserOnSuccess verifies the happy
+// path where the caller takes ownership of the agent connection via io.Closer.
+func TestAgentSignersWithKeyIndexReturnsOwnerCloserOnSuccess(t *testing.T) {
+	client, server := pipePair()
+	keyring := agent.NewKeyring()
+	newFakeAgent(t, server, keyring)
+
+	withDialAgent(t, func() (net.Conn, error) { return client, nil })
+
+	_, closer, err := AgentSignersWithKeyIndex(-1)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if closer == nil {
+		t.Fatalf("expected non-nil closer on success")
+	}
+	if got := client.closeCount.Load(); got != 0 {
+		t.Fatalf("underlying conn must stay open on success until caller closes, closeCount=%d", got)
+	}
+
+	if err := closer.Close(); err != nil {
+		t.Fatalf("closer.Close returned error: %v", err)
+	}
+	if got := client.closeCount.Load(); got < 1 {
+		t.Fatalf("expected closer to close the underlying agent conn, closeCount=%d", got)
+	}
+}
+
+// TestAgentSignersWithKeyIndexOutOfRangeClosesConn verifies that when the
+// requested key index exceeds the number of agent signers the connection is
+// released.
+func TestAgentSignersWithKeyIndexOutOfRangeClosesConn(t *testing.T) {
+	client, server := pipePair()
+	keyring := agent.NewKeyring() // empty keyring => no signers
+	newFakeAgent(t, server, keyring)
+
+	withDialAgent(t, func() (net.Conn, error) { return client, nil })
+
+	_, closer, err := AgentSignersWithKeyIndex(0)
+	if err == nil {
+		t.Fatalf("expected out-of-range error on empty keyring, got nil")
+	}
+	if closer == nil {
+		t.Fatalf("expected non-nil closer on out-of-range error")
+	}
+	_ = closer.Close()
+	if got := client.closeCount.Load(); got < 1 {
+		t.Fatalf("expected conn close on out-of-range error, closeCount=%d", got)
+	}
+}
+
+// Compile-time sanity: the returned closer must satisfy io.Closer.
+var _ io.Closer = (*countingConn)(nil)