From 9cd54baca4d608564ad49e4bf260d3c070c7f20b Mon Sep 17 00:00:00 2001
From: Julien Salleyron <julien.salleyron@gmail.com>
Date: Thu, 22 Sep 2022 10:00:09 +0200
Subject: [PATCH] Optimize websocket headers handling

Co-authored-by: Kevin Pollet <pollet.kevin@gmail.com>
---
 pkg/server/service/proxy.go | 31 +++++++++++++++++++++----------
 1 file changed, 21 insertions(+), 10 deletions(-)

diff --git a/pkg/server/service/proxy.go b/pkg/server/service/proxy.go
index 93cd1c033..12ce95ce7 100644
--- a/pkg/server/service/proxy.go
+++ b/pkg/server/service/proxy.go
@@ -15,6 +15,7 @@ import (
 	ptypes "github.com/traefik/paerser/types"
 	"github.com/traefik/traefik/v2/pkg/config/dynamic"
 	"github.com/traefik/traefik/v2/pkg/log"
+	"golang.org/x/net/http/httpguts"
 )
 
 // StatusClientClosedRequest non-standard HTTP status code for client disconnection.
@@ -67,16 +68,18 @@ func buildProxy(passHostHeader *bool, responseForwarding *dynamic.ResponseForwar
 			// some servers need Sec-WebSocket-Key, Sec-WebSocket-Extensions, Sec-WebSocket-Accept,
 			// Sec-WebSocket-Protocol and Sec-WebSocket-Version to be case-sensitive.
 			// https://tools.ietf.org/html/rfc6455#page-20
-			outReq.Header["Sec-WebSocket-Key"] = outReq.Header["Sec-Websocket-Key"]
-			outReq.Header["Sec-WebSocket-Extensions"] = outReq.Header["Sec-Websocket-Extensions"]
-			outReq.Header["Sec-WebSocket-Accept"] = outReq.Header["Sec-Websocket-Accept"]
-			outReq.Header["Sec-WebSocket-Protocol"] = outReq.Header["Sec-Websocket-Protocol"]
-			outReq.Header["Sec-WebSocket-Version"] = outReq.Header["Sec-Websocket-Version"]
-			delete(outReq.Header, "Sec-Websocket-Key")
-			delete(outReq.Header, "Sec-Websocket-Extensions")
-			delete(outReq.Header, "Sec-Websocket-Accept")
-			delete(outReq.Header, "Sec-Websocket-Protocol")
-			delete(outReq.Header, "Sec-Websocket-Version")
+			if isWebSocketUpgrade(outReq) {
+				outReq.Header["Sec-WebSocket-Key"] = outReq.Header["Sec-Websocket-Key"]
+				outReq.Header["Sec-WebSocket-Extensions"] = outReq.Header["Sec-Websocket-Extensions"]
+				outReq.Header["Sec-WebSocket-Accept"] = outReq.Header["Sec-Websocket-Accept"]
+				outReq.Header["Sec-WebSocket-Protocol"] = outReq.Header["Sec-Websocket-Protocol"]
+				outReq.Header["Sec-WebSocket-Version"] = outReq.Header["Sec-Websocket-Version"]
+				delete(outReq.Header, "Sec-Websocket-Key")
+				delete(outReq.Header, "Sec-Websocket-Extensions")
+				delete(outReq.Header, "Sec-Websocket-Accept")
+				delete(outReq.Header, "Sec-Websocket-Protocol")
+				delete(outReq.Header, "Sec-Websocket-Version")
+			}
 		},
 		Transport:     roundTripper,
 		FlushInterval: time.Duration(flushInterval),
@@ -112,6 +115,14 @@ func buildProxy(passHostHeader *bool, responseForwarding *dynamic.ResponseForwar
 	return proxy, nil
 }
 
+func isWebSocketUpgrade(req *http.Request) bool {
+	if !httpguts.HeaderValuesContainsToken(req.Header["Connection"], "Upgrade") {
+		return false
+	}
+
+	return strings.EqualFold(req.Header.Get("Upgrade"), "websocket")
+}
+
 func statusText(statusCode int) string {
 	if statusCode == StatusClientClosedRequest {
 		return StatusClientClosedRequestText