From 3b43c1507473179ed03d1cc904c9f7b740262a51 Mon Sep 17 00:00:00 2001
From: Para Dox <armhole.limeade_0l@icloud.com>
Date: Thu, 29 May 2025 01:08:24 +0700
Subject: [PATCH] more features

---
 benchmark-proxy/main.go | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/benchmark-proxy/main.go b/benchmark-proxy/main.go
index 7a4cbbe3..2ef872aa 100644
--- a/benchmark-proxy/main.go
+++ b/benchmark-proxy/main.go
@@ -965,6 +965,9 @@ func (sc *StatsCollector) printSummary() {
 	sc.totalRequests = 0
 	sc.totalWsConnections = 0
 
+	// Reset WebSocket connections to prevent memory leak
+	sc.wsConnections = sc.wsConnections[:0]
+
 	// Reset the interval start time for the next interval
 	sc.intervalStartTime = time.Now()
 }
@@ -1273,10 +1276,19 @@ func main() {
 func handleRequest(w http.ResponseWriter, r *http.Request, backends []Backend, client *http.Client, enableDetailedLogs bool, statsCollector *StatsCollector) {
 	startTime := time.Now()
 
+	// Limit request body size to 10MB to prevent memory exhaustion
+	const maxBodySize = 10 * 1024 * 1024 // 10MB
+	r.Body = http.MaxBytesReader(w, r.Body, maxBodySize)
+
 	// Read the entire request body
 	body, err := io.ReadAll(r.Body)
 	if err != nil {
-		http.Error(w, "Error reading request body", http.StatusBadRequest)
+		// Check if the error is due to body size limit
+		if strings.Contains(err.Error(), "request body too large") {
+			http.Error(w, "Request body too large (max 10MB)", http.StatusRequestEntityTooLarge)
+		} else {
+			http.Error(w, "Error reading request body", http.StatusBadRequest)
+		}
 		return
 	}
 	defer r.Body.Close()