beacon-server/internal/ws/handler_test.go

// Copyright 2026 Beacon Contributors
// SPDX-License-Identifier: AGPL-3.0-or-later

package ws

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/MeshCore-Beacon/beacon-server/internal/hub"
	"github.com/coder/websocket"
)

func TestAllowedOrigins(t *testing.T) {
	for _, tc := range []struct {
		name    string
		allowed []string
		origin  string
		wantOK  bool
	}{
		{"no origin header", nil, "", true},
		{"foreign origin by default", nil, "https://example.com", false},
		{"listed origin", []string{"https://example.com"}, "https://example.com", true},
		{"listed origin, other case", []string{"https://example.com"}, "https://Example.com", true},
		{"scheme mismatch", []string{"https://example.com"}, "http://example.com", false},
		{"port mismatch", []string{"https://example.com"}, "https://example.com:8443", false},
		{"unlisted origin", []string{"https://example.com"}, "https://other.example", false},
	} {
		t.Run(tc.name, func(t *testing.T) {
			server := httptest.NewServer(Handler(hub.New(), nil, 5, 100, tc.allowed))
			t.Cleanup(server.Close)
			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			header := http.Header{}
			if tc.origin != "" {
				header.Set("Origin", tc.origin)
			}
			conn, resp, err := websocket.Dial(ctx, "ws"+strings.TrimPrefix(server.URL, "http"), &websocket.DialOptions{HTTPHeader: header})
			if tc.wantOK {
				if err != nil {
					t.Fatalf("dial: %v", err)
				}
				conn.CloseNow()
				return
			}
			if err == nil {
				conn.CloseNow()
				t.Fatal("expected handshake to be refused")
			}
			if resp == nil || resp.StatusCode != http.StatusForbidden {
				t.Fatalf("expected 403, got %v (%v)", resp, err)
			}
		})
	}
}

func TestConfigureIncludeObserverKey(t *testing.T) {
	h := hub.New()
	go h.Run()
	server := httptest.NewServer(Handler(h, nil, 5, 100, nil))
	t.Cleanup(server.Close)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	read := func(conn *websocket.Conn) map[string]json.RawMessage {
		t.Helper()
		_, data, err := conn.Read(ctx)
		if err != nil {
			t.Fatal(err)
		}
		var msg map[string]json.RawMessage
		if err := json.Unmarshal(data, &msg); err != nil {
			t.Fatal(err)
		}
		return msg
	}
	send := func(conn *websocket.Conn, msg string) {
		t.Helper()
		if err := conn.Write(ctx, websocket.MessageText, []byte(msg)); err != nil {
			t.Fatal(err)
		}
	}
	connect := func(configure string) *websocket.Conn {
		t.Helper()
		conn, _, err := websocket.Dial(ctx, "ws"+strings.TrimPrefix(server.URL, "http"), nil)
		if err != nil {
			t.Fatal(err)
		}
		t.Cleanup(func() { conn.CloseNow() })
		read(conn) // hello
		if configure != "" {
			send(conn, configure)
			reply := read(conn)
			if string(reply["type"]) != `"configured"` || string(reply["includeObserverKey"]) != "true" || string(reply["resolvePath"]) != "false" {
				t.Fatalf("unexpected configured reply: %v", reply)
			}
		}
		send(conn, `{"v":1,"type":"subscribe","id":"s","scope":{}}`)
		read(conn) // subscribed
		return conn
	}

	plain := connect("")
	keyed := connect(`{"v":1,"type":"configure","id":"c","includeObserverKey":true}`)
	time.Sleep(20 * time.Millisecond) // let the hub apply both subscriptions
	h.Broadcast(hub.Event{
		Type:           hub.EventPacketObservation,
		Payload:        json.RawMessage(`{"variant":"base"}`),
		PayloadWithKey: json.RawMessage(`{"variant":"key"}`),
	})
	if got := string(read(plain)["data"]); got != `{"variant":"base"}` {
		t.Errorf("plain client got %s", got)
	}
	if got := string(read(keyed)["data"]); got != `{"variant":"key"}` {
		t.Errorf("opted-in client got %s", got)
	}
}