1package pipe
  2
  3import (
  4	"context"
  5	"crypto/ed25519"
  6	"crypto/rand"
  7	"errors"
  8	"io"
  9	"net"
 10	"runtime"
 11	"strings"
 12	"sync"
 13	"sync/atomic"
 14	"testing"
 15	"time"
 16
 17	"github.com/picosh/utils/pipe"
 18	"golang.org/x/crypto/ssh"
 19)
 20
 21// closers closes what a test server opened once the test ends.
 22type closers struct {
 23	mu sync.Mutex
 24	cs []io.Closer
 25}
 26
 27func newClosers(t *testing.T) *closers {
 28	c := &closers{}
 29	t.Cleanup(func() {
 30		c.mu.Lock()
 31		defer c.mu.Unlock()
 32		for _, cl := range c.cs {
 33			_ = cl.Close()
 34		}
 35	})
 36	return c
 37}
 38
 39func (c *closers) add(cl io.Closer) {
 40	c.mu.Lock()
 41	defer c.mu.Unlock()
 42	c.cs = append(c.cs, cl)
 43}
 44
 45// stalledConn is a pipe connection to a server that accepts connections and
 46// never answers.
 47func stalledConn(t *testing.T, timeout time.Duration) *pipeConn {
 48	t.Helper()
 49	ln, err := net.Listen("tcp", "127.0.0.1:0")
 50	if err != nil {
 51		t.Fatal(err)
 52	}
 53	open := newClosers(t)
 54	open.add(ln)
 55	go func() {
 56		for {
 57			conn, err := ln.Accept()
 58			if err != nil {
 59				return
 60			}
 61			open.add(conn)
 62		}
 63	}()
 64	conn := newPipeConn(&pipe.SSHClientInfo{RemoteHost: ln.Addr().String()})
 65	conn.timeout = timeout
 66	return conn
 67}
 68
 69func TestStalledHandshakeEndsWithRequest(t *testing.T) {
 70	conn := stalledConn(t, time.Minute)
 71
 72	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
 73	defer cancel()
 74	start := time.Now()
 75	if _, err := conn.run(ctx, "sub topic"); !errors.Is(err, context.DeadlineExceeded) {
 76		t.Fatalf("run = %v, want the request's deadline", err)
 77	}
 78	if took := time.Since(start); took > time.Second {
 79		t.Errorf("canceled request waited %v for the handshake", took)
 80	}
 81
 82	// The next request is not held up behind the canceled one.
 83	ctx, cancel = context.WithTimeout(context.Background(), 100*time.Millisecond)
 84	defer cancel()
 85	start = time.Now()
 86	if _, err := conn.run(ctx, "sub topic"); !errors.Is(err, context.DeadlineExceeded) {
 87		t.Fatalf("run = %v, want the request's deadline", err)
 88	}
 89	if took := time.Since(start); took > time.Second {
 90		t.Errorf("second request waited %v", took)
 91	}
 92}
 93
 94func TestStalledHandshakeTimesOut(t *testing.T) {
 95	conn := stalledConn(t, 200*time.Millisecond)
 96
 97	start := time.Now()
 98	_, err := conn.run(context.Background(), "sub topic")
 99	var netErr net.Error
100	if !errors.As(err, &netErr) || !netErr.Timeout() {
101		t.Fatalf("run = %v, want a timeout", err)
102	}
103	if took := time.Since(start); took > 2*time.Second {
104		t.Errorf("handshake took %v, want about 200ms", took)
105	}
106}
107
108// stalledSSH is a pipe connection to an ssh server that completes the
109// handshake and then stops answering: at opening a session when openStalls
110// is set, or else at starting the command. It counts the connections made
111// and the ones closed.
112func stalledSSH(t *testing.T, timeout time.Duration, openStalls bool) (conn *pipeConn, conns, closed *atomic.Int32) {
113	t.Helper()
114	_, key, err := ed25519.GenerateKey(rand.Reader)
115	if err != nil {
116		t.Fatal(err)
117	}
118	hostKey, err := ssh.NewSignerFromKey(key)
119	if err != nil {
120		t.Fatal(err)
121	}
122	cfg := &ssh.ServerConfig{NoClientAuth: true}
123	cfg.AddHostKey(hostKey)
124
125	ln, err := net.Listen("tcp", "127.0.0.1:0")
126	if err != nil {
127		t.Fatal(err)
128	}
129	open := newClosers(t)
130	open.add(ln)
131	conns, closed = &atomic.Int32{}, &atomic.Int32{}
132	go func() {
133		for {
134			nc, err := ln.Accept()
135			if err != nil {
136				return
137			}
138			conns.Add(1)
139			open.add(nc)
140			go func() {
141				sc, chans, reqs, err := ssh.NewServerConn(nc, cfg)
142				if err != nil {
143					return
144				}
145				go func() {
146					_ = sc.Wait()
147					closed.Add(1)
148				}()
149				go ssh.DiscardRequests(reqs)
150				for nch := range chans {
151					if openStalls {
152						continue
153					}
154					ch, chReqs, err := nch.Accept()
155					if err != nil {
156						return
157					}
158					open.add(ch)
159					// Never reply to the exec request.
160					go func() {
161						for range chReqs {
162						}
163					}()
164				}
165			}()
166		}
167	}()
168	conn = newPipeConn(&pipe.SSHClientInfo{RemoteHost: ln.Addr().String()})
169	conn.timeout = timeout
170	return conn, conns, closed
171}
172
173func TestStalledSessionEndsWithRequest(t *testing.T) {
174	for name, openStalls := range map[string]bool{"open": true, "exec": false} {
175		t.Run(name, func(t *testing.T) {
176			conn, _, _ := stalledSSH(t, time.Minute, openStalls)
177			ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
178			defer cancel()
179			start := time.Now()
180			if _, err := conn.run(ctx, "sub topic"); !errors.Is(err, context.DeadlineExceeded) {
181				t.Fatalf("run = %v, want the request's deadline", err)
182			}
183			if took := time.Since(start); took > time.Second {
184				t.Errorf("canceled request waited %v for the session", took)
185			}
186		})
187	}
188}
189
190func TestStalledSessionRedials(t *testing.T) {
191	for name, openStalls := range map[string]bool{"open": true, "exec": false} {
192		t.Run(name, func(t *testing.T) {
193			conn, conns, _ := stalledSSH(t, 200*time.Millisecond, openStalls)
194			_, err := conn.run(context.Background(), "sub topic")
195			if err == nil || !strings.Contains(err.Error(), "timed out") {
196				t.Fatalf("run = %v, want a timeout", err)
197			}
198			if n := conns.Load(); n != 2 {
199				t.Errorf("made %d connections, want the stalled one and a redial", n)
200			}
201		})
202	}
203}
204
205func TestCanceledStartStillTimesOut(t *testing.T) {
206	for name, openStalls := range map[string]bool{"open": true, "exec": false} {
207		t.Run(name, func(t *testing.T) {
208			conn, _, closed := stalledSSH(t, 300*time.Millisecond, openStalls)
209			before := runtime.NumGoroutine()
210			for range 5 {
211				ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
212				_, err := conn.run(ctx, "sub topic")
213				cancel()
214				if !errors.Is(err, context.DeadlineExceeded) {
215					t.Fatalf("run = %v, want the request's deadline", err)
216				}
217			}
218			for range 100 {
219				if closed.Load() == 1 && runtime.NumGoroutine() <= before {
220					return
221				}
222				time.Sleep(20 * time.Millisecond)
223			}
224			t.Errorf("closed %d connections and have %d goroutines, want 1 and at most %d", closed.Load(), runtime.NumGoroutine(), before)
225		})
226	}
227}