Fix websocket issue

main.go

@@ -28,7 +28,7 @@ import (
 	"golang.org/x/sys/windows/registry"
 )
 
-const AppVersion = "2.2"
+const AppVersion = "2.3"
 
 type Config struct {
 	DomainsOfInterest []string `toml:"domains_of_interest"`
@@ -332,6 +332,18 @@ func (p *ProxyServer) setupHandlers() {
 			}
 		}
 
+		if isWebSocketUpgradeResponse(r) {
+			if p.verbose || p.debugMode {
+				fmt.Printf(
+					"[%s][INFO][Interest=%v] WebSocket upgrade detected, leaving stream untouched: %s\n",
+					time.Now().Format("20060102T15:04:05.000000"),
+					p.isDomainOfInterest(ctx.Req.Host),
+					ctx.Req.URL.String(),
+				)
+			}
+			return r
+		}
+
 		if p.verbose || (p.debugMode && p.isDomainOfInterest(ctx.Req.Host)) {
 			fmt.Printf(
 				"[%s][INFO][Interest=%v] HTTP Response: %s %s\n",
@@ -541,6 +553,28 @@ func sanitizeFilename(filename string) string {
 	return s
 }
 
+func isWebSocketUpgradeResponse(resp *http.Response) bool {
+	if resp == nil {
+		return false
+	}
+	if resp.StatusCode != http.StatusSwitchingProtocols {
+		return false
+	}
+	return headerContainsValue(resp.Header, "Connection", "Upgrade") &&
+		headerContainsValue(resp.Header, "Upgrade", "websocket")
+}
+
+func headerContainsValue(header http.Header, name, value string) bool {
+	for _, v := range header.Values(name) {
+		for _, s := range strings.Split(v, ",") {
+			if strings.EqualFold(value, strings.TrimSpace(s)) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 func (p *ProxyServer) setSystemProxy() error {
 	// Get current proxy settings
 	k, err := registry.OpenKey(registry.CURRENT_USER, `Software\Microsoft\Windows\CurrentVersion\Internet Settings`, registry.QUERY_VALUE)