niko_trust/internal/server/ws_test.go

package server_test

import (
	"bytes"
	"context"
	"crypto/ed25519"
	"encoding/json"
	"io"
	"net/http"
	"strings"
	"testing"
	"time"

	"github.com/coder/websocket"

	"git.n1ko.dev/Niko/niko_trust/pkg/identity/signer"
	"git.n1ko.dev/Niko/niko_trust/pkg/protocol"
	"git.n1ko.dev/Niko/niko_trust/pkg/tce"
	"git.n1ko.dev/Niko/niko_trust/pkg/transport"
)

func wsURL(httpURL string) string {
	return "ws" + strings.TrimPrefix(httpURL, "http") + "/v1/ws"
}

func wsCtx(t *testing.T) context.Context {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	t.Cleanup(cancel)
	return ctx
}

func dialWS(t *testing.T, url, token string) *websocket.Conn {
	t.Helper()
	opts := &websocket.DialOptions{}
	if token != "" {
		opts.HTTPHeader = http.Header{"Authorization": []string{"Bearer " + token}}
	}
	conn, _, err := websocket.Dial(wsCtx(t), url, opts)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { conn.Close(websocket.StatusNormalClosure, "") })
	return conn
}

func wsWrite(t *testing.T, conn *websocket.Conn, v any) {
	t.Helper()
	raw, err := json.Marshal(v)
	if err != nil {
		t.Fatal(err)
	}
	if err := conn.Write(wsCtx(t), websocket.MessageText, raw); err != nil {
		t.Fatal(err)
	}
}

func wsRead(t *testing.T, conn *websocket.Conn) map[string]any {
	t.Helper()
	_, raw, err := conn.Read(wsCtx(t))
	if err != nil {
		t.Fatalf("read event: %v", err)
	}
	var ev map[string]any
	if json.Unmarshal(raw, &ev) != nil {
		t.Fatalf("bad event json: %s", raw)
	}
	return ev
}

func wsSubscribe(t *testing.T, conn *websocket.Conn, channel, key string) map[string]any {
	t.Helper()
	wsWrite(t, conn, map[string]string{"op": "subscribe", "channel": channel, "key": key})
	return wsRead(t, conn)
}

// wsAssert performs the challenge/assert handshake with an explicit scope.
func wsAssert(t *testing.T, baseURL string, key *signer.Signer, scope string) string {
	t.Helper()
	resp, err := http.Post(baseURL+"/v1/auth/challenge", "application/json", nil)
	if err != nil {
		t.Fatal(err)
	}
	var chRes struct {
		Challenge string `json:"challenge"`
	}
	json.NewDecoder(resp.Body).Decode(&chRes)
	resp.Body.Close()

	ch, err := hexDecode(chRes.Challenge)
	if err != nil {
		t.Fatal(err)
	}
	a := &protocol.AuthAssertion{
		PubKey:    key.Public(),
		Challenge: ch,
		Scope:     scope,
		Audience:  "trust.n1ko.dev",
		CreatedAt: uint64(time.Now().Unix()),
	}
	b, err := protocol.EncodeAuthAssertion(a)
	if err != nil {
		t.Fatal(err)
	}
	r2 := postEnvelope(t, baseURL+"/v1/auth/assert", &transport.Envelope{TCE: b, Signature: key.Sign(b)})
	defer r2.Body.Close()
	if r2.StatusCode != http.StatusOK {
		raw, _ := io.ReadAll(r2.Body)
		t.Fatalf("assert status %d: %s", r2.StatusCode, raw)
	}
	var out struct {
		SessionToken string `json:"session_token"`
	}
	json.NewDecoder(r2.Body).Decode(&out)
	if out.SessionToken == "" {
		t.Fatal("no session token")
	}
	return out.SessionToken
}

func wsPutClaim(t *testing.T, tsURL string, issuer *signer.Signer, subjectPub ed25519.PublicKey, nonce byte) {
	t.Helper()
	c := &protocol.Claim{
		Issuer:    issuer.Public(),
		Subject:   subjectPub,
		Claims:    map[string]tce.Value{"ws.test": tce.Bool(true)},
		CreatedAt: uint64(time.Now().Unix()),
		Serial:    1,
		Nonce:     bytes.Repeat([]byte{nonce}, tce.NonceSize),
	}
	b, err := protocol.EncodeClaim(c)
	if err != nil {
		t.Fatal(err)
	}
	resp := postEnvelope(t, tsURL+"/v1/objects", &transport.Envelope{TCE: b, Signature: issuer.Sign(b)})
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		raw, _ := io.ReadAll(resp.Body)
		t.Fatalf("put status %d: %s", resp.StatusCode, raw)
	}
}

func TestWSRequiresSession(t *testing.T) {
	ts := newTestServer(t)
	defer ts.Close()

	resp, err := http.Get(ts.URL + "/v1/ws")
	if err != nil {
		t.Fatal(err)
	}
	resp.Body.Close()
	if resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("status %d, want 401", resp.StatusCode)
	}
}

func TestWSRejectsUnknownChannel(t *testing.T) {
	ts := newTestServer(t)
	defer ts.Close()

	issuer, _ := signer.Generate()
	token := wsAssert(t, ts.URL, issuer, "read")
	conn := dialWS(t, wsURL(ts.URL), token)
	ev := wsSubscribe(t, conn, "gossip", "whatever")
	if !strings.Contains(ev["message"].(string), "unknown channel") {
		t.Fatalf("unexpected ack: %v", ev)
	}
}

func TestWSPushesMatchingClaimsOnly(t *testing.T) {
	ts := newTestServer(t)
	defer ts.Close()

	issuer, _ := signer.Generate()
	subject, _ := signer.Generate()
	token := wsAssert(t, ts.URL, issuer, "read")

	conn := dialWS(t, wsURL(ts.URL), token)
	ev := wsSubscribe(t, conn, "claims", subject.Identity().String())
	if ev["event"] != "subscribed" {
		t.Fatalf("no subscribed ack: %v", ev)
	}

	wsPutClaim(t, ts.URL, issuer, subject.Public(), 0x03)

	got := wsRead(t, conn)
	if got["event"] != "object" || got["channel"] != "claims" || got["key"] != subject.Identity().String() {
		t.Fatalf("unexpected event: %v", got)
	}
	env, ok := got["envelope"].(map[string]any)
	if !ok || env["tce"] == "" || env["signature"] == "" {
		t.Fatalf("event missing raw envelope: %v", got)
	}
	if id, _ := got["object_id"].(string); len(id) != 64 {
		t.Fatalf("event missing object id: %v", got)
	}

	// A claim about another subject must not arrive before a matching one.
	other, _ := signer.Generate()
	wsPutClaim(t, ts.URL, issuer, other.Public(), 0x04)
	wsPutClaim(t, ts.URL, issuer, subject.Public(), 0x05)

	got2 := wsRead(t, conn)
	if got2["key"] != subject.Identity().String() {
		t.Fatalf("non-matching claim leaked through: %v", got2)
	}
}

func TestWSScopeGatesChannels(t *testing.T) {
	ts := newTestServer(t)
	defer ts.Close()

	issuer, _ := signer.Generate()
	token := wsAssert(t, ts.URL, issuer, "read:claims") // no requests scope

	conn := dialWS(t, wsURL(ts.URL), token)
	ev := wsSubscribe(t, conn, "requests", issuer.Identity().String())
	msg, _ := ev["message"].(string)
	if !strings.Contains(msg, "insufficient scope") {
		t.Fatalf("scope not enforced: %v", ev)
	}

	// The granted channel still works.
	ev2 := wsSubscribe(t, conn, "claims", subjectAddrFor(t))
	if ev2["event"] != "subscribed" {
		t.Fatalf("granted channel refused: %v", ev2)
	}
}

func subjectAddrFor(t *testing.T) string {
	t.Helper()
	s, _ := signer.Generate()
	return s.Identity().String()
}

func hexDecode(s string) ([]byte, error) {
	if len(s)%2 != 0 {
		s = s[:len(s)-1]
	}
	out := make([]byte, len(s)/2)
	for i := 0; i < len(out); i++ {
		hi := hexNibble(s[2*i])
		lo := hexNibble(s[2*i+1])
		out[i] = hi<<4 | lo
	}
	return out, nil
}

func hexNibble(c byte) byte {
	switch {
	case c >= '0' && c <= '9':
		return c - '0'
	case c >= 'a' && c <= 'f':
		return c - 'a' + 10
	case c >= 'A' && c <= 'F':
		return c - 'A' + 10
	}
	return 0
}

func TestBatchFetchAndCursorPagination(t *testing.T) {
	ts := newTestServer(t)
	defer ts.Close()

	issuer, _ := signer.Generate()
	subject, _ := signer.Generate()
	var ids []string
	for i := byte(1); i <= 5; i++ {
		c := &protocol.Claim{
			Issuer:    issuer.Public(),
			Subject:   subject.Public(),
			Claims:    map[string]tce.Value{"batch.test": tce.Bool(true)},
			CreatedAt: uint64(time.Now().Unix()),
			Serial:    1,
			Nonce:     bytes.Repeat([]byte{i}, tce.NonceSize),
		}
		b, err := protocol.EncodeClaim(c)
		if err != nil {
			t.Fatal(err)
		}
		resp := postEnvelope(t, ts.URL+"/v1/objects", &transport.Envelope{TCE: b, Signature: issuer.Sign(b)})
		if resp.StatusCode != http.StatusOK {
			t.Fatalf("put %d status %d", i, resp.StatusCode)
		}
		var out struct {
			ObjectID string `json:"object_id"`
		}
		json.NewDecoder(resp.Body).Decode(&out)
		resp.Body.Close()
		ids = append(ids, out.ObjectID)
	}

	// Batch fetch returns exactly the requested objects.
	q := ids[0] + "," + ids[2] + ",ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"
	resp, err := http.Get(ts.URL + "/v1/objects?ids=" + q)
	if err != nil {
		t.Fatal(err)
	}
	var batch struct {
		Objects map[string]json.RawMessage `json:"objects"`
	}
	json.NewDecoder(resp.Body).Decode(&batch)
	resp.Body.Close()
	if len(batch.Objects) != 2 {
		t.Fatalf("batch returned %d objects, want 2", len(batch.Objects))
	}
	if _, ok := batch.Objects[ids[0]]; !ok {
		t.Fatal("batch missing first id")
	}
	if _, ok := batch.Objects[ids[2]]; !ok {
		t.Fatal("batch missing third id")
	}

	// Cursor pagination walks deterministically.
	token := wsAssert(t, ts.URL, issuer, "read")
	listPage := func(after string) []string {
		url := ts.URL + "/v1/claims?subject=" + subject.Identity().String() +
			"&limit=2&token=" + token
		if after != "" {
			url += "&after=" + after
		}
		resp, err := http.Get(url)
		if err != nil {
			t.Fatal(err)
		}
		defer resp.Body.Close()
		var out struct {
			Claims []struct {
				ObjectID string `json:"object_id"`
			} `json:"claims"`
		}
		json.NewDecoder(resp.Body).Decode(&out)
		var got []string
		for _, c := range out.Claims {
			got = append(got, c.ObjectID)
		}
		return got
	}

	page1 := listPage("")
	if len(page1) != 2 {
		t.Fatalf("page1 = %d entries", len(page1))
	}
	page2 := listPage(page1[len(page1)-1])
	if len(page2) == 0 || page2[0] <= page1[len(page1)-1] {
		t.Fatalf("cursor did not advance: page1 last %s, page2 %v", page1[len(page1)-1], page2)
	}
}