TrueCloudLab/neoneo-go
7
0
Fork 2
Code Issues Pull requests Projects Releases Packages Wiki Activity

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

Browse source 
rpc: avoid panic during WSClient request after WS connection is closed
This commit is contained in:
Roman Khimov 2022-04-27 17:02:45 +03:00 committed by GitHub
commit b8edf4ac73
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 24 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
pkg/rpc/client/wsclient.go
View file

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

15
pkg/rpc/client/wsclient_test.go
View file

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