Merge pull request #637 from elazarl/websocket-direct

Handle Websocket using standard library

examples/go.mod

@@ -3,9 +3,9 @@ module github.com/elazarl/goproxy/examples/goproxy-transparent
 go 1.20
 
 require (
-	github.com/elazarl/goproxy v1.3.0
-	github.com/elazarl/goproxy/ext v0.0.0-20250110140559-10fc34b80676
-	github.com/gorilla/websocket v1.5.3
+	github.com/coder/websocket v1.8.12
+	github.com/elazarl/goproxy v1.5.0
+	github.com/elazarl/goproxy/ext v0.0.0-20250117123040-e9229c451ab8
 	github.com/inconshreveable/go-vhost v1.0.0
 )
 

examples/go.sum

@@ -1,8 +1,10 @@
+github.com/coder/websocket v1.8.12 h1:5bUXkEPPIbewrnkU8LTCLVaxi4N4J8ahufH2vlo4NAo=
+github.com/coder/websocket v1.8.12/go.mod h1:LNVeNrXQZfe5qhS9ALED3uA+l5pPqvwXg3CKoDBB2gs=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/elazarl/goproxy/ext v0.0.0-20250110140559-10fc34b80676 h1:3bAtOWqImclW/5rXbhNyAcM122jafst+/+4J4vC8wZI=
 github.com/elazarl/goproxy/ext v0.0.0-20250110140559-10fc34b80676/go.mod h1:q2JQCFWg+AQfe6O2cbf7LJDB48R68w+q0pBU53v02iM=
-github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
-github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
+github.com/elazarl/goproxy/ext v0.0.0-20250117123040-e9229c451ab8 h1:rGxOExXmpBcmZc4ZnEXBGkcxSReZx7S9ECtuv6BtUYQ=
+github.com/elazarl/goproxy/ext v0.0.0-20250117123040-e9229c451ab8/go.mod h1:q2JQCFWg+AQfe6O2cbf7LJDB48R68w+q0pBU53v02iM=
 github.com/inconshreveable/go-vhost v1.0.0 h1:IK4VZTlXL4l9vz2IZoiSFbYaaqUW7dXJAiPriUN5Ur8=
 github.com/inconshreveable/go-vhost v1.0.0/go.mod h1:aA6DnFhALT3zH0y+A39we+zbrdMC2N0X/q21e6FI0LU=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=

examples/websockets/main.go

@@ -1,9 +1,10 @@
 package main
 
 import (
+	"context"
 	"crypto/tls"
+	"github.com/coder/websocket"
 	"github.com/elazarl/goproxy"
-	"github.com/gorilla/websocket"
 	"log"
 	"net/http"
 	"net/url"
@@ -12,26 +13,23 @@ import (
 	"time"
 )
 
-var _upgrader = websocket.Upgrader{
-	HandshakeTimeout: 10 * time.Second,
-}
-
 func echo(w http.ResponseWriter, r *http.Request) {
-	c, err := _upgrader.Upgrade(w, r, nil)
+	c, err := websocket.Accept(w, r, nil)
 	if err != nil {
 		log.Printf("upgrade: %v\n", err)
 		return
 	}
-	defer c.Close()
+	defer c.Close(websocket.StatusNormalClosure, "")
 
+	ctx := context.Background()
 	for {
-		mt, message, err := c.ReadMessage()
+		mt, message, err := c.Read(ctx)
 		if err != nil {
 			log.Printf("read: %v\n", err)
 			break
 		}
 		log.Printf("recv: %s\n", message)
-		if err := c.WriteMessage(mt, message); err != nil {
+		if err := c.Write(ctx, mt, message); err != nil {
 			log.Printf("write: %v\n", err)
 			break
 		}
@@ -74,27 +72,30 @@ func main() {
 		log.Fatal("unable to parse proxy URL")
 	}
 
-	dialer := websocket.Dialer{
-		Subprotocols: []string{"p1"},
-		TLSClientConfig: &tls.Config{
-			InsecureSkipVerify: true,
-		},
-		Proxy: http.ProxyURL(parsedProxy),
-	}
-
+	ctx := context.Background()
 	endpointUrl := "wss://localhost:12345"
-	c, _, err := dialer.Dial(endpointUrl, nil)
+
+	c, _, err := websocket.Dial(ctx, endpointUrl, &websocket.DialOptions{
+		HTTPClient: &http.Client{
+			Transport: &http.Transport{
+				TLSClientConfig: &tls.Config{
+					InsecureSkipVerify: true,
+				},
+				Proxy: http.ProxyURL(parsedProxy),
+			},
+		},
+		Subprotocols: []string{"p1"},
+	})
 	if err != nil {
 		log.Fatal("dial:", err)
 	}
-	defer c.Close()
 
 	done := make(chan struct{})
 
 	go func() {
 		defer close(done)
 		for {
-			_, message, err := c.ReadMessage()
+			_, message, err := c.Read(ctx)
 			if err != nil {
 				log.Println("read:", err)
 				return
@@ -110,15 +111,15 @@ func main() {
 		select {
 		case t := <-ticker.C: // Message send
 			// Write current time to the websocket client every 1 second
-			if err := c.WriteMessage(websocket.TextMessage, []byte(t.String())); err != nil {
+			if err := c.Write(ctx, websocket.MessageText, []byte(t.String())); err != nil {
 				log.Println("write:", err)
 				return
 			}
 		case <-interrupt: // Server shutdown
 			log.Println("interrupt")
 			// To cleanly close a connection, a client should send a close
 			// frame and wait for the server to close the connection.
-			err := c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
+			err := c.Close(websocket.StatusNormalClosure, "")
 			if err != nil {
 				log.Println("write close:", err)
 				return

http.go

@@ -18,12 +18,6 @@ func (proxy *ProxyHttpServer) handleHttp(w http.ResponseWriter, r *http.Request)
 	r, resp := proxy.filterRequest(r, ctx)
 
 	if resp == nil {
-		if isWebSocketRequest(r) {
-			ctx.Logf("Request looks like websocket upgrade.")
-			if conn, err := proxy.hijackConnection(ctx, w); err == nil {
-				proxy.serveWebsocket(ctx, conn, r)
-			}
-		}
 		if !proxy.KeepHeader {
 			RemoveProxyHeaders(ctx, r)
 		}
@@ -69,6 +63,23 @@ func (proxy *ProxyHttpServer) handleHttp(w http.ResponseWriter, r *http.Request)
 	}
 	copyHeaders(w.Header(), resp.Header, proxy.KeepDestinationHeaders)
 	w.WriteHeader(resp.StatusCode)
+
+	if isWebSocketHandshake(resp.Header) {
+		ctx.Logf("Response looks like websocket upgrade.")
+
+		// We have already written the "101 Switching Protocols" response,
+		// now we hijack the connection to send WebSocket data
+		if clientConn, err := proxy.hijackConnection(ctx, w); err == nil {
+			wsConn, ok := resp.Body.(io.ReadWriter)
+			if !ok {
+				ctx.Warnf("Unable to use Websocket connection")
+				return
+			}
+			proxy.proxyWebsocket(ctx, wsConn, clientConn)
+		}
+		return
+	}
+
 	var copyWriter io.Writer = w
 	// Content-Type header may also contain charset definition, so here we need to check the prefix.
 	// Transfer-Encoding can be a list of comma separated values, so we use Contains() for it.

https.go

@@ -251,8 +251,9 @@ func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request
 					}
 				}
 				resp = proxy.filterResponse(resp, ctx)
+				defer resp.Body.Close()
+
 				err = resp.Write(proxyClient)
-				_ = resp.Body.Close()
 				if err != nil {
 					httpError(proxyClient, ctx, err)
 					return false
@@ -353,16 +354,6 @@ func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request
 							}
 							return false
 						}
-						if isWebSocketRequest(req) {
-							ctx.Logf("Request looks like websocket upgrade.")
-							if req.URL.Scheme == "http" {
-								ctx.Logf("Enforced HTTP websocket forwarding over TLS")
-								proxy.serveWebsocket(ctx, rawClientTls, req)
-							} else {
-								proxy.serveWebsocketTLS(ctx, req, tlsConfig, rawClientTls)
-							}
-							return false
-						}
 						if err != nil {
 							if req.URL != nil {
 								ctx.Warnf("Illegal URL %s", "https://"+r.Host+req.URL.Path)
@@ -398,7 +389,8 @@ func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request
 						return false
 					}
 
-					if resp.Request.Method == http.MethodHead {
+					isWebsocket := isWebSocketHandshake(resp.Header)
+					if isWebsocket || resp.Request.Method == http.MethodHead {
 						// don't change Content-Length for HEAD request
 					} else if (resp.StatusCode >= 100 && resp.StatusCode < 200) ||
 						resp.StatusCode == http.StatusNoContent {
@@ -412,7 +404,9 @@ func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request
 						resp.Header.Set("Transfer-Encoding", "chunked")
 					}
 					// Force connection close otherwise chrome will keep CONNECT tunnel open forever
-					resp.Header.Set("Connection", "close")
+					if !isWebsocket {
+						resp.Header.Set("Connection", "close")
+					}
 					if err := resp.Header.Write(rawClientTls); err != nil {
 						ctx.Warnf("Cannot write TLS response header from mitm'd client: %v", err)
 						return false
@@ -422,6 +416,24 @@ func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request
 						return false
 					}
 
+					if isWebsocket {
+						ctx.Logf("Response looks like websocket upgrade.")
+
+						// According to resp.Body documentation:
+						// As of Go 1.12, the Body will also implement io.Writer
+						// on a successful "101 Switching Protocols" response,
+						// as used by WebSockets and HTTP/2's "h2c" mode.
+						wsConn, ok := resp.Body.(io.ReadWriter)
+						if !ok {
+							ctx.Warnf("Unable to use Websocket connection")
+							return false
+						}
+						proxy.proxyWebsocket(ctx, wsConn, rawClientTls)
+						// We can't reuse connection after WebSocket handshake,
+						// by returning false here, the underlying connection will be closed
+						return false
+					}
+
 					if resp.Request.Method == http.MethodHead ||
 						(resp.StatusCode >= 100 && resp.StatusCode < 200) ||
 						resp.StatusCode == http.StatusNoContent ||

proxy.go

@@ -110,7 +110,13 @@ func RemoveProxyHeaders(ctx *ProxyCtx, r *http.Request) {
 	//   options that are desired for that particular connection and MUST NOT
 	//   be communicated by proxies over further connections.
 
-	r.Header.Del("Connection")
+	// We need to keep "Connection: upgrade" header, since it's part of
+	// the WebSocket handshake, and it won't work without it.
+	// For all the other cases (close, keep-alive), we already handle them, by
+	// setting the r.Close variable in the previous lines.
+	if !isWebSocketHandshake(r.Header) {
+		r.Header.Del("Connection")
+	}
 }
 
 type flushWriter struct {