fix WebSocket serialization (#1425)

Unit test was updated to assert valid encoding of data. Hacky approach
to fix the issue is not working. Replacing normal JSON serialization
with JSON RPC serialization is required, but there still is an issue.

<!--
Please read and fill out this form before submitting your PR.

Please make sure you have reviewed our contributors guide before
submitting your
first PR.
-->

## Overview

Draft for visibility. CC: @Manav-Aggarwal 

Relates to: #1268 

<!-- 
Please provide an explanation of the PR, including the appropriate
context,
background, goal, and rationale. If there is an issue with this
information,
please provide a tl;dr and link the issue. 
-->

## Checklist

<!-- 
Please complete the checklist to ensure that the PR is ready to be
reviewed.

IMPORTANT:
PRs should be left in Draft until the below checklist is completed.
-->

- [ ] New and updated code has appropriate documentation
- [ ] New and updated code has new and/or updated testing
- [ ] Required CI checks are passing
- [ ] Visual proof for any user facing features like CLI or
documentation updates
- [ ] Linked issues closed with keywords


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
	- Enhanced WebSocket connection handling in JSON-RPC services.
  
- **Bug Fixes**
	- Improved subscription response encoding in JSON-RPC.

- **Refactor**
	- Updated WebSocket-related code to improve RPC request management.

- **Tests**
	- Altered test cases to reflect new JSON unmarshalling approach.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Ganesha Upadhyaya <ganeshrvce@gmail.com>
Co-authored-by: Ganesha Upadhyaya <gupadhyaya@Ganeshas-MacBook-Pro-2.local>

Author: 3 people

rpc/json/handler.go

@@ -57,6 +57,9 @@ func (h *handler) serveJSONRPC(w http.ResponseWriter, r *http.Request) {
 func (h *handler) serveJSONRPCforWS(w http.ResponseWriter, r *http.Request, wsConn *wsConn) {
 	// Create a new codec request.
 	codecReq := h.codec.NewRequest(r)
+	if wsConn != nil {
+		wsConn.codecReq = codecReq
+	}
 	// Get service method to be called.
 	method, err := codecReq.Method()
 	if err != nil {

rpc/json/service.go

@@ -1,15 +1,19 @@
 package json
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
 	"net/http"
 	"reflect"
 	"time"
 
+	cmjson "github.com/cometbft/cometbft/libs/json"
+
 	rpcclient "github.com/cometbft/cometbft/rpc/client"
 	ctypes "github.com/cometbft/cometbft/rpc/core/types"
+	"github.com/gorilla/rpc/v2"
 	"github.com/gorilla/rpc/v2/json2"
 
 	"github.com/rollkit/rollkit/third_party/log"
@@ -101,12 +105,25 @@ func (s *service) Subscribe(req *http.Request, args *subscribeArgs, wsConn *wsCo
 	}
 
 	go func() {
+		var codecReq rpc.CodecRequest
+		if wsConn != nil {
+			codecReq = wsConn.codecReq
+		} else {
+			codecReq = json2.NewCodec().NewRequest(req)
+		}
+
 		for msg := range sub {
-			data, err := json.Marshal(msg.Data)
+			var raw json.RawMessage
+			raw, err = cmjson.Marshal(msg.Data)
+			btz := new(bytes.Buffer)
+			w := newResponseWriter(btz)
 			if err != nil {
-				s.logger.Error("failed to marshal response data", "error", err)
-				continue
+				codecReq.WriteError(w, http.StatusInternalServerError, err)
+				return
 			}
+			codecReq.WriteResponse(w, raw)
+
+			data := btz.Bytes()
 			if wsConn != nil {
 				wsConn.queue <- data
 			}

rpc/json/ws.go

@@ -5,15 +5,18 @@ import (
 	"io"
 	"net/http"
 
+	"github.com/gorilla/rpc/v2"
+
 	"github.com/gorilla/websocket"
 
 	"github.com/rollkit/rollkit/third_party/log"
 )
 
 type wsConn struct {
-	conn   *websocket.Conn
-	queue  chan []byte
-	logger log.Logger
+	conn     *websocket.Conn
+	codecReq rpc.CodecRequest
+	queue    chan []byte
+	logger   log.Logger
 }
 
 func (wsc *wsConn) sendLoop() {

rpc/json/ws_test.go

@@ -9,6 +9,9 @@ import (
 	"testing"
 	"time"
 
+	cmjson "github.com/cometbft/cometbft/libs/json"
+	"github.com/go-kit/kit/transport/http/jsonrpc"
+
 	"github.com/cometbft/cometbft/libs/log"
 	cmtypes "github.com/cometbft/cometbft/types"
 	"github.com/gorilla/rpc/v2/json2"
@@ -63,8 +66,11 @@ func TestWebSockets(t *testing.T) {
 	require.NoError(err)
 	assert.Equal(websocket.TextMessage, typ)
 	assert.NotEmpty(msg)
+	var jsrpcResp jsonrpc.Response
+	err = json.Unmarshal(msg, &jsrpcResp)
+	require.NoError(err)
 	var payload cmtypes.EventDataNewBlock
-	err = json.Unmarshal(msg, &payload)
+	err = cmjson.Unmarshal(jsrpcResp.Result, &payload)
 	require.NoError(err)
 	assert.NotNil(payload.ResultFinalizeBlock)
 	assert.NotNil(payload.Block)