From b1f74b549dbf8a12267f4cf30bfe54da6fa097c5 Mon Sep 17 00:00:00 2001 From: Anna Shaleva Date: Tue, 26 Apr 2022 17:19:51 +0300 Subject: [PATCH] rpc: avoid panic during request after WS connection is closed Fix the following panic: ``` panic: assignment to entry in nil map goroutine 131 [running]: github.com/nspcc-dev/neo-go/pkg/rpc/client.(*WSClient).registerRespChannel(0xc00033c240, 0x0, 0xc00003e2a0) /home/denis/go/pkg/mod/github.com/nspcc-dev/neo-go@v0.98.2/pkg/rpc/client/wsclient.go:244 +0x96 github.com/nspcc-dev/neo-go/pkg/rpc/client.(*WSClient).makeWsRequest(0xc00033c240, 0xc002080000) /home/denis/go/pkg/mod/github.com/nspcc-dev/neo-go@v0.98.2/pkg/rpc/client/wsclient.go:264 +0x69 github.com/nspcc-dev/neo-go/pkg/rpc/client.(*Client).performRequest(0xc00033c240, {0xc9f173, 0xd}, {{0x13d09d0, 0x0, 0x0}}, {0xb44120, 0xc00147a000}) /home/denis/go/pkg/mod/github.com/nspcc-dev/neo-go@v0.98.2/pkg/rpc/client/client.go:186 +0x15d github.com/nspcc-dev/neo-go/pkg/rpc/client.(*Client).GetBlockCount(0xc001fb5440) /home/denis/go/pkg/mod/github.com/nspcc-dev/neo-go@v0.98.2/pkg/rpc/client/rpc.go:73 +0x69 ... ``` --- pkg/rpc/client/wsclient.go | 17 +++++++++-------- pkg/rpc/client/wsclient_test.go | 15 +++++++++++++++ 2 files changed, 24 insertions(+), 8 deletions(-) diff --git a/pkg/rpc/client/wsclient.go b/pkg/rpc/client/wsclient.go index 88dc1cbfb..9c4ff5272 100644 --- a/pkg/rpc/client/wsclient.go +++ b/pkg/rpc/client/wsclient.go @@ -243,12 +243,6 @@ func (c *WSClient) wsWriter() { } } -func (c *WSClient) registerRespChannel(id uint64, ch chan *response.Raw) { - c.respLock.Lock() - defer c.respLock.Unlock() - c.respChannels[id] = ch -} - func (c *WSClient) unregisterRespChannel(id uint64) { c.respLock.Lock() defer c.respLock.Unlock() @@ -266,8 +260,15 @@ func (c *WSClient) getResponseChannel(id uint64) chan *response.Raw { func (c *WSClient) makeWsRequest(r *request.Raw) (*response.Raw, error) { ch := make(chan *response.Raw) - c.registerRespChannel(r.ID, ch) - + c.respLock.Lock() + select { + case <-c.done: + c.respLock.Unlock() + return nil, errors.New("connection lost before registering response channel") + default: + c.respChannels[r.ID] = ch + c.respLock.Unlock() + } select { case <-c.done: return nil, errors.New("connection lost before sending the request") diff --git a/pkg/rpc/client/wsclient_test.go b/pkg/rpc/client/wsclient_test.go index 019dead39..531db6d5d 100644 --- a/pkg/rpc/client/wsclient_test.go +++ b/pkg/rpc/client/wsclient_test.go @@ -453,3 +453,18 @@ func TestWSDoubleClose(t *testing.T) { c.Close() }) } + +func TestWS_RequestAfterClose(t *testing.T) { + srv := initTestServer(t, "") + + c, err := NewWS(context.TODO(), httpURLtoWS(srv.URL), Options{}) + require.NoError(t, err) + + c.Close() + + require.NotPanics(t, func() { + _, err = c.GetBlockCount() + }) + require.Error(t, err) + require.True(t, strings.Contains(err.Error(), "connection lost before registering response channel")) +}