internal/mcpserver/sshrun_test.go
Ref: Size: 8.2 KiB History
package mcpserver
import (
"bytes"
"context"
"crypto/ed25519"
"crypto/rand"
"errors"
"io"
"net"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"golang.org/x/crypto/ssh"
)
// ── test SSH-CA scaffolding ──────────────────────────────────────────────────
// newSigner returns a fresh 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
}
// hostCertSigner builds a host-cert-backed signer for principal, signed by ca.
// A server AddHostKey'd with it presents that host cert during the handshake.
func hostCertSigner(t *testing.T, ca ssh.Signer, principal string) ssh.Signer {
t.Helper()
hostKey := newSigner(t)
cert := &ssh.Certificate{
Key: hostKey.PublicKey(),
CertType: ssh.HostCert,
ValidPrincipals: []string{principal},
ValidBefore: ssh.CertTimeInfinity,
}
require.NoError(t, cert.SignCert(rand.Reader, ca))
cs, err := ssh.NewCertSigner(cert, hostKey)
require.NoError(t, err)
return cs
}
// userCertSigner builds a user-cert-backed signer for principal, signed by ca —
// the credential a VM's sshd accepts.
func userCertSigner(t *testing.T, ca ssh.Signer, principal string) ssh.Signer {
t.Helper()
userKey := newSigner(t)
cert := &ssh.Certificate{
Key: userKey.PublicKey(),
CertType: ssh.UserCert,
ValidPrincipals: []string{principal},
ValidBefore: ssh.CertTimeInfinity,
}
require.NoError(t, cert.SignCert(rand.Reader, ca))
cs, err := ssh.NewCertSigner(cert, userKey)
require.NoError(t, err)
return cs
}
// caUserAuth accepts a client only if it presents a user cert signed by ca —
// mirroring the VM sshd's TrustedUserCAKeys policy.
func caUserAuth(ca ssh.PublicKey) func(ssh.ConnMetadata, ssh.PublicKey) (*ssh.Permissions, error) {
checker := &ssh.CertChecker{
IsUserAuthority: func(auth ssh.PublicKey) bool {
return auth != nil && ca != nil && bytes.Equal(auth.Marshal(), ca.Marshal())
},
}
return func(_ ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
cert, ok := key.(*ssh.Certificate)
if !ok {
return nil, errors.New("only certificate authentication is accepted")
}
if cert.CertType != ssh.UserCert {
return nil, errors.New("not a user certificate")
}
if !checker.IsUserAuthority(cert.SignatureKey) {
return nil, errors.New("certificate not signed by the test CA")
}
if len(cert.ValidPrincipals) == 0 {
return nil, errors.New("certificate has no principals")
}
// Feed CheckCert one of the cert's own principals so only CA-signature +
// validity gate authentication, not the arbitrary outer username.
if err := checker.CheckCert(cert.ValidPrincipals[0], cert); err != nil {
return nil, err
}
return &ssh.Permissions{}, nil
}
}
// startBackingVM runs a minimal VM sshd on a random loopback port: it presents
// hostSigner's host cert, accepts CA-signed user certs, and answers a single
// "exec" request with out + code. Returns its listen address.
func startBackingVM(t *testing.T, hostSigner ssh.Signer, userCA ssh.PublicKey, out string, code int) string {
t.Helper()
conf := &ssh.ServerConfig{PublicKeyCallback: caUserAuth(userCA)}
conf.AddHostKey(hostSigner)
ln, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err)
t.Cleanup(func() { ln.Close() })
go func() {
for {
nc, err := ln.Accept()
if err != nil {
return
}
go serveBackingVM(nc, conf, out, code)
}
}()
return ln.Addr().String()
}
func serveBackingVM(nc net.Conn, conf *ssh.ServerConfig, out string, code int) {
sc, chans, reqs, err := ssh.NewServerConn(nc, conf)
if err != nil {
nc.Close()
return
}
defer sc.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 handleExecSession(ch, chReqs, out, code)
}
}
func handleExecSession(ch ssh.Channel, reqs <-chan *ssh.Request, out string, code int) {
defer ch.Close()
for req := range reqs {
if req.Type == "exec" {
req.Reply(true, nil)
io.WriteString(ch, out)
ch.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{uint32(code)}))
return
}
if req.WantReply {
req.Reply(false, nil)
}
}
}
// testDialer stands in for whatever transport reached the VM. Runner is
// transport-blind, so its tests supply the plainest dialer there is: connect to
// a fixed address with a CA-signed user cert. Reaching the right VM and
// verifying its host certificate belong to the real dialer, and are pinned
// there (internal/server/vmssh).
type testDialer struct {
addr string
tenant string
signer ssh.Signer
hostCA ssh.PublicKey
}
func (d testDialer) ConnectName(_ context.Context, vmName string) (string, error) {
return d.connectName(vmName), nil
}
func (d testDialer) connectName(vmName string) string { return d.tenant + "." + vmName }
func (d testDialer) Dial(_ context.Context, vmName string) (*ssh.Client, error) {
name := d.connectName(vmName)
checker := &ssh.CertChecker{
IsHostAuthority: func(auth ssh.PublicKey, _ string) bool {
return bytes.Equal(auth.Marshal(), d.hostCA.Marshal())
},
}
return ssh.Dial("tcp", d.addr, &ssh.ClientConfig{
User: "ubuntu",
Auth: []ssh.AuthMethod{ssh.PublicKeys(d.signer)},
HostKeyCallback: func(_ string, remote net.Addr, key ssh.PublicKey) error {
return checker.CheckHostKey(net.JoinHostPort(name, "22"), remote, key)
},
Timeout: 10 * time.Second,
})
}
// newTestRunner wires a Runner to a backing VM that answers every exec with out
// and code.
func newTestRunner(t *testing.T, out string, code int) *Runner {
t.Helper()
ca := newSigner(t)
addr := startBackingVM(t, hostCertSigner(t, ca, "default.testvm"), ca.PublicKey(), out, code)
return NewRunner(testDialer{
addr: addr,
tenant: "default",
signer: userCertSigner(t, ca, "ubuntu"),
hostCA: ca.PublicKey(),
})
}
// ── tests ────────────────────────────────────────────────────────────────────
func TestExecReturnsOutput(t *testing.T) {
res, err := newTestRunner(t, "hi\n", 0).Exec(t.Context(), "testvm", "echo hi", 10*time.Second)
require.NoError(t, err)
assert.Equal(t, "hi\n", res.Stdout)
assert.Equal(t, 0, res.ExitCode)
assert.False(t, res.Truncated)
}
func TestExecNonZeroExitIsNotAnError(t *testing.T) {
res, err := newTestRunner(t, "nope\n", 3).Exec(t.Context(), "testvm", "false", 10*time.Second)
require.NoError(t, err, "a command that fails is a result, not a transport failure")
assert.Equal(t, 3, res.ExitCode)
assert.Equal(t, "nope\n", res.Stdout)
}
func TestExecReportsTruncationPastTheCap(t *testing.T) {
res, err := newTestRunner(t, string(bytes.Repeat([]byte("x"), outputCap+512)), 0).
Exec(t.Context(), "testvm", "cat big", 30*time.Second)
require.NoError(t, err)
// 1 MiB as a literal: a cap compared to itself is satisfied at any size,
// including one that fits the whole of a guest's syslog into one reply.
assert.Len(t, res.Stdout, 1<<20,
"captured output is capped at 1 MiB — the buffer is held in the plane's memory per concurrent exec, and the bytes "+
"go on to the model, so an unbounded `cat` on a guest is both a memory cost the plane never agreed to and a "+
"context window spent on one file")
assert.True(t, res.Truncated, "the model must be told the output was cut")
}
func TestExecSurfacesADialFailure(t *testing.T) {
// A dialer that cannot reach the VM: the error is the transport's, and the
// Runner passes it through rather than reporting an empty success.
r := NewRunner(testDialer{addr: "127.0.0.1:1", tenant: "default", signer: newSigner(t), hostCA: newSigner(t).PublicKey()})
_, err := r.Exec(t.Context(), "testvm", "echo hi", 5*time.Second)
require.Error(t, err)
}
func TestConnectNameComesFromTheDialer(t *testing.T) {
name, err := newTestRunner(t, "", 0).ConnectName(t.Context(), "web-1")
require.NoError(t, err)
assert.Equal(t, "default.web-1", name)
}