ssh_test.go

package proxy

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"net"
	"testing"
	"time"

	"github.com/nemirovsky/sluice/internal/vault"
	"golang.org/x/crypto/ssh"
)

// tcpConnPair creates a pair of connected TCP connections. Unlike net.Pipe(),
// TCP connections have kernel buffering so both sides can write concurrently
// without deadlocking (required for SSH version exchange).
func tcpConnPair(t *testing.T) (client, server net.Conn) {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = ln.Close() }()

	done := make(chan net.Conn, 1)
	go func() {
		c, err := ln.Accept()
		if err != nil {
			return
		}
		done <- c
	}()

	client, err = net.Dial("tcp", ln.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	server = <-done
	return client, server
}

// generateTestSSHKey creates an ECDSA key pair and returns the SSH
// public key and PEM-encoded private key suitable for vault storage.
func generateTestSSHKey(t *testing.T) (ssh.PublicKey, []byte) {
	t.Helper()
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	signer, err := ssh.NewSignerFromKey(key)
	if err != nil {
		t.Fatal(err)
	}
	privDER, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		t.Fatal(err)
	}
	privPEM := pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: privDER})
	return signer.PublicKey(), privPEM
}

// startTestSSHServer starts an in-process SSH server that only accepts
// connections authenticated with the given public key. It responds to
// "exec" requests by writing "hello from ssh" to stdout with exit code 0.
func startTestSSHServer(t *testing.T, authorizedKey ssh.PublicKey) net.Listener {
	t.Helper()

	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	hostSigner, err := ssh.NewSignerFromKey(key)
	if err != nil {
		t.Fatal(err)
	}

	config := &ssh.ServerConfig{
		PublicKeyCallback: func(conn ssh.ConnMetadata, pubKey ssh.PublicKey) (*ssh.Permissions, error) {
			if bytes.Equal(pubKey.Marshal(), authorizedKey.Marshal()) {
				return &ssh.Permissions{}, nil
			}
			return nil, fmt.Errorf("unknown public key for user %q", conn.User())
		},
	}
	config.AddHostKey(hostSigner)

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

	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				return
			}
			go serveTestSSHConn(conn, config)
		}
	}()

	return ln
}

func serveTestSSHConn(conn net.Conn, config *ssh.ServerConfig) {
	sshConn, chans, reqs, err := ssh.NewServerConn(conn, config)
	if err != nil {
		_ = conn.Close()
		return
	}
	defer func() { _ = sshConn.Close() }()

	go ssh.DiscardRequests(reqs)

	for newChan := range chans {
		if newChan.ChannelType() != "session" {
			_ = newChan.Reject(ssh.UnknownChannelType, "unsupported channel type")
			continue
		}
		ch, reqs, err := newChan.Accept()
		if err != nil {
			continue
		}
		go func(ch ssh.Channel, reqs <-chan *ssh.Request) {
			defer func() { _ = ch.Close() }()
			for req := range reqs {
				switch req.Type {
				case "exec":
					if req.WantReply {
						_ = req.Reply(true, nil)
					}
					_, _ = ch.Write([]byte("hello from ssh"))
					_, _ = ch.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{0}))
					_ = ch.CloseWrite()
					// Wait briefly before returning so the defer-close races
					// do not close the channel before the proxy's io.Copy
					// has drained the data buffer.
					time.Sleep(50 * time.Millisecond)
					return
				default:
					if req.WantReply {
						_ = req.Reply(false, nil)
					}
				}
			}
		}(ch, reqs)
	}
}

func TestSSHJumpHostInjectsKey(t *testing.T) {
	// Generate SSH key pair for upstream authentication.
	pubKey, privPEM := generateTestSSHKey(t)

	// Store private key in vault.
	dir := t.TempDir()
	store, err := vault.NewStore(dir)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := store.Add("ssh_key", string(privPEM)); err != nil {
		t.Fatal(err)
	}

	// Start test SSH server that only accepts our public key.
	sshServer := startTestSSHServer(t, pubKey)
	defer func() { _ = sshServer.Close() }()

	// Generate host key for the proxy's SSH server side.
	proxyHostKey, err := GenerateSSHHostKey()
	if err != nil {
		t.Fatal(err)
	}

	binding := vault.Binding{
		Credential: "ssh_key",
		Template:   "testuser",
		Protocols:  []string{"ssh"},
	}

	jumpHost := NewSSHJumpHost(store, proxyHostKey)
	jumpHost.HostKeyCallback = ssh.InsecureIgnoreHostKey()

	// Use a TCP connection pair (buffered, unlike net.Pipe).
	agentConn, proxyConn := tcpConnPair(t)

	ready := make(chan error, 1)
	errCh := make(chan error, 1)
	go func() {
		errCh <- jumpHost.HandleConnection(proxyConn, []string{sshServer.Addr().String()}, sshServer.Addr().String(), binding, ready)
	}()

	if setupErr := <-ready; setupErr != nil {
		t.Fatalf("handler setup: %v", setupErr)
	}

	// Agent SSH client connects with no credentials.
	agentSSH, agentChans, agentReqs, err := ssh.NewClientConn(agentConn, "proxy", &ssh.ClientConfig{
		User:            "ignored",
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	})
	if err != nil {
		t.Fatalf("agent SSH handshake: %v", err)
	}

	client := ssh.NewClient(agentSSH, agentChans, agentReqs)

	session, err := client.NewSession()
	if err != nil {
		t.Fatalf("open session: %v", err)
	}

	output, err := session.Output("echo test")
	if err != nil {
		t.Fatalf("exec: %v", err)
	}

	if string(output) != "hello from ssh" {
		t.Errorf("expected 'hello from ssh', got %q", string(output))
	}

	_ = client.Close()
	_ = agentSSH.Close()
}

func TestSSHJumpHostMissingCredential(t *testing.T) {
	dir := t.TempDir()
	store, err := vault.NewStore(dir)
	if err != nil {
		t.Fatal(err)
	}

	proxyHostKey, err := GenerateSSHHostKey()
	if err != nil {
		t.Fatal(err)
	}

	jumpHost := NewSSHJumpHost(store, proxyHostKey)
	jumpHost.HostKeyCallback = ssh.InsecureIgnoreHostKey()

	agentConn, proxyConn := tcpConnPair(t)
	defer func() { _ = agentConn.Close() }()

	binding := vault.Binding{
		Credential: "nonexistent",
		Template:   "testuser",
	}

	ready := make(chan error, 1)
	err = jumpHost.HandleConnection(proxyConn, []string{"127.0.0.1:22"}, "127.0.0.1:22", binding, ready)
	if err == nil {
		t.Fatal("expected error for missing credential")
	}
}

func TestSSHJumpHostBadKey(t *testing.T) {
	// Generate two different key pairs: one authorized, one not.
	pubKey, _ := generateTestSSHKey(t)
	_, wrongPEM := generateTestSSHKey(t)

	// Store the WRONG key in the vault.
	dir := t.TempDir()
	store, err := vault.NewStore(dir)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := store.Add("wrong_key", string(wrongPEM)); err != nil {
		t.Fatal(err)
	}

	// Start server that only accepts the first key.
	sshServer := startTestSSHServer(t, pubKey)
	defer func() { _ = sshServer.Close() }()

	proxyHostKey, err := GenerateSSHHostKey()
	if err != nil {
		t.Fatal(err)
	}

	binding := vault.Binding{
		Credential: "wrong_key",
		Template:   "testuser",
	}

	jumpHost := NewSSHJumpHost(store, proxyHostKey)
	jumpHost.HostKeyCallback = ssh.InsecureIgnoreHostKey()

	agentConn, proxyConn := tcpConnPair(t)
	defer func() { _ = agentConn.Close() }()

	// HandleConnection should fail because the upstream rejects our key.
	ready := make(chan error, 1)
	err = jumpHost.HandleConnection(proxyConn, []string{sshServer.Addr().String()}, sshServer.Addr().String(), binding, ready)
	if err == nil {
		t.Fatal("expected error when upstream rejects key")
	}
}

func TestSSHVaultIntegrityAfterHandshake(t *testing.T) {
	pubKey, privPEM := generateTestSSHKey(t)

	dir := t.TempDir()
	store, err := vault.NewStore(dir)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := store.Add("ssh_key_zero", string(privPEM)); err != nil {
		t.Fatal(err)
	}

	sshServer := startTestSSHServer(t, pubKey)
	defer func() { _ = sshServer.Close() }()

	proxyHostKey, err := GenerateSSHHostKey()
	if err != nil {
		t.Fatal(err)
	}

	binding := vault.Binding{
		Credential: "ssh_key_zero",
		Template:   "testuser",
	}

	jumpHost := NewSSHJumpHost(store, proxyHostKey)
	jumpHost.HostKeyCallback = ssh.InsecureIgnoreHostKey()

	agentConn, proxyConn := tcpConnPair(t)

	ready := make(chan error, 1)
	errCh := make(chan error, 1)
	go func() {
		errCh <- jumpHost.HandleConnection(proxyConn, []string{sshServer.Addr().String()}, sshServer.Addr().String(), binding, ready)
	}()

	if setupErr := <-ready; setupErr != nil {
		t.Fatalf("handler setup: %v", setupErr)
	}

	// Connect, run command, disconnect.
	agentSSH, agentChans, agentReqs, err := ssh.NewClientConn(agentConn, "proxy", &ssh.ClientConfig{
		User:            "ignored",
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	})
	if err != nil {
		t.Fatalf("agent SSH handshake: %v", err)
	}

	client := ssh.NewClient(agentSSH, agentChans, agentReqs)
	session, err := client.NewSession()
	if err != nil {
		t.Fatalf("open session: %v", err)
	}
	_, _ = session.Output("test")
	_ = client.Close()
	_ = agentSSH.Close()

	// Wait for HandleConnection to return.
	<-errCh

	// The vault's encrypted credential is unaffected by zeroing the
	// in-memory SecureBytes copy. Verify the vault still works.
	after, err := store.Get("ssh_key_zero")
	if err != nil {
		t.Fatalf("credential should still be readable from vault: %v", err)
	}
	if after.IsReleased() {
		t.Error("stored credential should not appear released")
	}
	after.Release()
}

// TestResolveHostKeyCallbackExplicit tests that an explicit callback is returned.
func TestResolveHostKeyCallbackExplicit(t *testing.T) {
	called := false
	cb := func(_ string, _ net.Addr, _ ssh.PublicKey) error {
		called = true
		return nil
	}

	h := &SSHJumpHost{HostKeyCallback: cb}
	got, err := h.resolveHostKeyCallback()
	if err != nil {
		t.Fatal(err)
	}

	// Call the returned callback to verify it's our custom one.
	_ = got("example.com:22", nil, nil)
	if !called {
		t.Error("expected custom callback to be called")
	}
}

// TestResolveHostKeyCallbackNoKnownHosts tests fallback when no known_hosts exists.
func TestResolveHostKeyCallbackNoKnownHosts(t *testing.T) {
	// Set HOME to a dir without .ssh/known_hosts.
	t.Setenv("HOME", t.TempDir())

	h := &SSHJumpHost{}
	_, err := h.resolveHostKeyCallback()
	if err == nil {
		t.Fatal("expected error when no known_hosts and no explicit callback")
	}
}

// TestSSHJumpHost_BurstCloseDoesNotDropExecReply is a focused regression
// test for the race where an upstream that replies + writes data + sends
// exit-status + closes in one burst (no sleep between exit-status and
// close) caused sluice to close srcChan before the agent-to-upstream
// forwarder finished writing the SUCCESS reply for the agent's
// session.SendRequest("exec", true, ...). gossh closes ch.msg on
// SSH_MSG_CHANNEL_CLOSE, so the blocked SendRequest returns io.EOF and
// session.Output(...) fails with "exec command via SSH: EOF".
//
// startTestSSHServer (used by other tests) papers over the race with a
// 50ms sleep before returning from the channel handler. This test
// spins up its own burst-close server with no such sleep, so the race
// is deterministically triggered without the inflightBarrier fix.
func TestSSHJumpHost_BurstCloseDoesNotDropExecReply(t *testing.T) {
	pubKey, privPEM := generateTestSSHKey(t)
	dir := t.TempDir()
	store, err := vault.NewStore(dir)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := store.Add("ssh_key", string(privPEM)); err != nil {
		t.Fatal(err)
	}

	sshServer := startBurstCloseSSHServer(t, pubKey)
	defer func() { _ = sshServer.Close() }()

	proxyHostKey, err := GenerateSSHHostKey()
	if err != nil {
		t.Fatal(err)
	}

	binding := vault.Binding{
		Credential: "ssh_key",
		Template:   "testuser",
		Protocols:  []string{"ssh"},
	}

	jumpHost := NewSSHJumpHost(store, proxyHostKey)
	jumpHost.HostKeyCallback = ssh.InsecureIgnoreHostKey()

	// Run the test many times in a single process to maximize the
	// chance the close race fires if the fix regresses. Each iteration
	// runs through a fresh proxy connection + fresh agent session.
	const iterations = 50
	for i := 0; i < iterations; i++ {
		agentConn, proxyConn := tcpConnPair(t)

		ready := make(chan error, 1)
		errCh := make(chan error, 1)
		go func() {
			errCh <- jumpHost.HandleConnection(proxyConn, []string{sshServer.Addr().String()}, sshServer.Addr().String(), binding, ready)
		}()

		if setupErr := <-ready; setupErr != nil {
			t.Fatalf("iter %d: handler setup: %v", i, setupErr)
		}

		agentSSH, agentChans, agentReqs, err := ssh.NewClientConn(agentConn, "proxy", &ssh.ClientConfig{
			User:            "ignored",
			HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		})
		if err != nil {
			t.Fatalf("iter %d: agent SSH handshake: %v", i, err)
		}

		client := ssh.NewClient(agentSSH, agentChans, agentReqs)
		session, err := client.NewSession()
		if err != nil {
			t.Fatalf("iter %d: open session: %v", i, err)
		}

		output, err := session.Output("whoami")
		if err != nil {
			t.Fatalf("iter %d: exec: %v (this is the EOF symptom of the close race)", i, err)
		}
		if string(output) != "ssh-injection-ok\n" {
			t.Errorf("iter %d: expected 'ssh-injection-ok', got %q", i, string(output))
		}
		_ = session.Close()
		_ = client.Close()
		_ = agentSSH.Close()
		_ = agentConn.Close()

		// Wait for HandleConnection to return so a leaked handler
		// goroutine (or a connection that fails to teardown after
		// close) surfaces as a test timeout rather than as silent
		// resource exhaustion on the next iteration.
		select {
		case <-errCh:
		case <-time.After(5 * time.Second):
			t.Fatalf("iter %d: HandleConnection did not return within 5s after close", i)
		}
	}
}

// startBurstCloseSSHServer is a test SSH server that, on the first exec
// request, replies + writes output + sends exit-status + closes the
// channel with no delay between exit-status and Close. The lack of any
// sleep is intentional: it deterministically triggers the close race
// in sluice's SSH jump host when the inflightBarrier fix is absent.
func startBurstCloseSSHServer(t *testing.T, authorizedKey ssh.PublicKey) net.Listener {
	t.Helper()

	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	hostSigner, err := ssh.NewSignerFromKey(key)
	if err != nil {
		t.Fatal(err)
	}

	config := &ssh.ServerConfig{
		PublicKeyCallback: func(_ ssh.ConnMetadata, pubKey ssh.PublicKey) (*ssh.Permissions, error) {
			if bytes.Equal(pubKey.Marshal(), authorizedKey.Marshal()) {
				return &ssh.Permissions{}, nil
			}
			return nil, fmt.Errorf("unknown public key")
		},
	}
	config.AddHostKey(hostSigner)

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

	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				return
			}
			go func(c net.Conn) {
				sshConn, chans, reqs, err := ssh.NewServerConn(c, config)
				if err != nil {
					_ = c.Close()
					return
				}
				defer func() { _ = sshConn.Close() }()
				go ssh.DiscardRequests(reqs)
				for newChan := range chans {
					if newChan.ChannelType() != "session" {
						_ = newChan.Reject(ssh.UnknownChannelType, "unsupported")
						continue
					}
					ch, chReqs, err := newChan.Accept()
					if err != nil {
						continue
					}
					go func(ch ssh.Channel, reqs <-chan *ssh.Request) {
						// Defer close so a request loop that exits without
						// hitting the exec path (early agent close,
						// non-exec request only) still releases the
						// server-side channel.
						defer func() { _ = ch.Close() }()
						for req := range reqs {
							if req.Type != "exec" {
								if req.WantReply {
									_ = req.Reply(false, nil)
								}
								continue
							}
							if req.WantReply {
								_ = req.Reply(true, nil)
							}
							_, _ = ch.Write([]byte("ssh-injection-ok\n"))
							_, _ = ch.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{0}))
							_ = ch.CloseWrite()
							// NO time.Sleep here. This is the whole point of
							// the test: close immediately after exit-status
							// to maximally tighten the race window in
							// sluice's sshHandleChannel.
							return
						}
					}(ch, chReqs)
				}
			}(conn)
		}
	}()
	return ln
}

// TestGenerateSSHHostKey tests SSH host key generation.
func TestGenerateSSHHostKey(t *testing.T) {
	signer, err := GenerateSSHHostKey()
	if err != nil {
		t.Fatal(err)
	}
	if signer == nil {
		t.Fatal("expected non-nil signer")
	}
	if signer.PublicKey() == nil {
		t.Error("expected non-nil public key")
	}
}