summaryrefslogtreecommitdiff
path: root/internal/ssh/ssh_agent_test.go
blob: b71c078c7e14f7611efda1779fb6088618f3a8ea (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
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)