Skip to content

Commit ddbd8c1

Browse files
elazarlelazarl
authored
Merge pull request #637 from elazarl/websocket-direct
Handle Websocket using standard library
2 parents af4d657 + a1a67aa commit ddbd8c1

7 files changed

Lines changed: 98 additions & 167 deletions

File tree

‎examples/go.mod‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,9 @@ module github.com/elazarl/goproxy/examples/goproxy-transparent
33
go 1.20
44

55
require (
6-
github.com/elazarl/goproxy v1.3.0
7-
github.com/elazarl/goproxy/ext v0.0.0-20250110140559-10fc34b80676
8-
github.com/gorilla/websocket v1.5.3
6+
github.com/coder/websocket v1.8.12
7+
github.com/elazarl/goproxy v1.5.0
8+
github.com/elazarl/goproxy/ext v0.0.0-20250117123040-e9229c451ab8
99
github.com/inconshreveable/go-vhost v1.0.0
1010
)
1111

‎examples/go.sum‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
1+
github.com/coder/websocket v1.8.12 h1:5bUXkEPPIbewrnkU8LTCLVaxi4N4J8ahufH2vlo4NAo=
2+
github.com/coder/websocket v1.8.12/go.mod h1:LNVeNrXQZfe5qhS9ALED3uA+l5pPqvwXg3CKoDBB2gs=
13
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
24
github.com/elazarl/goproxy/ext v0.0.0-20250110140559-10fc34b80676 h1:3bAtOWqImclW/5rXbhNyAcM122jafst+/+4J4vC8wZI=
35
github.com/elazarl/goproxy/ext v0.0.0-20250110140559-10fc34b80676/go.mod h1:q2JQCFWg+AQfe6O2cbf7LJDB48R68w+q0pBU53v02iM=
4-
github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
5-
github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
6+
github.com/elazarl/goproxy/ext v0.0.0-20250117123040-e9229c451ab8 h1:rGxOExXmpBcmZc4ZnEXBGkcxSReZx7S9ECtuv6BtUYQ=
7+
github.com/elazarl/goproxy/ext v0.0.0-20250117123040-e9229c451ab8/go.mod h1:q2JQCFWg+AQfe6O2cbf7LJDB48R68w+q0pBU53v02iM=
68
github.com/inconshreveable/go-vhost v1.0.0 h1:IK4VZTlXL4l9vz2IZoiSFbYaaqUW7dXJAiPriUN5Ur8=
79
github.com/inconshreveable/go-vhost v1.0.0/go.mod h1:aA6DnFhALT3zH0y+A39we+zbrdMC2N0X/q21e6FI0LU=
810
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=

‎examples/websockets/main.go‎

Lines changed: 23 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
package main
22

33
import (
4+
"context"
45
"crypto/tls"
6+
"github.com/coder/websocket"
57
"github.com/elazarl/goproxy"
6-
"github.com/gorilla/websocket"
78
"log"
89
"net/http"
910
"net/url"
@@ -12,26 +13,23 @@ import (
1213
"time"
1314
)
1415

15-
var _upgrader = websocket.Upgrader{
16-
HandshakeTimeout: 10 * time.Second,
17-
}
18-
1916
func echo(w http.ResponseWriter, r *http.Request) {
20-
c, err := _upgrader.Upgrade(w, r, nil)
17+
c, err := websocket.Accept(w, r, nil)
2118
if err != nil {
2219
log.Printf("upgrade: %v\n", err)
2320
return
2421
}
25-
defer c.Close()
22+
defer c.Close(websocket.StatusNormalClosure, "")
2623

24+
ctx := context.Background()
2725
for {
28-
mt, message, err := c.ReadMessage()
26+
mt, message, err := c.Read(ctx)
2927
if err != nil {
3028
log.Printf("read: %v\n", err)
3129
break
3230
}
3331
log.Printf("recv: %s\n", message)
34-
if err := c.WriteMessage(mt, message); err != nil {
32+
if err := c.Write(ctx, mt, message); err != nil {
3533
log.Printf("write: %v\n", err)
3634
break
3735
}
@@ -74,27 +72,30 @@ func main() {
7472
log.Fatal("unable to parse proxy URL")
7573
}
7674

77-
dialer := websocket.Dialer{
78-
Subprotocols: []string{"p1"},
79-
TLSClientConfig: &tls.Config{
80-
InsecureSkipVerify: true,
81-
},
82-
Proxy: http.ProxyURL(parsedProxy),
83-
}
84-
75+
ctx := context.Background()
8576
endpointUrl := "wss://localhost:12345"
86-
c, _, err := dialer.Dial(endpointUrl, nil)
77+
78+
c, _, err := websocket.Dial(ctx, endpointUrl, &websocket.DialOptions{
79+
HTTPClient: &http.Client{
80+
Transport: &http.Transport{
81+
TLSClientConfig: &tls.Config{
82+
InsecureSkipVerify: true,
83+
},
84+
Proxy: http.ProxyURL(parsedProxy),
85+
},
86+
},
87+
Subprotocols: []string{"p1"},
88+
})
8789
if err != nil {
8890
log.Fatal("dial:", err)
8991
}
90-
defer c.Close()
9192

9293
done := make(chan struct{})
9394

9495
go func() {
9596
defer close(done)
9697
for {
97-
_, message, err := c.ReadMessage()
98+
_, message, err := c.Read(ctx)
9899
if err != nil {
99100
log.Println("read:", err)
100101
return
@@ -110,15 +111,15 @@ func main() {
110111
select {
111112
case t := <-ticker.C: // Message send
112113
// Write current time to the websocket client every 1 second
113-
if err := c.WriteMessage(websocket.TextMessage, []byte(t.String())); err != nil {
114+
if err := c.Write(ctx, websocket.MessageText, []byte(t.String())); err != nil {
114115
log.Println("write:", err)
115116
return
116117
}
117118
case <-interrupt: // Server shutdown
118119
log.Println("interrupt")
119120
// To cleanly close a connection, a client should send a close
120121
// frame and wait for the server to close the connection.
121-
err := c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
122+
err := c.Close(websocket.StatusNormalClosure, "")
122123
if err != nil {
123124
log.Println("write close:", err)
124125
return

‎http.go‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,6 @@ func (proxy *ProxyHttpServer) handleHttp(w http.ResponseWriter, r *http.Request)
1818
r, resp := proxy.filterRequest(r, ctx)
1919

2020
if resp == nil {
21-
if isWebSocketRequest(r) {
22-
ctx.Logf("Request looks like websocket upgrade.")
23-
if conn, err := proxy.hijackConnection(ctx, w); err == nil {
24-
proxy.serveWebsocket(ctx, conn, r)
25-
}
26-
}
2721
if !proxy.KeepHeader {
2822
RemoveProxyHeaders(ctx, r)
2923
}
@@ -69,6 +63,23 @@ func (proxy *ProxyHttpServer) handleHttp(w http.ResponseWriter, r *http.Request)
6963
}
7064
copyHeaders(w.Header(), resp.Header, proxy.KeepDestinationHeaders)
7165
w.WriteHeader(resp.StatusCode)
66+
67+
if isWebSocketHandshake(resp.Header) {
68+
ctx.Logf("Response looks like websocket upgrade.")
69+
70+
// We have already written the "101 Switching Protocols" response,
71+
// now we hijack the connection to send WebSocket data
72+
if clientConn, err := proxy.hijackConnection(ctx, w); err == nil {
73+
wsConn, ok := resp.Body.(io.ReadWriter)
74+
if !ok {
75+
ctx.Warnf("Unable to use Websocket connection")
76+
return
77+
}
78+
proxy.proxyWebsocket(ctx, wsConn, clientConn)
79+
}
80+
return
81+
}
82+
7283
var copyWriter io.Writer = w
7384
// Content-Type header may also contain charset definition, so here we need to check the prefix.
7485
// Transfer-Encoding can be a list of comma separated values, so we use Contains() for it.

‎https.go‎

Lines changed: 25 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -251,8 +251,9 @@ func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request
251251
}
252252
}
253253
resp = proxy.filterResponse(resp, ctx)
254+
defer resp.Body.Close()
255+
254256
err = resp.Write(proxyClient)
255-
_ = resp.Body.Close()
256257
if err != nil {
257258
httpError(proxyClient, ctx, err)
258259
return false
@@ -353,16 +354,6 @@ func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request
353354
}
354355
return false
355356
}
356-
if isWebSocketRequest(req) {
357-
ctx.Logf("Request looks like websocket upgrade.")
358-
if req.URL.Scheme == "http" {
359-
ctx.Logf("Enforced HTTP websocket forwarding over TLS")
360-
proxy.serveWebsocket(ctx, rawClientTls, req)
361-
} else {
362-
proxy.serveWebsocketTLS(ctx, req, tlsConfig, rawClientTls)
363-
}
364-
return false
365-
}
366357
if err != nil {
367358
if req.URL != nil {
368359
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
398389
return false
399390
}
400391

401-
if resp.Request.Method == http.MethodHead {
392+
isWebsocket := isWebSocketHandshake(resp.Header)
393+
if isWebsocket || resp.Request.Method == http.MethodHead {
402394
// don't change Content-Length for HEAD request
403395
} else if (resp.StatusCode >= 100 && resp.StatusCode < 200) ||
404396
resp.StatusCode == http.StatusNoContent {
@@ -412,7 +404,9 @@ func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request
412404
resp.Header.Set("Transfer-Encoding", "chunked")
413405
}
414406
// Force connection close otherwise chrome will keep CONNECT tunnel open forever
415-
resp.Header.Set("Connection", "close")
407+
if !isWebsocket {
408+
resp.Header.Set("Connection", "close")
409+
}
416410
if err := resp.Header.Write(rawClientTls); err != nil {
417411
ctx.Warnf("Cannot write TLS response header from mitm'd client: %v", err)
418412
return false
@@ -422,6 +416,24 @@ func (proxy *ProxyHttpServer) handleHttps(w http.ResponseWriter, r *http.Request
422416
return false
423417
}
424418

419+
if isWebsocket {
420+
ctx.Logf("Response looks like websocket upgrade.")
421+
422+
// According to resp.Body documentation:
423+
// As of Go 1.12, the Body will also implement io.Writer
424+
// on a successful "101 Switching Protocols" response,
425+
// as used by WebSockets and HTTP/2's "h2c" mode.
426+
wsConn, ok := resp.Body.(io.ReadWriter)
427+
if !ok {
428+
ctx.Warnf("Unable to use Websocket connection")
429+
return false
430+
}
431+
proxy.proxyWebsocket(ctx, wsConn, rawClientTls)
432+
// We can't reuse connection after WebSocket handshake,
433+
// by returning false here, the underlying connection will be closed
434+
return false
435+
}
436+
425437
if resp.Request.Method == http.MethodHead ||
426438
(resp.StatusCode >= 100 && resp.StatusCode < 200) ||
427439
resp.StatusCode == http.StatusNoContent ||

‎proxy.go‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,13 @@ func RemoveProxyHeaders(ctx *ProxyCtx, r *http.Request) {
110110
// options that are desired for that particular connection and MUST NOT
111111
// be communicated by proxies over further connections.
112112

113-
r.Header.Del("Connection")
113+
// We need to keep "Connection: upgrade" header, since it's part of
114+
// the WebSocket handshake, and it won't work without it.
115+
// For all the other cases (close, keep-alive), we already handle them, by
116+
// setting the r.Close variable in the previous lines.
117+
if !isWebSocketHandshake(r.Header) {
118+
r.Header.Del("Connection")
119+
}
114120
}
115121

116122
type flushWriter struct {

0 commit comments

Comments
 (0)