puma-plus-core 0.1.0

data/go/internal/realtime/hub_test.go

@@ -0,0 +1,370 @@
+package realtime
+
+import (
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+// fakeTransport records what was sent, and can pretend to be either kind so
+// mixed-topic behaviour is testable without a real network.
+type fakeTransport struct {
+	kind string
+
+	mu        sync.Mutex
+	messages  [][]byte
+	datagrams [][]byte
+	streams   map[string][][]byte
+	closed    bool
+	closeCode int
+
+	// noDatagrams makes SendDatagram fail the way WebSocket does.
+	noDatagrams bool
+	// preframe enables the fan-out optimization path.
+	preframe bool
+
+	preframeCalls atomic.Int64
+}
+
+func newFake(kind string) *fakeTransport {
+	return &fakeTransport{
+		kind:        kind,
+		streams:     map[string][][]byte{},
+		noDatagrams: kind == "websocket",
+		preframe:    kind == "websocket",
+	}
+}
+
+func (f *fakeTransport) Kind() string       { return f.kind }
+func (f *fakeTransport) RemoteAddr() string { return "test" }
+
+func (f *fakeTransport) SendMessage(p []byte, _ bool) error {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.messages = append(f.messages, append([]byte(nil), p...))
+	return nil
+}
+
+func (f *fakeTransport) SendOnStream(id string, p []byte, _ bool) error {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.streams[id] = append(f.streams[id], append([]byte(nil), p...))
+	return nil
+}
+
+func (f *fakeTransport) OpenStream() (string, error) {
+	if f.kind == "websocket" {
+		return "", ErrNoStreams
+	}
+	return "s1", nil
+}
+
+func (f *fakeTransport) SendDatagram(p []byte) error {
+	if f.noDatagrams {
+		return ErrNoDatagrams
+	}
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.datagrams = append(f.datagrams, append([]byte(nil), p...))
+	return nil
+}
+
+func (f *fakeTransport) Close(code int, _ string) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.closed, f.closeCode = true, code
+}
+
+func (f *fakeTransport) PreframeMessage(p []byte, _ bool) []byte {
+	if !f.preframe {
+		return nil
+	}
+	f.preframeCalls.Add(1)
+	return append([]byte("FRAMED:"), p...)
+}
+
+func (f *fakeTransport) SendPreframed(frame []byte) error {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.messages = append(f.messages, append([]byte(nil), frame...))
+	return nil
+}
+
+func (f *fakeTransport) sent() int {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return len(f.messages)
+}
+
+func (f *fakeTransport) datagramCount() int {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return len(f.datagrams)
+}
+
+func (f *fakeTransport) isClosed() bool {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return f.closed
+}
+
+func add(t *testing.T, h *Hub, kind string, mode Mode) (*Conn, *fakeTransport) {
+	t.Helper()
+	tr := newFake(kind)
+	c := NewConn(h, tr, mode, "/ws", nil)
+	h.Add(c)
+	return c, tr
+}
+
+// TestPublishReachesBothTransports is the payoff of a shared hub: one topic may
+// hold WebSocket and WebTransport clients, and a single publish reaches both.
+// Without this the application would have to know who is on which protocol and
+// publish twice.
+func TestPublishReachesBothTransports(t *testing.T) {
+	h := NewHub()
+
+	var wsT, wtT []*fakeTransport
+	for i := 0; i < 3; i++ {
+		c, tr := add(t, h, "websocket", ModeFlexible)
+		h.Subscribe(c.ID, "room")
+		wsT = append(wsT, tr)
+	}
+	for i := 0; i < 2; i++ {
+		c, tr := add(t, h, "webtransport", ModeFlexible)
+		h.Subscribe(c.ID, "room")
+		wtT = append(wtT, tr)
+	}
+
+	if got := h.Publish("room", []byte("hello"), false); got != 5 {
+		t.Errorf("Publish delivered to %d, want 5", got)
+	}
+	for i, tr := range wsT {
+		if tr.sent() != 1 {
+			t.Errorf("websocket subscriber %d received %d messages", i, tr.sent())
+		}
+	}
+	for i, tr := range wtT {
+		if tr.sent() != 1 {
+			t.Errorf("webtransport subscriber %d received %d messages", i, tr.sent())
+		}
+	}
+}
+
+// TestPublishPreframesOncePerKind pins the fan-out optimization: the message is
+// encoded once per transport kind, not once per subscriber. That is most of why
+// fan-out belongs in Go.
+func TestPublishPreframesOncePerKind(t *testing.T) {
+	h := NewHub()
+	var all []*fakeTransport
+	for i := 0; i < 10; i++ {
+		c, tr := add(t, h, "websocket", ModeFlexible)
+		h.Subscribe(c.ID, "room")
+		all = append(all, tr)
+	}
+
+	h.Publish("room", []byte("x"), false)
+
+	var total int64
+	for _, tr := range all {
+		total += tr.preframeCalls.Load()
+	}
+	if total != 1 {
+		t.Errorf("encoded %d times for 10 subscribers; expected exactly once", total)
+	}
+}
+
+// TestPublishDatagramSkipsWebSocket: WebSocket has no unreliable channel, so
+// those subscribers are reported as unsupported rather than silently upgraded to
+// reliable delivery they did not ask for.
+func TestPublishDatagramSkipsWebSocket(t *testing.T) {
+	h := NewHub()
+	for i := 0; i < 3; i++ {
+		c, _ := add(t, h, "websocket", ModeFlexible)
+		h.Subscribe(c.ID, "game")
+	}
+	var wtT []*fakeTransport
+	for i := 0; i < 2; i++ {
+		c, tr := add(t, h, "webtransport", ModeFlexible)
+		h.Subscribe(c.ID, "game")
+		wtT = append(wtT, tr)
+	}
+
+	delivered, unsupported := h.PublishDatagram("game", []byte("pos"))
+	if delivered != 2 {
+		t.Errorf("delivered = %d, want 2 (the WebTransport clients)", delivered)
+	}
+	if unsupported != 3 {
+		t.Errorf("unsupported = %d, want 3 (the WebSocket clients)", unsupported)
+	}
+	for i, tr := range wtT {
+		if tr.datagramCount() != 1 {
+			t.Errorf("webtransport subscriber %d got %d datagrams", i, tr.datagramCount())
+		}
+	}
+}
+
+func TestUnsubscribeStopsDelivery(t *testing.T) {
+	h := NewHub()
+	c, _ := add(t, h, "websocket", ModeFlexible)
+	h.Subscribe(c.ID, "t")
+	if h.Subscribers("t") != 1 {
+		t.Fatal("subscribe did not register")
+	}
+	h.Unsubscribe(c.ID, "t")
+	if got := h.Subscribers("t"); got != 0 {
+		t.Errorf("Subscribers after unsubscribe = %d, want 0", got)
+	}
+	if n := h.Publish("t", []byte("x"), false); n != 0 {
+		t.Errorf("published to %d connections after unsubscribe", n)
+	}
+}
+
+// TestRemoveCleansUpTopics guards a leak: a closed connection must not linger in
+// every topic it ever joined.
+func TestRemoveCleansUpTopics(t *testing.T) {
+	h := NewHub()
+	c, _ := add(t, h, "websocket", ModeFlexible)
+	for _, topic := range []string{"a", "b", "c"} {
+		h.Subscribe(c.ID, topic)
+	}
+	h.Remove(c.ID)
+
+	for _, topic := range []string{"a", "b", "c"} {
+		if got := h.Subscribers(topic); got != 0 {
+			t.Errorf("topic %q still has %d subscribers", topic, got)
+		}
+	}
+	if h.Count() != 0 {
+		t.Errorf("Count = %d after removing the only connection", h.Count())
+	}
+}
+
+// TestDropWorkerClosesOnlyItsStickyConns pins the failover contract, and it must
+// behave identically for both transports.
+func TestDropWorkerClosesOnlyItsStickyConns(t *testing.T) {
+	h := NewHub()
+
+	mine, mineT := add(t, h, "websocket", ModeSticky)
+	mine.WorkerID, mine.Pinned = 7, true
+
+	mineWT, mineWTT := add(t, h, "webtransport", ModeSticky)
+	mineWT.WorkerID, mineWT.Pinned = 7, true
+
+	theirs, theirsT := add(t, h, "websocket", ModeSticky)
+	theirs.WorkerID, theirs.Pinned = 8, true
+
+	flexible, flexT := add(t, h, "websocket", ModeFlexible)
+	_ = flexible
+
+	unpinned, unpinnedT := add(t, h, "websocket", ModeSticky)
+	_ = unpinned
+
+	dropped := h.DropWorker(7, 1012, "worker retired")
+
+	if len(dropped) != 2 {
+		t.Errorf("DropWorker(7) closed %d connections, want 2", len(dropped))
+	}
+	if !mineT.isClosed() || !mineWTT.isClosed() {
+		t.Error("a connection pinned to worker 7 was not closed")
+	}
+	if theirsT.isClosed() {
+		t.Error("a connection pinned to a DIFFERENT worker was closed")
+	}
+	if flexT.isClosed() {
+		t.Error("a flexible connection was closed; it holds no per-worker state")
+	}
+	if unpinnedT.isClosed() {
+		t.Error("an unpinned sticky connection was closed")
+	}
+}
+
+func TestStatsCountsModesAndTransports(t *testing.T) {
+	h := NewHub()
+	for i := 0; i < 3; i++ {
+		add(t, h, "websocket", ModeSticky)
+	}
+	for i := 0; i < 2; i++ {
+		c, _ := add(t, h, "webtransport", ModeFlexible)
+		h.Subscribe(c.ID, "topic")
+	}
+
+	s := h.Stats()
+	if s.Connections != 5 || s.Sticky != 3 || s.Flexible != 2 {
+		t.Errorf("Stats = %+v, want 5 connections / 3 sticky / 2 flexible", s)
+	}
+	if s.WebSocket != 3 || s.WebTransport != 2 {
+		t.Errorf("Stats transports = ws %d / wt %d, want 3 / 2", s.WebSocket, s.WebTransport)
+	}
+	if s.Topics != 1 {
+		t.Errorf("Topics = %d, want 1", s.Topics)
+	}
+}
+
+// TestWebSocketRejectsWebTransportOnlyCalls: the asymmetry must surface as an
+// error rather than as silently different behaviour.
+func TestWebSocketRejectsWebTransportOnlyCalls(t *testing.T) {
+	h := NewHub()
+	c, _ := add(t, h, "websocket", ModeFlexible)
+
+	if err := c.SendDatagram([]byte("x")); err != ErrNoDatagrams {
+		t.Errorf("SendDatagram on a WebSocket = %v, want ErrNoDatagrams", err)
+	}
+	if _, err := c.OpenStream(); err != ErrNoStreams {
+		t.Errorf("OpenStream on a WebSocket = %v, want ErrNoStreams", err)
+	}
+}
+
+// TestConcurrentSubscribeAndPublish runs under -race: the hub is written from
+// every request goroutine and read by every fan-out.
+func TestConcurrentSubscribeAndPublish(t *testing.T) {
+	h := NewHub()
+	var wg sync.WaitGroup
+
+	for i := 0; i < 20; i++ {
+		kind := "websocket"
+		if i%2 == 0 {
+			kind = "webtransport"
+		}
+		c, _ := add(t, h, kind, ModeFlexible)
+
+		wg.Add(1)
+		go func(id string) {
+			defer wg.Done()
+			for j := 0; j < 20; j++ {
+				h.Subscribe(id, "hot")
+				h.Subscribers("hot")
+				h.Unsubscribe(id, "hot")
+			}
+		}(c.ID)
+	}
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for i := 0; i < 50; i++ {
+			h.Publish("hot", []byte("x"), false)
+			h.PublishDatagram("hot", []byte("y"))
+			h.Stats()
+		}
+	}()
+
+	wg.Wait()
+}
+
+func TestCloseAll(t *testing.T) {
+	h := NewHub()
+	var all []*fakeTransport
+	for i := 0; i < 5; i++ {
+		_, tr := add(t, h, "websocket", ModeFlexible)
+		all = append(all, tr)
+	}
+	h.CloseAll(1001, "going away")
+	for i, tr := range all {
+		if !tr.isClosed() {
+			t.Errorf("connection %d was not closed", i)
+		}
+	}
+}
+
+var _ = time.Now

data/go/internal/realtime/id.go

@@ -0,0 +1,69 @@
+package realtime
+
+import "crypto/rand"
+
+// base58Alphabet is the Bitcoin alphabet: no 0, O, I or l, so an id copied out
+// of a log or read aloud cannot be transcribed into a different valid id.
+const base58Alphabet = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"
+
+// idBytes is the entropy behind each connection id: 128 bits, about 22 base58
+// characters.
+const idBytes = 16
+
+// NewID mints a random base58 connection identifier.
+//
+// Random rather than sequential, and the reason is not just tidiness. A conn id
+// is the address the application uses to send to a connection, so it is
+// effectively a capability. A counter would let any client read its own id and
+// derive its neighbours' -- an app that accepts an id from a client and routes
+// on it would be trivially exploitable, and even one that does not still leaks
+// exactly how many connections the server has handled. 128 random bits make
+// that a non-question rather than a caveat.
+func NewID() string {
+	var b [idBytes]byte
+	if _, err := rand.Read(b[:]); err != nil {
+		// crypto/rand does not fail on any supported platform. If it somehow
+		// does, a predictable id is far worse than a crash.
+		panic("websocket: crypto/rand unavailable: " + err.Error())
+	}
+	return encodeBase58(b[:])
+}
+
+// encodeBase58 encodes bytes big-endian by repeated division, preserving leading
+// zero bytes as the alphabet's first character.
+func encodeBase58(input []byte) string {
+	zeros := 0
+	for zeros < len(input) && input[zeros] == 0 {
+		zeros++
+	}
+
+	// log(256)/log(58) is about 1.365, so 137/100 is a safe bound.
+	size := (len(input)-zeros)*137/100 + 1
+	buf := make([]byte, size)
+
+	high := size - 1
+	for i := zeros; i < len(input); i++ {
+		carry := int(input[i])
+		j := size - 1
+		for ; j > high || carry != 0; j-- {
+			carry += 256 * int(buf[j])
+			buf[j] = byte(carry % 58)
+			carry /= 58
+		}
+		high = j
+	}
+
+	j := 0
+	for j < size && buf[j] == 0 {
+		j++
+	}
+
+	out := make([]byte, zeros+size-j)
+	for i := 0; i < zeros; i++ {
+		out[i] = base58Alphabet[0]
+	}
+	for i := zeros; j < size; i, j = i+1, j+1 {
+		out[i] = base58Alphabet[buf[j]]
+	}
+	return string(out)
+}

data/go/internal/realtime/id_test.go

@@ -0,0 +1,64 @@
+package realtime
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestNewIDIsBase58AndUnique(t *testing.T) {
+	seen := map[string]bool{}
+	for i := 0; i < 20000; i++ {
+		id := NewID()
+		if seen[id] {
+			t.Fatalf("duplicate id after %d draws: %q", i, id)
+		}
+		seen[id] = true
+
+		if len(id) < 16 || len(id) > 24 {
+			t.Fatalf("id %q has length %d, want ~22", id, len(id))
+		}
+		// The excluded characters are the whole point of base58: 0/O and I/l
+		// are the pairs people mis-transcribe.
+		if strings.ContainsAny(id, "0OIl") {
+			t.Errorf("id %q contains an ambiguous character", id)
+		}
+		for _, r := range id {
+			if !strings.ContainsRune(base58Alphabet, r) {
+				t.Fatalf("id %q contains %q, which is not in the alphabet", id, r)
+			}
+		}
+	}
+}
+
+// TestNewIDIsNotSequential guards the property that motivated random ids: one
+// id must not reveal the next. A counter would make these adjacent.
+func TestNewIDIsNotSequential(t *testing.T) {
+	a, b := NewID(), NewID()
+	common := 0
+	for i := 0; i < len(a) && i < len(b) && a[i] == b[i]; i++ {
+		common++
+	}
+	if common > 4 {
+		t.Errorf("consecutive ids share a %d-character prefix (%q, %q); "+
+			"they should be independent", common, a, b)
+	}
+}
+
+func TestEncodeBase58KnownVectors(t *testing.T) {
+	for _, tc := range []struct {
+		in   []byte
+		want string
+	}{
+		{[]byte{}, ""},
+		{[]byte{0}, "1"},
+		{[]byte{0, 0}, "11"},
+		{[]byte{57}, "z"},
+		{[]byte{58}, "21"},
+		{[]byte{0, 1}, "12"},
+		{[]byte("hello world"), "StV1DL6CwTryKyV"},
+	} {
+		if got := encodeBase58(tc.in); got != tc.want {
+			t.Errorf("encodeBase58(%v) = %q, want %q", tc.in, got, tc.want)
+		}
+	}
+}

data/go/internal/websocket/frame.go

@@ -0,0 +1,186 @@
+// Package websocket implements RFC 6455 on the server side, so that Go can own
+// WebSocket connections on the application's behalf.
+//
+// The point of doing this here rather than in Ruby is capacity. In puma a
+// WebSocket occupies a thread for the life of the connection, so ten thousand
+// idle chat clients means ten thousand parked threads -- which is why Rails
+// pushes ActionCable into its own process. In falcon it is a fiber, much
+// cheaper, but still bound to the process that accepted it. Here the connection
+// is a goroutine and a few kilobytes, and Ruby is involved only when a message
+// actually arrives.
+package websocket
+
+import (
+	"crypto/sha1"
+	"encoding/base64"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"io"
+	"unicode/utf8"
+)
+
+// GUID is the RFC 6455 magic value used to derive Sec-WebSocket-Accept.
+const GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
+
+// Opcodes.
+const (
+	OpContinuation = 0x0
+	OpText         = 0x1
+	OpBinary       = 0x2
+	OpClose        = 0x8
+	OpPing         = 0x9
+	OpPong         = 0xA
+)
+
+// Close codes we originate.
+const (
+	CloseNormal            = 1000
+	CloseGoingAway         = 1001
+	CloseProtocolError     = 1002
+	CloseUnsupportedData   = 1003
+	CloseInvalidPayload    = 1007
+	ClosePolicyViolation   = 1008
+	CloseMessageTooBig     = 1009
+	CloseInternalError     = 1011
+	CloseServiceRestarting = 1012
+)
+
+// MaxMessageSize bounds a reassembled message. A peer that exceeds it gets a
+// 1009 close rather than being allowed to exhaust memory.
+const MaxMessageSize = 8 << 20
+
+var (
+	// ErrProtocol means the peer violated RFC 6455.
+	ErrProtocol = errors.New("websocket: protocol error")
+	// ErrTooLarge means a message exceeded MaxMessageSize.
+	ErrTooLarge = errors.New("websocket: message too large")
+)
+
+// AcceptKey derives the Sec-WebSocket-Accept response value.
+func AcceptKey(clientKey string) string {
+	h := sha1.New()
+	io.WriteString(h, clientKey)
+	io.WriteString(h, GUID)
+	return base64.StdEncoding.EncodeToString(h.Sum(nil))
+}
+
+// frame is one WebSocket frame as read off the wire.
+type frame struct {
+	fin     bool
+	opcode  byte
+	payload []byte
+}
+
+// readFrame reads a single frame.
+//
+// Every client-to-server frame must be masked (RFC 6455 §5.1); an unmasked one
+// is a protocol error rather than something to tolerate, because accepting it
+// is how cache-poisoning attacks against intermediaries work.
+func readFrame(r io.Reader, buf []byte) (frame, error) {
+	var hdr [2]byte
+	if _, err := io.ReadFull(r, hdr[:]); err != nil {
+		return frame{}, err
+	}
+
+	f := frame{
+		fin:    hdr[0]&0x80 != 0,
+		opcode: hdr[0] & 0x0F,
+	}
+	// RSV1-3 must be zero without a negotiated extension, and we negotiate none.
+	if hdr[0]&0x70 != 0 {
+		return f, fmt.Errorf("%w: reserved bits set", ErrProtocol)
+	}
+
+	masked := hdr[1]&0x80 != 0
+	if !masked {
+		return f, fmt.Errorf("%w: client frame was not masked", ErrProtocol)
+	}
+
+	length := uint64(hdr[1] & 0x7F)
+	switch length {
+	case 126:
+		var ext [2]byte
+		if _, err := io.ReadFull(r, ext[:]); err != nil {
+			return f, err
+		}
+		length = uint64(binary.BigEndian.Uint16(ext[:]))
+	case 127:
+		var ext [8]byte
+		if _, err := io.ReadFull(r, ext[:]); err != nil {
+			return f, err
+		}
+		length = binary.BigEndian.Uint64(ext[:])
+	}
+
+	// Control frames carry at most 125 bytes and can never be fragmented.
+	if f.opcode&0x8 != 0 {
+		if length > 125 {
+			return f, fmt.Errorf("%w: control frame payload of %d bytes", ErrProtocol, length)
+		}
+		if !f.fin {
+			return f, fmt.Errorf("%w: fragmented control frame", ErrProtocol)
+		}
+	}
+	if length > MaxMessageSize {
+		return f, ErrTooLarge
+	}
+
+	var mask [4]byte
+	if _, err := io.ReadFull(r, mask[:]); err != nil {
+		return f, err
+	}
+
+	if uint64(cap(buf)) < length {
+		buf = make([]byte, length)
+	}
+	buf = buf[:length]
+	if length > 0 {
+		if _, err := io.ReadFull(r, buf); err != nil {
+			return f, err
+		}
+		for i := range buf {
+			buf[i] ^= mask[i&3]
+		}
+	}
+	f.payload = buf
+	return f, nil
+}
+
+// appendFrame encodes a server-to-client frame. Server frames are never masked.
+func appendFrame(dst []byte, opcode byte, fin bool, payload []byte) []byte {
+	b0 := opcode
+	if fin {
+		b0 |= 0x80
+	}
+	dst = append(dst, b0)
+
+	n := len(payload)
+	switch {
+	case n < 126:
+		dst = append(dst, byte(n))
+	case n < 65536:
+		dst = append(dst, 126)
+		dst = binary.BigEndian.AppendUint16(dst, uint16(n))
+	default:
+		dst = append(dst, 127)
+		dst = binary.BigEndian.AppendUint64(dst, uint64(n))
+	}
+	return append(dst, payload...)
+}
+
+// appendClose encodes a close frame carrying a status code and reason.
+func appendClose(dst []byte, code int, reason string) []byte {
+	// A close frame's payload is capped at 125 bytes, of which two are the code.
+	if len(reason) > 123 {
+		reason = reason[:123]
+	}
+	body := make([]byte, 2, 2+len(reason))
+	binary.BigEndian.PutUint16(body, uint16(code))
+	body = append(body, reason...)
+	return appendFrame(dst, OpClose, true, body)
+}
+
+// validUTF8 reports whether a text message is valid UTF-8, which RFC 6455
+// requires. Sending invalid UTF-8 in a text frame earns a 1007 close.
+func validUTF8(b []byte) bool { return utf8.Valid(b) }