summercms/modules/pier/dispatch_hardening_test.go

package pier_test

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/url"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"git.golem15.com/golem15/summercms/modules/bouncer"
	"git.golem15.com/golem15/summercms/modules/pact"
	"github.com/modelcontextprotocol/go-sdk/mcp"
)

// dispatchGuard is a fakeGuard whose MCPDispatch answer is given per
// route method.
type dispatchGuard struct {
	*fakeGuard
	byMethod func(pact.MCPSubject, string) (pact.MCPDispatch, error)
}

func (g dispatchGuard) MCPDispatch(_ context.Context, s pact.MCPSubject, method string) (pact.MCPDispatch, error) {
	return g.byMethod(s, method)
}

type oneSegmentInput struct {
	ID string `json:"id" path:"id"`
}

type manySegmentsInput struct {
	Rest string `json:"rest" path:"rest"`
}

type numberInput struct {
	ID int64 `json:"id" path:"id"`
	N  int64 `json:"n,omitempty" query:"n"`
}

// TestDispatchHardening pins plan 06's dispatch hardening (A-WR-01, the
// runtime half of A-WR-02, A-WR-08, A-IN-04, A-IN-05, A-IN-08): path values
// the mux never delivers (., .., empty segments, a slash in a single
// segment) are refused before the route; a missing path value is refused;
// a provider error answers a fixed text and its message reaches only the
// log; a large integer keeps its digits in the path and the query; the
// inner request carries the outer Host and forwarding headers but never the
// outer credentials; and a provider whose answer contradicts the tool's
// guard is refused.
func TestDispatchHardening(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	var hits atomic.Int64
	seen := func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]any{
			"id":     r.PathValue("id"),
			"rest":   r.PathValue("rest"),
			"n":      r.URL.Query().Get("n"),
			"host":   r.Host,
			"proto":  r.Header.Get("X-Forwarded-Proto"),
			"fhost":  r.Header.Get("X-Forwarded-Host"),
			"auth":   r.Header.Get("Authorization"),
			"cookie": r.Header.Get("Cookie"),
		})
	}
	hard := routePlugin{id: "test.hard",
		routes: []testRoute{
			{method: http.MethodPost, path: "/hard/one/{id}", mw: []string{"acme.auth"}, h: seen},
			{method: http.MethodGet, path: "/hard/many/{rest...}", mw: []string{"acme.auth"}, h: seen},
			{method: http.MethodGet, path: "/hard/num/{id}", mw: []string{"acme.auth"}, h: seen},
		},
		tools: []pact.MCPTool{
			{Name: "one", Method: "POST", Path: "/hard/one/{id}", Input: oneSegmentInput{}},
			{Name: "many", Method: "GET", Path: "/hard/many/{rest...}", Input: manySegmentsInput{}},
			{Name: "num", Method: "GET", Path: "/hard/num/{id}", Input: numberInput{}},
		}}
	// One frontend guard provider (an OAuth grant acts under exactly one);
	// its dispatch answer depends on the route's method: PUT fails like a
	// database error, PATCH names the backend guard, DELETE answers a
	// backend principal.
	pgxLike := "ERROR: relation \"backend_users\" does not exist (SQLSTATE 42P01)"
	guard := dispatchGuard{fakeGuard: &fakeGuard{}}
	guard.byMethod = func(s pact.MCPSubject, method string) (pact.MCPDispatch, error) {
		switch method {
		case http.MethodPut:
			return pact.MCPDispatch{}, errors.New(pgxLike)
		case http.MethodPatch:
			return pact.MCPDispatch{Principal: &bouncer.Principal{ID: s.UserID}, Guards: []string{"backend"}}, nil
		case http.MethodDelete:
			return pact.MCPDispatch{Principal: &bouncer.Principal{ID: s.UserID, Backend: true}, Guards: []string{"acme"}}, nil
		}
		return pact.MCPDispatch{Principal: &bouncer.Principal{ID: s.UserID}, Guards: []string{"acme"}}, nil
	}
	hard.routes = append(hard.routes,
		testRoute{method: http.MethodPut, path: "/hard/broken", mw: []string{"acme.auth"}, h: seen},
		testRoute{method: http.MethodPatch, path: "/hard/claims-backend", mw: []string{"acme.auth"}, h: seen},
		testRoute{method: http.MethodDelete, path: "/hard/backend-principal", mw: []string{"acme.auth"}, h: seen},
	)
	hard.tools = append(hard.tools,
		pact.MCPTool{Name: "broken", Method: "PUT", Path: "/hard/broken", Input: struct{}{}},
		pact.MCPTool{Name: "claims_backend", Method: "PATCH", Path: "/hard/claims-backend", Input: struct{}{}},
		pact.MCPTool{Name: "backend_principal", Method: "DELETE", Path: "/hard/backend-principal", Input: struct{}{}},
	)
	app := newApp(t, appSpec{guards: []pact.MCPGuard{guard}})
	in := startInstance(t, app, hard)
	t.Cleanup(in.pier().StopReaper)
	outer, err := url.Parse(in.srv.URL)
	if err != nil {
		t.Fatal(err)
	}
	logs := captureLogs(t)
	cs := sdkClient(t, ctx, in.url(), "oas_at_test", clientOpts{headers: map[string]string{
		"Cookie":            "summer_session=outer",
		"X-Forwarded-Proto": "https",
		"X-Forwarded-Host":  "site.example",
	}})
	call := func(name string, args map[string]any) *mcp.CallToolResult {
		t.Helper()
		res, err := cs.CallTool(ctx, &mcp.CallToolParams{Name: name, Arguments: args})
		if err != nil {
			t.Fatalf("%s: %v", name, err)
		}
		return res
	}
	refused := func(t *testing.T, name string, args map[string]any, want string) {
		t.Helper()
		before := hits.Load()
		res := call(name, args)
		if !res.IsError || !strings.Contains(resultText(t, res), want) {
			t.Fatalf("%s %v = %+v, want a tool error with %q", name, args, res, want)
		}
		if hits.Load() != before {
			t.Fatalf("%s %v reached the route", name, args)
		}
	}

	t.Run("path values the mux never delivers are refused", func(t *testing.T) {
		for _, v := range []string{"..", "."} {
			refused(t, "hard_one", map[string]any{"id": v}, "must not contain empty, . or .. segments")
		}
		refused(t, "hard_one", map[string]any{"id": ""}, "one non-empty path segment")
		refused(t, "hard_one", map[string]any{"id": "a/b"}, "one non-empty path segment")
		for _, v := range []string{"a/../b", "a//b", "./a", "a/.."} {
			refused(t, "hard_many", map[string]any{"rest": v}, "must not contain empty, . or .. segments")
		}
		if res := call("hard_many", map[string]any{"rest": "a/b"}); res.IsError || structured(t, res)["rest"] != "a/b" {
			t.Fatalf("a/b on a {rest...} route = %+v", res)
		}
	})

	t.Run("a missing path value is refused", func(t *testing.T) {
		refused(t, "hard_one", map[string]any{}, "invalid arguments")
	})

	t.Run("a provider error answers the fixed text and is only logged", func(t *testing.T) {
		before := hits.Load()
		res := call("hard_broken", nil)
		if !res.IsError || resultText(t, res) != "pier: the tool could not be run for this caller" {
			t.Fatalf("provider error = %+v %q", res, resultText(t, res))
		}
		if strings.Contains(resultText(t, res), "SQLSTATE") {
			t.Fatal("the provider's message reached the agent")
		}
		if hits.Load() != before {
			t.Fatal("a failed dispatch reached the route")
		}
		if !strings.Contains(logs.String(), "42P01") || !strings.Contains(logs.String(), "guard dispatch failed") {
			t.Fatalf("the provider's message is not in the log: %s", logs.String())
		}
	})

	t.Run("a large integer keeps its digits in the path and the query", func(t *testing.T) {
		res := cs.CallTool
		raw := json.RawMessage(`{"id":9007199254740993,"n":9007199254740993}`)
		out, err := res(ctx, &mcp.CallToolParams{Name: "hard_num", Arguments: raw})
		if err != nil {
			t.Fatal(err)
		}
		got := structured(t, out)
		if out.IsError || got["id"] != "9007199254740993" || got["n"] != "9007199254740993" {
			t.Fatalf("large integer = %+v", got)
		}
	})

	t.Run("the inner request has the outer Host and forwarding headers, never credentials", func(t *testing.T) {
		got := structured(t, call("hard_num", map[string]any{"id": 5}))
		if got["host"] != outer.Host || got["proto"] != "https" || got["fhost"] != "site.example" {
			t.Fatalf("inner request host %v proto %v forwarded host %v, want %s https site.example", got["host"], got["proto"], got["fhost"], outer.Host)
		}
		if got["auth"] != "" || got["cookie"] != "" {
			t.Fatalf("outer credentials reached the route: auth %q cookie %q", got["auth"], got["cookie"])
		}
	})

	t.Run("a dispatch contradicting the tool's guard is refused", func(t *testing.T) {
		for _, name := range []string{"hard_claims_backend", "hard_backend_principal"} {
			refused(t, name, nil, "pier: the tool could not be run for this caller")
		}
		if !strings.Contains(logs.String(), "contradicts the tool's guard") {
			t.Fatalf("no contradiction line: %s", logs.String())
		}
	})
}