From 28d4373f673ccaf16db5b3690e87855067df1b81 Mon Sep 17 00:00:00 2001
From: yusing <yusing@6uo.me>
Date: Wed, 4 Dec 2024 06:09:52 +0800
Subject: [PATCH] fix potential issues with some websocket upstream servers

---
 internal/net/http/reverse_proxy_mod.go | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/internal/net/http/reverse_proxy_mod.go b/internal/net/http/reverse_proxy_mod.go
index 7070525..ed7a310 100644
--- a/internal/net/http/reverse_proxy_mod.go
+++ b/internal/net/http/reverse_proxy_mod.go
@@ -375,6 +375,10 @@ func (p *ReverseProxy) handler(rw http.ResponseWriter, req *http.Request) {
 	if reqUpType != "" {
 		outreq.Header.Set("Connection", "Upgrade")
 		outreq.Header.Set("Upgrade", reqUpType)
+
+		if strings.EqualFold(reqUpType, "websocket") {
+			cleanWebsocketHeaders(outreq)
+		}
 	}
 
 	if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil {
@@ -538,6 +542,25 @@ func RemoveHopByHopHeaders(h http.Header) {
 	}
 }
 
+// reference: https://github.com/traefik/traefik/blob/master/pkg/proxy/httputil/proxy.go
+// https://tools.ietf.org/html/rfc6455#page-20
+func cleanWebsocketHeaders(req *http.Request) {
+	req.Header["Sec-WebSocket-Key"] = req.Header["Sec-Websocket-Key"]
+	delete(req.Header, "Sec-Websocket-Key")
+
+	req.Header["Sec-WebSocket-Extensions"] = req.Header["Sec-Websocket-Extensions"]
+	delete(req.Header, "Sec-Websocket-Extensions")
+
+	req.Header["Sec-WebSocket-Accept"] = req.Header["Sec-Websocket-Accept"]
+	delete(req.Header, "Sec-Websocket-Accept")
+
+	req.Header["Sec-WebSocket-Protocol"] = req.Header["Sec-Websocket-Protocol"]
+	delete(req.Header, "Sec-Websocket-Protocol")
+
+	req.Header["Sec-WebSocket-Version"] = req.Header["Sec-Websocket-Version"]
+	delete(req.Header, "Sec-Websocket-Version")
+}
+
 func (p *ReverseProxy) handleUpgradeResponse(rw http.ResponseWriter, req *http.Request, res *http.Response) {
 	reqUpType := UpgradeType(req.Header)
 	resUpType := UpgradeType(res.Header)