msa-agent/internal/agent/agent_test.go

package agent

import (
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strconv"
	"testing"
	"time"

	"fun.chick/msa-agent/internal/policy"
	"fun.chick/msa-agent/internal/runtime"
	"fun.chick/msa-agent/internal/wire"
)

// These tests are the executable version of the ordering claim: they build real envelopes with real
// keys and check what the agent does with each one, rather than calling Handle with a hand-made
// struct that skips the parts under test.

type pair struct {
	appSign ed25519.PrivateKey
	appKx   []byte
}

func newPair(t testing.TB) pair {
	t.Helper()
	_, sign, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	kx := make([]byte, wire.KeyBytes)
	if _, err := rand.Read(kx); err != nil {
		t.Fatal(err)
	}
	return pair{appSign: sign, appKx: kx}
}

type harness struct {
	agent   *Agent
	app     pair
	appID   string
	agentKx []byte
	now     time.Time
	// The phone must keep the private half of the per-request ephemeral key until the response
	// arrives, because the response is encrypted to it. A test that forgot this would have to
	// weaken the protocol to pass, which is the failure this catches.
	sentEphemeral map[string][]byte
}

func newHarness(t testing.TB, handler http.Handler) *harness {
	t.Helper()
	server := httptest.NewServer(handler)
	t.Cleanup(server.Close)
	parsed, err := url.Parse(server.URL)
	if err != nil {
		t.Fatal(err)
	}
	port, err := strconv.Atoi(parsed.Port())
	if err != nil {
		t.Fatal(err)
	}
	pol, err := policy.Load()
	if err != nil {
		t.Fatal(err)
	}
	app := newPair(t)
	agentKx := make([]byte, wire.KeyBytes)
	if _, err := rand.Read(agentKx); err != nil {
		t.Fatal(err)
	}
	// The phone is given the public half derived from that private key, never the raw bytes. A test
	// that fed the raw private bytes in as a public key would derive a different shared secret on
	// each side and fail as an authentication error with no obvious cause.
	agentKxPublic, err := wire.PublicFrom(agentKx)
	if err != nil {
		t.Fatal(err)
	}
	_, agentSign, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	now := time.Unix(1_800_000_000, 0).UTC()
	exec, err := runtime.NewExecutor(pol, runtime.Config{
		Port:  port,
		Audit: runtime.NullAudit(),
		Clock: func() time.Time { return now },
	})
	if err != nil {
		t.Fatal(err)
	}
	agent, err := New("nd_test0001", Keys{Sign: agentSign, Kx: agentKx},
		StaticDevices{"ap_test0000": app.appSign.Public().(ed25519.PublicKey)}, pol, exec)
	if err != nil {
		t.Fatal(err)
	}
	agent.now = func() time.Time { return now }
	return &harness{
		agent: agent, app: app, appID: "ap_test0000", agentKx: agentKxPublic, now: now,
		sentEphemeral: map[string][]byte{},
	}
}

// seal builds a complete request envelope the way the phone does.
func (h *harness) seal(t testing.TB, reqID string, inner wire.InnerRequest) []byte {
	t.Helper()
	body, err := json.Marshal(inner)
	if err != nil {
		t.Fatal(err)
	}
	// The phone's side of the exchange, built from raw primitives rather than wire.NewEphemeral:
	// this direction needs the private half after the request is sent, which Ephemeral deliberately
	// does not hand out.
	phonePrivate := make([]byte, wire.KeyBytes)
	if _, err := rand.Read(phonePrivate); err != nil {
		t.Fatal(err)
	}
	phonePublic, err := wire.PublicFrom(phonePrivate)
	if err != nil {
		t.Fatal(err)
	}
	shared, err := wire.SharedSecret(phonePrivate, h.agentKx)
	if err != nil {
		t.Fatal(err)
	}
	key, iv, err := wire.DeriveKey(shared, "nd_test0001", reqID, false)
	if err != nil {
		t.Fatal(err)
	}
	sealed, err := wire.Seal(key, iv, body)
	if err != nil {
		t.Fatal(err)
	}
	envelope := wire.Envelope{
		V: wire.Version, Node: "nd_test0001", Req: reqID, App: h.appID,
		TS: uint64(h.now.Unix()), EK: phonePublic, Alg: wire.Suite, P: sealed,
	}
	if err := wire.SignRequest(&envelope, h.app.appSign); err != nil {
		t.Fatal(err)
	}
	raw, err := json.Marshal(envelope)
	if err != nil {
		t.Fatal(err)
	}
	h.sentEphemeral[reqID] = phonePrivate
	return raw
}

// open decrypts a response the way the phone does, and verifies it.
func (h *harness) open(t *testing.T, raw json.RawMessage, reqID string) wire.InnerResponse {
	t.Helper()
	var envelope wire.Envelope
	if err := json.Unmarshal(raw, &envelope); err != nil {
		t.Fatal(err)
	}
	if err := wire.VerifyResponse(envelope, h.agent.keys.Sign.Public().(ed25519.PublicKey)); err != nil {
		t.Fatalf("response signature: %v", err)
	}
	// The phone holds the private half of the ephemeral key it sent, so it can reproduce the
	// response key from the agent's ek.
	private, ok := h.sentEphemeral[reqID]
	if !ok {
		t.Fatalf("test never sent a request with id %s", reqID)
	}
	shared, err := wire.SharedSecret(private, envelope.EK)
	if err != nil {
		t.Fatal(err)
	}
	key, iv, err := wire.DeriveKey(shared, envelope.Node, reqID, true)
	if err != nil {
		t.Fatal(err)
	}
	opened, err := wire.Open(key, iv, envelope.P)
	if err != nil {
		t.Fatalf("open response: %v", err)
	}
	var inner wire.InnerResponse
	if err := json.Unmarshal(opened, &inner); err != nil {
		t.Fatal(err)
	}
	return inner
}

func okHandler(body string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(body))
	})
}

func TestARoundTripFromAPairedPhoneSucceeds(t *testing.T) {
	h := newHarness(t, okHandler(`{"status":"up"}`))
	raw, err := h.agent.Handle(context.Background(),
		h.seal(t, "11111111-1111-4111-8111-111111111111",
			wire.InnerRequest{Method: "GET", Path: "/api/health"}))
	if err != nil {
		t.Fatalf("Handle: %v", err)
	}
	inner := h.open(t, raw, "11111111-1111-4111-8111-111111111111")
	if inner.Status != 200 || string(inner.Body) != `{"status":"up"}` {
		t.Fatalf("got %d %q", inner.Status, inner.Body)
	}
	if inner.Outcome != "ok" {
		t.Fatalf("outcome is %q", inner.Outcome)
	}
}

// The relay cannot read why a command failed, so every refusal has to come back inside the
// ciphertext. If a refusal were returned as a transport error instead, this test would fail.
func TestRefusalsComeBackEncryptedNotAsTransportErrors(t *testing.T) {
	h := newHarness(t, okHandler(`{}`))
	_, refusal := h.agent.exec.Execute(context.Background(),
		wire.Envelope{TS: uint64(h.now.Unix())}, wire.InnerRequest{Method: "GET", Path: "/api/secrets"})
	if refusal == nil {
		t.Fatal("executor allowed a path outside the policy")
	}
	raw, err := h.agent.Handle(context.Background(),
		h.seal(t, "22222222-2222-4222-8222-222222222222",
			wire.InnerRequest{Method: "GET", Path: "/api/secrets"}))
	if err != nil {
		t.Fatalf("a policy refusal must not become a transport error: %v", err)
	}
	inner := h.open(t, raw, "22222222-2222-4222-8222-222222222222")
	if inner.Outcome != OutcomePolicy {
		t.Fatalf("outcome is %q, want %q", inner.Outcome, OutcomePolicy)
	}
}

func TestAnUnpairedDeviceIsRefusedBeforeThePayloadIsRead(t *testing.T) {
	h := newHarness(t, okHandler(`{}`))
	stranger := newPair(t)
	h.app.appSign = stranger.appSign
	raw, err := h.agent.Handle(context.Background(),
		h.seal(t, "33333333-3333-4333-8333-333333333333",
			wire.InnerRequest{Method: "GET", Path: "/api/health"}))
	if err != nil {
		t.Fatal(err)
	}
	// The signature check runs before the device lookup's counterpart, so an unpaired phone gets a
	// signature failure rather than a "you are not paired" oracle.
	inner := h.open(t, raw, "33333333-3333-4333-8333-333333333333")
	if inner.Outcome != OutcomeBadSig {
		t.Fatalf("outcome is %q, want %q", inner.Outcome, OutcomeBadSig)
	}
}

func TestAnEnvelopeForAnotherNodeIsNotAnswered(t *testing.T) {
	h := newHarness(t, okHandler(`{}`))
	raw := h.seal(t, "44444444-4444-4444-8444-444444444444",
		wire.InnerRequest{Method: "GET", Path: "/api/health"})
	var envelope map[string]any
	if err := json.Unmarshal(raw, &envelope); err != nil {
		t.Fatal(err)
	}
	envelope["node"] = "nd_other9999"
	// Re-marshalling breaks the signature, which is the point: the node check must fire first so the
	// agent does not confirm to a caller that it holds this node's keys.
	retargeted, err := json.Marshal(envelope)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := h.agent.Handle(context.Background(), retargeted); err == nil {
		t.Fatal("an envelope for another node was answered")
	}
}

func TestADangerousWriteStillNeedsItsIntentEndToEnd(t *testing.T) {
	h := newHarness(t, okHandler(`{"sent":true}`))
	body := []byte(`{"phone_number":"+8613800000000","content":"hello"}`)
	noIntent := h.seal(t, "55555555-5555-4555-8555-555555555555",
		wire.InnerRequest{Method: "POST", Path: "/api/sms/send", Body: body})
	raw, err := h.agent.Handle(context.Background(), noIntent)
	if err != nil {
		t.Fatal(err)
	}
	if inner := h.open(t, raw, "55555555-5555-4555-8555-555555555555"); inner.Outcome != OutcomeNoConfirm {
		t.Fatalf("outcome is %q, want %q", inner.Outcome, OutcomeNoConfirm)
	}

	intent, err := wire.IntentDigest("POST", "/api/sms/send", body)
	if err != nil {
		t.Fatal(err)
	}
	withIntent := h.seal(t, "66666666-6666-4666-8666-666666666666",
		wire.InnerRequest{Method: "POST", Path: "/api/sms/send", Body: body, Intent: intent})
	raw, err = h.agent.Handle(context.Background(), withIntent)
	if err != nil {
		t.Fatal(err)
	}
	inner := h.open(t, raw, "66666666-6666-4666-8666-666666666666")
	if inner.Outcome != "ok" || inner.Status != 200 {
		t.Fatalf("got %q %d", inner.Outcome, inner.Status)
	}
}

func TestNewRefusesToStartDegraded(t *testing.T) {
	h := newHarness(t, okHandler(`{}`))
	pol, err := policy.Load()
	if err != nil {
		t.Fatal(err)
	}
	if _, err := New("", h.agent.keys, StaticDevices{"a": nil}, pol, h.agent.exec); err == nil {
		t.Fatal("an empty node id was accepted")
	}
	if _, err := New("nd_x", Keys{Sign: h.agent.keys.Sign}, StaticDevices{"a": nil}, pol, h.agent.exec); err == nil {
		t.Fatal("a missing x25519 key was accepted")
	}
	if _, err := New("nd_x", h.agent.keys, nil, pol, h.agent.exec); err == nil {
		t.Fatal("a nil device list was accepted")
	}
	if _, err := New("nd_x", h.agent.keys, StaticDevices{}, pol, h.agent.exec); err == nil {
		t.Fatal("an empty device list was accepted; nothing could ever be authorised")
	}
}

// Status mapping exists so the phone can distinguish "try again later" from "this will never work"
// without parsing a message string.
func TestStatusMappingSeparatesRetryableFromTerminal(t *testing.T) {
	cases := map[string]int{
		runtime.ClassRateLimited: 429,
		runtime.ClassExpired:     408,
		runtime.ClassReplay:      408,
		runtime.ClassTooLarge:    413,
		runtime.ClassUpstream:    502,
		runtime.ClassPolicy:      403,
	}
	for class, want := range cases {
		if got := statusFor(class); got != want {
			t.Errorf("statusFor(%q) = %d, want %d", class, got, want)
		}
	}
}