github.com/e154/smart-home@v0.17.2-0.20240311175135-e530a6e5cd45/common/websocketproxy/websocketproxy.go (about) 1 // This file is part of the Smart Home 2 // Program complex distribution https://github.com/e154/smart-home 3 // Copyright (C) 2023, Filippov Alex 4 // 5 // This library is free software: you can redistribute it and/or 6 // modify it under the terms of the GNU Lesser General Public 7 // License as published by the Free Software Foundation; either 8 // version 3 of the License, or (at your option) any later version. 9 // 10 // This library is distributed in the hope that it will be useful, 11 // but WITHOUT ANY WARRANTY; without even the implied warranty of 12 // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU 13 // Library General Public License for more details. 14 // 15 // You should have received a copy of the GNU Lesser General Public 16 // License along with this library. If not, see 17 // <https://www.gnu.org/licenses/>. 18 19 package websocketproxy 20 21 import ( 22 "fmt" 23 "io" 24 "log" 25 "net" 26 "net/http" 27 "net/url" 28 "strings" 29 30 "github.com/gorilla/websocket" 31 ) 32 33 var ( 34 // DefaultUpgrader specifies the parameters for upgrading an HTTP 35 // connection to a WebSocket connection. 36 DefaultUpgrader = &websocket.Upgrader{ 37 ReadBufferSize: 1024, 38 WriteBufferSize: 1024, 39 } 40 41 // DefaultDialer is a dialer with all fields set to the default zero values. 42 DefaultDialer = websocket.DefaultDialer 43 ) 44 45 // WebsocketProxy is an HTTP Handler that takes an incoming WebSocket 46 // connection and proxies it to another server. 47 type WebsocketProxy struct { 48 // Director, if non-nil, is a function that may copy additional request 49 // headers from the incoming WebSocket connection into the output headers 50 // which will be forwarded to another server. 51 Director func(incoming *http.Request, out http.Header) 52 53 // Backend returns the backend URL which the proxy uses to reverse proxy 54 // the incoming WebSocket connection. Request is the initial incoming and 55 // unmodified request. 56 Backend func(*http.Request) *url.URL 57 58 // Upgrader specifies the parameters for upgrading a incoming HTTP 59 // connection to a WebSocket connection. If nil, DefaultUpgrader is used. 60 Upgrader *websocket.Upgrader 61 62 // Dialer contains options for connecting to the backend WebSocket server. 63 // If nil, DefaultDialer is used. 64 Dialer *websocket.Dialer 65 } 66 67 // ProxyHandler returns a new http.Handler interface that reverse proxies the 68 // request to the given target. 69 func ProxyHandler(target *url.URL) http.Handler { return NewProxy(target) } 70 71 // NewProxy returns a new Websocket reverse proxy that rewrites the 72 // URL's to the scheme, host and base path provider in target. 73 func NewProxy(target *url.URL) *WebsocketProxy { 74 backend := func(r *http.Request) *url.URL { 75 // Shallow copy 76 u := *target 77 u.Fragment = r.URL.Fragment 78 u.Path = r.URL.Path 79 u.RawQuery = r.URL.RawQuery 80 return &u 81 } 82 return &WebsocketProxy{Backend: backend} 83 } 84 85 // ServeHTTP implements the http.Handler that proxies WebSocket connections. 86 func (w *WebsocketProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) { 87 if w.Backend == nil { 88 log.Println("websocketproxy: backend function is not defined") 89 http.Error(rw, "internal server error (code: 1)", http.StatusInternalServerError) 90 return 91 } 92 93 backendURL := w.Backend(req) 94 if backendURL == nil { 95 log.Println("websocketproxy: backend URL is nil") 96 http.Error(rw, "internal server error (code: 2)", http.StatusInternalServerError) 97 return 98 } 99 100 dialer := w.Dialer 101 if w.Dialer == nil { 102 dialer = DefaultDialer 103 } 104 105 // Pass headers from the incoming request to the dialer to forward them to 106 // the final destinations. 107 requestHeader := http.Header{} 108 if origin := req.Header.Get("Origin"); origin != "" { 109 requestHeader.Add("Origin", origin) 110 } 111 for _, prot := range req.Header[http.CanonicalHeaderKey("Sec-WebSocket-Protocol")] { 112 requestHeader.Add("Sec-WebSocket-Protocol", prot) 113 } 114 for _, cookie := range req.Header[http.CanonicalHeaderKey("Cookie")] { 115 requestHeader.Add("Cookie", cookie) 116 } 117 if req.Host != "" { 118 requestHeader.Set("Host", req.Host) 119 } 120 121 // Pass X-Forwarded-For headers too, code below is a part of 122 // httputil.ReverseProxy. See http://en.wikipedia.org/wiki/X-Forwarded-For 123 // for more information 124 // TODO: use RFC7239 http://tools.ietf.org/html/rfc7239 125 if clientIP, _, err := net.SplitHostPort(req.RemoteAddr); err == nil { 126 // If we aren't the first proxy retain prior 127 // X-Forwarded-For information as a comma+space 128 // separated list and fold multiple headers into one. 129 if prior, ok := req.Header["X-Forwarded-For"]; ok { 130 clientIP = strings.Join(prior, ", ") + ", " + clientIP 131 } 132 requestHeader.Set("X-Forwarded-For", clientIP) 133 } 134 135 // Set the originating protocol of the incoming HTTP request. The SSL might 136 // be terminated on our site and because we doing proxy adding this would 137 // be helpful for applications on the backend. 138 requestHeader.Set("X-Forwarded-Proto", "http") 139 if req.TLS != nil { 140 requestHeader.Set("X-Forwarded-Proto", "https") 141 } 142 143 // Enable the director to copy any additional headers it desires for 144 // forwarding to the remote server. 145 if w.Director != nil { 146 w.Director(req, requestHeader) 147 } 148 149 // Connect to the backend URL, also pass the headers we get from the requst 150 // together with the Forwarded headers we prepared above. 151 // TODO: support multiplexing on the same backend connection instead of 152 // opening a new TCP connection time for each request. This should be 153 // optional: 154 // http://tools.ietf.org/html/draft-ietf-hybi-websocket-multiplexing-01 155 connBackend, resp, err := dialer.Dial(backendURL.String(), requestHeader) 156 if err != nil { 157 log.Printf("websocketproxy: couldn't dial to remote backend url %s", err) 158 if resp != nil { 159 // If the WebSocket handshake fails, ErrBadHandshake is returned 160 // along with a non-nil *http.Response so that callers can handle 161 // redirects, authentication, etcetera. 162 if err := copyResponse(rw, resp); err != nil { 163 log.Printf("websocketproxy: couldn't write response after failed remote backend handshake: %s", err) 164 } 165 } else { 166 http.Error(rw, http.StatusText(http.StatusServiceUnavailable), http.StatusServiceUnavailable) 167 } 168 return 169 } 170 defer connBackend.Close() 171 172 upgrader := w.Upgrader 173 if w.Upgrader == nil { 174 upgrader = DefaultUpgrader 175 } 176 177 // Only pass those headers to the upgrader. 178 upgradeHeader := http.Header{} 179 if hdr := resp.Header.Get("Sec-Websocket-Protocol"); hdr != "" { 180 upgradeHeader.Set("Sec-Websocket-Protocol", hdr) 181 } 182 if hdr := resp.Header.Get("Set-Cookie"); hdr != "" { 183 upgradeHeader.Set("Set-Cookie", hdr) 184 } 185 186 // Now upgrade the existing incoming request to a WebSocket connection. 187 // Also pass the header that we gathered from the Dial handshake. 188 connPub, err := upgrader.Upgrade(rw, req, upgradeHeader) 189 if err != nil { 190 log.Printf("websocketproxy: couldn't upgrade %s", err) 191 return 192 } 193 defer connPub.Close() 194 195 errClient := make(chan error, 1) 196 errBackend := make(chan error, 1) 197 replicateWebsocketConn := func(dst, src *websocket.Conn, errc chan error) { 198 for { 199 msgType, msg, err := src.ReadMessage() 200 if err != nil { 201 m := websocket.FormatCloseMessage(websocket.CloseNormalClosure, fmt.Sprintf("%v", err)) 202 if e, ok := err.(*websocket.CloseError); ok { 203 if e.Code != websocket.CloseNoStatusReceived { 204 m = websocket.FormatCloseMessage(e.Code, e.Text) 205 } 206 } 207 errc <- err 208 dst.WriteMessage(websocket.CloseMessage, m) 209 break 210 } 211 err = dst.WriteMessage(msgType, msg) 212 if err != nil { 213 errc <- err 214 break 215 } 216 } 217 } 218 219 go replicateWebsocketConn(connPub, connBackend, errClient) 220 go replicateWebsocketConn(connBackend, connPub, errBackend) 221 222 var message string 223 select { 224 case err = <-errClient: 225 message = "websocketproxy: Error when copying from backend to client: %v" 226 case err = <-errBackend: 227 message = "websocketproxy: Error when copying from client to backend: %v" 228 229 } 230 if e, ok := err.(*websocket.CloseError); !ok || e.Code == websocket.CloseAbnormalClosure { 231 log.Printf(message, err) 232 } 233 } 234 235 func copyHeader(dst, src http.Header) { 236 for k, vv := range src { 237 for _, v := range vv { 238 dst.Add(k, v) 239 } 240 } 241 } 242 243 func copyResponse(rw http.ResponseWriter, resp *http.Response) error { 244 copyHeader(rw.Header(), resp.Header) 245 rw.WriteHeader(resp.StatusCode) 246 defer resp.Body.Close() 247 248 _, err := io.Copy(rw, resp.Body) 249 return err 250 }