Implemented better testing and fixed skipepd frames

2026-07-01 16:39:34 +02:00
parent 7a326c5d78
commit 0bea41f866
46 changed files with 4358 additions and 1739 deletions
@@ -127,7 +127,12 @@ func (s SignalInfo) NumElements() int {
 	if c == 0 {
 		c = 1
 	}
-	return r * c
+	/* HI-2: cap at 1M to prevent integer overflow / OOM from crafted packets */
+	n := r * c
+	if n < 0 || n > 1024*1024 {
+		return 1024 * 1024
+	}
+	return n
 }

 // rawTypeSize returns the byte size for one element of the raw (unquantised) type.
@@ -227,6 +232,11 @@ func ParseConfig(payload []byte) ([]SignalInfo, uint8, error) {
 		return nil, 0, fmt.Errorf("config payload too short")
 	}
 	numSigs := binary.LittleEndian.Uint32(payload[0:4])
+	/* HI-2: validate numSigs against payload length before allocating */
+	maxSigs := uint32(len(payload) / SigDescSize)
+	if numSigs > maxSigs {
+		return nil, 0, fmt.Errorf("config claims %d signals but payload can hold at most %d", numSigs, maxSigs)
+	}
 	offset := 4
 	sigs := make([]SignalInfo, 0, numSigs)
 	for i := uint32(0); i < numSigs; i++ {
@@ -327,6 +337,10 @@ func ParseData(payload []byte, sigs []SignalInfo, publishMode uint8, arrivalTime
 		if numSamples == 0 {
 			return []DataSample{}, nil
 		}
+		/* HI-2: sanity-cap numSamples to prevent OOM from crafted packets */
+		if numSamples < 0 || numSamples > 1024*1024 {
+			return nil, fmt.Errorf("accumulate numSamples %d out of range", numSamples)
+		}

 		// Parse per-signal data blocks (all slots for a signal are contiguous).
 		accumVals := make(map[string][]float64, len(sigs)) // scalars: numSamples values
@@ -0,0 +1,97 @@
+package udpsprotocol
+
+import (
+	"encoding/binary"
+	"math"
+	"testing"
+	"time"
+)
+
+// TestParseConfig_HugeNumSigs_NoOOM — a CONFIG payload claiming 0xFFFFFFFF signals
+// must return an error, not panic/OOM.
+func TestParseConfig_HugeNumSigs_NoOOM(t *testing.T) {
+	// 4 bytes: numSigs = 0xFFFFFFFF, then nothing else
+	payload := make([]byte, 4)
+	binary.LittleEndian.PutUint32(payload[0:4], 0xFFFFFFFF)
+	sigs, _, err := ParseConfig(payload)
+	if err == nil {
+		t.Fatal("expected error for huge numSigs, got nil")
+	}
+	if sigs != nil {
+		t.Fatalf("expected nil sigs, got %d", len(sigs))
+	}
+}
+
+// TestParseConfig_ValidSmallConfig — a minimal valid CONFIG parses correctly.
+func TestParseConfig_ValidSmallConfig(t *testing.T) {
+	// 1 signal, then publish mode
+	payload := make([]byte, 4+SigDescSize+1)
+	binary.LittleEndian.PutUint32(payload[0:4], 1)
+	// Set typeCode to float32 (8) at offset 64
+	payload[4+64] = 8
+	// numRows=1, numCols=1 at offsets 67, 71
+	binary.LittleEndian.PutUint32(payload[4+67:4+71], 1)
+	binary.LittleEndian.PutUint32(payload[4+71:4+75], 1)
+	// publish mode = 0 (Strict)
+	payload[4+SigDescSize] = 0
+	sigs, pm, err := ParseConfig(payload)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if len(sigs) != 1 {
+		t.Fatalf("expected 1 signal, got %d", len(sigs))
+	}
+	if pm != PublishModeStrict {
+		t.Fatalf("expected Strict mode, got %d", pm)
+	}
+}
+
+// TestNumElements_OverflowCapped — huge numRows*numCols is capped, no panic.
+func TestNumElements_OverflowCapped(t *testing.T) {
+	s := SignalInfo{NumRows: 0xFFFFFFFF, NumCols: 0xFFFFFFFF}
+	n := s.NumElements()
+	if n <= 0 || n > 1024*1024 {
+		t.Fatalf("expected capped value 1M, got %d", n)
+	}
+}
+
+// TestNumElements_Normal — normal values work correctly.
+func TestNumElements_Normal(t *testing.T) {
+	s := SignalInfo{NumRows: 3, NumCols: 4}
+	if n := s.NumElements(); n != 12 {
+		t.Fatalf("expected 12, got %d", n)
+	}
+}
+
+// TestParseData_HugeNumSamples_NoOOM — an Accumulate DATA packet with
+// numSamples=0xFFFFFFFF must return an error, not OOM.
+func TestParseData_HugeNumSamples_NoOOM(t *testing.T) {
+	sigs := []SignalInfo{
+		{Name: "test", TypeCode: 8, NumRows: 1, NumCols: 1, QuantType: QuantNone},
+	}
+	payload := make([]byte, 12)
+	binary.LittleEndian.PutUint64(payload[0:8], 0) // HRT
+	binary.LittleEndian.PutUint32(payload[8:12], 0xFFFFFFFF)
+	_, err := ParseData(payload, sigs, PublishModeAccumulate, time.Now())
+	if err == nil {
+		t.Fatal("expected error for huge numSamples, got nil")
+	}
+}
+
+// TestParseData_ValidStrict — a valid Strict DATA packet parses without error.
+func TestParseData_ValidStrict(t *testing.T) {
+	sigs := []SignalInfo{
+		{Name: "test", TypeCode: 8, NumRows: 1, NumCols: 1, QuantType: QuantNone},
+	}
+	// 8 HRT + 4 bytes float32
+	payload := make([]byte, 12)
+	binary.LittleEndian.PutUint64(payload[0:8], 1000)
+	binary.LittleEndian.PutUint32(payload[8:12], math.Float32bits(3.14))
+	samples, err := ParseData(payload, sigs, PublishModeStrict, time.Now())
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if len(samples) != 1 {
+		t.Fatalf("expected 1 sample, got %d", len(samples))
+	}
+}
@@ -122,10 +122,48 @@ func (c *wsClient) readPump() {

 // ─── Hub ─────────────────────────────────────────────────────────────────────

+// allowedOrigins is the set of Origin values (scheme://host[:port]) that are
+// accepted for WebSocket upgrades.  If empty, same-origin is enforced by
+// comparing the Origin's host to the HTTP Host header.
+var allowedOrigins []string
+
+// SetAllowedOrigins configures the WebSocket Origin allowlist.  Pass an empty
+// slice to enforce same-origin only (the default).
+func SetAllowedOrigins(origins []string) {
+	allowedOrigins = origins
+}
+
+// checkOrigin validates the Origin header against the allowlist, falling back
+// to a same-origin check (Origin host == Host header) when no allowlist is
+// configured.  Requests with no Origin header (non-browser clients) are allowed.
+func checkOrigin(r *http.Request) bool {
+	origin := r.Header.Get("Origin")
+	if origin == "" {
+		return true // non-browser client
+	}
+	// Check explicit allowlist first.
+	for _, allowed := range allowedOrigins {
+		if origin == allowed {
+			return true
+		}
+	}
+	// Fall back to same-origin: compare the Origin's host to the Host header.
+	// Origin format: "scheme://host[:port]" — strip scheme.
+	host := origin
+	if idx := strings.Index(host, "://"); idx >= 0 {
+		host = host[idx+3:]
+	}
+	// Strip path if present.
+	if idx := strings.Index(host, "/"); idx >= 0 {
+		host = host[:idx]
+	}
+	return host == r.Host
+}
+
 var upgrader = websocket.Upgrader{
 	ReadBufferSize:  4096,
 	WriteBufferSize: 64 * 1024,
-	CheckOrigin:     func(r *http.Request) bool { return true },
+	CheckOrigin:     checkOrigin,
 }

 // sourceHubState holds all data for one active data source.
@@ -0,0 +1,72 @@
+package wshub
+
+import (
+	"net/http"
+	"testing"
+)
+
+// TestCheckOrigin_NoOriginHeader — non-browser clients (no Origin) are allowed.
+func TestCheckOrigin_NoOriginHeader(t *testing.T) {
+	r := &http.Request{Header: http.Header{}}
+	if !checkOrigin(r) {
+		t.Fatal("non-browser client (no Origin header) should be allowed")
+	}
+}
+
+// TestCheckOrigin_SameOrigin — Origin host matching Host header is allowed.
+func TestCheckOrigin_SameOrigin(t *testing.T) {
+	r := &http.Request{
+		Header: http.Header{
+			"Origin": []string{"http://localhost:8090"},
+		},
+		Host: "localhost:8090",
+	}
+	if !checkOrigin(r) {
+		t.Fatal("same-origin request should be allowed")
+	}
+}
+
+// TestCheckOrigin_CrossOriginBlocked — different Origin host is rejected.
+func TestCheckOrigin_CrossOriginBlocked(t *testing.T) {
+	r := &http.Request{
+		Header: http.Header{
+			"Origin": []string{"http://evil.example.com:8090"},
+		},
+		Host: "localhost:8090",
+	}
+	if checkOrigin(r) {
+		t.Fatal("cross-origin request should be blocked")
+	}
+}
+
+// TestCheckOrigin_Allowlist — explicitly allowed origins pass even if cross-origin.
+func TestCheckOrigin_Allowlist(t *testing.T) {
+	SetAllowedOrigins([]string{"http://evil.example.com:8090"})
+	defer SetAllowedOrigins(nil) // reset
+
+	r := &http.Request{
+		Header: http.Header{
+			"Origin": []string{"http://evil.example.com:8090"},
+		},
+		Host: "localhost:8090",
+	}
+	if !checkOrigin(r) {
+		t.Fatal("allowlisted origin should be allowed")
+	}
+}
+
+// TestCheckOrigin_AllowlistDoesNotMatch — non-allowlisted cross-origin is blocked.
+func TestCheckOrigin_AllowlistDoesNotMatch(t *testing.T) {
+	SetAllowedOrigins([]string{"http://good.example.com"})
+	defer SetAllowedOrigins(nil)
+
+	r := &http.Request{
+		Header: http.Header{
+			"Origin": []string{"http://evil.example.com"},
+		},
+		Host: "localhost:8090",
+	}
+	if checkOrigin(r) {
+		t.Fatal("non-allowlisted cross-origin should be blocked")
+	}
+}