unibus/pkg/client/client_test.go

package client_test

import (
	"net"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"sync"
	"testing"
	"time"

	cs "fn-registry/functions/cybersecurity"

	"github.com/enmanuel/unibus/pkg/blobstore"
	"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/enmanuel/unibus/pkg/room"
	server "github.com/nats-io/nats-server/v2/server"
)

// testHarness boots an embedded NATS server and an in-process membershipd HTTP
// server, returning their URLs and a cleanup func.
type testHarness struct {
	natsURL string
	ctrlURL string
	ns      *server.Server
	httpts  *httptest.Server
}

func freePort(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 newHarness(t *testing.T) *testHarness {
	t.Helper()
	dir := t.TempDir()

	ns, err := embeddednats.Start(filepath.Join(dir, "js"), freePort(t))
	if err != nil {
		t.Fatalf("embedded nats: %v", err)
	}

	store, err := membership.Open(filepath.Join(dir, "unibus.db"))
	if err != nil {
		ns.Shutdown()
		t.Fatalf("membership store: %v", err)
	}
	blobs, err := blobstore.New(filepath.Join(dir, "blobs"))
	if err != nil {
		ns.Shutdown()
		t.Fatalf("blob store: %v", err)
	}
	srv := membership.NewServer(store, blobs)
	httpts := httptest.NewServer(srv)

	h := &testHarness{natsURL: embeddednats.ClientURL(ns), ctrlURL: httpts.URL, ns: ns, httpts: httpts}
	t.Cleanup(func() {
		httpts.Close()
		store.Close()
		ns.Shutdown()
		ns.WaitForShutdown()
	})
	return h
}

func waitHealth(t *testing.T, ctrlURL string) {
	t.Helper()
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		resp, err := http.Get(ctrlURL + "/healthz")
		if err == nil && resp.StatusCode == 200 {
			resp.Body.Close()
			return
		}
		if resp != nil {
			resp.Body.Close()
		}
		time.Sleep(50 * time.Millisecond)
	}
	t.Fatalf("membershipd never became healthy")
}

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

// TestE2EEncryptedForwardSecrecy is the headline test: A creates an encrypted
// room, invites B, A publishes a message B decrypts, then A kicks B and
// publishes at the new epoch — B must NOT be able to decrypt the new message.
func TestE2EEncryptedForwardSecrecy(t *testing.T) {
	h := newHarness(t)
	waitHealth(t, h.ctrlURL)

	a, err := client.New(h.natsURL, h.ctrlURL, mustIdentity(t))
	if err != nil {
		t.Fatalf("connect A: %v", err)
	}
	defer a.Close()
	b, err := client.New(h.natsURL, h.ctrlURL, mustIdentity(t))
	if err != nil {
		t.Fatalf("connect B: %v", err)
	}
	defer b.Close()

	roomID, err := a.CreateRoom("room.test", room.ModeMatrix)
	if err != nil {
		t.Fatalf("A create room: %v", err)
	}
	if err := a.Invite(roomID, b.Endpoint()); err != nil {
		t.Fatalf("A invite B: %v", err)
	}
	if err := b.Join(roomID); err != nil {
		t.Fatalf("B join: %v", err)
	}

	var mu sync.Mutex
	var received []string
	sub, err := b.Subscribe(roomID, func(f frame.Frame, plaintext []byte) {
		mu.Lock()
		received = append(received, string(plaintext))
		mu.Unlock()
	})
	if err != nil {
		t.Fatalf("B subscribe: %v", err)
	}
	defer sub.Unsubscribe()
	time.Sleep(150 * time.Millisecond)

	const msg1 = "hola E2E"
	if err := a.Publish(roomID, []byte(msg1)); err != nil {
		t.Fatalf("A publish msg1: %v", err)
	}

	// Wait for B to receive and decrypt msg1.
	if !waitFor(&mu, &received, func(rs []string) bool {
		for _, r := range rs {
			if r == msg1 {
				return true
			}
		}
		return false
	}, 2*time.Second) {
		t.Fatalf("B did not decrypt pre-kick message %q; got %v", msg1, snapshot(&mu, &received))
	}

	// A kicks B (rotates K to a new epoch, re-sealed only for the owner).
	if err := a.Kick(roomID, b.Endpoint().ID); err != nil {
		t.Fatalf("A kick B: %v", err)
	}
	time.Sleep(150 * time.Millisecond)

	const msg2 = "secreto post-kick"
	if err := a.Publish(roomID, []byte(msg2)); err != nil {
		t.Fatalf("A publish msg2: %v", err)
	}

	// Give B a chance to (fail to) decrypt; assert it never sees msg2.
	time.Sleep(1 * time.Second)
	for _, r := range snapshot(&mu, &received) {
		if r == msg2 {
			t.Fatalf("forward secrecy broken: B decrypted post-kick message %q", msg2)
		}
	}

	// Sanity: A itself can still decrypt at the new epoch (self-loopback via a fresh subscriber).
	aSub := subscribeCollect(t, a, roomID)
	defer aSub.sub.Unsubscribe()
	time.Sleep(100 * time.Millisecond)
	const msg3 = "owner-still-works"
	if err := a.Publish(roomID, []byte(msg3)); err != nil {
		t.Fatalf("A publish msg3: %v", err)
	}
	if !waitFor(&aSub.mu, &aSub.msgs, func(rs []string) bool {
		for _, r := range rs {
			if r == msg3 {
				return true
			}
		}
		return false
	}, 2*time.Second) {
		t.Fatalf("owner could not decrypt own message at new epoch; got %v", snapshot(&aSub.mu, &aSub.msgs))
	}
}

// TestCleartextWorkerToChat validates the ModeNATS path: a publisher and a
// subscriber sharing a subject, no encryption, messages flow through verbatim.
func TestCleartextWorkerToChat(t *testing.T) {
	h := newHarness(t)
	waitHealth(t, h.ctrlURL)

	pub, err := client.New(h.natsURL, h.ctrlURL, mustIdentity(t))
	if err != nil {
		t.Fatalf("connect pub: %v", err)
	}
	defer pub.Close()
	subC, err := client.New(h.natsURL, h.ctrlURL, mustIdentity(t))
	if err != nil {
		t.Fatalf("connect sub: %v", err)
	}
	defer subC.Close()

	const subject = "proc.test.ticks"
	// Each peer owns a room mapped to the shared subject; NATS fans out by subject.
	pubRoom, err := pub.CreateRoom(subject, room.ModeNATS)
	if err != nil {
		t.Fatalf("pub create room: %v", err)
	}
	subRoom, err := subC.CreateRoom(subject, room.ModeNATS)
	if err != nil {
		t.Fatalf("sub create room: %v", err)
	}

	collector := subscribeCollect(t, subC, subRoom)
	defer collector.sub.Unsubscribe()
	time.Sleep(150 * time.Millisecond)

	const msg = "tick 1"
	if err := pub.Publish(pubRoom, []byte(msg)); err != nil {
		t.Fatalf("publish: %v", err)
	}
	if !waitFor(&collector.mu, &collector.msgs, func(rs []string) bool {
		for _, r := range rs {
			if r == msg {
				return true
			}
		}
		return false
	}, 2*time.Second) {
		t.Fatalf("subscriber did not receive cleartext message; got %v", snapshot(&collector.mu, &collector.msgs))
	}
}

// TestMediaBlobRoundTrip validates encrypted media via the object store.
func TestMediaBlobRoundTrip(t *testing.T) {
	h := newHarness(t)
	waitHealth(t, h.ctrlURL)

	a, err := client.New(h.natsURL, h.ctrlURL, mustIdentity(t))
	if err != nil {
		t.Fatalf("connect A: %v", err)
	}
	defer a.Close()
	b, err := client.New(h.natsURL, h.ctrlURL, mustIdentity(t))
	if err != nil {
		t.Fatalf("connect B: %v", err)
	}
	defer b.Close()

	roomID, err := a.CreateRoom("room.media", room.ModeMatrix)
	if err != nil {
		t.Fatalf("create room: %v", err)
	}
	if err := a.Invite(roomID, b.Endpoint()); err != nil {
		t.Fatalf("invite: %v", err)
	}
	if err := b.Join(roomID); err != nil {
		t.Fatalf("join: %v", err)
	}

	gotBlob := make(chan []byte, 1)
	sub, err := b.Subscribe(roomID, func(f frame.Frame, _ []byte) {
		if f.Blob == nil {
			return
		}
		data, err := b.FetchMedia(roomID, f)
		if err != nil {
			return
		}
		gotBlob <- data
	})
	if err != nil {
		t.Fatalf("subscribe: %v", err)
	}
	defer sub.Unsubscribe()
	time.Sleep(150 * time.Millisecond)

	payload := []byte("a fake image payload that should be encrypted in the store")
	if err := a.PublishMedia(roomID, payload); err != nil {
		t.Fatalf("publish media: %v", err)
	}

	select {
	case got := <-gotBlob:
		if string(got) != string(payload) {
			t.Fatalf("media mismatch: got %q want %q", got, payload)
		}
	case <-time.After(2 * time.Second):
		t.Fatalf("B never received/decrypted the media blob")
	}
}

// ---- test helpers ---------------------------------------------------------

type collector struct {
	mu   sync.Mutex
	msgs []string
	sub  interface{ Unsubscribe() error }
}

func subscribeCollect(t *testing.T, c *client.Client, roomID string) *collector {
	t.Helper()
	col := &collector{}
	sub, err := c.Subscribe(roomID, func(_ frame.Frame, plaintext []byte) {
		col.mu.Lock()
		col.msgs = append(col.msgs, string(plaintext))
		col.mu.Unlock()
	})
	if err != nil {
		t.Fatalf("subscribe: %v", err)
	}
	col.sub = sub
	return col
}

func waitFor(mu *sync.Mutex, slice *[]string, pred func([]string) bool, timeout time.Duration) bool {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		mu.Lock()
		cp := append([]string(nil), (*slice)...)
		mu.Unlock()
		if pred(cp) {
			return true
		}
		time.Sleep(25 * time.Millisecond)
	}
	return false
}

func snapshot(mu *sync.Mutex, slice *[]string) []string {
	mu.Lock()
	defer mu.Unlock()
	return append([]string(nil), (*slice)...)
}