AUTH-1070: added SSH/protocol forwarding

websocket/websocket.go

@@ -9,11 +9,24 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"time"
 
 	"github.com/gorilla/websocket"
+	"github.com/sirupsen/logrus"
 )
 
-var stripWebsocketHeaders = []string {
+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",
@@ -21,6 +34,32 @@ var stripWebsocketHeaders = []string {
 	"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)
@@ -36,7 +75,7 @@ func ClientConnect(req *http.Request, tlsClientConfig *tls.Config) (*websocket.C
 	d := &websocket.Dialer{TLSClientConfig: tlsClientConfig}
 	conn, response, err := d.Dial(req.URL.String(), wsHeaders)
 	if err != nil {
-		return nil, nil, err
+		return nil, response, err
 	}
 	response.Header.Set("Sec-WebSocket-Accept", generateAcceptKey(req))
 	return conn, response, err
@@ -74,16 +113,58 @@ func Stream(conn, backendConn io.ReadWriter) {
 	<-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()
+
+		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
+		wsHeaders[key] = val
 	}
 	// Assume the header keys are in canonical format.
-	for _,  header := range stripWebsocketHeaders {
+	for _, header := range stripWebsocketHeaders {
 		wsHeaders.Del(header)
 	}
 	return wsHeaders
@@ -115,3 +196,19 @@ func changeRequestScheme(req *http.Request) string {
 		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
+		}
+	}
+}

websocket/websocket_test.go

@@ -1,100 +1,138 @@
 package websocket
 
 import (
-  "crypto/tls"
-  "io"
-  "math/rand"
-  "net/http"
-  "testing"
+	"crypto/tls"
+	"io"
+	"math/rand"
+	"net/http"
+	"testing"
 
-  "github.com/sirupsen/logrus"
-  "github.com/stretchr/testify/assert"
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
 
-  "golang.org/x/net/websocket"
+	"golang.org/x/net/websocket"
 
-  "github.com/cloudflare/cloudflared/hello"
-  "github.com/cloudflare/cloudflared/tlsconfig"
+	"github.com/cloudflare/cloudflared/hello"
+	"github.com/cloudflare/cloudflared/tlsconfig"
 )
 
 const (
-  // example in Sec-Websocket-Key in rfc6455
-  testSecWebsocketKey    = "dGhlIHNhbXBsZSBub25jZQ=="
-  // example Sec-Websocket-Accept in rfc6455
-  testSecWebsocketAccept = "s3pPLMBiTxaQ9kYGzzhZRbK+xOo="
+	// example in Sec-Websocket-Key in rfc6455
+	testSecWebsocketKey = "dGhlIHNhbXBsZSBub25jZQ=="
+	// example Sec-Websocket-Accept in rfc6455
+	testSecWebsocketAccept = "s3pPLMBiTxaQ9kYGzzhZRbK+xOo="
 )
 
 func testRequest(t *testing.T, url string, stream io.ReadWriter) *http.Request {
-  req, err := http.NewRequest("GET", url, stream)
-  if err != nil {
-    t.Fatalf("testRequestHeader error")
-  }
+	req, err := http.NewRequest("GET", url, stream)
+	if err != nil {
+		t.Fatalf("testRequestHeader error")
+	}
 
-  req.Header.Add("Connection", "Upgrade")
-  req.Header.Add("Upgrade", "WebSocket")
-  req.Header.Add("Sec-Websocket-Key", testSecWebsocketKey)
-  req.Header.Add("Sec-Websocket-Protocol", "tunnel-protocol")
-  req.Header.Add("Sec-Websocket-Version", "13")
-  req.Header.Add("User-Agent", "curl/7.59.0")
+	req.Header.Add("Connection", "Upgrade")
+	req.Header.Add("Upgrade", "WebSocket")
+	req.Header.Add("Sec-Websocket-Key", testSecWebsocketKey)
+	req.Header.Add("Sec-Websocket-Protocol", "tunnel-protocol")
+	req.Header.Add("Sec-Websocket-Version", "13")
+	req.Header.Add("User-Agent", "curl/7.59.0")
 
-  return req
+	return req
 }
 
 func websocketClientTLSConfig(t *testing.T) *tls.Config {
-  certPool, err := tlsconfig.LoadOriginCertPool(nil)
-  assert.NoError(t, err)
-  assert.NotNil(t, certPool)
-  return &tls.Config{RootCAs: certPool}
+	certPool, err := tlsconfig.LoadOriginCertPool(nil)
+	assert.NoError(t, err)
+	assert.NotNil(t, certPool)
+	return &tls.Config{RootCAs: certPool}
 }
 
 func TestWebsocketHeaders(t *testing.T) {
-  req := testRequest(t, "http://example.com", nil)
-  wsHeaders := websocketHeaders(req)
-  for _, header := range stripWebsocketHeaders {
-    assert.Empty(t, wsHeaders[header])
-  }
-  assert.Equal(t, "curl/7.59.0", wsHeaders.Get("User-Agent"))
+	req := testRequest(t, "http://example.com", nil)
+	wsHeaders := websocketHeaders(req)
+	for _, header := range stripWebsocketHeaders {
+		assert.Empty(t, wsHeaders[header])
+	}
+	assert.Equal(t, "curl/7.59.0", wsHeaders.Get("User-Agent"))
 }
 
 func TestGenerateAcceptKey(t *testing.T) {
-  req := testRequest(t, "http://example.com", nil)
-  assert.Equal(t, testSecWebsocketAccept, generateAcceptKey(req))
+	req := testRequest(t, "http://example.com", nil)
+	assert.Equal(t, testSecWebsocketAccept, generateAcceptKey(req))
 }
 
 func TestServe(t *testing.T) {
-  logger := logrus.New()
-  shutdownC := make(chan struct{})
-  errC := make(chan error)
-  listener, err := hello.CreateTLSListener("localhost:1111")
-  assert.NoError(t, err)
-  defer listener.Close()
+	logger := logrus.New()
+	shutdownC := make(chan struct{})
+	errC := make(chan error)
+	listener, err := hello.CreateTLSListener("localhost:1111")
+	assert.NoError(t, err)
+	defer listener.Close()
 
-  go func() {
-    errC <- hello.StartHelloWorldServer(logger, listener, shutdownC)
-  }()
+	go func() {
+		errC <- hello.StartHelloWorldServer(logger, listener, shutdownC)
+	}()
 
-  req := testRequest(t, "https://localhost:1111/ws", nil)
+	req := testRequest(t, "https://localhost:1111/ws", nil)
 
-  tlsConfig := websocketClientTLSConfig(t)
-  assert.NotNil(t, tlsConfig)
-  conn, resp, err := ClientConnect(req, tlsConfig)
-  assert.NoError(t, err)
-  assert.Equal(t, testSecWebsocketAccept, resp.Header.Get("Sec-WebSocket-Accept"))
+	tlsConfig := websocketClientTLSConfig(t)
+	assert.NotNil(t, tlsConfig)
+	conn, resp, err := ClientConnect(req, tlsConfig)
+	assert.NoError(t, err)
+	assert.Equal(t, testSecWebsocketAccept, resp.Header.Get("Sec-WebSocket-Accept"))
 
-  for i := 0; i < 1000; i++ {
-    messageSize := rand.Int() % 2048 + 1
-    clientMessage := make([]byte, messageSize)
-    // rand.Read always returns len(clientMessage) and a nil error
-    rand.Read(clientMessage)
-    err = conn.WriteMessage(websocket.BinaryFrame, clientMessage)
-    assert.NoError(t, err)
+	for i := 0; i < 1000; i++ {
+		messageSize := rand.Int()%2048 + 1
+		clientMessage := make([]byte, messageSize)
+		// rand.Read always returns len(clientMessage) and a nil error
+		rand.Read(clientMessage)
+		err = conn.WriteMessage(websocket.BinaryFrame, clientMessage)
+		assert.NoError(t, err)
 
-    messageType, message, err := conn.ReadMessage()
-    assert.NoError(t, err)
-    assert.Equal(t, websocket.BinaryFrame, messageType)
-    assert.Equal(t, clientMessage, message)
-  }
+		messageType, message, err := conn.ReadMessage()
+		assert.NoError(t, err)
+		assert.Equal(t, websocket.BinaryFrame, messageType)
+		assert.Equal(t, clientMessage, message)
+	}
 
-  conn.Close()
-  close(shutdownC)
-  <-errC
+	conn.Close()
+	close(shutdownC)
+	<-errC
 }
+
+// func TestStartProxyServer(t *testing.T) {
+// 	var wg sync.WaitGroup
+// 	remoteAddress := "localhost:1113"
+// 	listenerAddress := "localhost:1112"
+// 	message := "Good morning Austin! Time for another sunny day in the great state of Texas."
+// 	logger := logrus.New()
+// 	shutdownC := make(chan struct{})
+
+// 	listener, err := net.Listen("tcp", listenerAddress)
+// 	assert.NoError(t, err)
+// 	defer listener.Close()
+
+// 	remoteListener, err := net.Listen("tcp", remoteAddress)
+// 	assert.NoError(t, err)
+// 	defer remoteListener.Close()
+
+// 	wg.Add(1)
+// 	go func() {
+// 		defer wg.Done()
+// 		conn, err := remoteListener.Accept()
+// 		assert.NoError(t, err)
+// 		buf := make([]byte, len(message))
+// 		conn.Read(buf)
+// 		assert.Equal(t, string(buf), message)
+// 	}()
+
+// 	go func() {
+// 		StartProxyServer(logger, listener, remoteAddress, shutdownC)
+// 	}()
+
+// 	req := testRequest(t, fmt.Sprintf("http://%s/", listenerAddress), nil)
+// 	conn, _, err := ClientConnect(req, nil)
+// 	assert.NoError(t, err)
+// 	err = conn.WriteMessage(1, []byte(message))
+// 	assert.NoError(t, err)
+// 	wg.Wait()
+// }