fix(websocket): guard Dispatcher's handler map with a RWMutex (#1065)

apps/backend/pkg/websocket/handler.go

@@ -1,6 +1,9 @@
 package websocket
 
-import "context"
+import (
+	"context"
+	"sync"
+)
 
 // Handler is the interface for WebSocket message handlers
 type Handler interface {
@@ -16,8 +19,15 @@ func (f HandlerFunc) Handle(ctx context.Context, msg *Message) (*Message, error)
 	return f(ctx, msg)
 }
 
-// Dispatcher routes messages to appropriate handlers based on action
+// Dispatcher routes messages to appropriate handlers based on action.
+//
+// Dispatcher must be constructed via NewDispatcher; the zero value has a
+// nil handlers map and panics on use. Once constructed it is safe for
+// concurrent use: Register/RegisterFunc may be called from any goroutine,
+// including while Dispatch or HasHandler are in flight. Typical usage
+// registers all handlers at startup, but the API does not require it.
 type Dispatcher struct {
+	mu       sync.RWMutex
 	handlers map[string]Handler
 }
 
@@ -30,17 +40,28 @@ func NewDispatcher() *Dispatcher {
 
 // Register registers a handler for an action
 func (d *Dispatcher) Register(action string, handler Handler) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
 	d.handlers[action] = handler
 }
 
 // RegisterFunc registers a handler function for an action
 func (d *Dispatcher) RegisterFunc(action string, handler HandlerFunc) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
 	d.handlers[action] = handler
 }
 
-// Dispatch routes a message to the appropriate handler
+// Dispatch routes a message to the appropriate handler.
+//
+// The handler is looked up under the read lock, then the lock is released
+// before invoking Handle: holding it across user-supplied I/O would
+// serialise every dispatched message and risk deadlock if a handler
+// re-entered Register.
 func (d *Dispatcher) Dispatch(ctx context.Context, msg *Message) (*Message, error) {
+	d.mu.RLock()
 	handler, ok := d.handlers[msg.Action]
+	d.mu.RUnlock()
 	if !ok {
 		return NewError(msg.ID, msg.Action, ErrorCodeUnknownAction,
 			"Unknown action: "+msg.Action, nil)
@@ -50,6 +71,8 @@ func (d *Dispatcher) Dispatch(ctx context.Context, msg *Message) (*Message, erro
 
 // HasHandler returns true if a handler is registered for the action
 func (d *Dispatcher) HasHandler(action string) bool {
+	d.mu.RLock()
+	defer d.mu.RUnlock()
 	_, ok := d.handlers[action]
 	return ok
 }

apps/backend/pkg/websocket/handler_race_test.go

@@ -0,0 +1,56 @@
+//go:build race
+
+package websocket
+
+import (
+	"context"
+	"sync"
+	"testing"
+)
+
+// TestDispatcher_ConcurrentRegisterAndDispatch exercises the Dispatcher
+// from many goroutines at once. Its purpose is to surface a data race on
+// the underlying handlers map when run under the race detector; the
+// //go:build race tag scopes the test to `go test -race` runs (which CI
+// uses for the backend) so the stress is not paid on plain `go test`
+// invocations where it would add wall time without observing the race.
+func TestDispatcher_ConcurrentRegisterAndDispatch(t *testing.T) {
+	d := NewDispatcher()
+
+	const goroutines = 16
+	const ops = 200
+
+	noop := HandlerFunc(func(_ context.Context, msg *Message) (*Message, error) {
+		return &Message{ID: msg.ID, Action: msg.Action}, nil
+	})
+
+	var wg sync.WaitGroup
+	wg.Add(goroutines * 2)
+
+	for g := 0; g < goroutines; g++ {
+		go func() {
+			defer wg.Done()
+			for i := 0; i < ops; i++ {
+				// Alternate between Register and RegisterFunc so both
+				// write paths are exercised under -race.
+				if i%2 == 0 {
+					d.Register("action", noop)
+				} else {
+					d.RegisterFunc("action", noop)
+				}
+			}
+		}()
+		go func() {
+			defer wg.Done()
+			for i := 0; i < ops; i++ {
+				_, _ = d.Dispatch(context.Background(), &Message{
+					ID:     "x",
+					Action: "action",
+				})
+				_ = d.HasHandler("action")
+			}
+		}()
+	}
+
+	wg.Wait()
+}

apps/backend/pkg/websocket/handler_test.go

@@ -0,0 +1,50 @@
+package websocket
+
+import (
+	"context"
+	"testing"
+)
+
+func TestDispatcher_RegisterAndDispatch(t *testing.T) {
+	d := NewDispatcher()
+
+	d.RegisterFunc("ping", func(_ context.Context, msg *Message) (*Message, error) {
+		return &Message{ID: msg.ID, Action: "pong"}, nil
+	})
+
+	if !d.HasHandler("ping") {
+		t.Fatal("expected HasHandler(\"ping\") to be true")
+	}
+
+	resp, err := d.Dispatch(context.Background(), &Message{ID: "1", Action: "ping"})
+	if err != nil {
+		t.Fatalf("dispatch returned err: %v", err)
+	}
+	if resp == nil || resp.Action != "pong" {
+		t.Fatalf("unexpected response: %#v", resp)
+	}
+}
+
+func TestDispatcher_UnknownActionReturnsError(t *testing.T) {
+	d := NewDispatcher()
+
+	resp, err := d.Dispatch(context.Background(), &Message{ID: "1", Action: "nope"})
+	if err != nil {
+		t.Fatalf("dispatch returned err: %v", err)
+	}
+	if resp == nil {
+		t.Fatal("expected error response, got nil")
+	}
+	if resp.Type != MessageTypeError {
+		t.Errorf("expected error message type, got %q", resp.Type)
+	}
+
+	var payload ErrorPayload
+	if err := resp.ParsePayload(&payload); err != nil {
+		t.Fatalf("parse error payload: %v", err)
+	}
+	if payload.Code != ErrorCodeUnknownAction {
+		t.Errorf("expected error code %q, got %q",
+			ErrorCodeUnknownAction, payload.Code)
+	}
+}