mesh-controller/internal/broker/state_test.go

package broker

import (
	"slices"
	"strings"
	"testing"

	"github.com/nats-io/nats.go"
)

// The grants measured against a running server (novox/hq research 024): an owner reads and writes
// its bucket, a reader only reads, and neither reaches any other bucket.
func TestAnOwnerWritesItsStateAndAReaderOnlyReads(t *testing.T) {
	owner, err := PermissionsFor(Principal{Kind: KindModule, Node: "one", Module: "claude-code",
		State: []string{"servers"}, PasswordHash: "x"})
	if err != nil {
		t.Fatal(err)
	}
	for _, s := range []string{
		"$KV.claude-code_servers.>",
		"$JS.API.STREAM.INFO.KV_claude-code_servers",
		"$JS.API.DIRECT.GET.KV_claude-code_servers.>",
		"$JS.API.CONSUMER.CREATE.KV_claude-code_servers.>",
		"$JS.API.CONSUMER.DELETE.KV_claude-code_servers.>",
		"$JS.FC.KV_claude-code_servers.>",
	} {
		has(t, owner.Publish, s)
	}
	hasNot(t, owner.Publish, "$KV.>")
	hasNot(t, owner.Publish, "$JS.API.>")

	reader, err := PermissionsFor(Principal{Kind: KindModule, Node: "one", Module: "console",
		Reads: []string{"claude-code.servers"}, PasswordHash: "x"})
	if err != nil {
		t.Fatal(err)
	}
	has(t, reader.Publish, "$JS.API.DIRECT.GET.KV_claude-code_servers.>")
	has(t, reader.Publish, "$JS.API.CONSUMER.CREATE.KV_claude-code_servers.>")
	hasNot(t, reader.Publish, "$KV.claude-code_servers.>")
	for _, s := range reader.Subscribe {
		if s == "$KV.claude-code_servers.>" {
			t.Fatalf("a reader subscribes the bucket's subjects directly: %v", reader.Subscribe)
		}
	}
}

// One runtime carries every module on its machine, so its grant is the union: the owner's write
// where an owner is carried, a read where only a reader is.
func TestTheRuntimeKeepsAndReadsStateForItsModules(t *testing.T) {
	perms, err := PermissionsFor(Principal{Kind: KindNodeTools, Node: "one", Module: RuntimeModule,
		Carries: []Declared{
			{Module: "claude-code", State: []Bucket{{Module: "claude-code", Name: "servers"}},
				Reads: []string{"licence-manager.bindings"}},
			{Module: "audit"},
		}, PasswordHash: "x"})
	if err != nil {
		t.Fatal(err)
	}
	has(t, perms.Publish, "$KV.claude-code_servers.>")
	has(t, perms.Publish, "$JS.API.DIRECT.GET.KV_licence-manager_bindings.>")
	hasNot(t, perms.Publish, "$KV.licence-manager_bindings.>")
}

// A module with no state is granted nothing of any bucket — the composition of every module that
// existed before this is unchanged.
func TestAModuleWithNoStateReachesNoBucket(t *testing.T) {
	perms, err := PermissionsFor(Principal{Kind: KindModule, Node: "one", Module: "billing",
		Emits: []string{"order.placed"}, PasswordHash: "x"})
	if err != nil {
		t.Fatal(err)
	}
	for _, s := range perms.Publish {
		if strings.HasPrefix(s, "$KV.") || strings.HasPrefix(s, "$JS.FC.") || strings.Contains(s, ".KV_") {
			t.Fatalf("granted %q without declaring state", s)
		}
	}
}

// A read that names no bucket grants nothing rather than something that happens to parse.
func TestAReadThatNamesNoBucketGrantsNothing(t *testing.T) {
	if got := stateGrants(stateAccess{Module: "a", Node: "one", Reads: []string{"nodot", "x.", ".y", "a.b>"},
		KeyedReads: []KeyedRead{{Read: "nodot", Key: "one"}, {Read: "a.b", Key: "x>"}}}); len(got) != 0 {
		t.Fatalf("granted %v for reads that name no bucket", got)
	}
}

// The membership lists every bucket the module's code may reach, by the name the module uses for
// it, and whether it may write it — the list the runtime refuses from.
func TestAMembershipListsTheStateItsModuleMayReach(t *testing.T) {
	m := MembershipFor("one", Declared{Module: "claude-code",
		State: []Bucket{{Module: "claude-code", Name: "servers"}},
		Reads: []string{"licence-manager.bindings"}}, Placements{})
	want := []StateIssued{
		{Name: "servers", Bucket: "claude-code_servers", Writes: true},
		{Name: "licence-manager.bindings", Bucket: "licence-manager_bindings"},
	}
	if !slices.Equal(m.State, want) {
		t.Fatalf("issued %+v, want %+v", m.State, want)
	}
	if none := MembershipFor("one", Declared{Module: "audit"}, Placements{}); none.State != nil {
		t.Fatalf("a module with no state was issued %+v", none.State)
	}
}

type buckets struct {
	ensured []string
	on      []string
}

func (b *buckets) EnsureBucket(x Bucket) error {
	b.ensured = append(b.ensured, x.Bucket())
	return nil
}
func (b *buckets) BucketNames() ([]string, error) { return b.on, nil }

// Every declared bucket is asserted; one on the server that nothing declares is said, not removed.
func TestRaisingStateReportsWhatNothingDeclares(t *testing.T) {
	b := &buckets{on: []string{"claude-code_servers", "gone_old", "ours_by_hand"}}
	undeclared, err := RaiseBuckets(b, []Bucket{{Module: "claude-code", Name: "servers"}, {Module: "a", Name: "b"}})
	if err != nil {
		t.Fatal(err)
	}
	if !slices.Equal(b.ensured, []string{"a_b", "claude-code_servers"}) {
		t.Fatalf("asserted %v", b.ensured)
	}
	if !slices.Equal(undeclared, []string{"gone_old", "ours_by_hand"}) {
		t.Fatalf("reported %v", undeclared)
	}
}

// Against a real server: a bucket is created with the owner's options and the mesh's caps,
// asserting it again changes nothing and keeps what it holds, and a changed option is brought to
// match in place.
func TestABucketIsAssertedInPlace(t *testing.T) {
	js := aLiveBus(t)
	b := Bucket{Module: "statetest", Name: "servers"}
	if _, err := RaiseBuckets(js, []Bucket{b}); err != nil {
		t.Fatalf("a real server refused a module's bucket: %v", err)
	}
	kv, err := js.Context().KeyValue(b.Bucket())
	if err != nil {
		t.Fatal(err)
	}
	if _, err := kv.Put("all.one", []byte(`{"kept":true}`)); err != nil {
		t.Fatal(err)
	}
	b.History = 3
	if _, err := RaiseBuckets(js, []Bucket{b}); err != nil {
		t.Fatalf("asserting the bucket again failed, so a restart would: %v", err)
	}
	got, err := kv.Get("all.one")
	if err != nil || string(got.Value()) != `{"kept":true}` {
		t.Fatalf("asserting again lost what the bucket held: %v %v", got, err)
	}
	status, err := kv.Status()
	if err != nil {
		t.Fatal(err)
	}
	if status.History() != 3 {
		t.Fatalf("history is %d, the owner declared 3", status.History())
	}
	if s, ok := status.(*nats.KeyValueBucketStatus); ok {
		if c := s.StreamInfo().Config; c.MaxMsgSize != StateMaxValueBytes || c.MaxBytes != StateMaxBytes {
			t.Fatalf("the mesh's caps are not on the bucket: value %d, bucket %d", c.MaxMsgSize, c.MaxBytes)
		}
	}
}

// Against a real server: the handle over a connection the control plane already holds asserts a
// bucket as Dial's does — what a push uses, so a module registered since the last start has its
// bucket before its membership names it.
func TestABucketIsAssertedOverAHeldConnection(t *testing.T) {
	js := aLiveBus(t)
	held := OnConn(js.Conn())
	if _, err := RaiseBuckets(held, []Bucket{{Module: "statetest", Name: "held"}}); err != nil {
		t.Fatalf("asserting over a held connection failed: %v", err)
	}
	if _, err := js.Context().KeyValue("statetest_held"); err != nil {
		t.Fatalf("the bucket is not there: %v", err)
	}
}

// **State keyed by machine reaches its machine's key alone** (novox/hq ADR 0260): the owner's instance on
// a machine writes and reads that key, and a bar granted a key reads that key — a direct get of its
// subject and a consumer filtered to it — and no other key of the bucket.
func TestPerMachineStateAndAKeyedReadReachOneKey(t *testing.T) {
	owner, err := PermissionsFor(Principal{Kind: KindModule, Node: "laptop", Module: "power", PasswordHash: "x",
		State: []string{"draw"}, PerMachine: []string{"draw"}})
	if err != nil {
		t.Fatal(err)
	}
	has(t, owner.Publish, "$KV.power_draw.laptop")
	has(t, owner.Publish, "$JS.API.DIRECT.GET.KV_power_draw.$KV.power_draw.laptop")
	has(t, owner.Publish, "$JS.API.CONSUMER.CREATE.KV_power_draw.*.$KV.power_draw.laptop")
	hasNot(t, owner.Publish, "$KV.power_draw.>")
	hasNot(t, owner.Publish, "$JS.API.DIRECT.GET.KV_power_draw.>")
	hasNot(t, owner.Publish, "$JS.API.CONSUMER.CREATE.KV_power_draw.>")

	bar, err := PermissionsFor(Principal{Kind: KindModule, Node: "laptop", Module: "bar", PasswordHash: "x",
		KeyedReads: []KeyedRead{{Read: "power.draw", Key: "laptop"}}})
	if err != nil {
		t.Fatal(err)
	}
	has(t, bar.Publish, "$JS.API.STREAM.INFO.KV_power_draw")
	has(t, bar.Publish, "$JS.API.DIRECT.GET.KV_power_draw.$KV.power_draw.laptop")
	has(t, bar.Publish, "$JS.API.CONSUMER.CREATE.KV_power_draw.*.$KV.power_draw.laptop")
	for _, s := range bar.Publish {
		if strings.HasPrefix(s, "$KV.") || strings.HasSuffix(s, "KV_power_draw.>") && !strings.Contains(s, "DELETE") && !strings.HasPrefix(s, "$JS.FC.") {
			t.Fatalf("a keyed read was granted %q", s)
		}
	}

	// And the machine's runtime, which carries both, is granted the union: still the one key.
	runtime, err := PermissionsFor(Principal{Kind: KindNodeTools, Node: "laptop", Module: RuntimeModule, PasswordHash: "x",
		Carries: []Declared{
			{Module: "power", State: []Bucket{{Module: "power", Name: "draw", PerMachine: true}}},
			{Module: "bar", KeyedReads: []KeyedRead{{Read: "power.draw", Key: "laptop"}}},
		}})
	if err != nil {
		t.Fatal(err)
	}
	has(t, runtime.Publish, "$KV.power_draw.laptop")
	hasNot(t, runtime.Publish, "$KV.power_draw.>")
	hasNot(t, runtime.Publish, "$JS.API.DIRECT.GET.KV_power_draw.>")
}

func TestAMembershipNamesTheOneKeyOfKeyedState(t *testing.T) {
	m := MembershipFor("laptop", Declared{Module: "bar", KeyedReads: []KeyedRead{{Read: "power.draw", Key: "laptop"}}}, Placements{})
	if len(m.State) != 1 || m.State[0] != (StateIssued{Name: "power.draw", Bucket: "power_draw", Key: "laptop"}) {
		t.Fatalf("%+v", m.State)
	}
	o := MembershipFor("laptop", Declared{Module: "power", State: []Bucket{{Module: "power", Name: "draw", PerMachine: true}}}, Placements{})
	if len(o.State) != 1 || o.State[0] != (StateIssued{Name: "draw", Bucket: "power_draw", Writes: true, Key: "laptop"}) {
		t.Fatalf("%+v", o.State)
	}
}