Merge pull request #2450 from nspcc-dev/fix-ws

rpc: avoid panic during WSClient request after WS connection is closed

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

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