From 64cd28b9dd4ee2fd80af26fb8bf4234705e791f5 Mon Sep 17 00:00:00 2001
From: Henrik Hautakoski <henrik@eossweden.org>
Date: Tue, 29 Nov 2022 05:22:14 +0100
Subject: [PATCH] client.go: In ShipClient.ReadRaw() when if error is an
 websocket.CloseError and return ErrSockClosed in that case.

---
 client.go      |  8 +++++++-
 client_test.go | 47 ++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/client.go b/client.go
index 8d7d19d..0517d70 100644
--- a/client.go
+++ b/client.go
@@ -225,7 +225,13 @@ func (c *ShipClient) ReadRaw() (int, []byte, error) {
 	// Read message from socket.
 	msg_type, data, err := c.sock.ReadMessage()
 	if err != nil {
-		return msg_type, data, ShipClientError{ErrSockRead, err.Error()}
+
+		errType := ErrSockRead
+		if _, ok := err.(*ws.CloseError); ok {
+			errType = ErrSockClosed
+		}
+
+		return msg_type, data, ShipClientError{errType, err.Error()}
 	}
 
 	// Check if we need to ack messages
diff --git a/client_test.go b/client_test.go
index dae6b35..c986eb9 100644
--- a/client_test.go
+++ b/client_test.go
@@ -25,6 +25,8 @@ type testHandler struct {
 
 	responses [][]byte
 
+	CloseError bool
+
 	ExpectedBlockRequest *ship.GetBlocksRequestV0
 	RespondBlocks        []ship.GetBlocksResultV0
 }
@@ -57,7 +59,12 @@ func (h testHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		h.t.Logf("Upgrade: %v", err)
 		return
 	}
-	defer wsock.Close()
+
+	if h.CloseError {
+		wsock.Close()
+	} else {
+		defer wsock.Close()
+	}
 
 	for {
 
@@ -109,6 +116,44 @@ func TestShipClient_ConnectFail(t *testing.T) {
 	assert.Error(t, err, "dial tcp :9999: connect: connection refused")
 }
 
+func TestShipClient_ReadFromNormalClosedSocket(t *testing.T) {
+	handler := testHandler{t: t}
+
+	s := newServerWithHandler(t, &handler)
+	defer s.Close()
+
+	client := NewClient(23617231, NULL_BLOCK_NUMBER, false)
+	err := client.ConnectURL(*s.URL)
+	assert.NilError(t, err)
+	err = client.SendCloseMessage()
+	assert.NilError(t, err)
+	err = client.Read()
+	assert.Error(t, err, "shipclient - socket closed: websocket: close 1000 (normal)")
+
+	shErr, ok := err.(ShipClientError)
+	assert.Equal(t, true, ok, "Failed to cast error to ShipClientError")
+	assert.Equal(t, shErr.Type, ErrSockClosed)
+}
+
+func TestShipClient_ReadFromAbnormalClosedSocket(t *testing.T) {
+	handler := testHandler{t: t, CloseError: true}
+
+	s := newServerWithHandler(t, &handler)
+	defer s.Close()
+
+	client := NewClient(72367186, NULL_BLOCK_NUMBER, false)
+	err := client.ConnectURL(*s.URL)
+	assert.NilError(t, err)
+	err = client.SendCloseMessage()
+	assert.NilError(t, err)
+	err = client.Read()
+	assert.Error(t, err, "shipclient - socket closed: websocket: close 1006 (abnormal closure): unexpected EOF")
+
+	shErr, ok := err.(ShipClientError)
+	assert.Equal(t, true, ok, "Failed to cast error to ShipClientError")
+	assert.Equal(t, shErr.Type, ErrSockClosed)
+}
+
 func TestShipClient_ReadBlockMessages(t *testing.T) {
 	called := false