@@ -2,14 +2,18 @@ package environment
22
33import (
44 "bytes"
5+ "context"
56 "encoding/json"
67 "io"
78 "net/http"
89 "net/http/httptest"
910 "net/url"
11+ "regexp"
12+ "time"
1013
1114 espressoTaggedBase64 "github.com/EspressoSystems/espresso-network/sdks/go/tagged-base64"
1215 espressoCommon "github.com/EspressoSystems/espresso-network/sdks/go/types"
16+ "github.com/coder/websocket"
1317 "github.com/ethereum-optimism/optimism/op-batcher/batcher"
1418 "github.com/ethereum-optimism/optimism/op-e2e/system/e2esys"
1519)
@@ -88,6 +92,13 @@ type proxyRequest struct {
8892
8993// ServeHTTP implements http.Handler
9094func (p * proxyRequest ) ServeHTTP (w http.ResponseWriter , r * http.Request ) {
95+ // Check if this is a websocket stream request
96+ if isWebSocketStreamRequest (r ) {
97+ p .proxyWebSocket (w , r )
98+ return
99+ }
100+
101+ // Handle regular HTTP requests
91102 defer r .Body .Close ()
92103 buf := new (bytes.Buffer )
93104 if _ , err := io .Copy (buf , r .Body ); err != nil && err != io .EOF {
@@ -135,6 +146,82 @@ func (p *proxyRequest) ServeHTTP(w http.ResponseWriter, r *http.Request) {
135146 }
136147}
137148
149+ // proxyWebSocket handles websocket upgrade and proxying
150+ func (p * proxyRequest ) proxyWebSocket (w http.ResponseWriter , r * http.Request ) {
151+ // Accept the websocket connection from the client
152+ clientConn , err := websocket .Accept (w , r , & websocket.AcceptOptions {
153+ InsecureSkipVerify : true ,
154+ })
155+ if err != nil {
156+ http .Error (w , err .Error (), http .StatusInternalServerError )
157+ return
158+ }
159+ defer clientConn .Close (websocket .StatusInternalError , "proxy error" )
160+
161+ // Create websocket URL for the backend
162+ backendURL := p .baseURL
163+ if backendURL .Scheme == "https" {
164+ backendURL .Scheme = "wss"
165+ } else {
166+ backendURL .Scheme = "ws"
167+ }
168+ backendURL .Path = r .URL .Path
169+ backendURL .RawQuery = r .URL .RawQuery
170+
171+ // Connect to the backend websocket
172+ ctx , cancel := context .WithTimeout (context .Background (), 30 * time .Second )
173+ defer cancel ()
174+
175+ //nolint:bodyclose // Not applicable to coder/websocket. From the websocket.Dial docs: "You never need to close resp.Body yourself."
176+ backendConn , _ , err := websocket .Dial (ctx , backendURL .String (), & websocket.DialOptions {})
177+ if err != nil {
178+ clientConn .Close (websocket .StatusInternalError , "backend connection failed" )
179+ return
180+ }
181+ defer backendConn .Close (websocket .StatusNormalClosure , "" )
182+
183+ // Proxy messages bidirectionally
184+ ctx , cancel = context .WithCancel (context .Background ())
185+ defer cancel ()
186+
187+ // Client to backend
188+ go func () {
189+ defer cancel ()
190+ for {
191+ msgType , data , err := clientConn .Read (ctx )
192+ if err != nil {
193+ return
194+ }
195+ err = backendConn .Write (ctx , msgType , data )
196+ if err != nil {
197+ return
198+ }
199+ }
200+ }()
201+
202+ // Backend to client
203+ go func () {
204+ defer cancel ()
205+ for {
206+ msgType , data , err := backendConn .Read (ctx )
207+ if err != nil {
208+ return
209+ }
210+ err = clientConn .Write (ctx , msgType , data )
211+ if err != nil {
212+ return
213+ }
214+ }
215+ }()
216+
217+ // Wait until context is cancelled (one of the goroutines finished)
218+ <- ctx .Done ()
219+
220+ // Close connections gracefully
221+ clientConn .Close (websocket .StatusNormalClosure , "" )
222+ backendConn .Close (websocket .StatusNormalClosure , "" )
223+ }
224+
138225// fakeSubmitTransactionSuccess is a simple HTTP handler that simulates a
139226// successful transaction submission by returning a fake commit hash.
140227type fakeSubmitTransactionSuccess struct {}
@@ -304,6 +391,13 @@ func isSubmitTransactionRequest(r *http.Request) bool {
304391 requestMatchesPath (r , http .MethodPost , stringEquals ("/v0/submit/submit" ))
305392}
306393
394+ // isWebSocketStreamRequest checks if the request is a websocket request
395+ // matching the pattern "vN/stream/*" where N is an integer.
396+ func isWebSocketStreamRequest (r * http.Request ) bool {
397+ matched , _ := regexp .MatchString (`/stream/` , r .URL .Path )
398+ return matched
399+ }
400+
307401// DecideHowToHandleRequest implements InterceptHandlerDecider
308402func (d * randomRollFakeSubmitTransactionSuccess ) DecideHowToHandleRequest (w http.ResponseWriter , r * http.Request ) InterceptHandleDecision {
309403 if isSubmitTransactionRequest (r ) {
0 commit comments