cloudflared-mirror/websocket/websocket_test.go

package websocket

import (
	"crypto/tls"
	"io"
	"math/rand"
	"net/http"
	"testing"

	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/assert"

	"golang.org/x/net/websocket"

	"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="
)

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.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
}

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}
}

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"))
}

func TestGenerateAcceptKey(t *testing.T) {
	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()

	go func() {
		errC <- hello.StartHelloWorldServer(logger, listener, shutdownC)
	}()

	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"))

	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)
	}

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