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  }