From 3f48703f265a36de5df80b1c94871f56704480a5 Mon Sep 17 00:00:00 2001
From: Felipe Andrade <felipe@oplabs.co>
Date: Thu, 27 Jul 2023 13:29:44 -0700
Subject: [PATCH] testz

---
 proxyd/proxyd/backend.go                   | 16 +------
 proxyd/proxyd/integration_tests/ws_test.go | 51 ++++++++++++++++++++--
 2 files changed, 49 insertions(+), 18 deletions(-)

diff --git a/proxyd/proxyd/backend.go b/proxyd/proxyd/backend.go
index 4314848..99f8b06 100644
--- a/proxyd/proxyd/backend.go
+++ b/proxyd/proxyd/backend.go
@@ -884,13 +884,6 @@ func (w *WSProxier) Proxy(ctx context.Context) error {
 
 func (w *WSProxier) clientPump(ctx context.Context, errC chan error) {
 	for {
-		err := w.clientConn.SetReadDeadline(time.Now().Add(w.readTimeout))
-		if err != nil {
-			log.Error("ws client read timeout", "err", err)
-			errC <- err
-			return
-		}
-
 		// Block until we get a message.
 		msgType, msg, err := w.clientConn.ReadMessage()
 		if err != nil {
@@ -974,13 +967,6 @@ func (w *WSProxier) clientPump(ctx context.Context, errC chan error) {
 
 func (w *WSProxier) backendPump(ctx context.Context, errC chan error) {
 	for {
-		err := w.backendConn.SetReadDeadline(time.Now().Add(w.readTimeout))
-		if err != nil {
-			log.Error("ws backend read timeout", "err", err)
-			errC <- err
-			return
-		}
-
 		// Block until we get a message.
 		msgType, msg, err := w.backendConn.ReadMessage()
 		if err != nil {
@@ -1085,7 +1071,7 @@ func (w *WSProxier) writeBackendConn(msgType int, msg []byte) error {
 		log.Error("ws backend write timeout", "err", err)
 		return err
 	}
-	err := w.writeBackendConn(msgType, msg)
+	err := w.backendConn.WriteMessage(msgType, msg)
 	return err
 }
 
diff --git a/proxyd/proxyd/integration_tests/ws_test.go b/proxyd/proxyd/integration_tests/ws_test.go
index c0907fe..fdaae0b 100644
--- a/proxyd/proxyd/integration_tests/ws_test.go
+++ b/proxyd/proxyd/integration_tests/ws_test.go
@@ -2,16 +2,17 @@ package integration_tests
 
 import (
 	"os"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
 
-	"github.com/ethereum/go-ethereum/log"
-
 	"github.com/ethereum-optimism/optimism/proxyd"
+	"github.com/ethereum/go-ethereum/log"
 	"github.com/gorilla/websocket"
 	"github.com/stretchr/testify/require"
+	"github.com/syndtr/goleveldb/leveldb/opt"
 )
 
 // TestConcurrentWSPanic tests for a panic in the websocket proxy
@@ -201,7 +202,7 @@ func TestWS(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			timeout := time.NewTicker(30 * time.Second)
+			timeout := time.NewTicker(10 * time.Second)
 			doneCh := make(chan struct{}, 1)
 			backendHdlr.SetMsgCB(func(conn *websocket.Conn, msgType int, data []byte) {
 				require.NoError(t, conn.WriteMessage(websocket.TextMessage, []byte(tt.backendRes)))
@@ -270,3 +271,47 @@ func TestWSClientClosure(t *testing.T) {
 		})
 	}
 }
+
+func TestWSClientExceedReadLimit(t *testing.T) {
+	backendHdlr := new(backendHandler)
+	clientHdlr := new(clientHandler)
+
+	backend := NewMockWSBackend(nil, func(conn *websocket.Conn, msgType int, data []byte) {
+		backendHdlr.MsgCB(conn, msgType, data)
+	}, func(conn *websocket.Conn, err error) {
+		backendHdlr.CloseCB(conn, err)
+	})
+	defer backend.Close()
+
+	require.NoError(t, os.Setenv("GOOD_BACKEND_RPC_URL", backend.URL()))
+
+	config := ReadConfig("ws")
+	_, shutdown, err := proxyd.Start(config)
+	require.NoError(t, err)
+	defer shutdown()
+
+	client, err := NewProxydWSClient("ws://127.0.0.1:8546", func(msgType int, data []byte) {
+		clientHdlr.MsgCB(msgType, data)
+	}, nil)
+	require.NoError(t, err)
+
+	closed := false
+	originalHandler := client.conn.CloseHandler()
+	client.conn.SetCloseHandler(func(code int, text string) error {
+		closed = true
+		return originalHandler(code, text)
+	})
+
+	backendHdlr.SetMsgCB(func(conn *websocket.Conn, msgType int, data []byte) {
+		t.Fatalf("backend should not get the large message")
+	})
+
+	clientReq := "{\"id\": 1, \"method\": \"eth_subscribe\", \"params\": [\"" + strings.Repeat("barf", 256*opt.KiB+1) + "\"]}"
+	err = client.WriteMessage(
+		websocket.TextMessage,
+		[]byte(clientReq),
+	)
+	require.Error(t, err)
+	require.True(t, closed)
+
+}