From 482c1f81f5adb44c450d836879150adfc083b3ac Mon Sep 17 00:00:00 2001 From: Bel LaPointe <153096461+breel-render@users.noreply.github.com> Date: Wed, 9 Sep 2026 07:24:03 -0700 Subject: [PATCH] try new ws proxy --- go.mod | 2 + go.sum | 2 + main.go | 24 +++-- websocketproxy.go | 233 ---------------------------------------------- 4 files changed, 20 insertions(+), 241 deletions(-) delete mode 100644 websocketproxy.go diff --git a/go.mod b/go.mod index 080ad92..584c382 100644 --- a/go.mod +++ b/go.mod @@ -8,3 +8,5 @@ require ( ) require github.com/viki-org/dnscache v0.0.0-20130720023526-c70c1f23c5d8 + +require github.com/pretty66/websocketproxy v0.0.0-20220507015215-930b3a686308 // indirect diff --git a/go.sum b/go.sum index 96f5fa5..5087a2a 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,7 @@ 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/pretty66/websocketproxy v0.0.0-20220507015215-930b3a686308 h1:JfSau4YABtkm5gRtFWuRWHT2Lsw4ZbyB4F/qORwf+BA= +github.com/pretty66/websocketproxy v0.0.0-20220507015215-930b3a686308/go.mod h1:hxhFuMswfNko9fAxYeqBapfUdJHAgDafBs/MzOZh0X8= github.com/viki-org/dnscache v0.0.0-20130720023526-c70c1f23c5d8 h1:EVObHAr8DqpoJCVv6KYTle8FEImKhtkfcZetNqxDoJQ= github.com/viki-org/dnscache v0.0.0-20130720023526-c70c1f23c5d8/go.mod h1:dniwbG03GafCjFohMDmz6Zc6oCuiqgH6tGNyXTkHzXE= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= diff --git a/main.go b/main.go index bd92349..ff766f3 100644 --- a/main.go +++ b/main.go @@ -12,6 +12,7 @@ import ( "strings" "time" + "github.com/pretty66/websocketproxy" "github.com/viki-org/dnscache" ) @@ -95,17 +96,24 @@ func (c Config) serveHTTPProxy(w http.ResponseWriter, r *http.Request, endpoint isWebsocket := r.Header.Get("Connection") == "Upgrade" || r.Header.Get("Se-Fetch-Mode") == "websocket" || r.Header.Get("Sec-WebSocket-Version") != "" || r.Header.Get("Upgrade") == "websocket" if isWebsocket { - if strings.HasPrefix(u.Scheme, "s") { - u.Scheme = "wss" + wsu := *u + if strings.HasPrefix(wsu.Scheme, "s") { + wsu.Scheme = "wss" } else { - u.Scheme = "ws" + wsu.Scheme = "ws" } - if DefaultUpgrader.CheckOrigin == nil { - DefaultUpgrader.CheckOrigin = func(*http.Request) bool { - return true - } + wp, err := websocketproxy.NewProxy(wsu.String(), func(r2 *http.Request) error { + r2.Header.Set("Cookie", r.Header.Get("Cookie")) + u3 := *u + u3.Path = "/" + r2.Header.Set("Origin", u3.String()) + return nil + }) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return } - NewProxy(u).ServeHTTP(w, r) + wp.ServeHTTP(w, r) return } diff --git a/websocketproxy.go b/websocketproxy.go deleted file mode 100644 index 4150cc5..0000000 --- a/websocketproxy.go +++ /dev/null @@ -1,233 +0,0 @@ -// Package websocketproxy is a reverse proxy for WebSocket connections. -package main - -import ( - "fmt" - "io" - "log" - "net" - "net/http" - "net/url" - "strings" - - "github.com/gorilla/websocket" -) - -var ( - // DefaultUpgrader specifies the parameters for upgrading an HTTP - // connection to a WebSocket connection. - DefaultUpgrader = &websocket.Upgrader{ - ReadBufferSize: 1024, - WriteBufferSize: 1024, - } - - // DefaultDialer is a dialer with all fields set to the default zero values. - DefaultDialer = websocket.DefaultDialer -) - -// WebsocketProxy is an HTTP Handler that takes an incoming WebSocket -// connection and proxies it to another server. -type WebsocketProxy struct { - // Director, if non-nil, is a function that may copy additional request - // headers from the incoming WebSocket connection into the output headers - // which will be forwarded to another server. - Director func(incoming *http.Request, out http.Header) - - // Backend returns the backend URL which the proxy uses to reverse proxy - // the incoming WebSocket connection. Request is the initial incoming and - // unmodified request. - Backend func(*http.Request) *url.URL - - // Upgrader specifies the parameters for upgrading a incoming HTTP - // connection to a WebSocket connection. If nil, DefaultUpgrader is used. - Upgrader *websocket.Upgrader - - // Dialer contains options for connecting to the backend WebSocket server. - // If nil, DefaultDialer is used. - Dialer *websocket.Dialer -} - -// ProxyHandler returns a new http.Handler interface that reverse proxies the -// request to the given target. -func ProxyHandler(target *url.URL) http.Handler { return NewProxy(target) } - -// NewProxy returns a new Websocket reverse proxy that rewrites the -// URL's to the scheme, host and base path provider in target. -func NewProxy(target *url.URL) *WebsocketProxy { - backend := func(r *http.Request) *url.URL { - // Shallow copy - u := *target - u.Fragment = r.URL.Fragment - u.Path = r.URL.Path - u.RawQuery = r.URL.RawQuery - return &u - } - return &WebsocketProxy{Backend: backend} -} - -// ServeHTTP implements the http.Handler that proxies WebSocket connections. -func (w *WebsocketProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) { - if w.Backend == nil { - log.Println("websocketproxy: backend function is not defined") - http.Error(rw, "internal server error (code: 1)", http.StatusInternalServerError) - return - } - - backendURL := w.Backend(req) - if backendURL == nil { - log.Println("websocketproxy: backend URL is nil") - http.Error(rw, "internal server error (code: 2)", http.StatusInternalServerError) - return - } - - dialer := w.Dialer - if w.Dialer == nil { - dialer = DefaultDialer - } - - // Pass headers from the incoming request to the dialer to forward them to - // the final destinations. - requestHeader := http.Header{} - if origin := req.Header.Get("Origin"); origin != "" { - requestHeader.Add("Origin", origin) - } - for _, prot := range req.Header[http.CanonicalHeaderKey("Sec-WebSocket-Protocol")] { - requestHeader.Add("Sec-WebSocket-Protocol", prot) - } - for _, cookie := range req.Header[http.CanonicalHeaderKey("Cookie")] { - requestHeader.Add("Cookie", cookie) - } - if req.Host != "" { - requestHeader.Set("Host", req.Host) - } - - // Pass X-Forwarded-For headers too, code below is a part of - // httputil.ReverseProxy. See http://en.wikipedia.org/wiki/X-Forwarded-For - // for more information - // TODO: use RFC7239 http://tools.ietf.org/html/rfc7239 - if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil { - // If we aren't the first proxy retain prior - // X-Forwarded-For information as a comma+space - // separated list and fold multiple headers into one. - if prior, ok := req.Header["X-Forwarded-For"]; ok { - clientIP = strings.Join(prior, ", ") + ", " + clientIP - } - requestHeader.Set("X-Forwarded-For", clientIP) - } - - // Set the originating protocol of the incoming HTTP request. The SSL might - // be terminated on our site and because we doing proxy adding this would - // be helpful for applications on the backend. - requestHeader.Set("X-Forwarded-Proto", "http") - if req.TLS != nil { - 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) - } - - // Connect to the backend URL, also pass the headers we get from the requst - // together with the Forwarded headers we prepared above. - // TODO: support multiplexing on the same backend connection instead of - // opening a new TCP connection time for each request. This should be - // optional: - // http://tools.ietf.org/html/draft-ietf-hybi-websocket-multiplexing-01 - connBackend, resp, err := dialer.Dial(backendURL.String(), requestHeader) - if err != nil { - log.Printf("websocketproxy: couldn't dial to remote backend url %s", err) - if resp != nil { - // If the WebSocket handshake fails, ErrBadHandshake is returned - // along with a non-nil *http.Response so that callers can handle - // redirects, authentication, etcetera. - if err := copyResponse(rw, resp); err != nil { - log.Printf("websocketproxy: couldn't write response after failed remote backend handshake: %s", err) - } - } else { - http.Error(rw, http.StatusText(http.StatusServiceUnavailable), http.StatusServiceUnavailable) - } - return - } - defer connBackend.Close() - - upgrader := w.Upgrader - if w.Upgrader == nil { - upgrader = DefaultUpgrader - } - - // Only pass those headers to the upgrader. - 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) - } - - // Now upgrade the existing incoming request to a WebSocket connection. - // Also pass the header that we gathered from the Dial handshake. - connPub, err := upgrader.Upgrade(rw, req, upgradeHeader) - if err != nil { - log.Printf("websocketproxy: couldn't upgrade %s", err) - return - } - defer connPub.Close() - - errClient := make(chan error, 1) - errBackend := make(chan error, 1) - replicateWebsocketConn := func(dst, src *websocket.Conn, errc chan error) { - 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) - } - } - errc <- err - dst.WriteMessage(websocket.CloseMessage, m) - break - } - err = dst.WriteMessage(msgType, msg) - if err != nil { - errc <- err - break - } - } - } - - go replicateWebsocketConn(connPub, connBackend, errClient) - go replicateWebsocketConn(connBackend, connPub, errBackend) - - var message string - select { - case err = <-errClient: - message = "websocketproxy: Error when copying from backend to client: %v" - case err = <-errBackend: - message = "websocketproxy: Error when copying from client to backend: %v" - - } - if e, ok := err.(*websocket.CloseError); !ok || e.Code == websocket.CloseAbnormalClosure { - log.Printf(message, err) - } -} - -func copyHeader(dst, src http.Header) { - for k, vv := range src { - for _, v := range vv { - dst.Add(k, v) - } - } -} - -func copyResponse(rw http.ResponseWriter, resp *http.Response) error { - copyHeader(rw.Header(), resp.Header) - rw.WriteHeader(resp.StatusCode) - defer resp.Body.Close() - - _, err := io.Copy(rw, resp.Body) - return err -}