AUTH-1070: added SSH/protocol forwarding

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()
+// }