internal/smoke/mcp_test.go
Ref: Size: 17.1 KiB History
package smoke
import (
"context"
"crypto/ed25519"
"crypto/rand"
"errors"
"net/http"
"net/http/httptest"
"strconv"
"strings"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"golang.org/x/crypto/ssh"
)
// fakeMCP scripts an MCP endpoint: it records the calls the leg makes, answers
// each from results, and can be made to fail one of them. delegate_complete is
// answered like the real one — the certificate is parsed and its signing CA
// reported back — so the leg's own checks are exercised rather than mocked past.
type fakeMCP struct {
tools []string
results map[string]map[string]any
failOn string
// trusted is the tenant's registered CA set, as the control plane would
// hold it: a certificate signed by anything else is refused.
trusted []ssh.PublicKey
pubLine string
// uploaded records the CA lines ca_upload received, and exposedAs the
// protocol each vm_expose asked for.
uploaded []string
exposedAs []string
// delegated is what tenant_info reports; set once a certificate lands.
delegated bool
calls []string
}
func newFakeMCP(t *testing.T, trusted ...ssh.Signer) *fakeMCP {
t.Helper()
_, priv, err := ed25519.GenerateKey(rand.Reader)
require.NoError(t, err)
eitriKey, err := ssh.NewSignerFromSigner(priv)
require.NoError(t, err)
f := &fakeMCP{
tools: []string{
"vm_create", "vm_list", "vm_info", "vm_exec", "vm_write_file", "vm_read_file",
"vm_expose", "vm_exposures", "vm_unexpose", "vm_destroy",
"ca_upload", "tenant_info", "delegate_begin", "delegate_complete",
},
pubLine: authorizedLine(eitriKey.PublicKey()),
results: map[string]map[string]any{
"vm_create": {"id": "v-1", "name": "smoke-mcp"},
"vm_exec": {"stdout": "", "exit_code": 0.0},
"vm_unexpose": {"id": "x-1"},
"vm_destroy": {"id": "v-1"},
},
}
for _, ca := range trusted {
f.trusted = append(f.trusted, ca.PublicKey())
}
f.results["delegate_begin"] = map[string]any{"public_key": f.pubLine, "principal": "ubuntu"}
return f
}
func (f *fakeMCP) List(context.Context) ([]string, error) { return f.tools, nil }
func (f *fakeMCP) Call(_ context.Context, name string, args map[string]any) (map[string]any, error) {
f.calls = append(f.calls, name)
if name == f.failOn {
return nil, errors.New(name + ": refused")
}
if name == "ca_upload" {
line, _ := args["public_key"].(string)
pub, _, _, _, err := ssh.ParseAuthorizedKey([]byte(line))
if err != nil {
return nil, errors.New("ca_upload: not a public key")
}
f.uploaded = append(f.uploaded, line)
return map[string]any{"fingerprint": ssh.FingerprintSHA256(pub)}, nil
}
if name == "tenant_info" {
cas := []any{}
for _, line := range f.uploaded {
if pub, _, _, _, err := ssh.ParseAuthorizedKey([]byte(line)); err == nil {
cas = append(cas, map[string]any{"fingerprint": ssh.FingerprintSHA256(pub), "label": "eitri-smoke"})
}
}
return map[string]any{"registered_cas": cas, "delegated": f.delegated}, nil
}
if name == "delegate_complete" {
return f.complete(args)
}
if name == "vm_expose" {
// The control plane answers with the exposure it made, protocol
// included, and an unnamed protocol is tcp. A scripted result stands in
// for a plane answering something else.
proto, _ := args["protocol"].(string)
if proto == "" {
proto = "tcp"
}
f.exposedAs = append(f.exposedAs, proto)
if scripted, ok := f.results[name]; ok {
return scripted, nil
}
return map[string]any{"exposure": map[string]any{
"id": "x-1", "address": "10.0.0.1:30001", "protocol": proto,
}}, nil
}
out := map[string]any{}
for k, v := range f.results[name] {
out[k] = v
}
return out, nil
}
// complete stands in for the control plane's own validation: the certificate
// must be over the key this fake handed out, and signed by a registered CA.
func (f *fakeMCP) complete(args map[string]any) (map[string]any, error) {
line, _ := args["certificate"].(string)
pub, _, _, _, err := ssh.ParseAuthorizedKey([]byte(line))
if err != nil {
return nil, errors.New("delegate_complete: not a certificate")
}
cert, ok := pub.(*ssh.Certificate)
if !ok {
return nil, errors.New("delegate_complete: not a certificate")
}
if authorizedLine(cert.Key) != f.pubLine {
return nil, errors.New("delegate_complete: that certificate is for a different key")
}
for _, ca := range f.trusted {
if string(ca.Marshal()) == string(cert.SignatureKey.Marshal()) {
f.delegated = true
return map[string]any{
"ca_fingerprint": ssh.FingerprintSHA256(cert.SignatureKey),
"principals": []any{"ubuntu"},
"expires_at": time.Unix(int64(cert.ValidBefore), 0).UTC().Format(time.RFC3339),
}, nil
}
}
return nil, errors.New("delegate_complete: not a CA registered to this tenant")
}
// echoingMCP makes vm_exec behave like a guest: it echoes the command's
// argument, which is the nonce the leg generated.
type echoingMCP struct{ *fakeMCP }
func (e echoingMCP) Call(ctx context.Context, name string, args map[string]any) (map[string]any, error) {
out, err := e.fakeMCP.Call(ctx, name, args)
if err != nil || name != "vm_exec" {
return out, err
}
cmd, _ := args["command"].(string)
out["stdout"] = strings.TrimPrefix(cmd, "echo ") + "\n"
return out, nil
}
func okBannerDial(context.Context, string) (string, error) { return "SSH-2.0-OpenSSH_9.6\r\n", nil }
// okEcho is a published UDP port with the guest's echo behind it: whatever went
// in comes back.
func okEcho(_ context.Context, _, payload string) (string, error) { return payload, nil }
// newCA returns a signer standing in for somebody's own SSH user CA.
func newCA(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
}
// testDelegator builds the out-of-band half: ca is the one the leg registers
// and signs with, stranger is the one nobody registered.
func testDelegator(ca, stranger ssh.Signer) delegator {
return delegator{
ca: ca,
stranger: stranger,
principal: "ubuntu",
now: func() time.Time { return time.Unix(1_800_000_000, 0) },
}
}
// TestProveMCPDrivesTheWholeCycle pins the order the leg must run in: the CA is
// registered BEFORE the VM is created, because a guest bakes its trusted CA set
// at create and would otherwise refuse every certificate. The delegation itself
// may fall on either side of the create, which is the improvement over holding
// a signing key.
func TestProveMCPDrivesTheWholeCycle(t *testing.T) {
ca, stranger := newCA(t), newCA(t)
f := echoingMCP{newFakeMCP(t, ca)}
clock := &fakeClock{}
require.NoError(t, proveMCP(t.Context(), f, "smoke-mcp", testDelegator(ca, stranger), clock.now, clock.sleep, okBannerDial, okEcho))
assert.Equal(t, []string{authorizedLine(ca.PublicKey())}, f.uploaded,
"the CA is registered through the tool, before anything is created")
assert.Equal(t,
[]string{"ca_upload", "delegate_begin", "delegate_complete", "delegate_complete",
"tenant_info", "vm_create", "vm_exec", "vm_expose", "vm_unexpose",
"vm_exec", "vm_expose", "vm_unexpose", "vm_destroy"},
f.calls, "the UDP leg starts its own listener in the guest before publishing a port for it")
assert.Equal(t, []string{"tcp", "udp"}, f.exposedAs)
}
// TestProveUDPExposureRequiresTheDatagramBack pins what the UDP leg is for: a
// published port that binds and answers nothing passes no gate.
func TestProveUDPExposureRequiresTheDatagramBack(t *testing.T) {
f := newFakeMCP(t)
clock := &fakeClock{}
// A port that answers something else entirely is a failure on the spot —
// something is listening, and it is not the guest's echo.
wrong := func(context.Context, string, string) (string, error) { return "not-the-nonce", nil }
err := proveUDPExposure(t.Context(), f, "smoke-mcp", clock.now, clock.sleep, wrong)
require.Error(t, err)
assert.Contains(t, err.Error(), "echoed")
// A port that answers nothing at all is retried to the deadline and then
// fails with what the last attempt said.
silent := func(context.Context, string, string) (string, error) {
return "", errors.New("i/o timeout")
}
err = proveUDPExposure(t.Context(), f, "smoke-mcp", clock.now, clock.sleep, silent)
require.Error(t, err)
assert.Contains(t, err.Error(), "no echo from the MCP-published UDP port")
assert.Contains(t, err.Error(), "i/o timeout")
}
// TestUDPEchoCommandDetachesFromTheExecSession pins the two things about the
// guest-side echo that a passing leg depends on: it binds the port the leg
// publishes, and it outlives the exec that started it.
func TestUDPEchoCommandDetachesFromTheExecSession(t *testing.T) {
cmd := udpEchoCommand()
assert.Contains(t, cmd, strconv.Itoa(udpEchoPort), "the echo binds the port the leg publishes")
assert.Contains(t, cmd, "nohup", "the echo must survive the session that started it")
assert.Contains(t, cmd, ">/dev/null 2>&1 &", "a background process holding stdout would hang the exec")
}
// TestProveMCPRequiresTheUnregisteredCAToBeRefused is the negative leg itself:
// if the control plane ever accepted a certificate from a CA nobody registered,
// the gate must fail.
func TestProveMCPRequiresTheUnregisteredCAToBeRefused(t *testing.T) {
ca, stranger := newCA(t), newCA(t)
f := echoingMCP{newFakeMCP(t, ca, stranger)} // a plane that trusts everyone
clock := &fakeClock{}
err := proveMCP(t.Context(), f, "smoke-mcp", testDelegator(ca, stranger), clock.now, clock.sleep, okBannerDial, okEcho)
require.Error(t, err)
assert.Contains(t, err.Error(), "accepted a certificate from an unregistered CA")
}
// TestProveMCPRejectsAWrongPrincipal: the challenge names the principal the
// certificate must carry, and a plane asking for a different one is a plane the
// gate does not understand.
func TestProveMCPRejectsAWrongPrincipal(t *testing.T) {
ca, stranger := newCA(t), newCA(t)
f := newFakeMCP(t, ca)
f.results["delegate_begin"] = map[string]any{"public_key": f.pubLine, "principal": "root"}
clock := &fakeClock{}
err := proveMCP(t.Context(), f, "smoke-mcp", testDelegator(ca, stranger), clock.now, clock.sleep, okBannerDial, okEcho)
require.Error(t, err)
assert.Contains(t, err.Error(), `asked for principal "root"`)
}
// TestProveMCPRejectsADelegationWithNoExpiry: a delegation that never ends is
// not one.
func TestProveMCPRejectsADelegationWithNoExpiry(t *testing.T) {
ca, stranger := newCA(t), newCA(t)
f := noExpiryMCP{newFakeMCP(t, ca)}
clock := &fakeClock{}
err := proveMCP(t.Context(), f, "smoke-mcp", testDelegator(ca, stranger), clock.now, clock.sleep, okBannerDial, okEcho)
require.Error(t, err)
assert.Contains(t, err.Error(), "reports no expiry")
}
type noExpiryMCP struct{ *fakeMCP }
func (n noExpiryMCP) Call(ctx context.Context, name string, args map[string]any) (map[string]any, error) {
out, err := n.fakeMCP.Call(ctx, name, args)
if err == nil && name == "delegate_complete" {
delete(out, "expires_at")
}
return out, err
}
// TestProveMCPDestroysItsVMOnFailure: a failed leg must not leave a VM running
// on the live fleet.
func TestProveMCPDestroysItsVMOnFailure(t *testing.T) {
ca, stranger := newCA(t), newCA(t)
f := newFakeMCP(t, ca)
f.failOn = "vm_exec"
clock := &fakeClock{}
err := proveMCP(t.Context(), f, "smoke-mcp", testDelegator(ca, stranger), clock.now, clock.sleep, okBannerDial, okEcho)
require.Error(t, err)
assert.Contains(t, err.Error(), "vm_exec over MCP")
assert.Contains(t, f.calls, "vm_destroy", "the leg must reap its own VM even when it fails")
}
// TestProveToolList covers the cheap leg every extra origin gets on its own: an
// origin advertising fewer tools than the toolset means that route reaches
// something else, and a missing tool is named rather than counted.
func TestProveToolList(t *testing.T) {
cases := []struct {
name string
tools func(*fakeMCP)
wantErr string
}{
{"the whole toolset passes", func(*fakeMCP) {}, ""},
{"a short toolset fails", func(f *fakeMCP) { f.tools = f.tools[:5] }, "advertises 5 tools"},
{
"a renamed tool is named",
func(f *fakeMCP) {
for i, name := range f.tools {
if name == "delegate_begin" {
f.tools[i] = "vm_something_else"
}
}
},
"does not advertise delegate_begin",
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
f := newFakeMCP(t, newCA(t))
tc.tools(f)
err := proveToolList(t.Context(), f)
if tc.wantErr == "" {
require.NoError(t, err)
return
}
require.Error(t, err)
assert.Contains(t, err.Error(), tc.wantErr)
})
}
}
// TestProveMCPRejectsAShortToolset: the full leg refuses to spend a boot on an
// endpoint whose toolset is already wrong.
func TestProveMCPRejectsAShortToolset(t *testing.T) {
ca, stranger := newCA(t), newCA(t)
f := newFakeMCP(t, ca)
f.tools = f.tools[:5]
clock := &fakeClock{}
err := proveMCP(t.Context(), f, "smoke-mcp", testDelegator(ca, stranger), clock.now, clock.sleep, okBannerDial, okEcho)
require.Error(t, err)
assert.Contains(t, err.Error(), "advertises 5 tools")
assert.Empty(t, f.calls, "a wrong toolset must fail before any tool is called")
}
// TestProveMCPRejectsAMissingDelegateTool guards the tools a remote caller
// cannot work without.
func TestProveMCPRejectsAMissingDelegateTool(t *testing.T) {
ca, stranger := newCA(t), newCA(t)
f := newFakeMCP(t, ca)
for i, name := range f.tools {
if name == "delegate_begin" {
f.tools[i] = "vm_something_else"
}
}
clock := &fakeClock{}
err := proveMCP(t.Context(), f, "smoke-mcp", testDelegator(ca, stranger), clock.now, clock.sleep, okBannerDial, okEcho)
require.Error(t, err)
assert.Contains(t, err.Error(), "does not advertise delegate_begin")
}
// TestProveMCPRejectsAnExposureWithNoAddress: a grant with nowhere to dial is a
// failure, not something to retry.
func TestProveMCPRejectsAnExposureWithNoAddress(t *testing.T) {
ca, stranger := newCA(t), newCA(t)
f := echoingMCP{newFakeMCP(t, ca)}
f.results["vm_expose"] = map[string]any{"exposure": map[string]any{"id": "x-1"}}
clock := &fakeClock{}
err := proveMCP(t.Context(), f, "smoke-mcp", testDelegator(ca, stranger), clock.now, clock.sleep, okBannerDial, okEcho)
require.Error(t, err)
assert.Contains(t, err.Error(), "no address to dial")
}
// TestProveMCPRejectsAWrongBanner: something is listening on the published port,
// but it is not the guest's sshd.
func TestProveMCPRejectsAWrongBanner(t *testing.T) {
ca, stranger := newCA(t), newCA(t)
f := echoingMCP{newFakeMCP(t, ca)}
clock := &fakeClock{}
err := proveMCP(t.Context(), f, "smoke-mcp", testDelegator(ca, stranger), clock.now, clock.sleep,
func(context.Context, string) (string, error) { return "HTTP/1.1 200 OK", nil }, okEcho)
require.Error(t, err)
assert.Contains(t, err.Error(), "want an SSH-2.0 banner")
}
// TestProveMCPRejectsAnEchoThatDoesNotComeBack: the nonce proves the command
// actually ran in the guest rather than being answered by the control plane.
func TestProveMCPRejectsAnEchoThatDoesNotComeBack(t *testing.T) {
ca, stranger := newCA(t), newCA(t)
f := newFakeMCP(t, ca) // plain fake: vm_exec returns an empty stdout
clock := &fakeClock{}
err := proveMCP(t.Context(), f, "smoke-mcp", testDelegator(ca, stranger), clock.now, clock.sleep, okBannerDial, okEcho)
require.Error(t, err)
assert.Contains(t, err.Error(), "want it to contain")
}
// TestSignDelegationProducesAUsableCertificate pins the shape the control plane
// checks: a user certificate over the key it handed out, naming the login user,
// with a real validity window.
func TestSignDelegationProducesAUsableCertificate(t *testing.T) {
ca := newCA(t)
f := newFakeMCP(t, ca)
now := time.Unix(1_800_000_000, 0)
line, err := signDelegation(ca, f.pubLine, "ubuntu", now)
require.NoError(t, err)
pub, _, _, _, err := ssh.ParseAuthorizedKey([]byte(line))
require.NoError(t, err)
cert, ok := pub.(*ssh.Certificate)
require.True(t, ok)
assert.Equal(t, uint32(ssh.UserCert), cert.CertType)
assert.Equal(t, []string{"ubuntu"}, cert.ValidPrincipals)
assert.Equal(t, f.pubLine, authorizedLine(cert.Key))
assert.Equal(t, uint64(now.Add(delegationTTL).Unix()), cert.ValidBefore)
assert.LessOrEqual(t, cert.ValidBefore, uint64(now.Add(time.Hour).Unix()),
"a delegated user cert is a credential the smoke mints and leaves behind: bounded against delegationTTL it is valid "+
"for however long that constant says, so the lifetime is held to an hour here as well")
checker := &ssh.CertChecker{
IsUserAuthority: func(auth ssh.PublicKey) bool {
return string(auth.Marshal()) == string(ca.PublicKey().Marshal())
},
Clock: func() time.Time { return now },
}
require.NoError(t, checker.CheckCert("ubuntu", cert))
}
// TestProveRemoteMCPNeedsACredential passes only on a 401 — an endpoint that
// answers anything else is either unprotected or not the MCP endpoint.
func TestProveRemoteMCPNeedsACredential(t *testing.T) {
unauthorized := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
http.Error(w, "sign in required", http.StatusUnauthorized)
}))
defer unauthorized.Close()
assert.NoError(t, proveRemoteMCPNeedsACredential(unauthorized.URL))
open := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusOK)
}))
defer open.Close()
err := proveRemoteMCPNeedsACredential(open.URL)
require.Error(t, err)
assert.Contains(t, err.Error(), "answered 200, want 401")
}