fix(dot/subscription): check websocket message from untrusted data

dot/rpc/http.go

@@ -233,7 +233,7 @@ func (h *HTTPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	wsc := NewWSConn(ws, h.serverConfig)
 	h.wsConns = append(h.wsConns, wsc)
 
-	go wsc.HandleComm()
+	go wsc.HandleConn()
 }
 
 // NewWSConn to create new WebSocket Connection struct

dot/rpc/subscription/websocket.go

@@ -52,28 +52,27 @@ type WSConn struct {
 }
 
 // readWebsocketMessage will read and parse the message data to a string->interface{} data
-func (c *WSConn) readWebsocketMessage() ([]byte, *websocketMessage, error) {
-	_, mbytes, err := c.Wsconn.ReadMessage()
+func (c *WSConn) readWebsocketMessage() (bytes []byte, msg websocketMessage, err error) {
+	_, bytes, err = c.Wsconn.ReadMessage()
 	if err != nil {
 		logger.Debugf("websocket failed to read message: %s", err)
-		return nil, nil, errCannotReadFromWebsocket
+		return bytes, msg, errCannotReadFromWebsocket
 	}
 
-	logger.Tracef("websocket message received: %s", string(mbytes))
+	logger.Tracef("websocket message received: %s", string(bytes))
 
-	msg := new(websocketMessage)
-	err = json.Unmarshal(mbytes, &msg)
+	err = json.Unmarshal(bytes, &msg)
 
 	if err != nil {
 		logger.Debugf("websocket failed to unmarshal request message: %s", err)
-		return nil, nil, errCannotUnmarshalMessage
+		return bytes, msg, errCannotUnmarshalMessage
 	}
 
-	return mbytes, msg, nil
+	return bytes, msg, nil
 }
 
-//HandleComm handles messages received on websocket connections
-func (c *WSConn) HandleComm() {
+//HandleConn handles messages received on websocket connections
+func (c *WSConn) HandleConn() {
 	for {
 		mbytes, msg, err := c.readWebsocketMessage()
 		if errors.Is(err, errCannotReadFromWebsocket) {

dot/rpc/subscription/websocket_test.go

@@ -26,7 +26,7 @@ func TestWSConn_CheckWebsocketInvalidData(t *testing.T) {
 	wsconn.Subscriptions = make(map[uint32]Listener)
 	defer cancel()
 
-	go wsconn.HandleComm()
+	go wsconn.HandleConn()
 
 	tests := []struct {
 		sentMessage []byte
@@ -68,12 +68,12 @@ func TestWSConn_CheckWebsocketInvalidData(t *testing.T) {
 	}
 }
 
-func TestWSConn_HandleComm(t *testing.T) {
+func TestWSConn_HandleConn(t *testing.T) {
 	wsconn, c, cancel := setupWSConn(t)
 	wsconn.Subscriptions = make(map[uint32]Listener)
 	defer cancel()
 
-	go wsconn.HandleComm()
+	go wsconn.HandleConn()
 	time.Sleep(time.Second * 2)
 
 	// test storageChangeListener
@@ -341,7 +341,7 @@ func TestSubscribeAllHeads(t *testing.T) {
 	wsconn.Subscriptions = make(map[uint32]Listener)
 	defer cancel()
 
-	go wsconn.HandleComm()
+	go wsconn.HandleConn()
 	time.Sleep(time.Second * 2)
 
 	_, err := wsconn.initAllBlocksListerner(1, nil)