From a024ad53aecbdda02a9d474ce8d01caeadf9f3fe Mon Sep 17 00:00:00 2001
From: TaltaM <68383301+TaltaM@users.noreply.github.com>
Date: Fri, 9 Jun 2023 01:24:33 +0200
Subject: [PATCH] websocket: Check for slice size to avoid panic when
 classifying empty binary message (#1219)

* Check for slice size to avoid panic when classifying empty resp

* Simplify according to suggestion
---
 exchanges/stream/websocket_connection.go | 35 +++++++-----------------
 exchanges/stream/websocket_test.go       |  7 ++++-
 2 files changed, 16 insertions(+), 26 deletions(-)

diff --git a/exchanges/stream/websocket_connection.go b/exchanges/stream/websocket_connection.go
index 0c372599..1601c7a1 100644
--- a/exchanges/stream/websocket_connection.go
+++ b/exchanges/stream/websocket_connection.go
@@ -268,36 +268,21 @@ func (w *WebsocketConnection) ReadMessage() Response {
 
 // parseBinaryResponse parses a websocket binary response into a usable byte array
 func (w *WebsocketConnection) parseBinaryResponse(resp []byte) ([]byte, error) {
-	var standardMessage []byte
+	var reader io.ReadCloser
 	var err error
-	// Detect GZIP
-	if resp[0] == 31 && resp[1] == 139 {
-		b := bytes.NewReader(resp)
-		var gReader *gzip.Reader
-		gReader, err = gzip.NewReader(b)
+	if len(resp) >= 2 && resp[0] == 31 && resp[1] == 139 { // Detect GZIP
+		reader, err = gzip.NewReader(bytes.NewReader(resp))
 		if err != nil {
-			return standardMessage, err
-		}
-		standardMessage, err = io.ReadAll(gReader)
-		if err != nil {
-			return standardMessage, err
-		}
-		err = gReader.Close()
-		if err != nil {
-			return standardMessage, err
+			return nil, err
 		}
 	} else {
-		reader := flate.NewReader(bytes.NewReader(resp))
-		standardMessage, err = io.ReadAll(reader)
-		if err != nil {
-			return standardMessage, err
-		}
-		err = reader.Close()
-		if err != nil {
-			return standardMessage, err
-		}
+		reader = flate.NewReader(bytes.NewReader(resp))
 	}
-	return standardMessage, nil
+	standardMessage, err := io.ReadAll(reader)
+	if err != nil {
+		return nil, err
+	}
+	return standardMessage, reader.Close()
 }
 
 // GenerateMessageID Creates a messageID to checkout
diff --git a/exchanges/stream/websocket_test.go b/exchanges/stream/websocket_test.go
index 5cdab00c..5d851b61 100644
--- a/exchanges/stream/websocket_test.go
+++ b/exchanges/stream/websocket_test.go
@@ -908,7 +908,12 @@ func TestParseBinaryResponse(t *testing.T) {
 		t.Error(err3)
 	}
 	if !strings.EqualFold(string(resp2), "hello") {
-		t.Errorf("GZip conversion failed. Received: '%v', Expected: 'hello'", string(resp2))
+		t.Errorf("Deflate conversion failed. Received: '%v', Expected: 'hello'", string(resp2))
+	}
+
+	_, err4 := wc.parseBinaryResponse([]byte{})
+	if err4 == nil || err4.Error() != "unexpected EOF" {
+		t.Error("Expected error 'unexpected EOF'")
 	}
 }