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")) +}