From 6719441dd07e4001bb6e2d84b6e2f4b8d3f4cb71 Mon Sep 17 00:00:00 2001 From: ABowlOfEleven <22452170+ABowlOfEleven@users.noreply.github.com> Date: Thu, 17 Sep 2026 10:24:11 -0500 Subject: [PATCH] Fix nil dereference in awaitResponse when the connection closes Close() closes the response channel of every pending request. A receive on a closed channel yields nil, and awaitResponse read response.Error from it without a check, so any request that was still waiting when the connection closed panicked the calling goroutine. Close() runs on an explicit call, on a read error (for example when PsyNet drops the socket on DuplicateLogin) and on a pong timeout. awaitResponse now returns ctx.Err() if the context ended, or the new ErrConnectionClosed otherwise. Add two tests that panic without the fix: Close() with a pending request, and the server dropping the connection. --- psynetrpc.go | 16 +++++++- psynetrpc_test.go | 93 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 108 insertions(+), 1 deletion(-) diff --git a/psynetrpc.go b/psynetrpc.go index 56c116b..1d952e9 100644 --- a/psynetrpc.go +++ b/psynetrpc.go @@ -14,6 +14,10 @@ import ( "github.com/gorilla/websocket" ) +// ErrConnectionClosed is returned by a request that was still waiting +// for its response when the WebSocket connection was closed. +var ErrConnectionClosed = errors.New("websocket connection closed") + type emptyRequest struct{} type EventType int @@ -293,7 +297,17 @@ func (p *PsyNetRPC) sendRequestAsync(ctx context.Context, service string, data i func (p *PsyNetRPC) awaitResponse(ctx context.Context, respCh <-chan *PsyResponse, result interface{}) error { select { - case response := <-respCh: + case response, ok := <-respCh: + if !ok || response == nil { + // The channel was closed without a response: either Close() + // ran while this request was pending, or the request's + // context ended and the cleanup goroutine got here first. + if err := ctx.Err(); err != nil { + return err + } + return ErrConnectionClosed + } + if response.Error != nil { return response.Error } diff --git a/psynetrpc_test.go b/psynetrpc_test.go index eea9422..07a0b93 100644 --- a/psynetrpc_test.go +++ b/psynetrpc_test.go @@ -3,6 +3,7 @@ package rlapi import ( "context" "encoding/json" + "errors" "fmt" "net/http" "net/http/httptest" @@ -22,6 +23,7 @@ type MockWSServer struct { responses map[string]*PsyResponse // Predefined responses pongResponse bool // Whether to respond to pings with pong dropPongs bool // Whether to drop pong responses + closeOnReq bool // Whether to close the connection when a request arrives } func NewMockWSServer() *MockWSServer { @@ -60,6 +62,10 @@ func (m *MockWSServer) SetDropPongs(drop bool) { m.dropPongs = drop } +func (m *MockWSServer) SetCloseOnRequest(closeOnReq bool) { + m.closeOnReq = closeOnReq +} + func (m *MockWSServer) handleWebSocket(w http.ResponseWriter, r *http.Request) { conn, err := m.upgrader.Upgrade(w, r, nil) if err != nil { @@ -85,6 +91,9 @@ func (m *MockWSServer) handleWebSocket(w http.ResponseWriter, r *http.Request) { // Parse PsyRequestID from the message if strings.Contains(msg, "PsyRequestID:") { + if m.closeOnReq { + return // drop the connection without answering + } lines := strings.Split(msg, "\r\n") for _, line := range lines { if strings.HasPrefix(line, "PsyRequestID:") { @@ -447,6 +456,90 @@ func TestPsyNetRPC_ConcurrentContextCancellation(t *testing.T) { } } +// A request that is still waiting when the connection goes away must +// get an error back. Close() closes every pending response channel, so +// the waiting side receives nil from it. +func TestPsyNetRPC_CloseWithPendingRequest(t *testing.T) { + // Setup mock server (no responses, the request will hang) + mockServer := NewMockWSServer() + defer mockServer.Close() + + psyNet := NewPsyNet() + rpc, err := psyNet.establishSocket(mockServer.URL(), "test-token", "test-session", "test-player") + if err != nil { + t.Fatalf("Failed to establish socket: %v", err) + } + + go rpc.readMessages() + rpc.schedulePing() + defer rpc.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + errCh := make(chan error, 1) + go func() { + var result map[string]interface{} + errCh <- rpc.sendRequestSync(ctx, "Test/Pending", map[string]interface{}{}, &result) + }() + + // Wait until the request is registered as pending + deadline := time.Now().Add(2 * time.Second) + for { + rpc.mu.Lock() + pending := len(rpc.pendingReqs) + rpc.mu.Unlock() + if pending == 1 { + break + } + if time.Now().After(deadline) { + t.Fatal("Request never became pending") + } + time.Sleep(5 * time.Millisecond) + } + + rpc.Close() + + select { + case err := <-errCh: + if !errors.Is(err, ErrConnectionClosed) { + t.Errorf("Expected ErrConnectionClosed, got %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("sendRequestSync did not return after Close") + } +} + +// Same as above, but the server drops the connection (what PsyNet does +// on DuplicateLogin), so Close() runs from readMessages. +func TestPsyNetRPC_ServerDropsConnectionWithPendingRequest(t *testing.T) { + mockServer := NewMockWSServer() + defer mockServer.Close() + mockServer.SetCloseOnRequest(true) + + psyNet := NewPsyNet() + rpc, err := psyNet.establishSocket(mockServer.URL(), "test-token", "test-session", "test-player") + if err != nil { + t.Fatalf("Failed to establish socket: %v", err) + } + + go rpc.readMessages() + rpc.schedulePing() + defer rpc.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + var result map[string]interface{} + err = rpc.sendRequestSync(ctx, "Test/Dropped", map[string]interface{}{}, &result) + if !errors.Is(err, ErrConnectionClosed) { + t.Errorf("Expected ErrConnectionClosed, got %v", err) + } + if rpc.IsConnected() { + t.Error("Expected connection to be marked as closed") + } +} + func TestPsyNetRPC_ParseMessage(t *testing.T) { rpc := &PsyNetRPC{}