internal/smoke/mcp.go
Ref: Size: 17.6 KiB History
package smoke
import (
"context"
"crypto/rand"
"encoding/binary"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"net/http"
"slices"
"strings"
"time"
"github.com/modelcontextprotocol/go-sdk/mcp"
"golang.org/x/crypto/ssh"
)
// mcpTools is what the remote-MCP leg needs from an MCP session: the advertised
// tool list, and calling one. Declaring it lets the leg's sequencing be tested
// without a control plane.
type mcpTools interface {
List(ctx context.Context) ([]string, error)
Call(ctx context.Context, name string, args map[string]any) (map[string]any, error)
}
// remoteToolCount is the toolset a bearer PAT gets over HTTP: the ten VM tools,
// ca_upload, tenant_info, and the two delegation tools.
const remoteToolCount = 14
// dialMCP opens a real MCP client against the control plane's /mcp endpoint,
// authenticating with a bearer PAT exactly as a remote LLM client does. The
// returned close function tears the session down.
func dialMCP(ctx context.Context, serverURL, pat string) (sdkTools, func(), error) {
tr := &mcp.StreamableClientTransport{
Endpoint: strings.TrimRight(serverURL, "/") + "/mcp",
HTTPClient: &http.Client{Transport: patTransport{pat: pat}, Timeout: 15 * time.Minute},
}
session, err := mcp.NewClient(&mcp.Implementation{Name: "eitri-smoke", Version: "1"}, nil).Connect(ctx, tr, nil)
if err != nil {
return sdkTools{}, nil, fmt.Errorf("connect to %s/mcp with a bearer PAT: %w", serverURL, err)
}
return sdkTools{session: session}, func() { session.Close() }, nil
}
// patTransport attaches the PAT the way a remote MCP client configuration does.
type patTransport struct{ pat string }
func (p patTransport) RoundTrip(r *http.Request) (*http.Response, error) {
r.Header.Set("Authorization", "Bearer "+p.pat)
return http.DefaultTransport.RoundTrip(r)
}
// sdkTools adapts an MCP session to mcpTools, folding a tool-level error into a
// Go error so the leg reads one way for both failure kinds.
type sdkTools struct{ session *mcp.ClientSession }
func (s sdkTools) List(ctx context.Context) ([]string, error) {
res, err := s.session.ListTools(ctx, nil)
if err != nil {
return nil, err
}
names := make([]string, 0, len(res.Tools))
for _, tool := range res.Tools {
names = append(names, tool.Name)
}
return names, nil
}
func (s sdkTools) Call(ctx context.Context, name string, args map[string]any) (map[string]any, error) {
res, err := s.session.CallTool(ctx, &mcp.CallToolParams{Name: name, Arguments: args})
if err != nil {
return nil, fmt.Errorf("%s: %w", name, err)
}
if res.IsError {
var b strings.Builder
for _, c := range res.Content {
if tc, ok := c.(*mcp.TextContent); ok {
b.WriteString(tc.Text)
}
}
return nil, fmt.Errorf("%s: %s", name, b.String())
}
raw, err := json.Marshal(res.StructuredContent)
if err != nil {
return nil, fmt.Errorf("%s: %w", name, err)
}
var out map[string]any
if err := json.Unmarshal(raw, &out); err != nil {
return nil, fmt.Errorf("%s: %w", name, err)
}
return out, nil
}
// proveRemoteMCPNeedsACredential is the cheap, fleet-independent half: the /mcp
// endpoint must refuse an unauthenticated caller. It runs before any VM exists,
// so a misrouted or unauthenticated endpoint fails the gate in a second rather
// than after a boot.
func proveRemoteMCPNeedsACredential(serverURL string) error {
body := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18",` +
`"capabilities":{},"clientInfo":{"name":"eitri-smoke","version":"1"}}}`
req, err := http.NewRequest(http.MethodPost, strings.TrimRight(serverURL, "/")+"/mcp", strings.NewReader(body))
if err != nil {
return err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json, text/event-stream")
resp, err := (&http.Client{Timeout: 30 * time.Second}).Do(req)
if err != nil {
return fmt.Errorf("POST %s/mcp: %w", serverURL, err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusUnauthorized {
return fmt.Errorf("FAIL: unauthenticated POST /mcp answered %d, want 401", resp.StatusCode)
}
return nil
}
// delegator is the out-of-band half of the exchange, as the smoke performs it:
// register a CA with the tenant, and sign eitri's ephemeral key with it. In a
// real session a human runs ssh-keygen; here the gate does it in-process,
// which is the same act.
type delegator struct {
ca ssh.Signer
// stranger is a CA nobody registered, for the negative leg.
stranger ssh.Signer
principal string
now func() time.Time
}
// proveToolList proves an origin advertises the whole remote toolset. It is the
// cheap half of the MCP leg, and all a second origin gets: a plane fronted two
// ways routes and authenticates /mcp on both, but only one of them needs to
// carry the long call that a proxy's silent-origin timeout would cut.
func proveToolList(ctx context.Context, c mcpTools) error {
names, err := c.List(ctx)
if err != nil {
return fmt.Errorf("list remote MCP tools: %w", err)
}
if len(names) != remoteToolCount {
return fmt.Errorf("FAIL: remote MCP advertises %d tools (%s), want %d", len(names), strings.Join(names, ", "), remoteToolCount)
}
for _, want := range []string{"ca_upload", "tenant_info", "delegate_begin", "delegate_complete", "vm_create", "vm_exec", "vm_expose", "vm_destroy"} {
if !slices.Contains(names, want) {
return fmt.Errorf("FAIL: remote MCP does not advertise %s (has: %s)", want, strings.Join(names, ", "))
}
}
return nil
}
// proveRemoteMCPToolset dials one origin with a bearer PAT, proves its toolset,
// and hangs up.
func proveRemoteMCPToolset(ctx context.Context, serverURL, pat string) error {
tools, closeMCP, err := dialMCP(ctx, serverURL, pat)
if err != nil {
return err
}
defer closeMCP()
if err := proveToolList(ctx, tools); err != nil {
return fmt.Errorf("%s: %w", serverURL, err)
}
return nil
}
// proveMCP drives one full cycle through the remote MCP endpoint and nothing
// else: delegate access to eitri, create a VM, run a command in it, publish a
// TCP port and read the guest's banner back through it, publish a UDP port and
// read a datagram back through that, and destroy the VM.
//
// The CA is registered BEFORE the VM is created, because a guest bakes its CA
// set at create and a CA registered afterwards is one it will never trust. The
// DELEGATION itself may happen on either side of the create — that is the whole
// improvement over holding a signing key, and it is worth stating plainly here
// so nobody reintroduces an ordering constraint that no longer exists.
func proveMCP(ctx context.Context, c mcpTools, vmName string, d delegator, now func() time.Time, sleep func(time.Duration), dial bannerFunc, echo echoFunc) error {
if err := proveToolList(ctx, c); err != nil {
return err
}
// Through the tool, not the endpoint behind it: registering a CA is the
// first step a bare token has to take, so the gate drives the same call an
// LLM client would.
uploaded, err := c.Call(ctx, "ca_upload", map[string]any{
"public_key": authorizedLine(d.ca.PublicKey()),
"label": "eitri-smoke",
})
if err != nil {
return fmt.Errorf("ca_upload: %w", err)
}
if got, _ := uploaded["fingerprint"].(string); got != ssh.FingerprintSHA256(d.ca.PublicKey()) {
return fmt.Errorf("FAIL: ca_upload registered %s, want %s", got, ssh.FingerprintSHA256(d.ca.PublicKey()))
}
begin, err := c.Call(ctx, "delegate_begin", nil)
if err != nil {
return fmt.Errorf("delegate_begin: %w", err)
}
pubLine, _ := begin["public_key"].(string)
if pubLine == "" {
return errors.New("FAIL: delegate_begin returned no public key to sign")
}
principal, _ := begin["principal"].(string)
if principal != d.principal {
return fmt.Errorf("FAIL: delegate_begin asked for principal %q, want %q", principal, d.principal)
}
// The negative leg: a certificate from a CA nobody registered must be
// refused. It costs one call and no VM, and it is the only live proof that
// the trust check is wired to the tenant's real CA set rather than to
// whatever signature happens to verify.
stranger, err := signDelegation(d.stranger, pubLine, principal, d.now())
if err != nil {
return err
}
if _, err := c.Call(ctx, "delegate_complete", map[string]any{"certificate": stranger}); err == nil {
return errors.New("FAIL: delegate_complete accepted a certificate from an unregistered CA")
}
cert, err := signDelegation(d.ca, pubLine, principal, d.now())
if err != nil {
return err
}
done, err := c.Call(ctx, "delegate_complete", map[string]any{"certificate": cert})
if err != nil {
return fmt.Errorf("delegate_complete: %w", err)
}
if got, _ := done["ca_fingerprint"].(string); got != ssh.FingerprintSHA256(d.ca.PublicKey()) {
return fmt.Errorf("FAIL: delegation reports CA %s, want %s", got, ssh.FingerprintSHA256(d.ca.PublicKey()))
}
expiresAt, _ := done["expires_at"].(string)
if expiresAt == "" {
return errors.New("FAIL: delegation reports no expiry; a delegation that never ends is not one")
}
// What a caller would read back before doing anything else: the CA it just
// registered is listed, and the delegation it just made is live.
info, err := c.Call(ctx, "tenant_info", nil)
if err != nil {
return fmt.Errorf("tenant_info: %w", err)
}
if delegated, _ := info["delegated"].(bool); !delegated {
return errors.New("FAIL: tenant_info reports no delegation immediately after delegate_complete")
}
if err := infoListsCA(info, ssh.FingerprintSHA256(d.ca.PublicKey())); err != nil {
return err
}
if _, err := c.Call(ctx, "vm_create", map[string]any{"name": vmName}); err != nil {
return fmt.Errorf("vm_create over MCP: %w", err)
}
defer func() {
// Destroy on the way out even when the leg failed: the smoke leaves no
// VM behind on the live fleet.
if _, err := c.Call(context.WithoutCancel(ctx), "vm_destroy", map[string]any{"vm": vmName}); err != nil {
fmt.Printf("eitri-smoke: MCP vm_destroy failed; vm %s left behind: %v\n", vmName, err)
}
}()
nonce, err := randNonce()
if err != nil {
return err
}
exec, err := c.Call(ctx, "vm_exec", map[string]any{"vm": vmName, "command": "echo " + nonce})
if err != nil {
return fmt.Errorf("vm_exec over MCP: %w", err)
}
if stdout, _ := exec["stdout"].(string); !strings.Contains(stdout, nonce) {
return fmt.Errorf("FAIL: vm_exec over MCP echoed %q, want it to contain %q", stdout, nonce)
}
exposed, err := c.Call(ctx, "vm_expose", map[string]any{"vm": vmName, "guest_port": 22})
if err != nil {
return fmt.Errorf("vm_expose over MCP: %w", err)
}
address, err := exposureAddress(exposed)
if err != nil {
return err
}
if err := proveBanner(ctx, address, now, sleep, dial); err != nil {
return err
}
if _, err := c.Call(ctx, "vm_unexpose", map[string]any{"vm": vmName, "guest_port": 22}); err != nil {
return fmt.Errorf("vm_unexpose over MCP: %w", err)
}
return proveUDPExposure(ctx, c, vmName, now, sleep, echo)
}
// udpEchoPort is where the smoke's own echo listens inside the guest. A TCP
// exposure can be proven against the sshd that is already there; UDP has no
// such standing service, so the leg brings its own — which is the honest test
// anyway, since it proves a port nothing else on the host or the guest is
// touching.
const udpEchoPort = 17007
// udpEchoCommand starts that echo and returns. python3 is what a guest that
// booted through cloud-init necessarily has, so this needs no package install
// and no image of its own; nohup and the redirects are what let the exec's
// session close while the echo keeps running.
func udpEchoCommand() string {
return fmt.Sprintf(`nohup python3 -c '
import socket
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
s.bind(("0.0.0.0", %d))
while True:
datagram, sender = s.recvfrom(2048)
s.sendto(datagram, sender)
' >/dev/null 2>&1 &`, udpEchoPort)
}
// proveUDPExposure publishes a UDP port of the MCP leg's own VM and proves a
// datagram makes the whole round trip: through the host's packet proxy, into
// the guest, and back out to the caller through the same published port. The
// grant is revoked before the leg returns; the echo dies with the VM.
func proveUDPExposure(ctx context.Context, c mcpTools, vmName string, now func() time.Time, sleep func(time.Duration), echo echoFunc) error {
if _, err := c.Call(ctx, "vm_exec", map[string]any{"vm": vmName, "command": udpEchoCommand()}); err != nil {
return fmt.Errorf("start the guest UDP echo over MCP: %w", err)
}
exposed, err := c.Call(ctx, "vm_expose", map[string]any{
"vm": vmName, "guest_port": udpEchoPort, "protocol": "udp",
})
if err != nil {
return fmt.Errorf("vm_expose udp over MCP: %w", err)
}
if got, _ := exposed["exposure"].(map[string]any); got != nil {
if proto, _ := got["protocol"].(string); proto != "udp" {
return fmt.Errorf("FAIL: vm_expose published %q, want a udp exposure", proto)
}
}
address, err := exposureAddress(exposed)
if err != nil {
return err
}
nonce, err := randNonce()
if err != nil {
return err
}
var lastErr error
err = pollLoop(ctx, now, sleep, 60*time.Second, 3*time.Second, func() (bool, error) {
// Everything transient lives in this retry: the host binds the socket on
// its next converge, the guest's echo takes a moment to come up, and a
// datagram is allowed to go missing on any hop between them.
back, derr := echo(ctx, address, nonce)
if derr != nil {
lastErr = derr
return false, nil
}
if back != nonce {
return false, fmt.Errorf("FAIL: published UDP port %s echoed %q, want %q", address, truncateBanner(back), nonce)
}
return true, nil
})
if errors.Is(err, errPollTimeout) {
return fmt.Errorf("FAIL: no echo from the MCP-published UDP port %s within 60s: %v", address, lastErr)
}
if err != nil {
return err
}
if _, err := c.Call(ctx, "vm_unexpose", map[string]any{
"vm": vmName, "guest_port": udpEchoPort, "protocol": "udp",
}); err != nil {
return fmt.Errorf("vm_unexpose udp over MCP: %w", err)
}
return nil
}
// delegationTTL is how long the smoke's delegation lives. Long enough for one
// gate run, short enough that a run which dies mid-way leaves nothing usable.
const delegationTTL = 30 * time.Minute
// signDelegation is `ssh-keygen -s`, performed in process: the caller's CA
// signs eitri's ephemeral public key for the guest login user.
func signDelegation(ca ssh.Signer, pubLine, principal string, now time.Time) (string, error) {
pub, _, _, _, err := ssh.ParseAuthorizedKey([]byte(pubLine))
if err != nil {
return "", fmt.Errorf("parse the key eitri asked us to sign: %w", err)
}
var serial uint64
if err := binary.Read(rand.Reader, binary.BigEndian, &serial); err != nil {
return "", fmt.Errorf("certificate serial: %w", err)
}
cert := &ssh.Certificate{
Key: pub,
Serial: serial,
CertType: ssh.UserCert,
KeyId: "eitri-smoke-delegation",
ValidPrincipals: []string{principal},
ValidAfter: uint64(now.Add(-time.Minute).Unix()),
ValidBefore: uint64(now.Add(delegationTTL).Unix()),
Permissions: ssh.Permissions{Extensions: map[string]string{
"permit-pty": "", "permit-port-forwarding": "",
}},
}
if err := cert.SignCert(rand.Reader, ca); err != nil {
return "", fmt.Errorf("sign the delegation certificate: %w", err)
}
return authorizedLine(cert), nil
}
// authorizedLine renders a key or certificate as the one-line form ssh-keygen
// writes and every endpoint here accepts.
func authorizedLine(k ssh.PublicKey) string {
return strings.TrimSpace(string(ssh.MarshalAuthorizedKey(k)))
}
// infoListsCA checks tenant_info named the CA the leg registered, and that it
// described it rather than reproducing it.
func infoListsCA(info map[string]any, want string) error {
cas, _ := info["registered_cas"].([]any)
for _, raw := range cas {
ca, _ := raw.(map[string]any)
if fp, _ := ca["fingerprint"].(string); fp == want {
if _, leaked := ca["pubkey"]; leaked {
return errors.New("FAIL: tenant_info returned CA key material; it must describe, not reproduce")
}
return nil
}
}
return fmt.Errorf("FAIL: tenant_info does not list the CA this run registered (%s)", want)
}
// exposureAddress pulls the dialable address out of a vm_expose result. The
// address is empty until the VM's host has reported its uplink, so an empty one
// is a real failure rather than something to retry.
func exposureAddress(out map[string]any) (string, error) {
exposure, ok := out["exposure"].(map[string]any)
if !ok {
return "", errors.New("FAIL: vm_expose returned no exposure")
}
address, _ := exposure["address"].(string)
if address == "" {
return "", errors.New("FAIL: vm_expose returned no address to dial")
}
return address, nil
}
// proveBanner dials a published address until the guest's sshd answers, the
// same proof proveExposure makes of the API path.
func proveBanner(ctx context.Context, address string, now func() time.Time, sleep func(time.Duration), dial bannerFunc) error {
var lastErr error
err := pollLoop(ctx, now, sleep, 60*time.Second, 3*time.Second, func() (bool, error) {
banner, derr := dial(ctx, address)
if derr != nil {
// The listener binds on the host's next converge, a tick away.
lastErr = derr
return false, nil
}
if !strings.HasPrefix(banner, sshBannerPrefix) {
return false, fmt.Errorf("FAIL: port published over MCP at %s answered %q, want an %s banner",
address, truncateBanner(banner), sshBannerPrefix)
}
return true, nil
})
if errors.Is(err, errPollTimeout) {
return fmt.Errorf("FAIL: no %s banner from the MCP-published port %s within 60s: %v", sshBannerPrefix, address, lastErr)
}
return err
}
// randNonce returns a value the guest cannot have echoed by accident.
func randNonce() (string, error) {
var b [6]byte
if _, err := rand.Read(b[:]); err != nil {
return "", fmt.Errorf("generate exec nonce: %w", err)
}
return "mcp-" + hex.EncodeToString(b[:]), nil
}