TUN-3868: Refactor singleTCPService and bridgeService to tcpOverWSService and rawTCPService

origin/proxy.go

@@ -13,7 +13,6 @@ import (
 	"github.com/cloudflare/cloudflared/connection"
 	"github.com/cloudflare/cloudflared/ingress"
 	tunnelpogs "github.com/cloudflare/cloudflared/tunnelrpc/pogs"
-	"github.com/cloudflare/cloudflared/websocket"
 	"github.com/pkg/errors"
 	"github.com/rs/zerolog"
 )
@@ -45,6 +44,7 @@ func NewOriginProxy(
 	}
 }
 
+// Caller is responsible for writing any error to ResponseWriter
 func (p *proxy) Proxy(w connection.ResponseWriter, req *http.Request, sourceConnectionType connection.Type) error {
 	incrementRequests()
 	defer decrementConcurrentRequests()
@@ -62,27 +62,31 @@ func (p *proxy) Proxy(w connection.ResponseWriter, req *http.Request, sourceConn
 			p.log.Error().Msg(err.Error())
 			return err
 		}
-		resp, err := p.proxyConnection(serveCtx, w, req, sourceConnectionType, p.warpRouting.Proxy)
-		if err != nil {
+		logFields := logFields{
+			cfRay:   cfRay,
+			lbProbe: lbProbe,
+			rule:    ingress.ServiceWarpRouting,
+		}
+		if err := p.proxyStreamRequest(serveCtx, w, req, sourceConnectionType, p.warpRouting.Proxy, logFields); err != nil {
 			p.logRequestError(err, cfRay, ingress.ServiceWarpRouting)
-			w.WriteErrorResponse()
 			return err
 		}
-		p.logOriginResponse(resp, cfRay, lbProbe, ingress.ServiceWarpRouting)
 		return nil
 	}
 
 	rule, ruleNum := p.ingressRules.FindMatchingRule(req.Host, req.URL.Path)
-	p.logRequest(req, cfRay, lbProbe, ruleNum)
+	logFields := logFields{
+		cfRay:   cfRay,
+		lbProbe: lbProbe,
+		rule:    ruleNum,
+	}
+	p.logRequest(req, logFields)
 
 	if sourceConnectionType == connection.TypeHTTP {
-		resp, err := p.proxyHTTP(w, req, rule)
-		if err != nil {
-			p.logErrorAndWriteResponse(w, err, cfRay, ruleNum)
+		if err := p.proxyHTTPRequest(w, req, rule, logFields); err != nil {
+			p.logRequestError(err, cfRay, ruleNum)
 			return err
 		}
-
-		p.logOriginResponse(resp, cfRay, lbProbe, ruleNum)
 		return nil
 	}
 
@@ -92,22 +96,14 @@ func (p *proxy) Proxy(w connection.ResponseWriter, req *http.Request, sourceConn
 		return fmt.Errorf("Not a connection-oriented service")
 	}
 
-	resp, err := p.proxyConnection(serveCtx, w, req, sourceConnectionType, connectionProxy)
-	if err != nil {
-		p.logErrorAndWriteResponse(w, err, cfRay, ruleNum)
+	if err := p.proxyStreamRequest(serveCtx, w, req, sourceConnectionType, connectionProxy, logFields); err != nil {
+		p.logRequestError(err, cfRay, ruleNum)
 		return err
 	}
-
-	p.logOriginResponse(resp, cfRay, lbProbe, ruleNum)
 	return nil
 }
 
-func (p *proxy) logErrorAndWriteResponse(w connection.ResponseWriter, err error, cfRay string, ruleNum int) {
-	p.logRequestError(err, cfRay, ruleNum)
-	w.WriteErrorResponse()
-}
-
-func (p *proxy) proxyHTTP(w connection.ResponseWriter, req *http.Request, rule *ingress.Rule) (*http.Response, error) {
+func (p *proxy) proxyHTTPRequest(w connection.ResponseWriter, req *http.Request, rule *ingress.Rule, fields logFields) error {
 	// Support for WSGI Servers by switching transfer encoding from chunked to gzip/deflate
 	if rule.Config.DisableChunkedEncoding {
 		req.TransferEncoding = []string{"gzip", "deflate"}
@@ -123,18 +119,18 @@ func (p *proxy) proxyHTTP(w connection.ResponseWriter, req *http.Request, rule *
 	httpService, ok := rule.Service.(ingress.HTTPOriginProxy)
 	if !ok {
 		p.log.Error().Msgf("%s is not a http service", rule.Service)
-		return nil, fmt.Errorf("Not a http service")
+		return fmt.Errorf("Not a http service")
 	}
 
 	resp, err := httpService.RoundTrip(req)
 	if err != nil {
-		return nil, errors.Wrap(err, "Error proxying request to origin")
+		return errors.Wrap(err, "Error proxying request to origin")
 	}
 	defer resp.Body.Close()
 
 	err = w.WriteRespHeaders(resp.StatusCode, resp.Header)
 	if err != nil {
-		return nil, errors.Wrap(err, "Error writing response header")
+		return errors.Wrap(err, "Error writing response header")
 	}
 	if connection.IsServerSentEvent(resp.Header) {
 		p.log.Debug().Msg("Detected Server-Side Events from Origin")
@@ -146,43 +142,30 @@ func (p *proxy) proxyHTTP(w connection.ResponseWriter, req *http.Request, rule *
 		defer p.bufferPool.Put(buf)
 		_, _ = io.CopyBuffer(w, resp.Body, buf)
 	}
-	return resp, nil
+	p.logOriginResponse(resp, fields)
+	return nil
 }
 
-func (p *proxy) proxyConnection(
+// proxyStreamRequest first establish a connection with origin, then it writes the status code and headers, and finally it streams data between
+// eyeball and origin.
+func (p *proxy) proxyStreamRequest(
 	serveCtx context.Context,
 	w connection.ResponseWriter,
 	req *http.Request,
 	sourceConnectionType connection.Type,
 	connectionProxy ingress.StreamBasedOriginProxy,
-) (*http.Response, error) {
-	originConn, connectionResp, err := connectionProxy.EstablishConnection(req)
+	fields logFields,
+) error {
+	originConn, resp, err := connectionProxy.EstablishConnection(req)
 	if err != nil {
-		return nil, err
+		return err
+	}
+	if resp.Body != nil {
+		defer resp.Body.Close()
 	}
 
-	var eyeballConn io.ReadWriter = w
-	respHeader := http.Header{}
-	if connectionResp != nil {
-		respHeader = connectionResp.Header
-	}
-	if sourceConnectionType == connection.TypeWebsocket {
-		wsReadWriter := websocket.NewConn(serveCtx, w, p.log)
-		// If cloudflared <-> origin is not websocket, we need to decode TCP data out of WS frames
-		if originConn.Type() != sourceConnectionType {
-			eyeballConn = wsReadWriter
-		}
-	}
-	status := http.StatusSwitchingProtocols
-	resp := &http.Response{
-		Status:        http.StatusText(status),
-		StatusCode:    status,
-		Header:        respHeader,
-		ContentLength: -1,
-	}
-	w.WriteRespHeaders(http.StatusSwitchingProtocols, respHeader)
-	if err != nil {
-		return nil, errors.Wrap(err, "Error writing response header")
+	if err = w.WriteRespHeaders(resp.StatusCode, resp.Header); err != nil {
+		return err
 	}
 
 	streamCtx, cancel := context.WithCancel(serveCtx)
@@ -194,8 +177,9 @@ func (p *proxy) proxyConnection(
 		originConn.Close()
 	}()
 
-	originConn.Stream(eyeballConn, p.log)
-	return resp, nil
+	originConn.Stream(serveCtx, w, p.log)
+	p.logOriginResponse(resp, fields)
+	return nil
 }
 
 func (p *proxy) writeEventStream(w connection.ResponseWriter, respBody io.ReadCloser) {
@@ -215,39 +199,45 @@ func (p *proxy) appendTagHeaders(r *http.Request) {
 	}
 }
 
-func (p *proxy) logRequest(r *http.Request, cfRay string, lbProbe bool, rule interface{}) {
-	if cfRay != "" {
-		p.log.Debug().Msgf("CF-RAY: %s %s %s %s", cfRay, r.Method, r.URL, r.Proto)
-	} else if lbProbe {
-		p.log.Debug().Msgf("CF-RAY: %s Load Balancer health check %s %s %s", cfRay, r.Method, r.URL, r.Proto)
+type logFields struct {
+	cfRay   string
+	lbProbe bool
+	rule    interface{}
+}
+
+func (p *proxy) logRequest(r *http.Request, fields logFields) {
+	if fields.cfRay != "" {
+		p.log.Debug().Msgf("CF-RAY: %s %s %s %s", fields.cfRay, r.Method, r.URL, r.Proto)
+	} else if fields.lbProbe {
+		p.log.Debug().Msgf("CF-RAY: %s Load Balancer health check %s %s %s", fields.cfRay, r.Method, r.URL, r.Proto)
 	} else {
 		p.log.Debug().Msgf("All requests should have a CF-RAY header. Please open a support ticket with Cloudflare. %s %s %s ", r.Method, r.URL, r.Proto)
 	}
-	p.log.Debug().Msgf("CF-RAY: %s Request Headers %+v", cfRay, r.Header)
-	p.log.Debug().Msgf("CF-RAY: %s Serving with ingress rule %v", cfRay, rule)
+	p.log.Debug().Msgf("CF-RAY: %s Request Headers %+v", fields.cfRay, r.Header)
+	p.log.Debug().Msgf("CF-RAY: %s Serving with ingress rule %v", fields.cfRay, fields.rule)
 
 	if contentLen := r.ContentLength; contentLen == -1 {
-		p.log.Debug().Msgf("CF-RAY: %s Request Content length unknown", cfRay)
+		p.log.Debug().Msgf("CF-RAY: %s Request Content length unknown", fields.cfRay)
 	} else {
-		p.log.Debug().Msgf("CF-RAY: %s Request content length %d", cfRay, contentLen)
+		p.log.Debug().Msgf("CF-RAY: %s Request content length %d", fields.cfRay, contentLen)
 	}
 }
 
-func (p *proxy) logOriginResponse(r *http.Response, cfRay string, lbProbe bool, rule interface{}) {
-	responseByCode.WithLabelValues(strconv.Itoa(r.StatusCode)).Inc()
-	if cfRay != "" {
-		p.log.Debug().Msgf("CF-RAY: %s Status: %s served by ingress %d", cfRay, r.Status, rule)
-	} else if lbProbe {
-		p.log.Debug().Msgf("Response to Load Balancer health check %s", r.Status)
+func (p *proxy) logOriginResponse(resp *http.Response, fields logFields) {
+	responseByCode.WithLabelValues(strconv.Itoa(resp.StatusCode)).Inc()
+	if fields.cfRay != "" {
+		p.log.Debug().Msgf("CF-RAY: %s Status: %s served by ingress %d", fields.cfRay, resp.Status, fields.rule)
+	} else if fields.lbProbe {
+		p.log.Debug().Msgf("Response to Load Balancer health check %s", resp.Status)
 	} else {
-		p.log.Debug().Msgf("Status: %s served by ingress %v", r.Status, rule)
+		p.log.Debug().Msgf("Status: %s served by ingress %v", resp.Status, fields.rule)
 	}
-	p.log.Debug().Msgf("CF-RAY: %s Response Headers %+v", cfRay, r.Header)
+	p.log.Debug().Msgf("CF-RAY: %s Response Headers %+v", fields.cfRay, resp.Header)
 
-	if contentLen := r.ContentLength; contentLen == -1 {
-		p.log.Debug().Msgf("CF-RAY: %s Response content length unknown", cfRay)
+	if contentLen := resp.ContentLength; contentLen == -1 {
+		p.log.Debug().Msgf("CF-RAY: %s Response content length unknown", fields.cfRay)
 	} else {
-		p.log.Debug().Msgf("CF-RAY: %s Response content length %d", cfRay, contentLen)
+		p.log.Debug().Msgf("CF-RAY: %s Response content length %d", fields.cfRay, contentLen)
 	}
 }
 

origin/proxy_test.go

@@ -347,10 +347,7 @@ func TestProxyError(t *testing.T) {
 	req, err := http.NewRequest(http.MethodGet, "http://127.0.0.1", nil)
 	assert.NoError(t, err)
 
-	err = proxy.Proxy(respWriter, req, connection.TypeHTTP)
-	assert.Error(t, err)
-	assert.Equal(t, http.StatusBadGateway, respWriter.Code)
-	assert.Equal(t, "http response error", respWriter.Body.String())
+	assert.Error(t, proxy.Proxy(respWriter, req, connection.TypeHTTP))
 }
 
 type replayer struct {
@@ -421,15 +418,17 @@ func TestConnections(t *testing.T) {
 			originService:        runEchoWSService,
 			eyeballService:       newWSRespWriter([]byte("test1"), replayer),
 			connectionType:       connection.TypeWebsocket,
-			requestHeaders: map[string][]string{
-				"Test-Cloudflared-Echo": []string{"Echo"},
+			requestHeaders: http.Header{
+				// Example key from https://tools.ietf.org/html/rfc6455#section-1.2
+				"Sec-Websocket-Key":     {"dGhlIHNhbXBsZSBub25jZQ=="},
+				"Test-Cloudflared-Echo": {"Echo"},
 			},
 			wantMessage: []byte("echo-test1"),
-			wantHeaders: map[string][]string{
-				"Connection":            []string{"Upgrade"},
-				"Sec-Websocket-Accept":  []string{"Kfh9QIsMVZcl6xEPYxPHzW8SZ8w="},
-				"Upgrade":               []string{"websocket"},
-				"Test-Cloudflared-Echo": []string{"Echo"},
+			wantHeaders: http.Header{
+				"Connection":            {"Upgrade"},
+				"Sec-Websocket-Accept":  {"s3pPLMBiTxaQ9kYGzzhZRbK+xOo="},
+				"Upgrade":               {"websocket"},
+				"Test-Cloudflared-Echo": {"Echo"},
 			},
 		},
 		{
@@ -441,25 +440,23 @@ func TestConnections(t *testing.T) {
 				replayer,
 			),
 			connectionType: connection.TypeTCP,
-			requestHeaders: map[string][]string{
-				"Cf-Cloudflared-Proxy-Src": []string{"non-blank-value"},
+			requestHeaders: http.Header{
+				"Cf-Cloudflared-Proxy-Src": {"non-blank-value"},
 			},
 			wantMessage: []byte("echo-test2"),
-			wantHeaders: http.Header{},
 		},
 		{
 			name:                 "tcp-ws proxy",
 			ingressServicePrefix: "ws://",
 			originService:        runEchoWSService,
 			eyeballService:       newPipedWSWriter(&mockTCPRespWriter{}, []byte("test3")),
-			requestHeaders: map[string][]string{
-				"Cf-Cloudflared-Proxy-Src": []string{"non-blank-value"},
+			requestHeaders: http.Header{
+				"Cf-Cloudflared-Proxy-Src": {"non-blank-value"},
 			},
 			connectionType: connection.TypeTCP,
 			wantMessage:    []byte("echo-test3"),
 			// We expect no headers here because they are sent back via
 			// the stream.
-			wantHeaders: http.Header{},
 		},
 		{
 			name:                 "ws-tcp proxy",
@@ -467,8 +464,16 @@ func TestConnections(t *testing.T) {
 			originService:        runEchoTCPService,
 			eyeballService:       newWSRespWriter([]byte("test4"), replayer),
 			connectionType:       connection.TypeWebsocket,
-			wantMessage:          []byte("echo-test4"),
-			wantHeaders:          http.Header{},
+			requestHeaders: http.Header{
+				// Example key from https://tools.ietf.org/html/rfc6455#section-1.2
+				"Sec-Websocket-Key": {"dGhlIHNhbXBsZSBub25jZQ=="},
+			},
+			wantMessage: []byte("echo-test4"),
+			wantHeaders: http.Header{
+				"Connection":           {"Upgrade"},
+				"Sec-Websocket-Accept": {"s3pPLMBiTxaQ9kYGzzhZRbK+xOo="},
+				"Upgrade":              {"websocket"},
+			},
 		},
 	}
 
@@ -477,19 +482,18 @@ func TestConnections(t *testing.T) {
 			ctx, cancel := context.WithCancel(context.Background())
 			ln, err := net.Listen("tcp", "127.0.0.1:0")
 			require.NoError(t, err)
+			// Starts origin service
 			test.originService(t, ln)
+
 			ingressRule := createSingleIngressConfig(t, test.ingressServicePrefix+ln.Addr().String())
 			var wg sync.WaitGroup
 			errC := make(chan error)
 			ingressRule.StartOrigins(&wg, logger, ctx.Done(), errC)
 			proxy := NewOriginProxy(ingressRule, ingress.NewWarpRoutingService(), testTags, logger)
+
 			req, err := http.NewRequest(http.MethodGet, test.ingressServicePrefix+ln.Addr().String(), nil)
 			require.NoError(t, err)
-			reqHeaders := make(http.Header)
-			for k, vs := range test.requestHeaders {
-				reqHeaders[k] = vs
-			}
-			req.Header = reqHeaders
+			req.Header = test.requestHeaders
 
 			if pipedWS, ok := test.eyeballService.(*pipedWSWriter); ok {
 				go func() {