Revert "Fix race condition in proxy test"

This reverts commit 1d68ad253f.

proxy/proxy_test.go

@@ -465,10 +465,11 @@ func (r *replayer) Bytes() []byte {
 // eyeball sends tcp packets wrapped in websockets. (E.g: cloudflared access).
 func TestConnections(t *testing.T) {
 	logger := logger.Create(nil)
+	replayer := &replayer{rw: &bytes.Buffer{}}
 	type args struct {
 		ingressServiceScheme  string
 		originService         func(*testing.T, net.Listener)
-		eyeballResponseWriter func(io.Writer) connection.ResponseWriter
+		eyeballResponseWriter connection.ResponseWriter
 		eyeballRequestBody    io.ReadCloser
 
 		// Can be set to nil to show warp routing is not enabled.
@@ -497,9 +498,7 @@ func TestConnections(t *testing.T) {
 			args: args{
 				ingressServiceScheme:  "ws://",
 				originService:         runEchoWSService,
-				eyeballResponseWriter: func(w io.Writer) connection.ResponseWriter {
+				eyeballResponseWriter: newWSRespWriter(replayer),
-					return newWSRespWriter(w)
-				},
 				eyeballRequestBody:    newWSRequestBody([]byte("test1")),
 				connectionType:        connection.TypeWebsocket,
 				requestHeaders: map[string][]string{
@@ -523,9 +522,7 @@ func TestConnections(t *testing.T) {
 			args: args{
 				ingressServiceScheme:  "tcp://",
 				originService:         runEchoTCPService,
-				eyeballResponseWriter: func(w io.Writer) connection.ResponseWriter {
+				eyeballResponseWriter: newTCPRespWriter(replayer),
-					return newTCPRespWriter(w)
-				},
 				eyeballRequestBody:    newTCPRequestBody([]byte("test2")),
 				warpRoutingService:    ingress.NewWarpRoutingService(testWarpRouting),
 				connectionType:        connection.TypeTCP,
@@ -563,9 +560,7 @@ func TestConnections(t *testing.T) {
 			args: args{
 				ingressServiceScheme:  "tcp://",
 				originService:         runEchoTCPService,
-				eyeballResponseWriter: func(w io.Writer) connection.ResponseWriter {
+				eyeballResponseWriter: newWSRespWriter(replayer),
-					return newWSRespWriter(w)
-				},
 				eyeballRequestBody:    newWSRequestBody([]byte("test4")),
 				connectionType:        connection.TypeWebsocket,
 				requestHeaders: map[string][]string{
@@ -588,9 +583,7 @@ func TestConnections(t *testing.T) {
 			args: args{
 				ingressServiceScheme:  "tcp://",
 				originService:         runEchoTCPService,
-				eyeballResponseWriter: func(_ io.Writer) connection.ResponseWriter {
+				eyeballResponseWriter: newMockHTTPRespWriter(),
-					return newMockHTTPRespWriter()
-				},
 				eyeballRequestBody:    http.NoBody,
 				connectionType:        connection.TypeHTTP,
 				requestHeaders: map[string][]string{
@@ -598,7 +591,7 @@ func TestConnections(t *testing.T) {
 				},
 			},
 			want: want{
-				message: nil,
+				message: []byte{},
 				headers: map[string][]string{},
 			},
 		},
@@ -607,9 +600,7 @@ func TestConnections(t *testing.T) {
 			args: args{
 				ingressServiceScheme:  "tcp://",
 				originService:         runEchoTCPService,
-				eyeballResponseWriter: func(w io.Writer) connection.ResponseWriter {
+				eyeballResponseWriter: newTCPRespWriter(replayer),
-					return newTCPRespWriter(w)
-				},
 				eyeballRequestBody:    newTCPRequestBody([]byte("test2")),
 				connectionType:        connection.TypeTCP,
 				requestHeaders: map[string][]string{
@@ -617,7 +608,7 @@ func TestConnections(t *testing.T) {
 				},
 			},
 			want: want{
-				message: nil,
+				message: []byte{},
 				err:     true,
 			},
 		},
@@ -626,9 +617,7 @@ func TestConnections(t *testing.T) {
 			args: args{
 				ingressServiceScheme:  "ws://",
 				originService:         runEchoWSService,
-				eyeballResponseWriter: func(w io.Writer) connection.ResponseWriter {
+				eyeballResponseWriter: newWSRespWriter(replayer),
-					return newWSRespWriter(w)
-				},
 				eyeballRequestBody:    newWSRequestBody([]byte("test1")),
 				connectionType:        connection.TypeWebsocket,
 				requestHeaders: map[string][]string{
@@ -656,9 +645,7 @@ func TestConnections(t *testing.T) {
 					// closing the listener created by the test.
 					ln.Close()
 				},
-				eyeballResponseWriter: func(w io.Writer) connection.ResponseWriter {
+				eyeballResponseWriter: newTCPRespWriter(replayer),
-					return newTCPRespWriter(w)
-				},
 				eyeballRequestBody:    newTCPRequestBody([]byte("test2")),
 				connectionType:        connection.TypeTCP,
 				requestHeaders: map[string][]string{
@@ -666,7 +653,7 @@ func TestConnections(t *testing.T) {
 				},
 			},
 			want: want{
-				message: nil,
+				message: []byte{},
 				err:     true,
 			},
 		},
@@ -674,7 +661,6 @@ func TestConnections(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			replayer := &replayer{rw: &bytes.Buffer{}}
 			ctx, cancel := context.WithCancel(context.Background())
 			ln, err := net.Listen("tcp", "127.0.0.1:0")
 			require.NoError(t, err)
@@ -695,18 +681,15 @@ func TestConnections(t *testing.T) {
 			require.NoError(t, err)
 
 			req.Header = test.args.requestHeaders
+			respWriter := test.args.eyeballResponseWriter
 
-			var respWriter connection.ResponseWriter
 			if pipedReqBody, ok := test.args.eyeballRequestBody.(*pipedRequestBody); ok {
 				respWriter = newTCPRespWriter(pipedReqBody.pipedConn)
 				go func() {
 					resp := pipedReqBody.roundtrip(test.args.ingressServiceScheme + ln.Addr().String())
 					replayer.Write(resp)
 				}()
-			} else {
-				respWriter = test.args.eyeballResponseWriter(replayer)
 			}
-
 			if test.args.connectionType == connection.TypeTCP {
 				rwa := connection.NewHTTPResponseReadWriterAcker(respWriter, req)
 				err = proxy.ProxyTCP(ctx, rwa, &connection.TCPRequest{Dest: dest})