TUN-3489: Add unit tests to cover proxy logic in connection package of cloudflared

Author: chungthuang

connection/connection.go

@@ -44,7 +44,7 @@ type OriginClient interface {
 
 type ResponseWriter interface {
 	WriteRespHeaders(*http.Response) error
-	WriteErrorResponse(error)
+	WriteErrorResponse()
 	io.ReadWriter
 }
 

connection/connection_test.go

@@ -0,0 +1,113 @@
+package connection
+
+import (
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+	"time"
+
+	"github.com/cloudflare/cloudflared/cmd/cloudflared/ui"
+	"github.com/cloudflare/cloudflared/logger"
+	"github.com/gobwas/ws/wsutil"
+)
+
+const (
+	largeFileSize = 2 * 1024 * 1024
+)
+
+var (
+	testConfig = &Config{
+		OriginClient: &mockOriginClient{},
+		GracePeriod:  time.Millisecond * 100,
+	}
+	testLogger, _ = logger.New()
+	testOriginURL = &url.URL{
+		Scheme: "https",
+		Host:   "connectiontest.argotunnel.com",
+	}
+	testTunnelEventChan = make(chan ui.TunnelEvent)
+	testObserver        = &Observer{
+		testLogger,
+		m,
+		testTunnelEventChan,
+	}
+	testLargeResp = make([]byte, largeFileSize)
+)
+
+type testRequest struct {
+	name           string
+	endpoint       string
+	expectedStatus int
+	expectedBody   []byte
+	isProxyError   bool
+}
+
+type mockOriginClient struct {
+}
+
+func (moc *mockOriginClient) Proxy(w ResponseWriter, r *http.Request, isWebsocket bool) error {
+	if isWebsocket {
+		return wsEndpoint(w, r)
+	}
+	switch r.URL.Path {
+	case "/ok":
+		originRespEndpoint(w, http.StatusOK, []byte(http.StatusText(http.StatusOK)))
+	case "/large_file":
+		originRespEndpoint(w, http.StatusOK, testLargeResp)
+	case "/400":
+		originRespEndpoint(w, http.StatusBadRequest, []byte(http.StatusText(http.StatusBadRequest)))
+	case "/500":
+		originRespEndpoint(w, http.StatusInternalServerError, []byte(http.StatusText(http.StatusInternalServerError)))
+	case "/error":
+		return fmt.Errorf("Failed to proxy to origin")
+	default:
+		originRespEndpoint(w, http.StatusNotFound, []byte("page not found"))
+	}
+	return nil
+}
+
+type nowriter struct {
+	io.Reader
+}
+
+func (nowriter) Write(p []byte) (int, error) {
+	return 0, fmt.Errorf("Writer not implemented")
+}
+
+func wsEndpoint(w ResponseWriter, r *http.Request) error {
+	resp := &http.Response{
+		StatusCode: http.StatusSwitchingProtocols,
+	}
+	w.WriteRespHeaders(resp)
+	clientReader := nowriter{r.Body}
+	go func() {
+		for {
+			data, err := wsutil.ReadClientText(clientReader)
+			if err != nil {
+				return
+			}
+			if err := wsutil.WriteServerText(w, data); err != nil {
+				return
+			}
+		}
+	}()
+	<-r.Context().Done()
+	return nil
+}
+
+func originRespEndpoint(w ResponseWriter, status int, data []byte) {
+	resp := &http.Response{
+		StatusCode: status,
+	}
+	w.WriteRespHeaders(resp)
+	w.Write(data)
+}
+
+type mockConnectedFuse struct{}
+
+func (mcf mockConnectedFuse) Connected() {}
+
+func (mcf mockConnectedFuse) IsConnected() bool {
+	return true
+}

connection/h2mux.go

@@ -88,9 +88,9 @@ func (h *h2muxConnection) ServeNamedTunnel(ctx context.Context, namedTunnel *Nam
 			return err
 		}
 		rpcClient := newRegistrationRPCClient(ctx, stream, h.observer)
-		defer rpcClient.close()
+		defer rpcClient.Close()
 
-		if err = registerConnection(serveCtx, rpcClient, namedTunnel, connOptions, h.connIndex, h.observer); err != nil {
+		if err = rpcClient.RegisterConnection(serveCtx, namedTunnel, connOptions, h.connIndex, h.observer); err != nil {
 			return err
 		}
 		connectedFuse.Connected()
@@ -177,11 +177,16 @@ func (h *h2muxConnection) ServeStream(stream *h2mux.MuxedStream) error {
 
 	req, reqErr := h.newRequest(stream)
 	if reqErr != nil {
-		respWriter.WriteErrorResponse(reqErr)
+		respWriter.WriteErrorResponse()
 		return reqErr
 	}
 
-	return h.config.OriginClient.Proxy(respWriter, req, websocket.IsWebSocketUpgrade(req))
+	err := h.config.OriginClient.Proxy(respWriter, req, websocket.IsWebSocketUpgrade(req))
+	if err != nil {
+		respWriter.WriteErrorResponse()
+		return err
+	}
+	return nil
 }
 
 func (h *h2muxConnection) newRequest(stream *h2mux.MuxedStream) (*http.Request, error) {
@@ -206,7 +211,7 @@ func (rp *h2muxRespWriter) WriteRespHeaders(resp *http.Response) error {
 	return rp.WriteHeaders(headers)
 }
 
-func (rp *h2muxRespWriter) WriteErrorResponse(err error) {
+func (rp *h2muxRespWriter) WriteErrorResponse() {
 	rp.WriteHeaders([]h2mux.Header{
 		{Name: ":status", Value: "502"},
 		{Name: responseMetaHeaderField, Value: responseMetaHeaderCfd},