Added manual toggle for websocket header copy

- Added setting for toggling websocket header copy
- Added close connection in TLS mode

src/mod/dynamicproxy/Server.go

@@ -222,10 +222,12 @@ func (h *ProxyHandler) handleRootRouting(w http.ResponseWriter, r *http.Request)
 		h.Parent.logRequest(r, false, 444, "root-noresponse", domainOnly)
 		hijacker, ok := w.(http.Hijacker)
 		if !ok {
+			w.Header().Set("Connection", "close")
 			return
 		}
 		conn, _, err := hijacker.Hijack()
 		if err != nil {
+			w.Header().Set("Connection", "close")
 			return
 		}
 		conn.Close()

src/mod/dynamicproxy/default.go

@@ -42,23 +42,24 @@ func GetDefaultHeaderRewriteRules() *HeaderRewriteRules {
 func GetDefaultProxyEndpoint() ProxyEndpoint {
 	randomPrefix := uuid.New().String()
 	return ProxyEndpoint{
-		ProxyType:              ProxyTypeHost,
-		RootOrMatchingDomain:   randomPrefix + ".internal",
-		MatchingDomainAlias:    []string{},
-		ActiveOrigins:          []*loadbalance.Upstream{},
-		InactiveOrigins:        []*loadbalance.Upstream{},
-		UseStickySession:       false,
-		UseActiveLoadBalance:   false,
-		Disabled:               false,
-		BypassGlobalTLS:        false,
-		VirtualDirectories:     []*VirtualDirectoryEndpoint{},
-		HeaderRewriteRules:     GetDefaultHeaderRewriteRules(),
-		AuthenticationProvider: GetDefaultAuthenticationProvider(),
-		RequireRateLimit:       false,
-		RateLimit:              0,
-		DisableUptimeMonitor:   false,
-		AccessFilterUUID:       "default",
-		DefaultSiteOption:      DefaultSite_InternalStaticWebServer,
-		DefaultSiteValue:       "",
+		ProxyType:                    ProxyTypeHost,
+		RootOrMatchingDomain:         randomPrefix + ".internal",
+		MatchingDomainAlias:          []string{},
+		ActiveOrigins:                []*loadbalance.Upstream{},
+		InactiveOrigins:              []*loadbalance.Upstream{},
+		UseStickySession:             false,
+		UseActiveLoadBalance:         false,
+		Disabled:                     false,
+		BypassGlobalTLS:              false,
+		VirtualDirectories:           []*VirtualDirectoryEndpoint{},
+		HeaderRewriteRules:           GetDefaultHeaderRewriteRules(),
+		EnableWebsocketCustomHeaders: false,
+		AuthenticationProvider:       GetDefaultAuthenticationProvider(),
+		RequireRateLimit:             false,
+		RateLimit:                    0,
+		DisableUptimeMonitor:         false,
+		AccessFilterUUID:             "default",
+		DefaultSiteOption:            DefaultSite_InternalStaticWebServer,
+		DefaultSiteValue:             "",
 	}
 }

src/mod/dynamicproxy/proxyRequestHandler.go

@@ -152,7 +152,7 @@ func (h *ProxyHandler) hostRequest(w http.ResponseWriter, r *http.Request, targe
 		wspHandler := websocketproxy.NewProxy(u, websocketproxy.Options{
 			SkipTLSValidation:  selectedUpstream.SkipCertValidations,
 			SkipOriginCheck:    selectedUpstream.SkipWebSocketOriginCheck,
-			CopyAllHeaders:     domainsniff.RequireWebsocketHeaderCopy(r),
+			CopyAllHeaders:     target.EnableWebsocketCustomHeaders,
 			UserDefinedHeaders: target.HeaderRewriteRules.UserDefinedHeaders,
 			Logger:             h.Parent.Option.Logger,
 		})
@@ -232,11 +232,16 @@ func (h *ProxyHandler) vdirRequest(w http.ResponseWriter, r *http.Request, targe
 		if target.RequireTLS {
 			u, _ = url.Parse("wss://" + wsRedirectionEndpoint + r.URL.String())
 		}
+
+		if target.parent.HeaderRewriteRules != nil {
+			target.parent.HeaderRewriteRules = GetDefaultHeaderRewriteRules()
+		}
+
 		h.Parent.logRequest(r, true, 101, "vdir-websocket", target.Domain)
 		wspHandler := websocketproxy.NewProxy(u, websocketproxy.Options{
 			SkipTLSValidation:  target.SkipCertValidations,
-			SkipOriginCheck:    true,                                      //You should not use websocket via virtual directory. But keep this to true for compatibility
-			CopyAllHeaders:     domainsniff.RequireWebsocketHeaderCopy(r), //Left this as default to prevent nginx user setting / as vdir
+			SkipOriginCheck:    target.parent.EnableWebsocketCustomHeaders, //You should not use websocket via virtual directory. But keep this to true for compatibility
+			CopyAllHeaders:     domainsniff.RequireWebsocketHeaderCopy(r),  //Left this as default to prevent nginx user setting / as vdir
 			UserDefinedHeaders: target.parent.HeaderRewriteRules.UserDefinedHeaders,
 			Logger:             h.Parent.Option.Logger,
 		})

src/mod/dynamicproxy/typedef.go

@@ -1,5 +1,12 @@
 package dynamicproxy
 
+/*
+	typdef.go
+
+	This script handle the type definition for dynamic proxy and endpoints
+
+	If you are looking for the default object initailization, please refer to default.go
+*/
 import (
 	_ "embed"
 	"net"
@@ -165,7 +172,8 @@ type ProxyEndpoint struct {
 	VirtualDirectories []*VirtualDirectoryEndpoint
 
 	//Custom Headers
-	HeaderRewriteRules *HeaderRewriteRules
+	HeaderRewriteRules           *HeaderRewriteRules
+	EnableWebsocketCustomHeaders bool //Enable custom headers for websocket connections as well (default only http reqiests)
 
 	//Authentication
 	AuthenticationProvider *AuthenticationProvider

src/mod/websocketproxy/websocketproxy.go

@@ -172,6 +172,11 @@ func (w *WebsocketProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	if req.Host != "" {
 		requestHeader.Set("Host", req.Host)
 	}
+	if userAgent := req.Header.Get("User-Agent"); userAgent != "" {
+		requestHeader.Set("User-Agent", userAgent)
+	} else {
+		requestHeader.Set("User-Agent", "zoraxy-wsproxy/1.1")
+	}
 
 	// Pass X-Forwarded-For headers too, code below is a part of
 	// httputil.ReverseProxy. See http://en.wikipedia.org/wiki/X-Forwarded-For
@@ -195,19 +200,29 @@ func (w *WebsocketProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 		requestHeader.Set("X-Forwarded-Proto", "https")
 	}
 
-	// Enable the director to copy any additional headers it desires for
-	// forwarding to the remote server.
-	if w.Director != nil {
-		w.Director(req, requestHeader)
-	}
-
 	// Replace header variables and copy user-defined headers
-	rewrittenUserDefinedHeaders := rewrite.PopulateRequestHeaderVariables(req, w.Options.UserDefinedHeaders)
-	upstreamHeaders, _ := rewrite.SplitUpDownStreamHeaders(&rewrite.HeaderRewriteOptions{
-		UserDefinedHeaders: rewrittenUserDefinedHeaders,
-	})
-	for _, headerValuePair := range upstreamHeaders {
-		requestHeader.Set(headerValuePair[0], headerValuePair[1])
+	if w.Options.CopyAllHeaders {
+		// Rewrite the user defined headers
+		// This is reported to be not compatible with Proxmox and Home Assistant
+		// but required by some other projects like MeshCentral
+		// we will make this optional
+		rewrittenUserDefinedHeaders := rewrite.PopulateRequestHeaderVariables(req, w.Options.UserDefinedHeaders)
+		upstreamHeaders, _ := rewrite.SplitUpDownStreamHeaders(&rewrite.HeaderRewriteOptions{
+			UserDefinedHeaders: rewrittenUserDefinedHeaders,
+		})
+		for _, headerValuePair := range upstreamHeaders {
+			//Do not copy Upgrade and Connection headers, it will be handled by the upgrader
+			if strings.EqualFold(headerValuePair[0], "Upgrade") || strings.EqualFold(headerValuePair[0], "Connection") {
+				continue
+			}
+			requestHeader.Set(headerValuePair[0], headerValuePair[1])
+		}
+
+		// Enable the director to copy any additional headers it desires for
+		// forwarding to the remote server.
+		if w.Director != nil {
+			w.Director(req, requestHeader)
+		}
 	}
 
 	// Connect to the backend URL, also pass the headers we get from the requst