snapshot of https://github.com/centrifugal/centrifugo.git (master, 2026-09-16; blobs > 131072B omitted)

internal/websocket/close_code_test.go

@@ -0,0 +1,72 @@
+package websocket
+
+import (
+	"bytes"
+	"testing"
+	"time"
+)
+
+// TestCloseCode_Unset returns zero values before any close frame is observed.
+func TestCloseCode_Unset(t *testing.T) {
+	c := newTestConn(nil, &bytes.Buffer{}, true)
+	if code, incoming := c.CloseCode(); code != 0 || incoming {
+		t.Fatalf("CloseCode() = (%d, %v), want (0, false)", code, incoming)
+	}
+}
+
+// TestCloseCode_Outgoing records a close code we write as outgoing.
+func TestCloseCode_Outgoing(t *testing.T) {
+	c := newTestConn(nil, &bytes.Buffer{}, true)
+	_ = c.WriteControl(CloseMessage, FormatCloseMessage(CloseMessageTooBig, ""), time.Now().Add(time.Second))
+	if code, incoming := c.CloseCode(); code != CloseMessageTooBig || incoming {
+		t.Fatalf("CloseCode() = (%d, %v), want (%d, false)", code, incoming, CloseMessageTooBig)
+	}
+}
+
+// TestCloseCode_Incoming records a close code received from the peer as incoming
+// (and the subsequent RFC echo must not flip it to outgoing).
+func TestCloseCode_Incoming(t *testing.T) {
+	var buf bytes.Buffer
+	// Client writes a close frame (masked) into buf.
+	wc := newTestConn(nil, &buf, false)
+	_ = wc.WriteControl(CloseMessage, FormatCloseMessage(CloseGoingAway, ""), time.Now().Add(time.Second))
+
+	// Server reads it; advanceFrame records it as incoming and echoes a close.
+	rc := newTestConn(&buf, &bytes.Buffer{}, true)
+	if _, _, err := rc.NextReader(); !IsCloseError(err, CloseGoingAway) {
+		t.Fatalf("NextReader err = %v, want CloseError(%d)", err, CloseGoingAway)
+	}
+	if code, incoming := rc.CloseCode(); code != CloseGoingAway || !incoming {
+		t.Fatalf("CloseCode() = (%d, %v), want (%d, true)", code, incoming, CloseGoingAway)
+	}
+}
+
+// TestCloseCode_FirstWins keeps the initiating close, ignoring later frames.
+func TestCloseCode_FirstWins(t *testing.T) {
+	c := newTestConn(nil, &bytes.Buffer{}, true)
+	deadline := time.Now().Add(time.Second)
+	_ = c.WriteControl(CloseMessage, FormatCloseMessage(CloseMessageTooBig, ""), deadline)
+	_ = c.WriteControl(CloseMessage, FormatCloseMessage(CloseNormalClosure, ""), deadline)
+	if code, incoming := c.CloseCode(); code != CloseMessageTooBig || incoming {
+		t.Fatalf("CloseCode() = (%d, %v), want (%d, false)", code, incoming, CloseMessageTooBig)
+	}
+}
+
+// TestCloseCode_IncomingAppRangeStillIncoming proves an attacker-supplied
+// application-range close code (3000-4999) is recorded as incoming, so the
+// Centrifuge layer (which only records !incoming) will not turn it into a label.
+func TestCloseCode_IncomingAppRangeStillIncoming(t *testing.T) {
+	const appCode = 4321
+	var buf bytes.Buffer
+	wc := newTestConn(nil, &buf, false)
+	_ = wc.WriteControl(CloseMessage, FormatCloseMessage(appCode, ""), time.Now().Add(time.Second))
+
+	rc := newTestConn(&buf, &bytes.Buffer{}, true)
+	if _, _, err := rc.NextReader(); !IsCloseError(err, appCode) {
+		t.Fatalf("NextReader err = %v, want CloseError(%d)", err, appCode)
+	}
+	code, incoming := rc.CloseCode()
+	if code != appCode || !incoming {
+		t.Fatalf("CloseCode() = (%d, %v), want (%d, true)", code, incoming, appCode)
+	}
+}