a73x

internal/server/mcphttp/mcphttp_test.go

Ref:   Size: 20.5 KiB   History

package mcphttp

import (
	"bytes"
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/a73x/eitri/internal/server/api"
	"github.com/a73x/eitri/internal/server/api/types"
	"github.com/a73x/eitri/internal/server/delegation"
	"github.com/a73x/eitri/internal/server/hub"
	"github.com/a73x/eitri/internal/server/registry"
	"github.com/a73x/eitri/internal/server/store"
	"github.com/a73x/eitri/internal/server/vmssh"
	"github.com/modelcontextprotocol/go-sdk/mcp"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/crypto/ssh"
)

// fixture is a control plane with two tenants, each holding a PAT, fronted by
// the same root mux the server binary builds: /api/ and /mcp behind one
// authentication.
type fixture struct {
	ts      *httptest.Server
	st      *store.Store
	reg     *registry.Registry
	tenantA string
	patA    string
	tenantB string
	patB    string
	keyring *delegation.Keyring
}

// testCreds is the pairing boot makes: the in-memory keyring plus the store.
type testCreds struct {
	k  *delegation.Keyring
	st *store.Store
}

func (c testCreds) Delegated(tenant string) (ssh.Signer, bool)  { return c.k.Signer(tenant) }
func (c testCreds) TenantHasUserCA(tenant string) (bool, error) { return c.st.TenantHasUserCA(tenant) }

func newFixture(t *testing.T) fixture {
	t.Helper()
	st, err := store.Open(t.TempDir()+"/db", "10.77.0.0/16")
	require.NoError(t, err)
	t.Cleanup(func() { st.Close() })

	reg := registry.New(time.Now)
	a := api.New(api.Config{
		HostSecret:       []byte("hostsecret"),
		AdvertiseHTTP:    "http://127.0.0.1:8080",
		AdvertiseQUIC:    "127.0.0.1:8443",
		ServerCertSHA256: strings.Repeat("c", 64),
	}, st, reg, hub.New())
	t.Cleanup(a.Close)

	f := fixture{st: st, reg: reg}
	f.tenantA, f.patA = newTenant(t, st, "alpha")
	f.tenantB, f.patB = newTenant(t, st, "beta")

	// A host CA stands in for a configured jump gate, so the exec path gets past
	// "no CA configured" and reaches the tenant's own CA situation.
	_, hostCAPriv, err := ed25519.GenerateKey(rand.Reader)
	require.NoError(t, err)
	hostCA, err := ssh.NewSignerFromSigner(hostCAPriv)
	require.NoError(t, err)

	// The delegation keyring the control plane wires: shared between the API
	// (which fills it) and the exec path (which reads it), exactly as in boot.
	f.keyring = delegation.New(time.Now, "ubuntu")
	a.SetDelegations(f.keyring)

	root := http.NewServeMux()
	root.Handle("/api/", a.Handler())
	mcpHandler := a.UserAuth(New(Deps{
		Handler: a.Handler(), Creds: testCreds{k: f.keyring, st: st}, VMUser: "ubuntu", HostCA: hostCA.PublicKey(),
		Lookup: func(string, string) (vmssh.VM, bool) { return vmssh.VM{}, false },
	}))
	root.Handle("/mcp", mcpHandler)
	root.Handle("/mcp/", mcpHandler)

	f.ts = httptest.NewServer(root)
	t.Cleanup(f.ts.Close)
	return f
}

func newTenant(t *testing.T, st *store.Store, handle string) (string, string) {
	t.Helper()
	tn, err := st.CreateTenantForIdentity("https://issuer.example", "sub-"+handle, handle+"@example.com")
	require.NoError(t, err)
	pat, _, err := st.CreateAPIToken(tn.ID, handle, 0)
	require.NoError(t, err)
	return tn.ID, pat
}

// seedVM gives a tenant one host and one ready VM, so vm_list has something to
// filter and the exec tools get as far as trying to reach a guest.
//
// The report at the end is what makes the VM ready, not the row: a VM whose
// host is not reporting reads `unreachable` however settled its columns look
// (see deriveLifecycle), and every exec tool gates on `ready`. reg may be nil
// for a caller that only needs the rows.
func seedVM(t *testing.T, st *store.Store, reg *registry.Registry, tenant, hostName, vmName string) {
	t.Helper()
	tok, err := st.CreateEnrollmentToken(tenant)
	require.NoError(t, err)
	h, err := st.RedeemEnrollmentToken(tok, store.EnrollFacts{
		Name: hostName, OS: "linux", Arch: "amd64", Provisioner: "cloudhv"})
	require.NoError(t, err)
	require.NoError(t, st.CreateVM(store.VM{
		ID: vmName + "-id", HostID: h.ID, Name: vmName,
		ImageURL: "https://images.example/x.img", ImageSHA256: strings.Repeat("a", 64),
		VCPUs: 1, MemMB: 512, DiskGB: 5, PowerState: "running"}))
	_, err = st.RecordVMStatus(vmName+"-id", h.ID, "ready", "", "10.77.0.5")
	require.NoError(t, err)
	if reg != nil {
		reg.UpdateReport(h.ID, registry.Report{VMs: []registry.VMStatus{{
			VMID: vmName + "-id", PowerState: "running", Phase: "ready", IP: "10.77.0.5"}}})
	}
}

// connect opens a real MCP client against /mcp with pat as the bearer token.
func connect(t *testing.T, f fixture, pat string) *mcp.ClientSession {
	t.Helper()
	tr := &mcp.StreamableClientTransport{
		Endpoint:   f.ts.URL + "/mcp",
		HTTPClient: &http.Client{Transport: bearerTransport{pat: pat}},
	}
	cs, err := mcp.NewClient(&mcp.Implementation{Name: "test", Version: "0"}, nil).Connect(t.Context(), tr, nil)
	require.NoError(t, err)
	t.Cleanup(func() { cs.Close() })
	return cs
}

// bearerTransport attaches the PAT the way a remote MCP client would.
type bearerTransport struct{ pat string }

func (b bearerTransport) RoundTrip(r *http.Request) (*http.Response, error) {
	if b.pat != "" {
		r.Header.Set("Authorization", "Bearer "+b.pat)
	}
	return http.DefaultTransport.RoundTrip(r)
}

// post sends a raw JSON-RPC body, for the cases an MCP client cannot express
// (no credential, a bad one, the wrong method).
func post(t *testing.T, f fixture, method, pat, body string) *http.Response {
	t.Helper()
	req, err := http.NewRequest(method, f.ts.URL+"/mcp", strings.NewReader(body))
	require.NoError(t, err)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json, text/event-stream")
	if pat != "" {
		req.Header.Set("Authorization", "Bearer "+pat)
	}
	resp, err := http.DefaultClient.Do(req)
	require.NoError(t, err)
	t.Cleanup(func() { resp.Body.Close() })
	return resp
}

const initBody = `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"t","version":"0"}}}`

// ── tests ────────────────────────────────────────────────────────────────────

// TestMCPRequiresACredential: the endpoint lives outside the REST route table
// but not outside its authentication.
func TestMCPRequiresACredential(t *testing.T) {
	f := newFixture(t)
	assert.Equal(t, http.StatusUnauthorized, post(t, f, "POST", "", initBody).StatusCode)
}

// TestMCPRejectsABadPAT: an invalid token is indistinguishable from none.
func TestMCPRejectsABadPAT(t *testing.T) {
	f := newFixture(t)
	assert.Equal(t, http.StatusUnauthorized, post(t, f, "POST", "eitri_pat_notreal", initBody).StatusCode)
}

// TestMCPGetIsNotAllowed: stateless streamable HTTP has no server-initiated
// stream to open, so a GET is answered 405 rather than hanging.
func TestMCPGetIsNotAllowed(t *testing.T) {
	f := newFixture(t)
	resp := post(t, f, "GET", f.patA, "")
	assert.Equal(t, http.StatusMethodNotAllowed, resp.StatusCode)
}

// TestMCPExposesTheFullToolset: a remote caller gets the ten VM tools plus the
// two delegation tools, which is what makes a bare PAT enough.
func TestMCPExposesTheFullToolset(t *testing.T) {
	f := newFixture(t)
	cs := connect(t, f, f.patA)

	res, err := cs.ListTools(t.Context(), nil)
	require.NoError(t, err)
	names := make([]string, 0, len(res.Tools))
	for _, tool := range res.Tools {
		names = append(names, tool.Name)
	}
	assert.Len(t, names, 14)
	assert.Contains(t, names, "delegate_begin")
	assert.Contains(t, names, "delegate_complete")
	assert.Contains(t, names, "ca_upload")
	assert.Contains(t, names, "tenant_info")
	assert.Contains(t, names, "vm_exec")
}

// TestVMListIsTenantIsolated is the property the whole per-request-identity
// design exists for: two PATs, two tenants, and neither sees the other's VMs.
func TestVMListIsTenantIsolated(t *testing.T) {
	f := newFixture(t)
	seedVM(t, f.st, f.reg, f.tenantA, "host-alpha", "alpha-vm")
	seedVM(t, f.st, f.reg, f.tenantB, "host-beta", "beta-vm")

	assert.Equal(t, []string{"alpha-vm"}, listVMNames(t, connect(t, f, f.patA)))
	assert.Equal(t, []string{"beta-vm"}, listVMNames(t, connect(t, f, f.patB)))
}

func listVMNames(t *testing.T, cs *mcp.ClientSession) []string {
	t.Helper()
	res, err := cs.CallTool(t.Context(), &mcp.CallToolParams{Name: "vm_list"})
	require.NoError(t, err)
	require.False(t, res.IsError, "vm_list failed: %+v", res.Content)

	var out struct {
		VMs []types.VM `json:"vms"`
	}
	raw, err := json.Marshal(res.StructuredContent)
	require.NoError(t, err)
	require.NoError(t, json.Unmarshal(raw, &out))
	names := make([]string, 0, len(out.VMs))
	for _, vm := range out.VMs {
		names = append(names, vm.Name)
	}
	return names
}

// TestDelegationRoundTripThroughTheTransport drives the whole exchange the way
// a model would: ask for a key, sign it out of band with the tenant's own CA,
// hand the certificate back, and find eitri holding a credential it could not
// have made for itself.
func TestDelegationRoundTripThroughTheTransport(t *testing.T) {
	f := newFixture(t)
	ca := registerCA(t, f, f.patA)
	cs := connect(t, f, f.patA)

	begin := callBegin(t, cs)
	require.NotEmpty(t, begin.PublicKey)
	assert.Equal(t, "ubuntu", begin.Principal)
	assert.Contains(t, begin.Instructions, "ssh-keygen -s")
	assert.Contains(t, begin.Note, "cannot sign this itself")

	done := callComplete(t, cs, signDelegation(t, ca, begin.PublicKey, "ubuntu", time.Now().Add(time.Hour)))
	assert.Equal(t, ssh.FingerprintSHA256(ca.PublicKey()), done.CAFingerprint)
	assert.Equal(t, []string{"ubuntu"}, done.Principals)
	assert.Contains(t, done.Note, "It holds no signing key")

	// eitri can now authenticate as this tenant — and only as this tenant.
	_, ok := f.keyring.Signer(f.tenantA)
	assert.True(t, ok)
	_, ok = f.keyring.Signer(f.tenantB)
	assert.False(t, ok, "one tenant's delegation is nobody else's")
}

// TestDelegateCompleteRefusesAnUnregisteredCA is the live proof that the trust
// check is wired to the tenant's real CA set, and that its refusal says so.
func TestDelegateCompleteRefusesAnUnregisteredCA(t *testing.T) {
	f := newFixture(t)
	registerCA(t, f, f.patA)
	cs := connect(t, f, f.patA)

	begin := callBegin(t, cs)
	stranger := newCA(t)
	res, err := cs.CallTool(t.Context(), &mcp.CallToolParams{
		Name: "delegate_complete",
		Arguments: map[string]any{
			"certificate": signDelegation(t, stranger, begin.PublicKey, "ubuntu", time.Now().Add(time.Hour)),
		},
	})
	require.NoError(t, err)
	require.True(t, res.IsError)
	assert.Contains(t, toolErrorText(res), "not a CA registered to this tenant")

	_, ok := f.keyring.Signer(f.tenantA)
	assert.False(t, ok, "a refused certificate must leave eitri with nothing")
}

// TestDelegateCompleteRefusesTheWrongPrincipal: the most likely user mistake
// has to be caught here, or it resurfaces as an opaque SSH failure later.
func TestDelegateCompleteRefusesTheWrongPrincipal(t *testing.T) {
	f := newFixture(t)
	ca := registerCA(t, f, f.patA)
	cs := connect(t, f, f.patA)

	begin := callBegin(t, cs)
	res, err := cs.CallTool(t.Context(), &mcp.CallToolParams{
		Name: "delegate_complete",
		Arguments: map[string]any{
			"certificate": signDelegation(t, ca, begin.PublicKey, "alex", time.Now().Add(time.Hour)),
		},
	})
	require.NoError(t, err)
	require.True(t, res.IsError)
	assert.Contains(t, toolErrorText(res), `must include "ubuntu"`)
}

// newCA returns a signer standing in for someone'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
}

// registerCA uploads a CA to the caller's tenant through the real endpoint,
// which is what makes a delegation signed by it acceptable.
func registerCA(t *testing.T, f fixture, pat string) ssh.Signer {
	t.Helper()
	ca := newCA(t)
	body, err := json.Marshal(types.UserCARequest{
		PublicKey: string(ssh.MarshalAuthorizedKey(ca.PublicKey())), Label: "mine"})
	require.NoError(t, err)
	req, err := http.NewRequest(http.MethodPost, f.ts.URL+"/api/v1/user-cas", bytes.NewReader(body))
	require.NoError(t, err)
	req.Header.Set("Authorization", "Bearer "+pat)
	req.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(req)
	require.NoError(t, err)
	defer resp.Body.Close()
	require.Equal(t, http.StatusCreated, resp.StatusCode)
	return ca
}

// signDelegation is what the human runs: `ssh-keygen -s`, in code.
func signDelegation(t *testing.T, ca ssh.Signer, pubLine, principal string, expiry time.Time) string {
	t.Helper()
	pub, _, _, _, err := ssh.ParseAuthorizedKey([]byte(pubLine))
	require.NoError(t, err)
	cert := &ssh.Certificate{
		Key:             pub,
		Serial:          7,
		CertType:        ssh.UserCert,
		KeyId:           "eitri-delegation",
		ValidPrincipals: []string{principal},
		ValidAfter:      uint64(time.Now().Add(-time.Minute).Unix()),
		ValidBefore:     uint64(expiry.Unix()),
	}
	require.NoError(t, cert.SignCert(rand.Reader, ca))
	return strings.TrimSpace(string(ssh.MarshalAuthorizedKey(cert)))
}

type beginOut struct {
	PublicKey    string `json:"public_key"`
	Principal    string `json:"principal"`
	Instructions string `json:"instructions"`
	Note         string `json:"note"`
}

type delegationOut struct {
	ExpiresAt     string   `json:"expires_at"`
	CAFingerprint string   `json:"ca_fingerprint"`
	Principals    []string `json:"principals"`
	Note          string   `json:"note"`
}

func callBegin(t *testing.T, cs *mcp.ClientSession) beginOut {
	t.Helper()
	var out beginOut
	callToolInto(t, cs, &mcp.CallToolParams{Name: "delegate_begin"}, &out)
	return out
}

func callComplete(t *testing.T, cs *mcp.ClientSession, cert string) delegationOut {
	t.Helper()
	var out delegationOut
	callToolInto(t, cs, &mcp.CallToolParams{
		Name: "delegate_complete", Arguments: map[string]any{"certificate": cert}}, &out)
	return out
}

func callToolInto(t *testing.T, cs *mcp.ClientSession, params *mcp.CallToolParams, into any) {
	t.Helper()
	res, err := cs.CallTool(t.Context(), params)
	require.NoError(t, err)
	require.False(t, res.IsError, "%s failed: %+v", params.Name, res.Content)
	raw, err := json.Marshal(res.StructuredContent)
	require.NoError(t, err)
	require.NoError(t, json.Unmarshal(raw, into))
}

// TestExecWithoutADelegationRefusesInWords: a tool failure comes back as an MCP
// error the model can read and act on, not a 500 — and eitri cannot get access
// by itself, so the message has to be a request with a recipe in it.
func TestExecWithoutADelegationRefusesInWords(t *testing.T) {
	f := newFixture(t)
	seedVM(t, f.st, f.reg, f.tenantA, "host-alpha", "alpha-vm")
	cs := connect(t, f, f.patA)

	res, err := cs.CallTool(t.Context(), &mcp.CallToolParams{
		Name:      "vm_exec",
		Arguments: map[string]any{"vm": "alpha-vm", "command": "true"},
	})
	require.NoError(t, err, "a refused tool must not fail the transport")
	require.True(t, res.IsError)
	assert.Contains(t, toolErrorText(res), "no live SSH delegation")
	assert.Contains(t, toolErrorText(res), "delegate_begin")
}

// TestExecWithoutAJumpGateSaysSo: on a control plane with no SSH CA at all,
// guests carry no host certificate and remote exec cannot be made safe. The
// refusal names that, rather than sending the caller off to delegate.
func TestExecWithoutAJumpGateSaysSo(t *testing.T) {
	st, err := store.Open(t.TempDir()+"/db", "10.77.0.0/16")
	require.NoError(t, err)
	t.Cleanup(func() { st.Close() })
	reg := registry.New(time.Now)
	a := api.New(api.Config{HostSecret: []byte("hostsecret"), ServerCertSHA256: strings.Repeat("c", 64)},
		st, reg, hub.New())
	t.Cleanup(a.Close)
	tenant, pat := newTenant(t, st, "alpha")
	seedVM(t, st, reg, tenant, "host-alpha", "alpha-vm")

	root := http.NewServeMux()
	root.Handle("/api/", a.Handler())
	root.Handle("/mcp", a.UserAuth(New(Deps{
		Handler: a.Handler(), Creds: testCreds{k: delegation.New(time.Now, "ubuntu"), st: st}, VMUser: "ubuntu"})))
	ts := httptest.NewServer(root)
	t.Cleanup(ts.Close)

	cs := connect(t, fixture{ts: ts}, pat)
	res, err := cs.CallTool(t.Context(), &mcp.CallToolParams{
		Name:      "vm_exec",
		Arguments: map[string]any{"vm": "alpha-vm", "command": "true"},
	})
	require.NoError(t, err)
	require.True(t, res.IsError)
	assert.Contains(t, toolErrorText(res), "no SSH CA configured")
}

// TestToolErrorsNeverEchoTheCredential: the PAT rides every request, so it must
// never come back out in a message a model will repeat.
func TestToolErrorsNeverEchoTheCredential(t *testing.T) {
	f := newFixture(t)
	cs := connect(t, f, f.patA)

	res, err := cs.CallTool(t.Context(), &mcp.CallToolParams{
		Name:      "vm_exec",
		Arguments: map[string]any{"vm": "nonexistent", "command": "true"},
	})
	require.NoError(t, err)
	require.True(t, res.IsError)
	assert.NotContains(t, toolErrorText(res), f.patA)
}

func toolErrorText(res *mcp.CallToolResult) string {
	var b strings.Builder
	for _, c := range res.Content {
		if tc, ok := c.(*mcp.TextContent); ok {
			b.WriteString(tc.Text)
		}
	}
	return b.String()
}

// TestInprocPassesTheCallersIdentity pins the in-process transport's contract:
// it re-enters the API's own handler carrying the request unchanged, so the
// answer is the one a remote client with that PAT would get.
func TestInprocPassesTheCallersIdentity(t *testing.T) {
	var gotAuth string
	tr := inproc{h: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotAuth = r.Header.Get("Authorization")
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusTeapot)
		w.Write([]byte(`{"ok":true}`))
	})}

	req, err := http.NewRequestWithContext(context.Background(), "GET", "http://eitri.internal/api/v1/me", nil)
	require.NoError(t, err)
	req.Header.Set("Authorization", "Bearer eitri_pat_x")

	resp, err := tr.RoundTrip(req)
	require.NoError(t, err)
	defer resp.Body.Close()
	assert.Equal(t, "Bearer eitri_pat_x", gotAuth)
	assert.Equal(t, http.StatusTeapot, resp.StatusCode)
	assert.Equal(t, "application/json", resp.Header.Get("Content-Type"))
	assert.Equal(t, int64(len(`{"ok":true}`)), resp.ContentLength)
}

// TestInprocDefaultsToOK covers the handler that writes a body without ever
// naming a status.
func TestInprocDefaultsToOK(t *testing.T) {
	tr := inproc{h: http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Write([]byte("hi"))
	})}
	req, err := http.NewRequest("GET", "http://eitri.internal/x", nil)
	require.NoError(t, err)
	resp, err := tr.RoundTrip(req)
	require.NoError(t, err)
	defer resp.Body.Close()
	assert.Equal(t, http.StatusOK, resp.StatusCode)
}

// TestCAUploadRegistersWithTheCallersOwnTenant drives the tool end to end
// through the transport and out into the store: the CA lands on the tenant the
// PAT names, with its label, and on nobody else's. Registering a CA is the
// first step of the journey a bare token has to make, so it has to work with
// nothing but that token.
func TestCAUploadRegistersWithTheCallersOwnTenant(t *testing.T) {
	f := newFixture(t)
	ca := newCA(t)
	line := strings.TrimSpace(string(ssh.MarshalAuthorizedKey(ca.PublicKey())))

	var out struct {
		Fingerprint string `json:"fingerprint"`
		Label       string `json:"label"`
		Note        string `json:"note"`
	}
	callToolInto(t, connect(t, f, f.patA), &mcp.CallToolParams{
		Name:      "ca_upload",
		Arguments: map[string]any{"public_key": line, "label": "from-mcp"},
	}, &out)

	assert.Equal(t, ssh.FingerprintSHA256(ca.PublicKey()), out.Fingerprint)
	assert.Equal(t, "from-mcp", out.Label)

	cas, err := f.st.ListTenantUserCAs(f.tenantA)
	require.NoError(t, err)
	require.Len(t, cas, 1)
	assert.Equal(t, line, cas[0].Pubkey)
	assert.Equal(t, "from-mcp", cas[0].Label)

	other, err := f.st.ListTenantUserCAs(f.tenantB)
	require.NoError(t, err)
	assert.Empty(t, other, "a CA belongs to the tenant whose token uploaded it")
}

// TestCAUploadThenDelegateIsTheWholeJourney: upload a CA with a bare PAT, then
// delegate against it. The second step only works because the first registered
// the CA the delegation is signed by — which is the point of having both tools.
func TestCAUploadThenDelegateIsTheWholeJourney(t *testing.T) {
	f := newFixture(t)
	ca := newCA(t)
	cs := connect(t, f, f.patA)

	var uploaded struct {
		Fingerprint string `json:"fingerprint"`
	}
	callToolInto(t, cs, &mcp.CallToolParams{
		Name:      "ca_upload",
		Arguments: map[string]any{"public_key": strings.TrimSpace(string(ssh.MarshalAuthorizedKey(ca.PublicKey())))},
	}, &uploaded)

	begin := callBegin(t, cs)
	done := callComplete(t, cs, signDelegation(t, ca, begin.PublicKey, "ubuntu", time.Now().Add(time.Hour)))
	assert.Equal(t, uploaded.Fingerprint, done.CAFingerprint,
		"the delegation must chain to the CA the same session just uploaded")

	_, ok := f.keyring.Signer(f.tenantA)
	assert.True(t, ok)
}