repos / pico

pico services mono repo
git clone https://github.com/picosh/pico.git

pico / pkg / pssh
Antonio Mika  ·  2025-03-12

server_test.go

  1package pssh_test
  2
  3import (
  4	"context"
  5	"errors"
  6	"log/slog"
  7	"net"
  8	"testing"
  9	"time"
 10
 11	"github.com/picosh/pico/pkg/pssh"
 12	"golang.org/x/crypto/ssh"
 13)
 14
 15func TestNewSSHServer(t *testing.T) {
 16	ctx := context.Background()
 17	logger := slog.Default()
 18	server := pssh.NewSSHServer(ctx, logger, &pssh.SSHServerConfig{})
 19
 20	if server == nil {
 21		t.Fatal("expected non-nil server")
 22	}
 23
 24	if server.Ctx == nil {
 25		t.Error("expected non-nil context")
 26	}
 27
 28	if server.Logger == nil {
 29		t.Error("expected non-nil logger")
 30	}
 31
 32	if server.Config == nil {
 33		t.Error("expected non-nil config")
 34	}
 35
 36	if server.Conns == nil {
 37		t.Error("expected non-nil connections map")
 38	}
 39}
 40
 41func TestNewSSHServerConn(t *testing.T) {
 42	ctx := context.Background()
 43	logger := slog.Default()
 44	server := pssh.NewSSHServer(ctx, logger, &pssh.SSHServerConfig{})
 45	conn := &ssh.ServerConn{}
 46
 47	serverConn := pssh.NewSSHServerConn(ctx, logger, conn, server)
 48
 49	if serverConn == nil {
 50		t.Fatal("expected non-nil server connection")
 51	}
 52
 53	if serverConn.Ctx == nil {
 54		t.Error("expected non-nil context")
 55	}
 56
 57	if serverConn.Logger == nil {
 58		t.Error("expected non-nil logger")
 59	}
 60
 61	if serverConn.Conn != conn {
 62		t.Error("expected conn to match")
 63	}
 64
 65	if serverConn.SSHServer != server {
 66		t.Error("expected server to match")
 67	}
 68}
 69
 70func TestSSHServerConnClose(t *testing.T) {
 71	ctx := context.Background()
 72	logger := slog.Default()
 73	server := pssh.NewSSHServer(ctx, logger, &pssh.SSHServerConfig{})
 74	conn := &ssh.ServerConn{}
 75
 76	serverConn := pssh.NewSSHServerConn(ctx, logger, conn, server)
 77	err := serverConn.Close()
 78
 79	if err != nil {
 80		t.Errorf("unexpected error: %v", err)
 81	}
 82
 83	// Should be canceled after close
 84	select {
 85	case <-serverConn.Ctx.Done():
 86		// Context was canceled as expected
 87	default:
 88		t.Error("context was not canceled after Close()")
 89	}
 90}
 91
 92func TestSSHServerClose(t *testing.T) {
 93	ctx := context.Background()
 94	logger := slog.Default()
 95	server := pssh.NewSSHServer(ctx, logger, &pssh.SSHServerConfig{})
 96
 97	// Create a mock listener to test Close()
 98	listener, err := net.Listen("tcp", "127.0.0.1:0")
 99	if err != nil {
100		t.Fatalf("failed to create listener: %v", err)
101	}
102
103	server.Listener = listener
104	err = server.Close()
105
106	if err != nil {
107		t.Errorf("unexpected error: %v", err)
108	}
109
110	// Should be canceled after close
111	select {
112	case <-server.Ctx.Done():
113		// Context was canceled as expected
114	default:
115		t.Error("context was not canceled after Close()")
116	}
117}
118
119func TestSSHServerNilParams(t *testing.T) {
120	// Test with nil context and logger
121	//nolint:staticcheck // SA1012 ignores nil check
122	//lint:ignore SA1012 ignores nil check
123	server := pssh.NewSSHServer(nil, nil, nil)
124
125	if server == nil {
126		t.Fatal("expected non-nil server")
127	}
128
129	if server.Ctx == nil {
130		t.Error("expected non-nil context even when nil is passed")
131	}
132
133	if server.Logger == nil {
134		t.Error("expected non-nil logger even when nil is passed")
135	}
136
137	// Test with nil context and logger for connection
138	//nolint:staticcheck // SA1012 ignores nil check
139	//lint:ignore SA1012 ignores nil check
140	conn := pssh.NewSSHServerConn(nil, nil, &ssh.ServerConn{}, server)
141
142	if conn == nil {
143		t.Fatal("expected non-nil server connection")
144	}
145
146	if conn.Ctx == nil {
147		t.Error("expected non-nil context even when nil is passed")
148	}
149
150	if conn.Logger == nil {
151		t.Error("expected non-nil logger even when nil is passed")
152	}
153}
154
155func TestSSHServerHandleConn(t *testing.T) {
156	ctx, cancel := context.WithCancel(context.Background())
157	defer cancel()
158	logger := slog.Default()
159	server := pssh.NewSSHServer(ctx, logger, &pssh.SSHServerConfig{})
160
161	// Setup a basic SSH server config
162	config := &ssh.ServerConfig{
163		NoClientAuth: true,
164	}
165
166	server.Config.ServerConfig = config
167
168	// Create a mock connection
169	client, server_conn := net.Pipe()
170	defer client.Close()
171
172	// Start HandleConn in a goroutine
173	errChan := make(chan error, 1)
174	go func() {
175		errChan <- server.HandleConn(server_conn)
176	}()
177
178	// Configure SSH client
179	clientConfig := &ssh.ClientConfig{
180		User:            "testuser",
181		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
182	}
183
184	// Try to establish SSH connection
185	_, _, _, err := ssh.NewClientConn(client, "", clientConfig)
186
187	// It should fail since we're using a pipe and not a proper SSH handshake
188	if err == nil {
189		t.Error("expected SSH handshake to fail with test pipe")
190	}
191
192	// Close connections to ensure HandleConn returns
193	client.Close()
194	server_conn.Close()
195
196	// Wait for HandleConn to return
197	select {
198	case <-errChan:
199		// Expected HandleConn to return
200	case <-time.After(2 * time.Second):
201		t.Error("HandleConn did not return after connection closed")
202	}
203}
204
205func TestSSHServerListenAndServe(t *testing.T) {
206	ctx, cancel := context.WithCancel(context.Background())
207	defer cancel()
208	logger := slog.Default()
209	server := pssh.NewSSHServer(ctx, logger, &pssh.SSHServerConfig{})
210
211	config := &ssh.ServerConfig{
212		NoClientAuth: true,
213	}
214
215	// Set a random port
216	port := "127.0.0.1:0"
217	server.Config.ListenAddr = port
218	server.Config.ServerConfig = config
219
220	// Start server in a goroutine
221	errChan := make(chan error, 1)
222	go func() {
223		err := server.ListenAndServe()
224		errChan <- err
225	}()
226
227	// Wait a bit for the server to start
228	time.Sleep(100 * time.Millisecond)
229
230	// Trigger cancellation to stop the server
231	cancel()
232
233	// Wait for server to stop
234	select {
235	case err := <-errChan:
236		if err != nil && !errors.Is(err, net.ErrClosed) {
237			t.Errorf("unexpected error: %v", err)
238		}
239	case <-time.After(2 * time.Second):
240		t.Error("server did not shut down in time")
241	}
242}
243
244func TestSSHServerConnHandle(t *testing.T) {
245	ctx, cancel := context.WithCancel(context.Background())
246	defer cancel()
247	logger := slog.Default()
248	server := pssh.NewSSHServer(ctx, logger, &pssh.SSHServerConfig{})
249	conn := &ssh.ServerConn{}
250
251	serverConn := pssh.NewSSHServerConn(ctx, logger, conn, server)
252
253	// Create channels for testing
254	chans := make(chan ssh.NewChannel)
255	reqs := make(chan *ssh.Request)
256
257	// Start handle in a goroutine
258	errChan := make(chan error, 1)
259	go func() {
260		errChan <- serverConn.Handle(chans, reqs)
261	}()
262
263	// Ensure handle returns when context is canceled
264	cancel()
265
266	// Wait for handle to return
267	select {
268	case err := <-errChan:
269		if err != nil {
270			t.Errorf("unexpected error: %v", err)
271		}
272	case <-time.After(2 * time.Second):
273		t.Error("Handle did not return after context canceled")
274	}
275}