a73x

internal/server/sshgate/gate_test.go

Ref:   Size: 31.9 KiB   History

package sshgate

import (
	"bytes"
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"io"
	"net"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/crypto/ssh"
)

// keysEqual reports whether two SSH public keys are byte-identical. Only tests
// compare keys this way; production matches via the certificate chain.
func keysEqual(a, b ssh.PublicKey) bool {
	return a != nil && b != nil && bytes.Equal(a.Marshal(), b.Marshal())
}

// newSigner returns a throwaway ed25519 ssh.Signer.
func newSigner(t *testing.T) ssh.Signer {
	t.Helper()
	_, priv, err := ed25519.GenerateKey(rand.Reader)
	require.NoError(t, err)
	s, err := ssh.NewSignerFromSigner(priv)
	require.NoError(t, err)
	return s
}

// mintCertSigner signs clientKey with ca into a user cert (principal ubuntu) and
// returns a cert signer usable as an SSH auth method.
func mintCertSigner(t *testing.T, ca, clientKey ssh.Signer) ssh.Signer {
	t.Helper()
	cert := &ssh.Certificate{
		Key:             clientKey.PublicKey(),
		Serial:          1,
		CertType:        ssh.UserCert,
		KeyId:           "ubuntu",
		ValidPrincipals: []string{"ubuntu"},
		ValidAfter:      uint64(time.Now().Add(-time.Minute).Unix()),
		ValidBefore:     uint64(time.Now().Add(time.Hour).Unix()),
	}
	require.NoError(t, cert.SignCert(rand.Reader, ca))
	cs, err := ssh.NewCertSigner(cert, clientKey)
	require.NoError(t, err)
	return cs
}

// mintCertSignerNoPrincipals signs clientKey with ca into a user cert with an
// EMPTY principal set. Under CheckCert's wildcard rule such a cert is valid for
// ANY principal, so the gate must reject it outright.
func mintCertSignerNoPrincipals(t *testing.T, ca, clientKey ssh.Signer) ssh.Signer {
	t.Helper()
	cert := &ssh.Certificate{
		Key:         clientKey.PublicKey(),
		Serial:      1,
		CertType:    ssh.UserCert,
		KeyId:       "ubuntu",
		ValidAfter:  uint64(time.Now().Add(-time.Minute).Unix()),
		ValidBefore: uint64(time.Now().Add(time.Hour).Unix()),
	}
	require.NoError(t, cert.SignCert(rand.Reader, ca))
	cs, err := ssh.NewCertSigner(cert, clientKey)
	require.NoError(t, err)
	return cs
}

// testGate wires a gate over a loopback listener and returns the client-side
// dial address plus the wired fakes' observed state.
type testGate struct {
	gate       *Gate
	addr       string
	hostKey    ssh.Signer
	dialCalls  chan [3]string // hostID, vmID, port-as-string per dial
	authorized bool
}

// startGate builds a gate with an echoing dialer and a single known VM "vm1",
// serving on 127.0.0.1:0. authorize returns the given result.
func startGate(t *testing.T, userCA ssh.PublicKey, authorized bool) *testGate {
	t.Helper()
	tg := &testGate{
		hostKey:    newSigner(t),
		dialCalls:  make(chan [3]string, 4),
		authorized: authorized,
	}
	resolve := func(tenant, name string) (string, string, bool) {
		if tenant == "default" && name == "vm1" {
			return "host-1", "vm-1", true
		}
		return "", "", false
	}
	authorize := func(tenant, vmID string) bool { return tg.authorized }
	dial := func(_ context.Context, hostID, vmID string, port uint32) (io.ReadWriteCloser, error) {
		tg.dialCalls <- [3]string{hostID, vmID, "22"}
		a, b := net.Pipe()
		go func() { _, _ = io.Copy(b, b); b.Close() }() // echo server = fake VM sshd
		return a, nil
	}
	g := New(tg.hostKey, singleCALookup(userCA, "default"), resolve, authorize, dial, nil, nil)

	tg.gate = g
	l, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	tg.addr = l.Addr().String()
	go func() { _ = g.Serve(l) }()
	t.Cleanup(func() { _ = l.Close() })
	return tg
}

// gateStarting exposes the gate's startup semaphore for the MaxStartups test.
func (tg *testGate) gateStarting() chan struct{} { return tg.gate.starting }

// singleCALookup builds a UserCALookup that trusts exactly one CA public key,
// stamping its connections with tenant. Any other signing key resolves to
// ok=false (rejected) — the test-side analogue of the DB-backed lookup.
func singleCALookup(ca ssh.PublicKey, tenant string) UserCALookup {
	return func(pub ssh.PublicKey) (string, bool) {
		if keysEqual(pub, ca) {
			return tenant, true
		}
		return "", false
	}
}

// startGateRevoked is startGate with an explicit revocation predicate wired, so
// a test can assert a revoked serial fails auth.
func startGateRevoked(t *testing.T, userCA ssh.PublicKey, isRevoked Revoker) *testGate {
	t.Helper()
	tg := &testGate{hostKey: newSigner(t), dialCalls: make(chan [3]string, 4), authorized: true}
	resolve := func(tenant, name string) (string, string, bool) {
		if tenant == "default" && name == "vm1" {
			return "host-1", "vm-1", true
		}
		return "", "", false
	}
	authorize := func(tenant, vmID string) bool { return tg.authorized }
	dial := func(_ context.Context, hostID, vmID string, port uint32) (io.ReadWriteCloser, error) {
		tg.dialCalls <- [3]string{hostID, vmID, "22"}
		a, b := net.Pipe()
		go func() { _, _ = io.Copy(b, b); b.Close() }()
		return a, nil
	}
	g := New(tg.hostKey, singleCALookup(userCA, "default"), resolve, authorize, dial, isRevoked, nil)
	tg.gate = g
	l, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	tg.addr = l.Addr().String()
	go func() { _ = g.Serve(l) }()
	t.Cleanup(func() { _ = l.Close() })
	return tg
}

// mintCertSignerSerial signs clientKey with ca into a user cert carrying an
// explicit serial, so a revocation test can target that serial.
func mintCertSignerSerial(t *testing.T, ca, clientKey ssh.Signer, serial uint64) ssh.Signer {
	t.Helper()
	cert := &ssh.Certificate{
		Key:             clientKey.PublicKey(),
		Serial:          serial,
		CertType:        ssh.UserCert,
		KeyId:           "ubuntu",
		ValidPrincipals: []string{"ubuntu"},
		ValidAfter:      uint64(time.Now().Add(-time.Minute).Unix()),
		ValidBefore:     uint64(time.Now().Add(time.Hour).Unix()),
	}
	require.NoError(t, cert.SignCert(rand.Reader, ca))
	cs, err := ssh.NewCertSigner(cert, clientKey)
	require.NoError(t, err)
	return cs
}

// dialClient connects an SSH client to the gate using certSigner.
func dialClient(t *testing.T, tg *testGate, certSigner ssh.Signer) *ssh.Client {
	t.Helper()
	cfg := &ssh.ClientConfig{
		User:            "some-random-outer-name", // must NOT matter on the gate hop
		Auth:            []ssh.AuthMethod{ssh.PublicKeys(certSigner)},
		HostKeyCallback: ssh.FixedHostKey(tg.hostKey.PublicKey()),
		Timeout:         5 * time.Second,
	}
	c, err := ssh.Dial("tcp", tg.addr, cfg)
	require.NoError(t, err)
	t.Cleanup(func() { _ = c.Close() })
	return c
}

// startGateWithHostKey is startGate with an explicit host-key signer, so a test
// can present a CA-signed host certificate (via ssh.NewCertSigner) rather than a
// bare host key.
func startGateWithHostKey(t *testing.T, hostKey ssh.Signer, userCA ssh.PublicKey) *testGate {
	t.Helper()
	tg := &testGate{hostKey: hostKey, dialCalls: make(chan [3]string, 4), authorized: true}
	resolve := func(tenant, name string) (string, string, bool) {
		if tenant == "default" && name == "vm1" {
			return "host-1", "vm-1", true
		}
		return "", "", false
	}
	authorize := func(tenant, vmID string) bool { return tg.authorized }
	dial := func(_ context.Context, hostID, vmID string, port uint32) (io.ReadWriteCloser, error) {
		tg.dialCalls <- [3]string{hostID, vmID, "22"}
		a, b := net.Pipe()
		go func() { _, _ = io.Copy(b, b); b.Close() }()
		return a, nil
	}
	g := New(hostKey, singleCALookup(userCA, "default"), resolve, authorize, dial, nil, nil)
	tg.gate = g
	l, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	tg.addr = l.Addr().String()
	go func() { _ = g.Serve(l) }()
	t.Cleanup(func() { _ = l.Close() })
	return tg
}

// TestGatePresentsCASignedHostCert verifies the gate presents a host key whose
// certificate is signed by the eitri CA, and that a client doing
// `@cert-authority`-style verification (CertChecker.IsHostAuthority) accepts it
// WITHOUT any prior TOFU pin — the whole point of host-cert signing.
func TestGatePresentsCASignedHostCert(t *testing.T) {
	ca := newSigner(t)
	hostKey := newSigner(t)

	// Sign a HOST cert for the gate's host key, scoped to the name the client
	// dials ("gate"), and present it via a cert signer.
	cert := &ssh.Certificate{
		Key:             hostKey.PublicKey(),
		Serial:          1,
		CertType:        ssh.HostCert,
		KeyId:           "eitri-gate",
		ValidPrincipals: []string{"gate"},
		ValidAfter:      uint64(time.Now().Add(-time.Minute).Unix()),
		ValidBefore:     uint64(time.Now().Add(time.Hour).Unix()),
	}
	require.NoError(t, cert.SignCert(rand.Reader, ca))
	hostCertSigner, err := ssh.NewCertSigner(cert, hostKey)
	require.NoError(t, err)

	tg := startGateWithHostKey(t, hostCertSigner, ca.PublicKey())

	checker := &ssh.CertChecker{
		IsHostAuthority: func(k ssh.PublicKey, _ string) bool {
			return keysEqual(k, ca.PublicKey())
		},
	}
	cfg := &ssh.ClientConfig{
		User:            "ubuntu",
		Auth:            []ssh.AuthMethod{ssh.PublicKeys(mintCertSigner(t, ca, newSigner(t)))},
		HostKeyCallback: checker.CheckHostKey,
		Timeout:         5 * time.Second,
	}
	// The client dials the gate as host "gate" (the cert principal) so principal
	// scoping is exercised, not just the CA signature.
	nc, err := net.Dial("tcp", tg.addr)
	require.NoError(t, err)
	c, chans, reqs, err := ssh.NewClientConn(nc, "gate:22", cfg)
	require.NoError(t, err, "client must accept a CA-signed host cert without a TOFU pin")
	client := ssh.NewClient(c, chans, reqs)
	t.Cleanup(func() { _ = client.Close() })

	// And the tunnel still works end-to-end over the cert-authenticated host.
	conn, err := client.Dial("tcp", "default.vm1:22")
	require.NoError(t, err)
	_ = conn.Close()
}

func TestGateRejectsSessionChannel(t *testing.T) {
	ca := newSigner(t)
	tg := startGate(t, ca.PublicKey(), true)
	client := dialClient(t, tg, mintCertSigner(t, ca, newSigner(t)))

	_, _, err := client.OpenChannel("session", nil)
	require.Error(t, err, "session channel must be rejected")
	var oce *ssh.OpenChannelError
	require.ErrorAs(t, err, &oce)
	assert.Equal(t, ssh.UnknownChannelType, oce.Reason)
}

func TestGateRejectsTCPIPForwardGlobalRequest(t *testing.T) {
	ca := newSigner(t)
	tg := startGate(t, ca.PublicKey(), true)
	client := dialClient(t, tg, mintCertSigner(t, ca, newSigner(t)))

	// tcpip-forward (ssh -R) must NOT be honored: the bastion is not an ingress relay.
	ok, _, err := client.SendRequest("tcpip-forward", true, ssh.Marshal(struct {
		Addr string
		Port uint32
	}{"0.0.0.0", 0}))
	require.NoError(t, err)
	assert.False(t, ok, "tcpip-forward must get a false reply")
}

func TestGateDirectTCPIPToPort22RoundTrips(t *testing.T) {
	ca := newSigner(t)
	tg := startGate(t, ca.PublicKey(), true)
	client := dialClient(t, tg, mintCertSigner(t, ca, newSigner(t)))

	conn, err := client.Dial("tcp", "default.vm1:22")
	require.NoError(t, err)
	defer conn.Close()

	// The dialer must have been reached with the resolved host/vm.
	select {
	case got := <-tg.dialCalls:
		assert.Equal(t, [3]string{"host-1", "vm-1", "22"}, got)
	case <-time.After(2 * time.Second):
		t.Fatal("dialer was never called")
	}

	// Bytes must round-trip through the echoing fake VM.
	want := []byte("hello-vm")
	_, err = conn.Write(want)
	require.NoError(t, err)
	got := make([]byte, len(want))
	_, err = io.ReadFull(conn, got)
	require.NoError(t, err)
	assert.Equal(t, want, got)
}

// TestGateHandshakeGraceDropsStalledConn pins the pre-auth DoS guard: a client
// that connects and then never completes the SSH handshake must be dropped by
// the handshake grace, not parked forever holding a goroutine + fd.
func TestGateHandshakeGraceDropsStalledConn(t *testing.T) {
	orig := handshakeGrace
	handshakeGrace = 150 * time.Millisecond
	t.Cleanup(func() { handshakeGrace = orig })

	ca := newSigner(t)
	tg := startGate(t, ca.PublicKey(), true)

	conn, err := net.Dial("tcp", tg.addr)
	require.NoError(t, err)
	defer conn.Close()

	// Never send a client identification string. The gate emits its banner then
	// blocks reading ours; the grace deadline must fire and close the connection,
	// so our read drains the banner and hits a clean EOF well before this generous
	// client deadline. Without the grace the server would block forever and we'd
	// instead trip our own read timeout.
	start := time.Now()
	require.NoError(t, conn.SetReadDeadline(time.Now().Add(5*time.Second)))
	_, err = io.ReadAll(conn)
	require.NoError(t, err, "server must close the stalled conn (EOF), not leave us to time out")
	assert.Less(t, time.Since(start), 3*time.Second, "stalled conn must drop near the handshake grace")
}

func TestGateDirectTCPIPToNonSSHPortRejected(t *testing.T) {
	ca := newSigner(t)
	tg := startGate(t, ca.PublicKey(), true)
	client := dialClient(t, tg, mintCertSigner(t, ca, newSigner(t)))

	_, err := client.Dial("tcp", "default.vm1:2222")
	require.Error(t, err, "only port 22 may be tunnelled")
}

func TestGateUnknownVMRejected(t *testing.T) {
	ca := newSigner(t)
	tg := startGate(t, ca.PublicKey(), true)
	client := dialClient(t, tg, mintCertSigner(t, ca, newSigner(t)))

	_, err := client.Dial("tcp", "default.nope:22")
	require.Error(t, err, "unknown VM name must be rejected")
}

func TestGateAuthzDenyRejected(t *testing.T) {
	ca := newSigner(t)
	tg := startGate(t, ca.PublicKey(), false) // authorize → deny
	client := dialClient(t, tg, mintCertSigner(t, ca, newSigner(t)))

	_, err := client.Dial("tcp", "default.vm1:22")
	require.Error(t, err, "authz denial must reject the channel")
}

// TestGateBareNameResolvesInCertTenant: a connect name with NO tenant prefix is
// the bare VM name within the connection's own tenant — the cert's signing CA
// already names the tenant — so it resolves and tunnels exactly like the
// explicit <tenant>.<vm> form.
func TestGateBareNameResolvesInCertTenant(t *testing.T) {
	ca := newSigner(t)
	tg := startGate(t, ca.PublicKey(), true)
	client := dialClient(t, tg, mintCertSigner(t, ca, newSigner(t)))

	conn, err := client.Dial("tcp", "vm1:22")
	require.NoError(t, err, "a bare VM name must resolve within the cert's tenant")
	defer conn.Close()

	// The dialer must have been reached with the VM resolved inside "default".
	select {
	case got := <-tg.dialCalls:
		assert.Equal(t, [3]string{"host-1", "vm-1", "22"}, got)
	case <-time.After(2 * time.Second):
		t.Fatal("dialer was never called")
	}
}

// TestGateBareNameForeignTenantRejected: resolution is tenant-scoped, so a bare
// name that belongs to ANOTHER tenant does not resolve within the connection's
// tenant and is rejected as an unknown VM.
func TestGateBareNameForeignTenantRejected(t *testing.T) {
	ca := newSigner(t)
	// The only known VM ("victim") lives in tenant "other", while the cert (and
	// thus the connection) is tenant "default": a tenant-scoped resolver never
	// hands it back for "default".
	tg := &testGate{hostKey: newSigner(t), dialCalls: make(chan [3]string, 4), authorized: true}
	resolve := func(tenant, name string) (string, string, bool) {
		if tenant == "other" && name == "victim" {
			return "host-x", "vm-x", true
		}
		return "", "", false
	}
	authorize := func(string, string) bool { return true }
	dial := func(_ context.Context, hostID, vmID string, port uint32) (io.ReadWriteCloser, error) {
		tg.dialCalls <- [3]string{hostID, vmID, "22"}
		a, b := net.Pipe()
		go func() { _, _ = io.Copy(b, b); b.Close() }()
		return a, nil
	}
	g := New(tg.hostKey, singleCALookup(ca.PublicKey(), "default"), resolve, authorize, dial, nil, nil)
	tg.gate = g
	l, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	tg.addr = l.Addr().String()
	go func() { _ = g.Serve(l) }()
	t.Cleanup(func() { _ = l.Close() })

	client := dialClient(t, tg, mintCertSigner(t, ca, newSigner(t)))
	_, err = client.Dial("tcp", "victim:22")
	require.Error(t, err, "a bare name for another tenant's VM must be rejected as unknown")
}

// TestGateForeignTenantNameIndistinguishable: an explicit <tenant>.<vm> whose
// prefix is a FOREIGN tenant is rejected with the SAME message as a nonexistent
// VM in the connection's own tenant, so tenancy structure is not probeable from
// the gate.
func TestGateForeignTenantNameIndistinguishable(t *testing.T) {
	ca := newSigner(t)
	tg := startGate(t, ca.PublicKey(), true)
	client := dialClient(t, tg, mintCertSigner(t, ca, newSigner(t)))

	// rejectMsg dials target (which must be rejected) and returns the gate's
	// channel-open reject message.
	rejectMsg := func(target string) string {
		_, err := client.Dial("tcp", target+":22")
		require.Error(t, err, "dialing %q must be rejected", target)
		var oce *ssh.OpenChannelError
		require.ErrorAs(t, err, &oce)
		return oce.Message
	}

	// A foreign tenant prefix must be rejected with the SAME message as a
	// nonexistent VM in our own tenant: the two are indistinguishable. Test-lock
	// the equal messages so a future refactor cannot split them silently.
	assert.Equal(t, rejectMsg("default.nope"), rejectMsg("other.vm1"),
		"a foreign tenant prefix must be indistinguishable from a nonexistent VM")

	// Control: the correctly namespaced name resolves and tunnels.
	conn, err := client.Dial("tcp", "default.vm1:22")
	require.NoError(t, err, "correctly namespaced name must succeed")
	_ = conn.Close()
}

func TestGateRejectsCertFromForeignCA(t *testing.T) {
	ca := newSigner(t)
	foreignCA := newSigner(t)
	tg := startGate(t, ca.PublicKey(), true)

	// A cert signed by a CA the gate does not trust must fail auth outright.
	certSigner := mintCertSigner(t, foreignCA, newSigner(t))
	cfg := &ssh.ClientConfig{
		User:            "ubuntu",
		Auth:            []ssh.AuthMethod{ssh.PublicKeys(certSigner)},
		HostKeyCallback: ssh.FixedHostKey(tg.hostKey.PublicKey()),
		Timeout:         5 * time.Second,
	}
	_, err := ssh.Dial("tcp", tg.addr, cfg)
	require.Error(t, err, "cert not signed by the eitri CA must fail auth")
}

func TestGateRejectsCertWithNoPrincipals(t *testing.T) {
	ca := newSigner(t)
	tg := startGate(t, ca.PublicKey(), true)

	// A cert with an empty principal set is valid for ANY principal under
	// CheckCert's wildcard rule — the gate must refuse it rather than let the
	// wildcard (and the ValidPrincipals[0] index) through.
	certSigner := mintCertSignerNoPrincipals(t, ca, newSigner(t))
	cfg := &ssh.ClientConfig{
		User:            "ubuntu",
		Auth:            []ssh.AuthMethod{ssh.PublicKeys(certSigner)},
		HostKeyCallback: ssh.FixedHostKey(tg.hostKey.PublicKey()),
		Timeout:         5 * time.Second,
	}
	_, err := ssh.Dial("tcp", tg.addr, cfg)
	require.Error(t, err, "cert with no principals must fail auth")
}

// TestGateRejectsRevokedCert mints a user cert with a known serial and asserts
// that with isRevoked true for that serial the client's auth FAILS, and with
// isRevoked false the same cert authenticates and tunnels through.
func TestGateRejectsRevokedCert(t *testing.T) {
	ca := newSigner(t)
	const serial = uint64(0xDEADBEEFCAFEF00D)

	// Revoked ⇒ auth fails.
	revoked := startGateRevoked(t, ca.PublicKey(), func(_ string, s uint64) bool { return s == serial })
	cfg := &ssh.ClientConfig{
		User:            "ubuntu",
		Auth:            []ssh.AuthMethod{ssh.PublicKeys(mintCertSignerSerial(t, ca, newSigner(t), serial))},
		HostKeyCallback: ssh.FixedHostKey(revoked.hostKey.PublicKey()),
		Timeout:         5 * time.Second,
	}
	_, err := ssh.Dial("tcp", revoked.addr, cfg)
	require.Error(t, err, "a revoked cert must fail auth at the gate")

	// Not revoked ⇒ the same serial authenticates and tunnels.
	allowed := startGateRevoked(t, ca.PublicKey(), func(string, uint64) bool { return false })
	client := dialClient(t, allowed, mintCertSignerSerial(t, ca, newSigner(t), serial))
	conn, err := client.Dial("tcp", "default.vm1:22")
	require.NoError(t, err, "a non-revoked cert must still tunnel")
	_ = conn.Close()
}

func TestGateRejectsBarePublicKey(t *testing.T) {
	ca := newSigner(t)
	tg := startGate(t, ca.PublicKey(), true)

	// A raw (non-certificate) key must be rejected — the gate is cert-only.
	cfg := &ssh.ClientConfig{
		User:            "ubuntu",
		Auth:            []ssh.AuthMethod{ssh.PublicKeys(newSigner(t))},
		HostKeyCallback: ssh.FixedHostKey(tg.hostKey.PublicKey()),
		Timeout:         5 * time.Second,
	}
	_, err := ssh.Dial("tcp", tg.addr, cfg)
	require.Error(t, err, "bare public key (no cert) must fail auth")
}

// TestGateTrustsRegisteredCARejectsUnknown proves the multi-CA boundary: a cert
// signed by a REGISTERED tenant CA authenticates and is stamped with THAT CA's
// tenant, while a cert signed by an UNREGISTERED CA is rejected. This drives the
// gate's PublicKeyCallback directly so the stamped tenant (an unexported
// Permissions extension) is asserted at the source, not merely inferred.
func TestGateTrustsRegisteredCARejectsUnknown(t *testing.T) {
	caAcme := newSigner(t) // registered as tenant "acme"
	caEvil := newSigner(t) // not registered

	lookup := func(pub ssh.PublicKey) (string, bool) {
		if keysEqual(pub, caAcme.PublicKey()) {
			return "acme", true
		}
		return "", false
	}
	resolve := func(string, string) (string, string, bool) { return "", "", false }
	authorize := func(string, string) bool { return true }
	dial := func(context.Context, string, string, uint32) (io.ReadWriteCloser, error) { return nil, nil }
	g := New(newSigner(t), lookup, resolve, authorize, dial, nil, nil)
	cb := g.cfg.PublicKeyCallback

	// A cert signed by the registered CA authenticates and is stamped "acme".
	acmeCert := mintCertSigner(t, caAcme, newSigner(t)).PublicKey().(*ssh.Certificate)
	perms, err := cb(nil, acmeCert)
	require.NoError(t, err, "a cert signed by a registered tenant CA must authenticate")
	require.NotNil(t, perms)
	assert.Equal(t, "acme", perms.Extensions[tenantExt],
		"the connection's tenant must be stamped from the signing CA")

	// A cert signed by an unregistered CA is rejected — the crypto boundary.
	evilCert := mintCertSigner(t, caEvil, newSigner(t)).PublicKey().(*ssh.Certificate)
	_, err = cb(nil, evilCert)
	require.Error(t, err, "a cert signed by an unregistered CA must be rejected")
}

// auditLog is a concurrency-safe recorder standing in for the store-backed
// audit sink. The gate calls it from each connection's own goroutine.
type auditLog struct {
	mu     sync.Mutex
	events []auditEvent
}

type auditEvent struct {
	tenant, action string
	detail         map[string]string
}

func (a *auditLog) record() Audit {
	return func(tenant, action string, detail map[string]string) {
		a.mu.Lock()
		defer a.mu.Unlock()
		a.events = append(a.events, auditEvent{tenant, action, detail})
	}
}

// waitFor returns the first recorded event with the given action, polling until
// it appears — the gate audits on the connection goroutine, so a client call
// can return before the row lands.
func (a *auditLog) waitFor(t *testing.T, action string) auditEvent {
	t.Helper()
	for deadline := time.Now().Add(2 * time.Second); time.Now().Before(deadline); {
		a.mu.Lock()
		for _, e := range a.events {
			if e.action == action {
				a.mu.Unlock()
				return e
			}
		}
		a.mu.Unlock()
		time.Sleep(5 * time.Millisecond)
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	t.Fatalf("no %q event; recorded %v", action, a.events)
	return auditEvent{}
}

// startGateAudited is startGate with an audit sink wired, so a test can assert
// the trail the gate leaves rather than only the bytes it moves.
func startGateAudited(t *testing.T, userCA ssh.PublicKey, authorized bool) (*testGate, *auditLog) {
	t.Helper()
	log := &auditLog{}
	tg := &testGate{hostKey: newSigner(t), dialCalls: make(chan [3]string, 4), authorized: authorized}
	resolve := func(tenant, name string) (string, string, bool) {
		if tenant == "default" && name == "vm1" {
			return "host-1", "vm-1", true
		}
		return "", "", false
	}
	authorize := func(tenant, vmID string) bool { return tg.authorized }
	dial := func(_ context.Context, hostID, vmID string, port uint32) (io.ReadWriteCloser, error) {
		tg.dialCalls <- [3]string{hostID, vmID, "22"}
		a, b := net.Pipe()
		go func() { _, _ = io.Copy(b, b); b.Close() }()
		return a, nil
	}
	g := New(tg.hostKey, singleCALookup(userCA, "default"), resolve, authorize, dial, nil, log.record())
	tg.gate = g
	l, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	tg.addr = l.Addr().String()
	go func() { _ = g.Serve(l) }()
	t.Cleanup(func() { _ = l.Close() })
	return tg, log
}

// TestGateAuditsAuthAndTunnel is the gate's answer to "who reached which VM".
// Both halves must be recorded: the certificate that opened the connection, and
// the VM it was pointed at.
func TestGateAuditsAuthAndTunnel(t *testing.T) {
	ca := newSigner(t)
	tg, log := startGateAudited(t, ca.PublicKey(), true)
	client := dialClient(t, tg, mintCertSigner(t, ca, newSigner(t)))

	auth := log.waitFor(t, "gate.auth")
	assert.Equal(t, "default", auth.tenant, "tenant comes from the signing CA")
	// The serial is the field a revocation names, so it is what ties a session
	// back to a certificate after the fact.
	assert.Equal(t, "1", auth.detail["serial"])
	assert.Equal(t, "ubuntu", auth.detail["key_id"])
	assert.NotEmpty(t, auth.detail["remote"])

	conn, err := client.Dial("tcp", "default.vm1:22")
	require.NoError(t, err)
	defer conn.Close()

	tun := log.waitFor(t, "gate.tunnel")
	assert.Equal(t, "default", tun.tenant)
	assert.Equal(t, "vm-1", tun.detail["vm_id"])
	assert.Equal(t, "host-1", tun.detail["host_id"])
}

// TestGateAuditsRefusedLogin pins that a login the gate turns away is recorded
// too — an unattributed event, since no tenant was ever established.
func TestGateAuditsRefusedLogin(t *testing.T) {
	ca, foreign := newSigner(t), newSigner(t)
	tg, log := startGateAudited(t, ca.PublicKey(), true)

	cfg := &ssh.ClientConfig{
		User:            "probe",
		Auth:            []ssh.AuthMethod{ssh.PublicKeys(mintCertSigner(t, foreign, newSigner(t)))},
		HostKeyCallback: ssh.FixedHostKey(tg.hostKey.PublicKey()),
		Timeout:         5 * time.Second,
	}
	c, err := ssh.Dial("tcp", tg.addr, cfg)
	require.Error(t, err, "a cert from an unregistered CA must not authenticate")
	if c != nil {
		_ = c.Close()
	}

	ev := log.waitFor(t, "gate.auth.denied")
	assert.Empty(t, ev.tenant, "a refused login has no tenant to attribute")
	assert.NotEmpty(t, ev.detail["remote"])
	assert.NotEmpty(t, ev.detail["reason"])
}

// TestGateAuditsRefusedTunnel pins the other refusal: an authenticated client
// that asks for something it may not have still leaves a row naming what it
// asked for.
func TestGateAuditsRefusedTunnel(t *testing.T) {
	ca := newSigner(t)
	tg, log := startGateAudited(t, ca.PublicKey(), true)
	client := dialClient(t, tg, mintCertSigner(t, ca, newSigner(t)))

	_, err := client.Dial("tcp", "default.nope:22")
	require.Error(t, err)

	ev := log.waitFor(t, "gate.tunnel.denied")
	assert.Equal(t, "default", ev.tenant)
	assert.Equal(t, "default.nope", ev.detail["target"])
	assert.Equal(t, "unknown VM", ev.detail["reason"])
}

// TestGateShedsBeyondMaxStartups pins the aggregate pre-auth bound. The grace
// deadline caps how long one stalled client holds a slot; this caps how many,
// which is what stops a flood of merely-slow clients from crowding out logins.
func TestGateShedsBeyondMaxStartups(t *testing.T) {
	prev := maxStartups
	maxStartups = 2
	t.Cleanup(func() { maxStartups = prev })

	ca := newSigner(t)
	tg, log := startGateAudited(t, ca.PublicKey(), true)

	// Occupy every slot with clients that connect and then say nothing, so each
	// sits in the handshake holding its slot until the grace expires.
	var stalled []net.Conn
	for range maxStartups {
		c, err := net.Dial("tcp", tg.addr)
		require.NoError(t, err)
		stalled = append(stalled, c)
	}
	t.Cleanup(func() {
		for _, c := range stalled {
			_ = c.Close()
		}
	})

	// The gate accepts asynchronously, so wait until the slots are actually held.
	require.Eventually(t, func() bool { return len(tg.gateStarting()) == maxStartups },
		2*time.Second, 5*time.Millisecond, "slots never filled")

	// One more must be shed rather than queued. A shed connection is closed
	// without a handshake, so it never sees the version banner an accepted one
	// gets — and it happens at once, rather than after the grace window.
	over, err := net.Dial("tcp", tg.addr)
	require.NoError(t, err, "TCP still accepts; the gate sheds above the SSH layer")
	defer over.Close()

	require.NoError(t, over.SetReadDeadline(time.Now().Add(2*time.Second)))
	_, err = over.Read(make([]byte, 1))
	require.ErrorIs(t, err, io.EOF, "a shed connection is closed, not left in the handshake")

	// Shedding must NOT write an audit row: it is the cheapest event to provoke
	// on the gate, and a row per refusal would make the defence amplify the
	// flood it exists to absorb.
	log.mu.Lock()
	defer log.mu.Unlock()
	for _, e := range log.events {
		assert.NotEqual(t, "gate.auth.denied", e.action, "the shed path must not audit")
	}
}

// TestGateAuthenticatedConnDoesNotHoldStartupSlot is the other half of the
// bound: a tunnel may last hours, and must not occupy the pre-auth window while
// it does — or a busy gate would refuse logins it has every reason to take.
func TestGateAuthenticatedConnDoesNotHoldStartupSlot(t *testing.T) {
	prev := maxStartups
	maxStartups = 1
	t.Cleanup(func() { maxStartups = prev })

	ca := newSigner(t)
	tg, _ := startGateAudited(t, ca.PublicKey(), true)

	// Hold one authenticated connection open with a live tunnel.
	client := dialClient(t, tg, mintCertSigner(t, ca, newSigner(t)))
	conn, err := client.Dial("tcp", "default.vm1:22")
	require.NoError(t, err)
	defer conn.Close()

	// With the only slot released at auth, a second login must still succeed.
	second := dialClient(t, tg, mintCertSigner(t, ca, newSigner(t)))
	_, err = second.Dial("tcp", "default.vm1:22")
	require.NoError(t, err, "an established tunnel must not consume a startup slot")
}

// mintCertSignerCriticalOptions signs clientKey with ca into a user cert
// carrying the given critical options, so a test can present a cert whose
// restrictions the gate is expected to honour or refuse.
func mintCertSignerCriticalOptions(t *testing.T, ca, clientKey ssh.Signer, opts map[string]string) ssh.Signer {
	t.Helper()
	cert := &ssh.Certificate{
		Key:             clientKey.PublicKey(),
		Serial:          1,
		CertType:        ssh.UserCert,
		KeyId:           "ubuntu",
		ValidPrincipals: []string{"ubuntu"},
		ValidAfter:      uint64(time.Now().Add(-time.Minute).Unix()),
		ValidBefore:     uint64(time.Now().Add(time.Hour).Unix()),
		Permissions:     ssh.Permissions{CriticalOptions: opts},
	}
	require.NoError(t, cert.SignCert(rand.Reader, ca))
	cs, err := ssh.NewCertSigner(cert, clientKey)
	require.NoError(t, err)
	return cs
}

// TestGateEnforcesSourceAddress proves the gate honours a tenant CA's
// source-address restriction: a cert scoped to a network the client is not on
// fails auth, and the same cert scoped to the client's own network tunnels.
// Tests dial from 127.0.0.1, so 203.0.113.0/24 is out and 127.0.0.0/8 is in.
func TestGateEnforcesSourceAddress(t *testing.T) {
	ca := newSigner(t)

	outside := startGate(t, ca.PublicKey(), true)
	cfg := &ssh.ClientConfig{
		User: "ubuntu",
		Auth: []ssh.AuthMethod{ssh.PublicKeys(mintCertSignerCriticalOptions(t, ca, newSigner(t),
			map[string]string{"source-address": "203.0.113.0/24"}))},
		HostKeyCallback: ssh.FixedHostKey(outside.hostKey.PublicKey()),
		Timeout:         5 * time.Second,
	}
	_, err := ssh.Dial("tcp", outside.addr, cfg)
	require.Error(t, err, "a cert restricted to a network the client is not on must fail auth")

	inside := startGate(t, ca.PublicKey(), true)
	client := dialClient(t, inside, mintCertSignerCriticalOptions(t, ca, newSigner(t),
		map[string]string{"source-address": "127.0.0.0/8"}))
	conn, err := client.Dial("tcp", "default.vm1:22")
	require.NoError(t, err, "a cert restricted to the client's own network must still tunnel")
	_ = conn.Close()
}

// TestGateRejectsUnknownCriticalOption pins that source-address is the only
// restriction the gate can honour: any other critical option is a promise it
// cannot keep, so the cert fails auth rather than being silently relaxed.
func TestGateRejectsUnknownCriticalOption(t *testing.T) {
	ca := newSigner(t)
	tg := startGate(t, ca.PublicKey(), true)

	cfg := &ssh.ClientConfig{
		User: "ubuntu",
		Auth: []ssh.AuthMethod{ssh.PublicKeys(mintCertSignerCriticalOptions(t, ca, newSigner(t),
			map[string]string{"force-command": "/bin/false"}))},
		HostKeyCallback: ssh.FixedHostKey(tg.hostKey.PublicKey()),
		Timeout:         5 * time.Second,
	}
	_, err := ssh.Dial("tcp", tg.addr, cfg)
	require.Error(t, err, "a cert carrying a critical option the gate cannot enforce must fail auth")
}