From be752525387db89417a1eea4258795977140c736 Mon Sep 17 00:00:00 2001 From: "sudipta.deb" Date: Mon, 30 Mar 2020 16:36:02 +0530 Subject: [PATCH 1/4] Added On-demand websocket proxy close. Useful when implementing authenticated proxy. Signed-off-by: sudipta.deb --- websocketproxy.go | 56 +++++++++++++++++++++++++++++++++-------------- 1 file changed, 39 insertions(+), 17 deletions(-) diff --git a/websocketproxy.go b/websocketproxy.go index 63d39ba..b0f0f8a 100644 --- a/websocketproxy.go +++ b/websocketproxy.go @@ -45,6 +45,10 @@ type WebsocketProxy struct { // Dialer contains options for connecting to the backend WebSocket server. // If nil, DefaultDialer is used. Dialer *websocket.Dialer + + // Stop channels to close the websocket on demand + stopClientChan chan struct{} + stopBackendChan chan struct{} } // ProxyHandler returns a new http.Handler interface that reverse proxies the @@ -65,6 +69,12 @@ func NewProxy(target *url.URL) *WebsocketProxy { return &WebsocketProxy{Backend: backend} } +// Stop websocket proxy on demand +func (w *WebsocketProxy) CloseProxy() { + close(w.stopBackendChan) + close(w.stopClientChan) +} + // ServeHTTP implements the http.Handler that proxies WebSocket connections. func (w *WebsocketProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) { if w.Backend == nil { @@ -177,30 +187,42 @@ func (w *WebsocketProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) { errClient := make(chan error, 1) errBackend := make(chan error, 1) - replicateWebsocketConn := func(dst, src *websocket.Conn, errc chan error) { + replicateWebsocketConn := func(dst, src *websocket.Conn, errc chan error, stopChan chan struct{}) { for { - msgType, msg, err := src.ReadMessage() - if err != nil { - m := websocket.FormatCloseMessage(websocket.CloseNormalClosure, fmt.Sprintf("%v", err)) - if e, ok := err.(*websocket.CloseError); ok { - if e.Code != websocket.CloseNoStatusReceived { - m = websocket.FormatCloseMessage(e.Code, e.Text) + // do until stopChan gets any message + select { + default: + msgType, msg, err := src.ReadMessage() + if err != nil { + m := websocket.FormatCloseMessage(websocket.CloseNormalClosure, fmt.Sprintf("%v", err)) + if e, ok := err.(*websocket.CloseError); ok { + if e.Code != websocket.CloseNoStatusReceived { + m = websocket.FormatCloseMessage(e.Code, e.Text) + } } + errc <- err + dst.WriteMessage(websocket.CloseMessage, m) + break } - errc <- err - dst.WriteMessage(websocket.CloseMessage, m) - break - } - err = dst.WriteMessage(msgType, msg) - if err != nil { - errc <- err - break + err = dst.WriteMessage(msgType, msg) + if err != nil { + errc <- err + break + } + case <-stopChan: + dst.WriteMessage(websocket.CloseMessage, []byte("Closed by proxy")) + return } + } } - go replicateWebsocketConn(connPub, connBackend, errClient) - go replicateWebsocketConn(connBackend, connPub, errBackend) + // initiate the stop channels + w.stopClientChan = make(chan struct{}) + w.stopBackendChan = make(chan struct{}) + + go replicateWebsocketConn(connPub, connBackend, errClient, w.stopClientChan) + go replicateWebsocketConn(connBackend, connPub, errBackend, w.stopBackendChan) var message string select { From a99a5d6b3002d980aae6b9d7375027ee100dde92 Mon Sep 17 00:00:00 2001 From: Sudipta Deb Date: Mon, 30 Mar 2020 16:51:01 +0530 Subject: [PATCH 2/4] Renaming links in readme --- README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 526bb43..83cf081 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -# WebsocketProxy [![GoDoc](https://godoc.org/github.com/koding/websocketproxy?status.svg)](https://godoc.org/github.com/koding/websocketproxy) [![Build Status](https://travis-ci.org/koding/websocketproxy.svg)](https://travis-ci.org/koding/websocketproxy) +# WebsocketProxy [![GoDoc](https://godoc.org/github.com/sudiptadeb/websocketproxy?status.svg)](https://godoc.org/github.com/sudiptadeb/websocketproxy) [![Build Status](https://travis-ci.org/sudiptadeb/websocketproxy.svg)](https://travis-ci.org/sudiptadeb/websocketproxy) WebsocketProxy is an http.Handler interface build on top of [gorilla/websocket](https://github.com/gorilla/websocket) that you can plug @@ -7,7 +7,7 @@ into your existing Go webserver to provide WebSocket reverse proxy. ## Install ```bash -go get github.com/koding/websocketproxy +go get github.com/sudiptadeb/websocketproxy ``` ## Example @@ -22,7 +22,7 @@ import ( "net/http" "net/url" - "github.com/koding/websocketproxy" + "github.com/sudiptadeb/websocketproxy" ) var ( From 09ac27f37b9d5b875001aa30e2f773d10a04572f Mon Sep 17 00:00:00 2001 From: "sudipta.deb" Date: Fri, 15 May 2020 19:24:33 +0530 Subject: [PATCH 3/4] Forwarding headers --- websocketproxy.go | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/websocketproxy.go b/websocketproxy.go index b0f0f8a..0677ab3 100644 --- a/websocketproxy.go +++ b/websocketproxy.go @@ -167,14 +167,14 @@ func (w *WebsocketProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) { upgrader = DefaultUpgrader } - // Only pass those headers to the upgrader. + // passing all headers except those which can become duplicate upgradeHeader := http.Header{} - if hdr := resp.Header.Get("Sec-Websocket-Protocol"); hdr != "" { - upgradeHeader.Set("Sec-Websocket-Protocol", hdr) - } - if hdr := resp.Header.Get("Set-Cookie"); hdr != "" { - upgradeHeader.Set("Set-Cookie", hdr) - } + copyHeader(upgradeHeader, resp.Header) + + // These are extra header which the upgrader actually sets itself so need to remove these to avoid duplicate headers + upgradeHeader.Del("Connection") + upgradeHeader.Del("Upgrade") + upgradeHeader.Del("Sec-WebSocket-Accept") // Now upgrade the existing incoming request to a WebSocket connection. // Also pass the header that we gathered from the Dial handshake. From db1f908f29d3590f3ddd25d3cd82446467b7e341 Mon Sep 17 00:00:00 2001 From: "goutham.r" Date: Wed, 16 Aug 2023 12:39:09 +0530 Subject: [PATCH 4/4] [FIX] Handle nil channel when closing WebSocketProxy --- websocketproxy.go | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/websocketproxy.go b/websocketproxy.go index 0677ab3..28cf7c3 100644 --- a/websocketproxy.go +++ b/websocketproxy.go @@ -2,6 +2,7 @@ package websocketproxy import ( + "errors" "fmt" "io" "log" @@ -23,6 +24,8 @@ var ( // DefaultDialer is a dialer with all fields set to the default zero values. DefaultDialer = websocket.DefaultDialer + + errNilChannelClose = errors.New("trying to close nil channel") ) // WebsocketProxy is an HTTP Handler that takes an incoming WebSocket @@ -70,9 +73,19 @@ func NewProxy(target *url.URL) *WebsocketProxy { } // Stop websocket proxy on demand -func (w *WebsocketProxy) CloseProxy() { - close(w.stopBackendChan) - close(w.stopClientChan) +func (w *WebsocketProxy) CloseProxy() (err error) { + err = nil + if w.stopBackendChan != nil { + close(w.stopBackendChan) + } else { + err = errNilChannelClose + } + if w.stopClientChan != nil { + close(w.stopClientChan) + } else { + err = errNilChannelClose + } + return err } // ServeHTTP implements the http.Handler that proxies WebSocket connections.