properly increase read limit

websocket.go

@@ -34,7 +34,7 @@ const (
 	wsPingInterval     = 60 * time.Second
 	wsPingWriteTimeout = 5 * time.Second
 	wsPongTimeout      = 30 * time.Second
-	wsMessageSizeLimit = 32 * 1024 * 1024
+	wsMessageSizeLimit = 128 * 1024 * 1024
 )
 
 // WebsocketHandler returns a handler that serves JSON-RPC to WebSocket connections.
@@ -154,11 +154,11 @@ func heartbeat(ctx context.Context, c *websocket.Conn, d time.Duration) {
 }
 
 func newWebsocketCodec(ctx context.Context, c *websocket.Conn, host string, req http.Header) ServerCodec {
-	c.SetReadLimit(wsMessageSizeLimit)
 	jsonWriter := func(v any) error {
 		return wsjson.Write(context.Background(), c, v)
 	}
 	jsonReader := func(v any) error {
+		c.SetReadLimit(wsMessageSizeLimit)
 		return wsjson.Read(context.Background(), c, v)
 	}
 	conn := websocket.NetConn(ctx, c, websocket.MessageText)

wsjson/wsjson.go

@@ -70,18 +70,17 @@ func Write(ctx context.Context, c *websocket.Conn, v interface{}) error {
 }
 
 func write(ctx context.Context, c *websocket.Conn, v interface{}) (err error) {
-
 	w, err := c.Writer(ctx, websocket.MessageText)
 	if err != nil {
 		return err
 	}
-
 	// json.Marshal cannot reuse buffers between calls as it has to return
 	// a copy of the byte slice but Encoder does as it directly writes to w.
-	err = jzon.NewEncoder(w).Encode(v)
+	st := jsoniter.NewStream(jzon, w, 1024)
+	st.WriteVal(v)
+	err = st.Flush()
 	if err != nil {
 		return fmt.Errorf("failed to marshal JSON: %w", err)
 	}
-
 	return w.Close()
 }