1package pipe
  2
  3import (
  4	"context"
  5	"fmt"
  6	"io"
  7	"net"
  8	"os"
  9	"strings"
 10	"sync"
 11	"time"
 12
 13	"github.com/picosh/utils/pipe"
 14	"golang.org/x/crypto/ssh"
 15)
 16
 17// dialTimeout bounds connecting to pipe, handshake included, and starting a
 18// command on it, so a server that stops answering can't hold up requests.
 19const dialTimeout = 10 * time.Second
 20
 21// pipeConn is pipe-web's ssh connection to pipe. Each request runs its
 22// command in a session of its own, which is never reopened, so the command
 23// ends when the request does.
 24type pipeConn struct {
 25	info    *pipe.SSHClientInfo
 26	timeout time.Duration
 27
 28	mu      sync.Mutex
 29	client  *ssh.Client
 30	dialing *dial
 31}
 32
 33// dial is a connection being opened, shared by the requests waiting on it.
 34type dial struct {
 35	done   chan struct{}
 36	client *ssh.Client
 37	err    error
 38}
 39
 40func newPipeConn(info *pipe.SSHClientInfo) *pipeConn {
 41	return &pipeConn{info: info, timeout: dialTimeout}
 42}
 43
 44// connect returns the open connection, dialing a new one when there is none
 45// or when the open one is dead. A request that ends stops waiting, and the
 46// dial carries on for the requests after it.
 47func (c *pipeConn) connect(ctx context.Context, dead *ssh.Client) (*ssh.Client, error) {
 48	c.mu.Lock()
 49	if c.client != nil && c.client != dead {
 50		client := c.client
 51		c.mu.Unlock()
 52		return client, nil
 53	}
 54	if c.client != nil {
 55		_ = c.client.Close()
 56		c.client = nil
 57	}
 58	d := c.dialing
 59	if d == nil {
 60		d = &dial{done: make(chan struct{})}
 61		c.dialing = d
 62		go c.dial(d)
 63	}
 64	c.mu.Unlock()
 65
 66	select {
 67	case <-d.done:
 68		return d.client, d.err
 69	case <-ctx.Done():
 70		return nil, ctx.Err()
 71	}
 72}
 73
 74func (c *pipeConn) dial(d *dial) {
 75	d.client, d.err = dialSSH(c.info, c.timeout)
 76	c.mu.Lock()
 77	c.dialing = nil
 78	if d.err == nil {
 79		c.client = d.client
 80	}
 81	c.mu.Unlock()
 82	close(d.done)
 83}
 84
 85// run starts cmd on pipe. A connection that can't open a session is dialed
 86// again once.
 87func (c *pipeConn) run(ctx context.Context, cmd string) (*cmdSession, error) {
 88	client, err := c.connect(ctx, nil)
 89	if err != nil {
 90		return nil, err
 91	}
 92	s, err := c.start(ctx, client, cmd)
 93	if err == nil || ctx.Err() != nil {
 94		return s, err
 95	}
 96	client, err = c.connect(ctx, client)
 97	if err != nil {
 98		return nil, err
 99	}
100	return c.start(ctx, client, cmd)
101}
102
103type started struct {
104	s   *cmdSession
105	err error
106}
107
108// start runs startCmd until the request ends. A connection that can't
109// start a command in time is closed, which also ends the startCmd waiting
110// on it, and that deadline holds after the request has gone.
111func (c *pipeConn) start(ctx context.Context, client *ssh.Client, cmd string) (*cmdSession, error) {
112	done := make(chan started, 1)
113	go func() {
114		s, err := startCmd(client, cmd)
115		done <- started{s, err}
116	}()
117	timer := time.NewTimer(c.timeout)
118	select {
119	case res := <-done:
120		timer.Stop()
121		return res.s, res.err
122	case <-ctx.Done():
123		go abandon(done, timer, client)
124		return nil, ctx.Err()
125	case <-timer.C:
126		_ = client.Close()
127		go closeStarted(done)
128		return nil, fmt.Errorf("starting %q on pipe timed out", cmd)
129	}
130}
131
132// abandon closes a session no request is waiting for, or the connection if
133// the session doesn't start in time.
134func abandon(done <-chan started, timer *time.Timer, client *ssh.Client) {
135	select {
136	case res := <-done:
137		timer.Stop()
138		if res.s != nil {
139			_ = res.s.Close()
140		}
141	case <-timer.C:
142		_ = client.Close()
143		closeStarted(done)
144	}
145}
146
147func closeStarted(done <-chan started) {
148	if res := <-done; res.s != nil {
149		_ = res.s.Close()
150	}
151}
152
153func dialSSH(info *pipe.SSHClientInfo, timeout time.Duration) (*ssh.Client, error) {
154	signer, err := clientSigner(info)
155	if err != nil {
156		return nil, err
157	}
158	var auth []ssh.AuthMethod
159	if signer != nil {
160		auth = append(auth, ssh.PublicKeys(signer))
161	}
162	addr := info.RemoteHost
163	if !strings.Contains(addr, ":") {
164		addr += ":22"
165	}
166
167	deadline := time.Now().Add(timeout)
168	conn, err := (&net.Dialer{Deadline: deadline}).Dial("tcp", addr)
169	if err != nil {
170		return nil, err
171	}
172	_ = conn.SetDeadline(deadline)
173	sshConn, chans, reqs, err := ssh.NewClientConn(conn, info.RemoteHostname, &ssh.ClientConfig{
174		Auth:            auth,
175		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
176		User:            info.RemoteUser,
177	})
178	if err != nil {
179		_ = conn.Close()
180		return nil, err
181	}
182	_ = conn.SetDeadline(time.Time{})
183	return ssh.NewClient(sshConn, chans, reqs), nil
184}
185
186// clientSigner loads the key, and the certificate for it if there is one.
187// It returns nil when there is no key.
188func clientSigner(info *pipe.SSHClientInfo) (ssh.Signer, error) {
189	if info.KeyLocation == "" {
190		return nil, nil
191	}
192	data, err := os.ReadFile(info.KeyLocation)
193	if err != nil {
194		return nil, err
195	}
196	var signer ssh.Signer
197	if info.KeyPassphrase != "" {
198		signer, err = ssh.ParsePrivateKeyWithPassphrase(data, []byte(info.KeyPassphrase))
199	} else {
200		signer, err = ssh.ParsePrivateKey(data)
201	}
202	if err != nil || info.CertificateLocation == "" {
203		return signer, err
204	}
205	certData, err := os.ReadFile(info.CertificateLocation)
206	if err != nil {
207		return nil, err
208	}
209	pub, _, _, _, err := ssh.ParseAuthorizedKey(certData)
210	if err != nil {
211		return nil, fmt.Errorf("failed to parse certificate: %w", err)
212	}
213	cert, ok := pub.(*ssh.Certificate)
214	if !ok {
215		return nil, fmt.Errorf("%s is not an ssh certificate", info.CertificateLocation)
216	}
217	return ssh.NewCertSigner(cert, signer)
218}
219
220// cmdSession is a command running on pipe, read from its stdout and written
221// to its stdin.
222type cmdSession struct {
223	*ssh.Session
224	stdin  io.WriteCloser
225	stdout io.Reader
226}
227
228func startCmd(client *ssh.Client, cmd string) (*cmdSession, error) {
229	session, err := client.NewSession()
230	if err != nil {
231		return nil, err
232	}
233	s := &cmdSession{Session: session}
234	if s.stdin, err = session.StdinPipe(); err == nil {
235		if s.stdout, err = session.StdoutPipe(); err == nil {
236			err = session.Start(cmd)
237		}
238	}
239	if err != nil {
240		_ = session.Close()
241		return nil, err
242	}
243	return s, nil
244}
245
246func (s *cmdSession) Read(p []byte) (int, error) {
247	return s.stdout.Read(p)
248}
249
250func (s *cmdSession) Write(p []byte) (int, error) {
251	return s.stdin.Write(p)
252}
253
254// CloseWrite sends the command EOF on its stdin.
255func (s *cmdSession) CloseWrite() error {
256	return s.stdin.Close()
257}