From 2a252795c9c657ef11de0e12e3e02940da49ca2c Mon Sep 17 00:00:00 2001
From: kaituo <kaituo@google.com>
Date: Fri, 25 Jul 2025 17:20:57 +0000
Subject: [PATCH] fix: Redact websocket chunks when record and replay them from
 logs

---
 internal/record/recording_https_proxy.go | 11 ++++++-----
 internal/replay/replay_http_server.go    |  6 +++---
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/internal/record/recording_https_proxy.go b/internal/record/recording_https_proxy.go
index b0c180a..105188a 100644
--- a/internal/record/recording_https_proxy.go
+++ b/internal/record/recording_https_proxy.go
@@ -251,8 +251,8 @@ func (r *RecordingHTTPSProxy) proxyWebsocket(w http.ResponseWriter, req *http.Re
 	c := make(chan []byte)
 	quit := make(chan int)
 
-	go pumpWebsocket(clientConn, conn, c, quit, ">")
-	go pumpWebsocket(conn, clientConn, c, quit, "<")
+	go r.pumpWebsocket(clientConn, conn, c, quit, ">")
+	go r.pumpWebsocket(conn, clientConn, c, quit, "<")
 
 	recordPath := filepath.Join(r.recordingDir, fileName+".websocket.log")
 	f, err := os.Create(recordPath)
@@ -280,7 +280,7 @@ func (r *RecordingHTTPSProxy) proxyWebsocket(w http.ResponseWriter, req *http.Re
 	}
 }
 
-func pumpWebsocket(src, dst *websocket.Conn, c chan []byte, quit chan int, prepend string) {
+func (r *RecordingHTTPSProxy) pumpWebsocket(src, dst *websocket.Conn, c chan []byte, quit chan int, prepend string) {
 	for {
 		msgType, buf, err := src.ReadMessage()
 		if err != nil {
@@ -293,8 +293,9 @@ func pumpWebsocket(src, dst *websocket.Conn, c chan []byte, quit chan int, prepe
 			return
 		}
 		buf = append(buf, '\n')
-		prefix := fmt.Sprintf("%s%d ", prepend, len(buf))
-		c <- append([]byte(prefix), buf...)
+		redactedBuf := r.redactor.Bytes(buf)
+		prefix := fmt.Sprintf("%s%d ", prepend, len(redactedBuf))
+		c <- append([]byte(prefix), redactedBuf...)
 		err = dst.WriteMessage(msgType, buf)
 		if err != nil {
 			fmt.Printf("Error writing to websocket: %v\n", err)
diff --git a/internal/replay/replay_http_server.go b/internal/replay/replay_http_server.go
index f799915..1eff395 100644
--- a/internal/replay/replay_http_server.go
+++ b/internal/replay/replay_http_server.go
@@ -230,7 +230,7 @@ func (r *ReplayHTTPServer) proxyWebsocket(w http.ResponseWriter, req *http.Reque
 		return
 	}
 	defer clientConn.Close()
-	replayWebsocket(clientConn, chunks)
+	r.replayWebsocket(clientConn, chunks)
 }
 
 func (r *ReplayHTTPServer) loadWebsocketChunks(fileName string) ([]string, error) {
@@ -273,11 +273,11 @@ func (r *ReplayHTTPServer) loadWebsocketChunks(fileName string) ([]string, error
 	return chunks, nil
 }
 
-func replayWebsocket(conn *websocket.Conn, chunks []string) {
+func (r *ReplayHTTPServer) replayWebsocket(conn *websocket.Conn, chunks []string) {
 	for _, chunk := range chunks {
 		if strings.HasPrefix(chunk, ">") {
 			_, buf, err := conn.ReadMessage()
-			reqChunk := string(buf)
+			reqChunk := r.redactor.String(string(buf))
 			if err != nil {
 				fmt.Printf("Error reading from websocket: %v\n", err)
 				return