cloudflared-mirror/websocket/websocket.go

220 lines
5.8 KiB
Go

package websocket
import (
"bufio"
"crypto/sha1"
"crypto/tls"
"encoding/base64"
"errors"
"io"
"net"
"net/http"
"time"
"github.com/gorilla/websocket"
"github.com/sirupsen/logrus"
)
const (
// Time allowed to write a message to the peer.
writeWait = 10 * time.Second
// Time allowed to read the next pong message from the peer.
pongWait = 60 * time.Second
// Send pings to peer with this period. Must be less than pongWait.
pingPeriod = (pongWait * 9) / 10
)
var stripWebsocketHeaders = []string{
"Upgrade",
"Connection",
"Sec-Websocket-Key",
"Sec-Websocket-Version",
"Sec-Websocket-Extensions",
}
// Conn is a wrapper around the standard gorilla websocket
// but implements a ReadWriter
type Conn struct {
*websocket.Conn
}
// Read will read messages from the websocket connection
func (c *Conn) Read(p []byte) (int, error) {
_, message, err := c.Conn.ReadMessage()
if err != nil {
return 0, err
}
return copy(p, message), nil
}
// Write will write messages to the websocket connection
func (c *Conn) Write(p []byte) (int, error) {
if err := c.Conn.WriteMessage(websocket.BinaryMessage, p); err != nil {
return 0, err
}
return len(p), nil
}
// IsWebSocketUpgrade checks to see if the request is a WebSocket connection.
func IsWebSocketUpgrade(req *http.Request) bool {
return websocket.IsWebSocketUpgrade(req)
}
// ClientConnect creates a WebSocket client connection for provided request. Caller is responsible for closing
// the connection. The response body may not contain the entire response and does
// not need to be closed by the application.
func ClientConnect(req *http.Request, tlsClientConfig *tls.Config) (*websocket.Conn, *http.Response, error) {
req.URL.Scheme = changeRequestScheme(req)
wsHeaders := websocketHeaders(req)
d := &websocket.Dialer{TLSClientConfig: tlsClientConfig}
conn, response, err := d.Dial(req.URL.String(), wsHeaders)
if err != nil {
return nil, response, err
}
response.Header.Set("Sec-WebSocket-Accept", generateAcceptKey(req))
return conn, response, err
}
// HijackConnection takes over an HTTP connection. Caller is responsible for closing connection.
func HijackConnection(w http.ResponseWriter) (net.Conn, *bufio.ReadWriter, error) {
hj, ok := w.(http.Hijacker)
if !ok {
return nil, nil, errors.New("hijack error")
}
conn, brw, err := hj.Hijack()
if err != nil {
return nil, nil, err
}
return conn, brw, nil
}
// Stream copies copy data to & from provided io.ReadWriters.
func Stream(conn, backendConn io.ReadWriter) {
proxyDone := make(chan struct{}, 2)
go func() {
io.Copy(conn, backendConn)
proxyDone <- struct{}{}
}()
go func() {
io.Copy(backendConn, conn)
proxyDone <- struct{}{}
}()
// If one side is done, we are done.
<-proxyDone
}
// StartProxyServer will start a websocket server that will decode
// the websocket data and write the resulting data to the provided
// address
func StartProxyServer(logger *logrus.Logger, listener net.Listener, remote string, shutdownC <-chan struct{}) error {
upgrader := websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
}
httpServer := &http.Server{Addr: listener.Addr().String(), Handler: nil}
go func() {
<-shutdownC
httpServer.Close()
}()
http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
stream, err := net.Dial("tcp", remote)
if err != nil {
logger.WithError(err).Error("Cannot connect to remote.")
return
}
defer stream.Close()
if !websocket.IsWebSocketUpgrade(r) {
w.Write(nonWebSocketRequestPage())
return
}
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
logger.WithError(err).Error("failed to upgrade")
return
}
conn.SetReadDeadline(time.Now().Add(pongWait))
conn.SetPongHandler(func(string) error { conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
done := make(chan struct{})
go pinger(logger, conn, done)
defer func() {
<-done
conn.Close()
}()
Stream(&Conn{conn}, stream)
})
return httpServer.Serve(listener)
}
// the gorilla websocket library sets its own Upgrade, Connection, Sec-WebSocket-Key,
// Sec-WebSocket-Version and Sec-Websocket-Extensions headers.
// https://github.com/gorilla/websocket/blob/master/client.go#L189-L194.
func websocketHeaders(req *http.Request) http.Header {
wsHeaders := make(http.Header)
for key, val := range req.Header {
wsHeaders[key] = val
}
// Assume the header keys are in canonical format.
for _, header := range stripWebsocketHeaders {
wsHeaders.Del(header)
}
wsHeaders.Set("Host", req.Host) // See TUN-1097
return wsHeaders
}
// sha1Base64 sha1 and then base64 encodes str.
func sha1Base64(str string) string {
hasher := sha1.New()
io.WriteString(hasher, str)
hash := hasher.Sum(nil)
return base64.StdEncoding.EncodeToString(hash)
}
// generateAcceptKey returns the string needed for the Sec-WebSocket-Accept header.
// https://tools.ietf.org/html/rfc6455#section-1.3 describes this process in more detail.
func generateAcceptKey(req *http.Request) string {
return sha1Base64(req.Header.Get("Sec-WebSocket-Key") + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11")
}
// changeRequestScheme is needed as the gorilla websocket library requires the ws scheme.
// (even though it changes it back to http/https, but ¯\_(ツ)_/¯.)
func changeRequestScheme(req *http.Request) string {
switch req.URL.Scheme {
case "https":
return "wss"
case "http":
return "ws"
default:
return req.URL.Scheme
}
}
// pinger simulates the websocket connection to keep it alive
func pinger(logger *logrus.Logger, ws *websocket.Conn, done chan struct{}) {
ticker := time.NewTicker(pingPeriod)
defer ticker.Stop()
for {
select {
case <-ticker.C:
if err := ws.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(writeWait)); err != nil {
logger.WithError(err).Debug("failed to send ping message")
}
case <-done:
return
}
}
}