ssh_test.go

package xshellz

import (
	"bytes"
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"net"
	"testing"

	"golang.org/x/crypto/ssh"
)

// startSSHServer runs a minimal in-process SSH server that accepts any public
// key and hands each exec request to handler. Returns the listen host/port.
func startSSHServer(t *testing.T, handler func(cmd string, ch ssh.Channel)) (string, int) {
	t.Helper()

	_, hostPriv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	hostSigner, err := ssh.NewSignerFromKey(hostPriv)
	if err != nil {
		t.Fatal(err)
	}

	config := &ssh.ServerConfig{
		PublicKeyCallback: func(ssh.ConnMetadata, ssh.PublicKey) (*ssh.Permissions, error) {
			return nil, nil
		},
	}
	config.AddHostKey(hostSigner)

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = ln.Close() })

	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				return
			}
			go func(conn net.Conn) {
				sconn, chans, reqs, err := ssh.NewServerConn(conn, config)
				if err != nil {
					return
				}
				defer sconn.Close()
				go ssh.DiscardRequests(reqs)
				for newCh := range chans {
					if newCh.ChannelType() != "session" {
						_ = newCh.Reject(ssh.UnknownChannelType, "only session")
						continue
					}
					ch, chReqs, err := newCh.Accept()
					if err != nil {
						continue
					}
					go func(ch ssh.Channel, chReqs <-chan *ssh.Request) {
						for req := range chReqs {
							if req.Type != "exec" {
								_ = req.Reply(false, nil)
								continue
							}
							var payload struct{ Command string }
							_ = ssh.Unmarshal(req.Payload, &payload)
							_ = req.Reply(true, nil)
							handler(payload.Command, ch)
							return
						}
					}(ch, chReqs)
				}
			}(conn)
		}
	}()

	addr := ln.Addr().(*net.TCPAddr)
	return "127.0.0.1", addr.Port
}

func exitStatus(code uint32) []byte {
	return ssh.Marshal(struct{ Status uint32 }{code})
}

func TestDialSSHRunsCommandsAgainstRealServer(t *testing.T) {
	host, port := startSSHServer(t, func(cmd string, ch ssh.Channel) {
		defer ch.Close()
		switch cmd {
		case "boom":
			_, _ = ch.Stderr().Write([]byte("kaboom\n"))
			_, _ = ch.SendRequest("exit-status", false, exitStatus(3))
		default:
			_, _ = ch.Write([]byte("ran: " + cmd + "\n"))
			_, _ = ch.SendRequest("exit-status", false, exitStatus(0))
		}
	})

	keys, err := generateKeypair()
	if err != nil {
		t.Fatal(err)
	}
	runner, err := dialSSH(context.Background(), host, port, keys.signer)
	if err != nil {
		t.Fatalf("dialSSH: %v", err)
	}

	var out, errBuf bytes.Buffer
	code, err := runner.run(context.Background(), "echo hi", nil, &out, &errBuf)
	if err != nil {
		t.Fatalf("run: %v", err)
	}
	if code != 0 || out.String() != "ran: echo hi\n" {
		t.Errorf("code=%d out=%q", code, out.String())
	}

	out.Reset()
	errBuf.Reset()
	code, err = runner.run(context.Background(), "boom", nil, &out, &errBuf)
	if err != nil {
		t.Fatalf("run(boom): %v", err)
	}
	if code != 3 || errBuf.String() != "kaboom\n" {
		t.Errorf("code=%d stderr=%q, want exit 3 via ssh.ExitError", code, errBuf.String())
	}

	if err := runner.close(); err != nil {
		t.Errorf("close: %v", err)
	}
	if _, err := runner.run(context.Background(), "true", nil, nil, nil); err == nil {
		t.Error("run after close should fail to open a session")
	}
}

func TestDialSSHConnectionRefused(t *testing.T) {
	// Grab a port and close it so nothing is listening.
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	port := ln.Addr().(*net.TCPAddr).Port
	_ = ln.Close()

	keys, err := generateKeypair()
	if err != nil {
		t.Fatal(err)
	}
	if _, err := dialSSH(context.Background(), "127.0.0.1", port, keys.signer); err == nil {
		t.Fatal("expected a dial error")
	}
}

func TestDialSSHHandshakeFailure(t *testing.T) {
	// A TCP listener that immediately closes connections: dial succeeds, the
	// SSH handshake does not.
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = ln.Close() })
	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				return
			}
			_ = conn.Close()
		}
	}()

	keys, err := generateKeypair()
	if err != nil {
		t.Fatal(err)
	}
	port := ln.Addr().(*net.TCPAddr).Port
	if _, err := dialSSH(context.Background(), "127.0.0.1", port, keys.signer); err == nil {
		t.Fatal("expected an SSH handshake error")
	}
}

func TestRunCancelledContextKillsCommand(t *testing.T) {
	host, port := startSSHServer(t, func(cmd string, ch ssh.Channel) {
		// Never send exit-status: simulates a hung command until the channel
		// is torn down by the client's cancellation path.
		<-make(chan struct{})
	})

	keys, err := generateKeypair()
	if err != nil {
		t.Fatal(err)
	}
	runner, err := dialSSH(context.Background(), host, port, keys.signer)
	if err != nil {
		t.Fatalf("dialSSH: %v", err)
	}
	defer func() { _ = runner.close() }()

	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, err := runner.run(ctx, "sleep 999", nil, nil, nil); err == nil {
		t.Fatal("expected a context cancellation error")
	}
}