From 850f56b36793239604deeeb1dc165eaef25d7d78 Mon Sep 17 00:00:00 2001
From: Anna Shaleva <shaleva.ann@nspcc.ru>
Date: Tue, 5 Apr 2022 11:28:26 +0300
Subject: [PATCH] rpc: avoid panic on double-call to *WSClient.Close()

Although it's the caller's duty to avoid WSClient re-closing, we
still can handle it.

Fixes the following neofs-node error:
```
panic: close of closed channel

goroutine 98 [running]:
github.com/nspcc-dev/neo-go/pkg/rpc/client.(*WSClient).Close(...)
        github.com/nspcc-dev/neo-go@v0.98.3-pre.0.20220321144433-3b639f518ebb/pkg/rpc/client/wsclient.go:120
github.com/nspcc-dev/neofs-node/pkg/morph/subscriber.(*subscriber).Close(0x13)
        github.com/nspcc-dev/neofs-node/pkg/morph/subscriber/subscriber.go:108 +0x29
github.com/nspcc-dev/neofs-node/pkg/morph/event.listener.Stop(...)
        github.com/nspcc-dev/neofs-node/pkg/morph/event/listener.go:573
created by github.com/nspcc-dev/neofs-node/pkg/innerring.(*Server).Stop
        github.com/nspcc-dev/neofs-node/pkg/innerring/innerring.go:285 +0x12f
```
---
 pkg/rpc/client/wsclient.go      | 23 ++++++++++++++---------
 pkg/rpc/client/wsclient_test.go | 12 ++++++++++++
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/pkg/rpc/client/wsclient.go b/pkg/rpc/client/wsclient.go
index 89bd4b934..88dc1cbfb 100644
--- a/pkg/rpc/client/wsclient.go
+++ b/pkg/rpc/client/wsclient.go
@@ -16,6 +16,7 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/rpc/response"
 	"github.com/nspcc-dev/neo-go/pkg/rpc/response/result/subscriptions"
 	"github.com/nspcc-dev/neo-go/pkg/util"
+	"go.uber.org/atomic"
 )
 
 // WSClient is a websocket-enabled RPC client that can be used with appropriate
@@ -34,10 +35,11 @@ type WSClient struct {
 	// be closed, so make sure to handle this.
 	Notifications chan Notification
 
-	ws       *websocket.Conn
-	done     chan struct{}
-	requests chan *request.Raw
-	shutdown chan struct{}
+	ws          *websocket.Conn
+	done        chan struct{}
+	requests    chan *request.Raw
+	shutdown    chan struct{}
+	closeCalled atomic.Bool
 
 	subscriptionsLock sync.RWMutex
 	subscriptions     map[string]bool
@@ -93,6 +95,7 @@ func NewWS(ctx context.Context, endpoint string, opts Options) (*WSClient, error
 		ws:            ws,
 		shutdown:      make(chan struct{}),
 		done:          make(chan struct{}),
+		closeCalled:   *atomic.NewBool(false),
 		respChannels:  make(map[uint64]chan *response.Raw),
 		requests:      make(chan *request.Raw),
 		subscriptions: make(map[string]bool),
@@ -113,11 +116,13 @@ func NewWS(ctx context.Context, endpoint string, opts Options) (*WSClient, error
 // Close closes connection to the remote side rendering this client instance
 // unusable.
 func (c *WSClient) Close() {
-	// Closing shutdown channel send signal to wsWriter to break out of the
-	// loop. In doing so it does ws.Close() closing the network connection
-	// which in turn makes wsReader receieve err from ws,ReadJSON() and also
-	// break out of the loop closing c.done channel in its shutdown sequence.
-	close(c.shutdown)
+	if c.closeCalled.CAS(false, true) {
+		// Closing shutdown channel send signal to wsWriter to break out of the
+		// loop. In doing so it does ws.Close() closing the network connection
+		// which in turn makes wsReader receieve err from ws,ReadJSON() and also
+		// break out of the loop closing c.done channel in its shutdown sequence.
+		close(c.shutdown)
+	}
 	<-c.done
 }
 
diff --git a/pkg/rpc/client/wsclient_test.go b/pkg/rpc/client/wsclient_test.go
index 0945aff77..019dead39 100644
--- a/pkg/rpc/client/wsclient_test.go
+++ b/pkg/rpc/client/wsclient_test.go
@@ -441,3 +441,15 @@ func TestWSConcurrentAccess(t *testing.T) {
 		batchCount*3+1) // batchCount*requestsPerBatch+1
 	wsc.Close()
 }
+
+func TestWSDoubleClose(t *testing.T) {
+	srv := initTestServer(t, "")
+
+	c, err := NewWS(context.TODO(), httpURLtoWS(srv.URL), Options{})
+	require.NoError(t, err)
+
+	require.NotPanics(t, func() {
+		c.Close()
+		c.Close()
+	})
+}