TUN-3403: Unit test for origin/proxy to test serving HTTP and Websocket

Author: chungthuang

connection/header.go

@@ -8,16 +8,14 @@ import (
 )
 
 const (
-	responseMetaHeaderField   = "cf-cloudflared-response-meta"
-	responseSourceCloudflared = "cloudflared"
-	responseSourceOrigin      = "origin"
+	responseMetaHeaderField = "cf-cloudflared-response-meta"
 )
 
 var (
 	canonicalResponseUserHeadersField = http.CanonicalHeaderKey(h2mux.ResponseUserHeadersField)
 	canonicalResponseMetaHeaderField  = http.CanonicalHeaderKey(responseMetaHeaderField)
-	responseMetaHeaderCfd             = mustInitRespMetaHeader(responseSourceCloudflared)
-	responseMetaHeaderOrigin          = mustInitRespMetaHeader(responseSourceOrigin)
+	responseMetaHeaderCfd             = mustInitRespMetaHeader("cloudflared")
+	responseMetaHeaderOrigin          = mustInitRespMetaHeader("origin")
 )
 
 type responseMetaHeader struct {

go.mod

@@ -27,6 +27,9 @@ require (
 	github.com/getsentry/raven-go v0.0.0-20180517221441-ed7bcb39ff10
 	github.com/gliderlabs/ssh v0.0.0-20191009160644-63518b5243e0
 	github.com/go-sql-driver/mysql v1.5.0
+	github.com/gobwas/httphead v0.0.0-20200921212729-da3d93bc3c58 // indirect
+	github.com/gobwas/pool v0.2.1 // indirect
+	github.com/gobwas/ws v1.0.4
 	github.com/golang-collections/collections v0.0.0-20130729185459-604e922904d3
 	github.com/google/go-cmp v0.5.2 // indirect
 	github.com/google/uuid v1.1.2

go.sum

@@ -233,6 +233,12 @@ github.com/go-sql-driver/mysql v1.4.0/go.mod h1:zAC/RDZ24gD3HViQzih4MyKcchzm+sOG
 github.com/go-sql-driver/mysql v1.5.0 h1:ozyZYNQW3x3HtqT1jira07DN2PArx2v7/mN66gGcHOs=
 github.com/go-sql-driver/mysql v1.5.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg=
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
+github.com/gobwas/httphead v0.0.0-20200921212729-da3d93bc3c58 h1:YyrUZvJaU8Q0QsoVo+xLFBgWDTam29PKea6GYmwvSiQ=
+github.com/gobwas/httphead v0.0.0-20200921212729-da3d93bc3c58/go.mod h1:L0fX3K22YWvt/FAX9NnzrNzcI4wNYi9Yku4O0LKYflo=
+github.com/gobwas/pool v0.2.1 h1:xfeeEhW7pwmX8nuLVlqbzVc7udMDrwetjEv+TZIz1og=
+github.com/gobwas/pool v0.2.1/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6WezmKEw=
+github.com/gobwas/ws v1.0.4 h1:5eXU1CZhpQdq5kXbKb+sECH5Ia5KiO6CYzIzdlVx6Bs=
+github.com/gobwas/ws v1.0.4/go.mod h1:szmBTxLgaFppYjEmNtny/v3w89xOydFnnZMcgRRu/EM=
 github.com/godbus/dbus/v5 v5.0.3/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
 github.com/gofrs/uuid v3.2.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
 github.com/gogo/googleapis v1.1.0/go.mod h1:gf4bu3Q80BeJ6H1S1vYPm8/ELATdvryBaNFGgqEef3s=

hello/hello.go

@@ -18,6 +18,11 @@ import (
 	"github.com/cloudflare/cloudflared/tlsconfig"
 )
 
+const (
+	UptimeRoute = "/uptime"
+	WSRoute     = "/ws"
+)
+
 type templateData struct {
 	ServerName string
 	Request    *http.Request
@@ -104,8 +109,8 @@ func StartHelloWorldServer(logger logger.Service, listener net.Listener, shutdow
 	}
 
 	muxer := http.NewServeMux()
-	muxer.HandleFunc("/uptime", uptimeHandler(time.Now()))
-	muxer.HandleFunc("/ws", websocketHandler(logger, upgrader))
+	muxer.HandleFunc(UptimeRoute, uptimeHandler(time.Now()))
+	muxer.HandleFunc(WSRoute, websocketHandler(logger, upgrader))
 	muxer.HandleFunc("/", rootHandler(serverName))
 	httpServer := &http.Server{Addr: listener.Addr().String(), Handler: muxer}
 	go func() {

origin/proxy.go

@@ -2,6 +2,7 @@ package origin
 
 import (
 	"bufio"
+	"context"
 	"crypto/tls"
 	"io"
 	"net/http"
@@ -112,20 +113,24 @@ func (c *client) proxyHTTP(w connection.ResponseWriter, req *http.Request) (*htt
 
 func (c *client) proxyWebsocket(w connection.ResponseWriter, req *http.Request) (*http.Response, error) {
 	c.setHostHeader(req)
-
 	conn, resp, err := websocket.ClientConnect(req, c.config.TLSConfig)
 	if err != nil {
 		return nil, err
 	}
-	defer conn.Close()
+
+	serveCtx, cancel := context.WithCancel(req.Context())
+	defer cancel()
+	go func() {
+		<-serveCtx.Done()
+		conn.Close()
+	}()
 	err = w.WriteRespHeaders(resp)
 	if err != nil {
 		return nil, errors.Wrap(err, "Error writing response header")
 	}
 	// Copy to/from stream to the undelying connection. Use the underlying
 	// connection because cloudflared doesn't operate on the message themselves
 	websocket.Stream(conn.UnderlyingConn(), w)
-
 	return resp, nil
 }
 

origin/proxy_test.go

@@ -0,0 +1,169 @@
+package origin
+
+import (
+	"bytes"
+	"context"
+	"crypto/tls"
+	"crypto/x509"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"sync"
+	"testing"
+
+	"github.com/cloudflare/cloudflared/connection"
+	"github.com/cloudflare/cloudflared/hello"
+	"github.com/cloudflare/cloudflared/logger"
+	"github.com/cloudflare/cloudflared/tlsconfig"
+
+	"github.com/gobwas/ws/wsutil"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type mockHTTPRespWriter struct {
+	*httptest.ResponseRecorder
+}
+
+func newMockHTTPRespWriter() *mockHTTPRespWriter {
+	return &mockHTTPRespWriter{
+		httptest.NewRecorder(),
+	}
+}
+
+func (w *mockHTTPRespWriter) WriteRespHeaders(resp *http.Response) error {
+	w.WriteHeader(resp.StatusCode)
+	for header, val := range resp.Header {
+		w.Header()[header] = val
+	}
+	return nil
+}
+
+func (w *mockHTTPRespWriter) WriteErrorResponse(err error) {
+	w.WriteHeader(http.StatusBadGateway)
+}
+
+func (w *mockHTTPRespWriter) Read(data []byte) (int, error) {
+	return 0, fmt.Errorf("mockHTTPRespWriter doesn't implement io.Reader")
+}
+
+type mockWSRespWriter struct {
+	*mockHTTPRespWriter
+	writeNotification chan []byte
+	reader            io.Reader
+}
+
+func newMockWSRespWriter(httpRespWriter *mockHTTPRespWriter, reader io.Reader) *mockWSRespWriter {
+	return &mockWSRespWriter{
+		httpRespWriter,
+		make(chan []byte),
+		reader,
+	}
+}
+
+func (w *mockWSRespWriter) Write(data []byte) (int, error) {
+	w.writeNotification <- data
+	return len(data), nil
+}
+
+func (w *mockWSRespWriter) respBody() io.ReadWriter {
+	data := <-w.writeNotification
+	return bytes.NewBuffer(data)
+}
+
+func (w *mockWSRespWriter) Read(data []byte) (int, error) {
+	return w.reader.Read(data)
+}
+
+func TestProxy(t *testing.T) {
+	logger, err := logger.New()
+	require.NoError(t, err)
+	// let runtime pick an available port
+	listener, err := hello.CreateTLSListener("127.0.0.1:0")
+	require.NoError(t, err)
+
+	originURL := &url.URL{
+		Scheme: "https",
+		Host:   listener.Addr().String(),
+	}
+	originCA := x509.NewCertPool()
+	helloCert, err := tlsconfig.GetHelloCertificateX509()
+	require.NoError(t, err)
+	originCA.AddCert(helloCert)
+	clientTLS := &tls.Config{
+		RootCAs: originCA,
+	}
+	proxyConfig := &ProxyConfig{
+		Client: &http.Transport{
+			TLSClientConfig: clientTLS,
+		},
+		URL:       originURL,
+		TLSConfig: clientTLS,
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	go func() {
+		hello.StartHelloWorldServer(logger, listener, ctx.Done())
+	}()
+
+	client := NewClient(proxyConfig, logger)
+	t.Run("testProxyHTTP", testProxyHTTP(t, client, originURL))
+	t.Run("testProxyWebsocket", testProxyWebsocket(t, client, originURL, clientTLS))
+	cancel()
+}
+
+func testProxyHTTP(t *testing.T, client connection.OriginClient, originURL *url.URL) func(t *testing.T) {
+	return func(t *testing.T) {
+		respWriter := newMockHTTPRespWriter()
+		req, err := http.NewRequest(http.MethodGet, originURL.String(), nil)
+		require.NoError(t, err)
+
+		err = client.Proxy(respWriter, req, false)
+		require.NoError(t, err)
+
+		assert.Equal(t, http.StatusOK, respWriter.Code)
+	}
+}
+
+func testProxyWebsocket(t *testing.T, client connection.OriginClient, originURL *url.URL, tlsConfig *tls.Config) func(t *testing.T) {
+	return func(t *testing.T) {
+		// WSRoute is a websocket echo handler
+		ctx, cancel := context.WithCancel(context.Background())
+		req, err := http.NewRequestWithContext(ctx, http.MethodGet, fmt.Sprintf("%s%s", originURL, hello.WSRoute), nil)
+
+		readPipe, writePipe := io.Pipe()
+		respWriter := newMockWSRespWriter(newMockHTTPRespWriter(), readPipe)
+
+		var wg sync.WaitGroup
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			err = client.Proxy(respWriter, req, true)
+			require.NoError(t, err)
+
+			require.Equal(t, http.StatusSwitchingProtocols, respWriter.Code)
+		}()
+
+		msg := []byte("test websocket")
+		err = wsutil.WriteClientText(writePipe, msg)
+		require.NoError(t, err)
+
+		// ReadServerText reads next data message from rw, considering that caller represents client side.
+		returnedMsg, err := wsutil.ReadServerText(respWriter.respBody())
+		require.NoError(t, err)
+		require.Equal(t, msg, returnedMsg)
+
+		err = wsutil.WriteClientBinary(writePipe, msg)
+		require.NoError(t, err)
+
+		returnedMsg, err = wsutil.ReadServerBinary(respWriter.respBody())
+		require.NoError(t, err)
+		require.Equal(t, msg, returnedMsg)
+
+		cancel()
+		wg.Wait()
+	}
+}