nixos-router/pkgs/router-webui/mcp_test.go

package main

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
)

// rpc posts a JSON-RPC request to /mcp.
func rpc(h http.Handler, token, method string, params any) (*httptest.ResponseRecorder, map[string]any) {
	body, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": 1, "method": method, "params": params})
	req := httptest.NewRequest("POST", "/mcp", strings.NewReader(string(body)))
	if token != "" {
		req.Header.Set("Authorization", "Bearer "+token)
	}
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	var out map[string]any
	_ = json.Unmarshal(rec.Body.Bytes(), &out)
	return rec, out
}

// call runs a tool and returns its text and whether it failed.
func call(t *testing.T, h http.Handler, token, tool string, args any) (string, bool) {
	t.Helper()
	rec, out := rpc(h, token, "tools/call", map[string]any{"name": tool, "arguments": args})
	if rec.Code != 200 {
		t.Fatalf("%s: HTTP %d", tool, rec.Code)
	}
	if out["error"] != nil {
		return out["error"].(map[string]any)["message"].(string), true
	}
	res := out["result"].(map[string]any)
	text := res["content"].([]any)[0].(map[string]any)["text"].(string)
	failed, _ := res["isError"].(bool)
	return text, failed
}

func TestMCPAuth(t *testing.T) {
	a, _ := newTestApp(t)
	h := a.routes()
	if rec, _ := rpc(h, "", "ping", nil); rec.Code != http.StatusUnauthorized || rec.Header().Get("WWW-Authenticate") == "" {
		t.Fatalf("no token: %d", rec.Code)
	}
	if rec, _ := rpc(h, "rtr_nonsense", "ping", nil); rec.Code != http.StatusUnauthorized {
		t.Fatalf("wrong token: %d", rec.Code)
	}
	if rec := do(h, "GET", "/mcp", nil, nil); rec.Code != http.StatusMethodNotAllowed {
		t.Fatalf("GET: %d", rec.Code)
	}
	_, secret, err := a.tokens.Create("test", "read")
	if err != nil {
		t.Fatal(err)
	}
	req := httptest.NewRequest("POST", "/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"ping"}`))
	req.Header.Set("Authorization", "Bearer "+secret)
	req.Header.Set("Origin", "https://evil.example")
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusForbidden {
		t.Fatalf("foreign origin: %d", rec.Code)
	}
	// A session cookie is no token.
	cookie := login(t, h)
	req = httptest.NewRequest("POST", "/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"ping"}`))
	req.AddCookie(cookie)
	rec = httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != http.StatusUnauthorized {
		t.Fatalf("session cookie: %d", rec.Code)
	}
}

func TestMCPProtocol(t *testing.T) {
	a, sys := newTestApp(t)
	h := a.routes()
	_, read, _ := a.tokens.Create("reader", "read")
	_, change, _ := a.tokens.Create("writer", "change")

	rec, out := rpc(h, read, "initialize", map[string]any{"protocolVersion": "2025-03-26"})
	if rec.Code != 200 || out["result"].(map[string]any)["protocolVersion"] != "2025-03-26" {
		t.Fatalf("initialize: %s", rec.Body)
	}
	// A notification gets no body.
	req := httptest.NewRequest("POST", "/mcp", strings.NewReader(`{"jsonrpc":"2.0","method":"notifications/initialized"}`))
	req.Header.Set("Authorization", "Bearer "+read)
	nrec := httptest.NewRecorder()
	h.ServeHTTP(nrec, req)
	if nrec.Code != http.StatusAccepted || nrec.Body.Len() != 0 {
		t.Fatalf("notification: %d %q", nrec.Code, nrec.Body)
	}
	if _, out := rpc(h, read, "nonsense", nil); out["error"] == nil {
		t.Fatal("unknown method accepted")
	}

	names := func(token string) string {
		_, out := rpc(h, token, "tools/list", nil)
		var n []string
		for _, t := range out["result"].(map[string]any)["tools"].([]any) {
			n = append(n, t.(map[string]any)["name"].(string))
		}
		return strings.Join(n, " ")
	}
	if got := names(read); strings.Contains(got, "update_settings") || !strings.Contains(got, "get_settings") {
		t.Fatalf("read tools: %s", got)
	}
	if got := names(change); !strings.Contains(got, "update_settings") || !strings.Contains(got, "discard_changes") {
		t.Fatalf("change tools: %s", got)
	}
	if text, failed := call(t, h, read, "update_settings", map[string]any{}); !failed || !strings.Contains(text, "unknown tool") {
		t.Fatalf("a read token edited settings: %s", text)
	}

	// Reading.
	text, failed := call(t, h, read, "get_settings", nil)
	if failed {
		t.Fatal(text)
	}
	var got struct {
		Etag     string
		Settings map[string]any
	}
	if err := json.Unmarshal([]byte(text), &got); err != nil || got.Etag == "" || got.Settings["router"] == nil {
		t.Fatalf("get_settings: %v %s", err, text)
	}
	if _, failed := call(t, h, read, "get_status", nil); failed {
		t.Fatal("get_status failed")
	}
	if text, failed := call(t, h, read, "list_interfaces", nil); failed || !strings.Contains(text, "WAN") {
		t.Fatalf("list_interfaces: %s", text)
	}
	if _, failed := call(t, h, read, "list_dhcp_leases", nil); failed {
		t.Fatal("list_dhcp_leases failed")
	}

	// Editing: a patch merges, a stale etag is refused, nothing is applied.
	patch := map[string]any{"etag": got.Etag, "patch": map[string]any{"system": map[string]any{"hostName": "assistant-was-here"}}}
	text, failed = call(t, h, change, "update_settings", patch)
	if failed || !strings.Contains(text, "assistant-was-here") || !strings.Contains(text, `"applied": false`) {
		t.Fatalf("update_settings: %s", text)
	}
	if data, _ := os.ReadFile(sys.settingsPath()); !strings.Contains(string(data), "assistant-was-here") {
		t.Fatal("not saved")
	}
	if text, failed := call(t, h, change, "update_settings", patch); !failed || !strings.Contains(text, "etag") {
		t.Fatalf("a stale etag was accepted: %s", text)
	}
	// A change shows up for review (the confirmed copy is the sample).
	os.WriteFile(sys.confirmedPath(), []byte(sample), 0o600)
	if text, _ := call(t, h, read, "get_pending_changes", nil); !strings.Contains(text, "+") || !strings.Contains(text, "assistant-was-here") {
		t.Fatalf("pending: %s", text)
	}
	if text, failed := call(t, h, change, "discard_changes", nil); failed {
		t.Fatal(text)
	}
	if data, _ := os.ReadFile(sys.settingsPath()); strings.Contains(string(data), "assistant-was-here") {
		t.Fatal("discard did not restore the settings")
	}

	// Every call was logged, with the token's name.
	audit := sys.RecentAudit(50)
	if len(audit) < 5 || audit[0].Tool != "discard_changes" || audit[0].Token != "writer" {
		t.Fatalf("audit: %+v", audit)
	}
}

func TestMCPEditLimits(t *testing.T) {
	a, sys := newTestApp(t)
	h := a.routes()
	_, change, _ := a.tokens.Create("writer", "change")
	etagNow := func() string {
		data, _ := os.ReadFile(sys.settingsPath())
		return etag(data)
	}
	for name, patch := range map[string]map[string]any{
		"ssh keys":      {"system": map[string]any{"sshKeys": []string{"ssh-ed25519 AAAA attacker"}}},
		"state version": {"system": map[string]any{"stateVersion": "20.09"}},
		"schema":        {"version": 99},
		"top level":     {"users": map[string]any{"x": 1}},
		"secret path":   {"router": map[string]any{"wan": map[string]any{"ipv4": map[string]any{"pppoe": map[string]any{"passwordFile": "/etc/shadow"}}}}},
		"path escape":   {"router": map[string]any{"wan": map[string]any{"ipv4": map[string]any{"pppoe": map[string]any{"passwordFile": sys.StateDir + "/../etc/shadow"}}}}},
		"nested file":   {"router": map[string]any{"wireguard": map[string]any{"wg0": map[string]any{"privateKeyFile": "/root/.ssh/id_rsa"}}}},
	} {
		before, _ := os.ReadFile(sys.settingsPath())
		if text, failed := call(t, h, change, "update_settings", map[string]any{"etag": etagNow(), "patch": patch}); !failed {
			t.Errorf("%s accepted: %s", name, text)
		}
		if after, _ := os.ReadFile(sys.settingsPath()); string(after) != string(before) {
			t.Errorf("%s changed the file", name)
		}
	}
	// Files in the state directory, and null to remove a key, are fine.
	ok := map[string]any{"router": map[string]any{"wan": map[string]any{"ipv4": map[string]any{"pppoe": map[string]any{"passwordFile": sys.StateDir + "/pppoe-password"}}}}}
	if text, failed := call(t, h, change, "update_settings", map[string]any{"etag": etagNow(), "patch": ok}); failed {
		t.Fatalf("a file in the state directory refused: %s", text)
	}
	rm := map[string]any{"router": map[string]any{"wan": map[string]any{"ipv4": map[string]any{"pppoe": nil}}}}
	if text, failed := call(t, h, change, "update_settings", map[string]any{"etag": etagNow(), "patch": rm}); failed || strings.Contains(text, "pppoe-password\"") && !strings.Contains(text, "-") {
		t.Fatalf("null did not remove: %s", text)
	}
}

func TestMergePatchAndDiff(t *testing.T) {
	got, _ := json.Marshal(mergePatch(
		map[string]any{"a": map[string]any{"b": 1, "c": 2}, "d": []any{1, 2}, "e": "x"},
		map[string]any{"a": map[string]any{"b": nil, "z": 3}, "d": []any{9}}))
	if string(got) != `{"a":{"c":2,"z":3},"d":[9],"e":"x"}` {
		t.Fatalf("merge: %s", got)
	}
	if diffLines("a\nb", "a\nb") != "" {
		t.Fatal("equal texts differ")
	}
	d := diffLines("1\n2\n3\n4\n5\n6\n7\n8\n9", "1\n2\n3\nfour\n5\n6\n7\n8\n9")
	if !strings.Contains(d, "- 4\n+ four\n") || !strings.Contains(d, "…") {
		t.Fatalf("diff:\n%s", d)
	}
}

func TestTokensPage(t *testing.T) {
	a, _ := newTestApp(t)
	h := a.routes()
	cookie := login(t, h)
	token := csrfFrom(t, do(h, "GET", "/tokens", nil, cookie).Body.String())
	rec := do(h, "POST", "/tokens", map[string][]string{"csrf": {token}, "name": {"laptop"}, "scope": {"read"}}, cookie)
	body := rec.Body.String()
	i := strings.Index(body, "rtr_")
	if rec.Code != 200 || i < 0 {
		t.Fatalf("create: %d\n%s", rec.Code, body)
	}
	secret := body[i : i+4+43]
	if _, ok := a.tokens.Lookup(secret); !ok {
		t.Fatal("the shown token does not work")
	}
	// The secret is not stored or shown again.
	if data, _ := os.ReadFile(a.sys.tokensPath()); strings.Contains(string(data), secret) {
		t.Fatal("the secret is stored")
	}
	if strings.Contains(do(h, "GET", "/tokens", nil, cookie).Body.String(), secret) {
		t.Fatal("the secret is shown again")
	}
	toks := a.tokens.List()
	do(h, "POST", "/tokens/revoke", map[string][]string{"csrf": {token}, "id": {toks[0].ID}}, cookie)
	if _, ok := a.tokens.Lookup(secret); ok {
		t.Fatal("a revoked token works")
	}
	if rec := do(h, "GET", "/changes", nil, cookie); rec.Code != 200 {
		t.Fatalf("changes page: %d", rec.Code)
	}
}