From 652371c965781050903f6ac8f59fd492e1e0ab4e Mon Sep 17 00:00:00 2001 From: vimzh Date: Mon, 25 May 2026 04:09:44 +0530 Subject: [PATCH] fix(websocket): guard Dispatcher's handler map with a RWMutex (#1065) --- apps/backend/pkg/websocket/handler.go | 29 +++++++++- .../pkg/websocket/handler_race_test.go | 56 +++++++++++++++++++ apps/backend/pkg/websocket/handler_test.go | 50 +++++++++++++++++ 3 files changed, 132 insertions(+), 3 deletions(-) create mode 100644 apps/backend/pkg/websocket/handler_race_test.go create mode 100644 apps/backend/pkg/websocket/handler_test.go diff --git a/apps/backend/pkg/websocket/handler.go b/apps/backend/pkg/websocket/handler.go index b864587bb4..862328d3f1 100644 --- a/apps/backend/pkg/websocket/handler.go +++ b/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 } diff --git a/apps/backend/pkg/websocket/handler_race_test.go b/apps/backend/pkg/websocket/handler_race_test.go new file mode 100644 index 0000000000..f69e009866 --- /dev/null +++ b/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() +} diff --git a/apps/backend/pkg/websocket/handler_test.go b/apps/backend/pkg/websocket/handler_test.go new file mode 100644 index 0000000000..1d1923cea2 --- /dev/null +++ b/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) + } +}