From ad4721820bcf9081ce57055e1baaee959fc9a896 Mon Sep 17 00:00:00 2001
From: Toby Chui <tobychui@users.noreply.github.com>
Date: Mon, 30 Dec 2024 21:01:45 +0800
Subject: [PATCH] Added websocket header test and benchmark tool

---
 tools/benchmark/go.mod         |  3 ++
 tools/benchmark/main.go        | 51 ++++++++++++++++++++++++++
 tools/benchmark/server.go      | 42 +++++++++++++++++++++
 tools/websocket_echo/go.mod    |  5 +++
 tools/websocket_echo/go.sum    |  2 +
 tools/websocket_echo/main.go   | 44 ++++++++++++++++++++++
 tools/websocket_echo/test.html | 67 ++++++++++++++++++++++++++++++++++
 7 files changed, 214 insertions(+)
 create mode 100644 tools/benchmark/go.mod
 create mode 100644 tools/benchmark/main.go
 create mode 100644 tools/benchmark/server.go
 create mode 100644 tools/websocket_echo/go.mod
 create mode 100644 tools/websocket_echo/go.sum
 create mode 100644 tools/websocket_echo/main.go
 create mode 100644 tools/websocket_echo/test.html

diff --git a/tools/benchmark/go.mod b/tools/benchmark/go.mod
new file mode 100644
index 0000000..76393e1
--- /dev/null
+++ b/tools/benchmark/go.mod
@@ -0,0 +1,3 @@
+module imuslab.com/zoraxy/benchmark
+
+go 1.23.2
diff --git a/tools/benchmark/main.go b/tools/benchmark/main.go
new file mode 100644
index 0000000..4f55c9b
--- /dev/null
+++ b/tools/benchmark/main.go
@@ -0,0 +1,51 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+)
+
+var (
+	//Global variables
+	stopchan chan bool
+
+	//Runtime flags
+	benchmarkWebserverListeningPort int
+)
+
+func init() {
+	flag.IntVar(&benchmarkWebserverListeningPort, "port", 8123, "Port to listen on")
+	flag.Parse()
+}
+
+/* SIGTERM handler, do shutdown sequences before closing */
+func SetupCloseHandler() {
+	c := make(chan os.Signal, 2)
+	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
+	go func() {
+		<-c
+		//Stop all request loops
+		fmt.Println("Stopping request generators")
+		if stopchan != nil {
+			stopchan <- true
+		}
+
+		// Wait for all goroutines to finish
+		time.Sleep(1 * time.Second)
+		os.Exit(0)
+	}()
+}
+
+func main() {
+	//Setup the SIGTERM handler
+	SetupCloseHandler()
+	//Start the web server
+	fmt.Println("Starting web server on port", benchmarkWebserverListeningPort)
+	fmt.Println("In Zoraxy, point your test proxy rule to this server at the given port")
+	startWebServer()
+	select {}
+}
diff --git a/tools/benchmark/server.go b/tools/benchmark/server.go
new file mode 100644
index 0000000..dd7d13a
--- /dev/null
+++ b/tools/benchmark/server.go
@@ -0,0 +1,42 @@
+package main
+
+import (
+	"fmt"
+	"net/http"
+	"time"
+)
+
+// Start the web server for reciving test request
+// in Zoraxy, point test.localhost to this server at the given port in the start variables
+func startWebServer() {
+	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+		// Print the request details to console
+		fmt.Printf("Timestamp: %s\n", time.Now().Format(time.RFC1123))
+		fmt.Printf("Request type: %s\n", r.Method)
+		fmt.Printf("Payload size: %d bytes\n", r.ContentLength)
+		fmt.Printf("Request URI: %s\n", r.RequestURI)
+		fmt.Printf("User Agent: %s\n", r.UserAgent())
+		fmt.Printf("Remote Address: %s\n", r.RemoteAddr)
+		fmt.Println("----------------------------------------")
+
+		//Set header to text
+		w.Header().Set("Content-Type", "text/plain")
+		// Send response, print the request details to web page
+		w.Write([]byte("----------------------------------------\n"))
+		w.Write([]byte("Request type: " + r.Method + "\n"))
+		w.Write([]byte(fmt.Sprintf("Payload size: %d bytes\n", r.ContentLength)))
+		w.Write([]byte("Request URI: " + r.RequestURI + "\n"))
+		w.Write([]byte("User Agent: " + r.UserAgent() + "\n"))
+		w.Write([]byte("Remote Address: " + r.RemoteAddr + "\n"))
+		w.Write([]byte("----------------------------------------\n"))
+	})
+
+	go func() {
+		err := http.ListenAndServe(fmt.Sprintf(":%d", benchmarkWebserverListeningPort), nil)
+		if err != nil {
+			fmt.Printf("Failed to start server: %v\n", err)
+			stopchan <- true
+		}
+	}()
+
+}
diff --git a/tools/websocket_echo/go.mod b/tools/websocket_echo/go.mod
new file mode 100644
index 0000000..2e76ab3
--- /dev/null
+++ b/tools/websocket_echo/go.mod
@@ -0,0 +1,5 @@
+module aroz.org/zoraxy/websocket-echo
+
+go 1.23.2
+
+require github.com/gorilla/websocket v1.5.3
diff --git a/tools/websocket_echo/go.sum b/tools/websocket_echo/go.sum
new file mode 100644
index 0000000..25a9fc4
--- /dev/null
+++ b/tools/websocket_echo/go.sum
@@ -0,0 +1,2 @@
+github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
+github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
diff --git a/tools/websocket_echo/main.go b/tools/websocket_echo/main.go
new file mode 100644
index 0000000..c7f10d1
--- /dev/null
+++ b/tools/websocket_echo/main.go
@@ -0,0 +1,44 @@
+package main
+
+import (
+	"fmt"
+	"log"
+	"net/http"
+
+	"github.com/gorilla/websocket"
+)
+
+var upgrader = websocket.Upgrader{
+	CheckOrigin: func(r *http.Request) bool {
+		return true
+	},
+}
+
+func echo(w http.ResponseWriter, r *http.Request) {
+	conn, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		log.Println("Upgrade error:", err)
+		return
+	}
+	defer conn.Close()
+
+	for key, values := range r.Header {
+		for _, value := range values {
+			message := fmt.Sprintf("%s: %s", key, value)
+			if err := conn.WriteMessage(websocket.TextMessage, []byte(message)); err != nil {
+				log.Println("WriteMessage error:", err)
+				return
+			}
+		}
+	}
+
+	if err := conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")); err != nil {
+		log.Println("CloseMessage error:", err)
+		return
+	}
+}
+
+func main() {
+	http.HandleFunc("/echo", echo)
+	log.Fatal(http.ListenAndServe(":8888", nil))
+}
diff --git a/tools/websocket_echo/test.html b/tools/websocket_echo/test.html
new file mode 100644
index 0000000..86c3445
--- /dev/null
+++ b/tools/websocket_echo/test.html
@@ -0,0 +1,67 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>WebSocket Echo Test</title>
+</head>
+<body>
+    <h1>WebSocket Echo Test</h1>
+    <p>1. Go run main.go</p>
+    <p>2. Create a Zoraxy proxy rule (and add to hosts file) from ws.localhost (port 80) to 127.0.0.1:8888</p>
+    <p>3. Click the Connect button below to test if headers are correctly sent over</p>
+    <button id="connectBtn">Connect</button>
+    <button id="disconnectBtn" disabled>Disconnect</button>
+    <input type="text" id="messageInput" placeholder="Enter message">
+    <button id="sendBtn" disabled>Send</button>
+    <div id="output"></div>
+
+    <script>
+        let socket;
+        const connectBtn = document.getElementById('connectBtn');
+        const disconnectBtn = document.getElementById('disconnectBtn');
+        const sendBtn = document.getElementById('sendBtn');
+        const messageInput = document.getElementById('messageInput');
+        const output = document.getElementById('output');
+
+        connectBtn.addEventListener('click', () => {
+            output.innerHTML = '';
+            //socket = new WebSocket('ws://localhost:8888/echo');
+            socket = new WebSocket('ws://ws.localhost/echo');
+
+            socket.onopen = () => {
+                output.innerHTML += '<p>Connected to WebSocket server</p>';
+                connectBtn.disabled = true;
+                disconnectBtn.disabled = false;
+                sendBtn.disabled = false;
+            };
+
+            socket.onmessage = (event) => {
+                output.innerHTML += `<p>Received: ${event.data}</p>`;
+            };
+
+            socket.onclose = () => {
+                output.innerHTML += '<p>Disconnected from WebSocket server</p>';
+                connectBtn.disabled = false;
+                disconnectBtn.disabled = true;
+                sendBtn.disabled = true;
+            };
+
+            socket.onerror = (error) => {
+                output.innerHTML += `<p>Error: ${error.message}</p>`;
+            };
+        });
+
+        disconnectBtn.addEventListener('click', () => {
+            socket.close();
+        });
+
+        sendBtn.addEventListener('click', () => {
+            const message = messageInput.value;
+            socket.send(message);
+            output.innerHTML += `<p>Sent: ${message}</p>`;
+            messageInput.value = '';
+        });
+    </script>
+</body>
+</html>
\ No newline at end of file