unibus/pkg/membership/acl_test.go

package membership_test

import (
	"encoding/hex"
	"net"
	"net/http/httptest"
	"path/filepath"
	"testing"
	"time"

	cs "fn-registry/functions/cybersecurity"

	"github.com/enmanuel/unibus/pkg/blobstore"
	"github.com/enmanuel/unibus/pkg/busauth"
	"github.com/enmanuel/unibus/pkg/client"
	"github.com/enmanuel/unibus/pkg/embeddednats"
	"github.com/enmanuel/unibus/pkg/frame"
	"github.com/enmanuel/unibus/pkg/membership"
	"github.com/nats-io/nats.go"
	server "github.com/nats-io/nats-server/v2/server"
)

func aclFreePort(t *testing.T) int {
	t.Helper()
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("free port: %v", err)
	}
	defer l.Close()
	return l.Addr().(*net.TCPAddr).Port
}

func mustID(t *testing.T) cs.Identity {
	t.Helper()
	id, err := cs.GenerateIdentity()
	if err != nil {
		t.Fatalf("identity: %v", err)
	}
	return id
}

// aclPermsFunc adapts membership.SubjectACLFor into the busauth.PermissionsFunc
// the ACL authenticator expects (same Allow set for publish and subscribe).
func aclPermsFunc(store membership.Store) busauth.PermissionsFunc {
	derive := membership.SubjectACLFor(store)
	return func(signPubHex string) (*server.Permissions, error) {
		subs, err := derive(signPubHex)
		if err != nil {
			return nil, err
		}
		sp := &server.SubjectPermission{Allow: subs}
		return &server.Permissions{Publish: sp, Subscribe: sp}, nil
	}
}

// startACLNats boots an embedded NATS whose authenticator confines each peer to
// the subjects of the rooms it belongs to (audit H4 residual).
func startACLNats(t *testing.T, store membership.Store) *server.Server {
	t.Helper()
	auth := busauth.NewNkeyAuthenticatorACL(store.IsAuthorized, aclPermsFunc(store))
	ns, err := embeddednats.StartServer(embeddednats.ServerConfig{
		StoreDir: t.TempDir(), Host: "127.0.0.1", Port: aclFreePort(t), Auth: auth,
	})
	if err != nil {
		t.Fatalf("acl nats: %v", err)
	}
	t.Cleanup(func() { ns.Shutdown(); ns.WaitForShutdown() })
	return ns
}

func nkeyConn(t *testing.T, natsURL string, id cs.Identity, errCh chan error) *nats.Conn {
	t.Helper()
	pub, sign, err := busauth.ClientNkey(id.SignPriv)
	if err != nil {
		t.Fatalf("nkey: %v", err)
	}
	nc, err := nats.Connect(natsURL,
		nats.Nkey(pub, sign),
		nats.ErrorHandler(func(_ *nats.Conn, _ *nats.Subscription, e error) {
			select {
			case errCh <- e:
			default:
			}
		}),
	)
	if err != nil {
		t.Fatalf("connect nkey: %v", err)
	}
	t.Cleanup(nc.Close)
	return nc
}

func mustAddUser(t *testing.T, store membership.Store, id cs.Identity, handle string) {
	t.Helper()
	if err := store.AddUser(hex.EncodeToString(id.SignPub), handle, membership.RoleMember); err != nil {
		t.Fatalf("add user %s: %v", handle, err)
	}
}

func mustCreateRoom(t *testing.T, store membership.Store, roomID, subject, ownerEP string, owner cs.Identity) {
	t.Helper()
	info := membership.RoomInfo{RoomID: roomID, Subject: subject, OwnerEndpoint: ownerEP}
	if err := store.CreateRoom(info, owner.SignPub, owner.KexPub, nil); err != nil {
		t.Fatalf("create room %s: %v", roomID, err)
	}
}

func newCtrl(t *testing.T, store membership.Store, blobs blobstore.Store) string {
	t.Helper()
	ts := httptest.NewServer(membership.NewServer(store, blobs, membership.AuthOff))
	t.Cleanup(ts.Close)
	return ts.URL
}

func waitErr(ch chan error, d time.Duration) error {
	select {
	case e := <-ch:
		return e
	case <-time.After(d):
		return nil
	}
}

func drain(ch chan error) {
	for {
		select {
		case <-ch:
		default:
			return
		}
	}
}

// TestSubjectACLIsolation closes the audit H4 residual: a registered peer is
// confined to the subjects of the rooms it belongs to. alice (member of room.A)
// may sub/pub room.A but is DENIED sub/pub on room.B, and never reads what bob
// (member of room.B) publishes there.
func TestSubjectACLIsolation(t *testing.T) {
	dir := t.TempDir()
	store, err := membership.Open(filepath.Join(dir, "unibus.db"))
	if err != nil {
		t.Fatalf("store: %v", err)
	}
	t.Cleanup(func() { store.Close() })

	alice, bob := mustID(t), mustID(t)
	aliceEP, bobEP := frame.EndpointID(alice.SignPub), frame.EndpointID(bob.SignPub)
	mustAddUser(t, store, alice, "alice")
	mustAddUser(t, store, bob, "bob")
	const subjA, subjB = "room.acl.a", "room.acl.b"
	mustCreateRoom(t, store, "ROOMA", subjA, aliceEP, alice)
	mustCreateRoom(t, store, "ROOMB", subjB, bobEP, bob)

	srv := startACLNats(t, store)
	url := srv.ClientURL()
	aliceErr := make(chan error, 4)
	bobErr := make(chan error, 4)
	aliceNC := nkeyConn(t, url, alice, aliceErr)
	bobNC := nkeyConn(t, url, bob, bobErr)

	// alice may subscribe to her own room (no error).
	aliceGot := make(chan string, 4)
	if _, err := aliceNC.Subscribe(subjA, func(m *nats.Msg) { aliceGot <- string(m.Data) }); err != nil {
		t.Fatalf("alice sub A: %v", err)
	}
	_ = aliceNC.Flush()
	if e := waitErr(aliceErr, 300*time.Millisecond); e != nil {
		t.Fatalf("alice sub to her OWN room raised an error: %v", e)
	}

	// alice subscribing to bob's room is a permissions violation.
	if _, err := aliceNC.Subscribe(subjB, func(m *nats.Msg) { aliceGot <- "LEAK:" + string(m.Data) }); err != nil {
		t.Fatalf("alice sub B (queue): %v", err)
	}
	_ = aliceNC.Flush()
	if e := waitErr(aliceErr, 1*time.Second); e == nil {
		t.Fatalf("alice subscribing to bob's room should raise a permissions violation")
	}

	// bob publishes in his room; alice (denied) must not receive it.
	bobGot := make(chan string, 4)
	if _, err := bobNC.Subscribe(subjB, func(m *nats.Msg) { bobGot <- string(m.Data) }); err != nil {
		t.Fatalf("bob sub B: %v", err)
	}
	_ = bobNC.Flush()
	if err := bobNC.Publish(subjB, []byte("internal-bob")); err != nil {
		t.Fatalf("bob pub B: %v", err)
	}
	_ = bobNC.Flush()
	select {
	case got := <-bobGot:
		if got != "internal-bob" {
			t.Fatalf("bob got %q", got)
		}
	case <-time.After(2 * time.Second):
		t.Fatalf("bob did not receive his own message")
	}
	select {
	case leak := <-aliceGot:
		t.Fatalf("alice received bob's room traffic despite the ACL: %q", leak)
	case <-time.After(500 * time.Millisecond):
		// good: alice never got it
	}

	// alice publishing into bob's room is denied; bob must not receive it.
	drain(aliceErr)
	if err := aliceNC.Publish(subjB, []byte("intruder")); err != nil {
		t.Fatalf("alice pub B (queue): %v", err)
	}
	_ = aliceNC.Flush()
	if e := waitErr(aliceErr, 1*time.Second); e == nil {
		t.Fatalf("alice publishing into bob's room should raise a permissions violation")
	}
	select {
	case got := <-bobGot:
		t.Fatalf("bob received alice's cross-room publish despite the ACL: %q", got)
	case <-time.After(500 * time.Millisecond):
		// good
	}
}

// TestRefreshSessionGainsNewRoom is the "permissions refreshed on join" path:
// alice is not in room B, so her connection has no permission for its subject;
// after she is added to room B and calls RefreshSession, the reconnect
// re-derives her permissions and she gains the room's subject.
func TestRefreshSessionGainsNewRoom(t *testing.T) {
	dir := t.TempDir()
	store, err := membership.Open(filepath.Join(dir, "unibus.db"))
	if err != nil {
		t.Fatalf("store: %v", err)
	}
	t.Cleanup(func() { store.Close() })

	alice, bob := mustID(t), mustID(t)
	aliceEP, bobEP := frame.EndpointID(alice.SignPub), frame.EndpointID(bob.SignPub)
	mustAddUser(t, store, alice, "alice")
	mustAddUser(t, store, bob, "bob")
	const subjB = "room.refresh.b"
	mustCreateRoom(t, store, "ROOMB", subjB, bobEP, bob)

	srv := startACLNats(t, store)
	blobs, _ := blobstore.New(filepath.Join(dir, "blobs"))
	ctrl := newCtrl(t, store, blobs)

	aliceC, err := client.NewWithOptions(srv.ClientURL(), ctrl, alice, client.Options{UseNkey: true})
	if err != nil {
		t.Fatalf("connect alice: %v", err)
	}
	defer aliceC.Close()

	// Add alice to room B (as if invited), then RefreshSession so the
	// authenticator re-derives her permissions on reconnect.
	if _, err := store.GetMember("ROOMB", aliceEP); err == nil {
		t.Fatalf("alice should not be a member yet")
	}
	if err := store.AddMember("ROOMB", membership.Member{Endpoint: aliceEP, Role: "member", SignPub: alice.SignPub, KexPub: alice.KexPub}, 1, nil); err != nil {
		t.Fatalf("add alice to room B: %v", err)
	}
	if err := aliceC.RefreshSession(); err != nil {
		t.Fatalf("refresh session: %v", err)
	}

	bobErr := make(chan error, 2)
	bobNC := nkeyConn(t, srv.ClientURL(), bob, bobErr)

	got := make(chan string, 2)
	sub, err := aliceC.Subscribe("ROOMB", func(_ frame.Frame, plaintext []byte) { got <- string(plaintext) })
	if err != nil {
		t.Fatalf("alice subscribe room B after refresh: %v", err)
	}
	defer sub.Unsubscribe()
	time.Sleep(200 * time.Millisecond)

	// bob publishes a minimal cleartext frame on subjB.
	f := frame.Frame{Type: frame.PUB, Subject: subjB, Sender: bobEP, MsgID: "m1", Payload: []byte("hello-after-join")}
	b, _ := f.Marshal()
	if err := bobNC.Publish(subjB, b); err != nil {
		t.Fatalf("bob publish: %v", err)
	}
	_ = bobNC.Flush()

	select {
	case msg := <-got:
		if msg != "hello-after-join" {
			t.Fatalf("alice got %q", msg)
		}
	case <-time.After(3 * time.Second):
		t.Fatalf("alice did not receive room B traffic after RefreshSession (permissions not refreshed)")
	}
}