Fix race condition in proxy test
This commit is contained in:
		
							parent
							
								
									301b2baca3
								
							
						
					
					
						commit
						1d68ad253f
					
				|  | @ -465,11 +465,10 @@ func (r *replayer) Bytes() []byte { | ||||||
| // eyeball sends tcp packets wrapped in websockets. (E.g: cloudflared access).
 | // eyeball sends tcp packets wrapped in websockets. (E.g: cloudflared access).
 | ||||||
| func TestConnections(t *testing.T) { | func TestConnections(t *testing.T) { | ||||||
| 	logger := logger.Create(nil) | 	logger := logger.Create(nil) | ||||||
| 	replayer := &replayer{rw: &bytes.Buffer{}} |  | ||||||
| 	type args struct { | 	type args struct { | ||||||
| 		ingressServiceScheme  string | 		ingressServiceScheme  string | ||||||
| 		originService         func(*testing.T, net.Listener) | 		originService         func(*testing.T, net.Listener) | ||||||
| 		eyeballResponseWriter connection.ResponseWriter | 		eyeballResponseWriter func(io.Writer) connection.ResponseWriter | ||||||
| 		eyeballRequestBody    io.ReadCloser | 		eyeballRequestBody    io.ReadCloser | ||||||
| 
 | 
 | ||||||
| 		// Can be set to nil to show warp routing is not enabled.
 | 		// Can be set to nil to show warp routing is not enabled.
 | ||||||
|  | @ -496,11 +495,13 @@ func TestConnections(t *testing.T) { | ||||||
| 		{ | 		{ | ||||||
| 			name: "ws-ws proxy", | 			name: "ws-ws proxy", | ||||||
| 			args: args{ | 			args: args{ | ||||||
| 				ingressServiceScheme:  "ws://", | 				ingressServiceScheme: "ws://", | ||||||
| 				originService:         runEchoWSService, | 				originService:        runEchoWSService, | ||||||
| 				eyeballResponseWriter: newWSRespWriter(replayer), | 				eyeballResponseWriter: func(w io.Writer) connection.ResponseWriter { | ||||||
| 				eyeballRequestBody:    newWSRequestBody([]byte("test1")), | 					return newWSRespWriter(w) | ||||||
| 				connectionType:        connection.TypeWebsocket, | 				}, | ||||||
|  | 				eyeballRequestBody: newWSRequestBody([]byte("test1")), | ||||||
|  | 				connectionType:     connection.TypeWebsocket, | ||||||
| 				requestHeaders: map[string][]string{ | 				requestHeaders: map[string][]string{ | ||||||
| 					// Example key from https://tools.ietf.org/html/rfc6455#section-1.2
 | 					// Example key from https://tools.ietf.org/html/rfc6455#section-1.2
 | ||||||
| 					"Sec-Websocket-Key":     {"dGhlIHNhbXBsZSBub25jZQ=="}, | 					"Sec-Websocket-Key":     {"dGhlIHNhbXBsZSBub25jZQ=="}, | ||||||
|  | @ -520,12 +521,14 @@ func TestConnections(t *testing.T) { | ||||||
| 		{ | 		{ | ||||||
| 			name: "tcp-tcp proxy", | 			name: "tcp-tcp proxy", | ||||||
| 			args: args{ | 			args: args{ | ||||||
| 				ingressServiceScheme:  "tcp://", | 				ingressServiceScheme: "tcp://", | ||||||
| 				originService:         runEchoTCPService, | 				originService:        runEchoTCPService, | ||||||
| 				eyeballResponseWriter: newTCPRespWriter(replayer), | 				eyeballResponseWriter: func(w io.Writer) connection.ResponseWriter { | ||||||
| 				eyeballRequestBody:    newTCPRequestBody([]byte("test2")), | 					return newTCPRespWriter(w) | ||||||
| 				warpRoutingService:    ingress.NewWarpRoutingService(testWarpRouting), | 				}, | ||||||
| 				connectionType:        connection.TypeTCP, | 				eyeballRequestBody: newTCPRequestBody([]byte("test2")), | ||||||
|  | 				warpRoutingService: ingress.NewWarpRoutingService(testWarpRouting), | ||||||
|  | 				connectionType:     connection.TypeTCP, | ||||||
| 				requestHeaders: map[string][]string{ | 				requestHeaders: map[string][]string{ | ||||||
| 					"Cf-Cloudflared-Proxy-Src": {"non-blank-value"}, | 					"Cf-Cloudflared-Proxy-Src": {"non-blank-value"}, | ||||||
| 				}, | 				}, | ||||||
|  | @ -558,11 +561,13 @@ func TestConnections(t *testing.T) { | ||||||
| 		{ | 		{ | ||||||
| 			name: "ws-tcp proxy", | 			name: "ws-tcp proxy", | ||||||
| 			args: args{ | 			args: args{ | ||||||
| 				ingressServiceScheme:  "tcp://", | 				ingressServiceScheme: "tcp://", | ||||||
| 				originService:         runEchoTCPService, | 				originService:        runEchoTCPService, | ||||||
| 				eyeballResponseWriter: newWSRespWriter(replayer), | 				eyeballResponseWriter: func(w io.Writer) connection.ResponseWriter { | ||||||
| 				eyeballRequestBody:    newWSRequestBody([]byte("test4")), | 					return newWSRespWriter(w) | ||||||
| 				connectionType:        connection.TypeWebsocket, | 				}, | ||||||
|  | 				eyeballRequestBody: newWSRequestBody([]byte("test4")), | ||||||
|  | 				connectionType:     connection.TypeWebsocket, | ||||||
| 				requestHeaders: map[string][]string{ | 				requestHeaders: map[string][]string{ | ||||||
| 					// Example key from https://tools.ietf.org/html/rfc6455#section-1.2
 | 					// Example key from https://tools.ietf.org/html/rfc6455#section-1.2
 | ||||||
| 					"Sec-Websocket-Key": {"dGhlIHNhbXBsZSBub25jZQ=="}, | 					"Sec-Websocket-Key": {"dGhlIHNhbXBsZSBub25jZQ=="}, | ||||||
|  | @ -581,45 +586,51 @@ func TestConnections(t *testing.T) { | ||||||
| 			// Send (unexpected) HTTP when origin expects WS (to unwrap for raw TCP)
 | 			// Send (unexpected) HTTP when origin expects WS (to unwrap for raw TCP)
 | ||||||
| 			name: "http-(ws)tcp proxy", | 			name: "http-(ws)tcp proxy", | ||||||
| 			args: args{ | 			args: args{ | ||||||
| 				ingressServiceScheme:  "tcp://", | 				ingressServiceScheme: "tcp://", | ||||||
| 				originService:         runEchoTCPService, | 				originService:        runEchoTCPService, | ||||||
| 				eyeballResponseWriter: newMockHTTPRespWriter(), | 				eyeballResponseWriter: func(_ io.Writer) connection.ResponseWriter { | ||||||
| 				eyeballRequestBody:    http.NoBody, | 					return newMockHTTPRespWriter() | ||||||
| 				connectionType:        connection.TypeHTTP, | 				}, | ||||||
|  | 				eyeballRequestBody: http.NoBody, | ||||||
|  | 				connectionType:     connection.TypeHTTP, | ||||||
| 				requestHeaders: map[string][]string{ | 				requestHeaders: map[string][]string{ | ||||||
| 					"Cf-Cloudflared-Proxy-Src": {"non-blank-value"}, | 					"Cf-Cloudflared-Proxy-Src": {"non-blank-value"}, | ||||||
| 				}, | 				}, | ||||||
| 			}, | 			}, | ||||||
| 			want: want{ | 			want: want{ | ||||||
| 				message: []byte{}, | 				message: nil, | ||||||
| 				headers: map[string][]string{}, | 				headers: map[string][]string{}, | ||||||
| 			}, | 			}, | ||||||
| 		}, | 		}, | ||||||
| 		{ | 		{ | ||||||
| 			name: "tcp-tcp proxy without warpRoutingService enabled", | 			name: "tcp-tcp proxy without warpRoutingService enabled", | ||||||
| 			args: args{ | 			args: args{ | ||||||
| 				ingressServiceScheme:  "tcp://", | 				ingressServiceScheme: "tcp://", | ||||||
| 				originService:         runEchoTCPService, | 				originService:        runEchoTCPService, | ||||||
| 				eyeballResponseWriter: newTCPRespWriter(replayer), | 				eyeballResponseWriter: func(w io.Writer) connection.ResponseWriter { | ||||||
| 				eyeballRequestBody:    newTCPRequestBody([]byte("test2")), | 					return newTCPRespWriter(w) | ||||||
| 				connectionType:        connection.TypeTCP, | 				}, | ||||||
|  | 				eyeballRequestBody: newTCPRequestBody([]byte("test2")), | ||||||
|  | 				connectionType:     connection.TypeTCP, | ||||||
| 				requestHeaders: map[string][]string{ | 				requestHeaders: map[string][]string{ | ||||||
| 					"Cf-Cloudflared-Proxy-Src": {"non-blank-value"}, | 					"Cf-Cloudflared-Proxy-Src": {"non-blank-value"}, | ||||||
| 				}, | 				}, | ||||||
| 			}, | 			}, | ||||||
| 			want: want{ | 			want: want{ | ||||||
| 				message: []byte{}, | 				message: nil, | ||||||
| 				err:     true, | 				err:     true, | ||||||
| 			}, | 			}, | ||||||
| 		}, | 		}, | ||||||
| 		{ | 		{ | ||||||
| 			name: "ws-ws proxy when origin is different", | 			name: "ws-ws proxy when origin is different", | ||||||
| 			args: args{ | 			args: args{ | ||||||
| 				ingressServiceScheme:  "ws://", | 				ingressServiceScheme: "ws://", | ||||||
| 				originService:         runEchoWSService, | 				originService:        runEchoWSService, | ||||||
| 				eyeballResponseWriter: newWSRespWriter(replayer), | 				eyeballResponseWriter: func(w io.Writer) connection.ResponseWriter { | ||||||
| 				eyeballRequestBody:    newWSRequestBody([]byte("test1")), | 					return newWSRespWriter(w) | ||||||
| 				connectionType:        connection.TypeWebsocket, | 				}, | ||||||
|  | 				eyeballRequestBody: newWSRequestBody([]byte("test1")), | ||||||
|  | 				connectionType:     connection.TypeWebsocket, | ||||||
| 				requestHeaders: map[string][]string{ | 				requestHeaders: map[string][]string{ | ||||||
| 					// Example key from https://tools.ietf.org/html/rfc6455#section-1.2
 | 					// Example key from https://tools.ietf.org/html/rfc6455#section-1.2
 | ||||||
| 					"Sec-Websocket-Key": {"dGhlIHNhbXBsZSBub25jZQ=="}, | 					"Sec-Websocket-Key": {"dGhlIHNhbXBsZSBub25jZQ=="}, | ||||||
|  | @ -645,15 +656,17 @@ func TestConnections(t *testing.T) { | ||||||
| 					// closing the listener created by the test.
 | 					// closing the listener created by the test.
 | ||||||
| 					ln.Close() | 					ln.Close() | ||||||
| 				}, | 				}, | ||||||
| 				eyeballResponseWriter: newTCPRespWriter(replayer), | 				eyeballResponseWriter: func(w io.Writer) connection.ResponseWriter { | ||||||
| 				eyeballRequestBody:    newTCPRequestBody([]byte("test2")), | 					return newTCPRespWriter(w) | ||||||
| 				connectionType:        connection.TypeTCP, | 				}, | ||||||
|  | 				eyeballRequestBody: newTCPRequestBody([]byte("test2")), | ||||||
|  | 				connectionType:     connection.TypeTCP, | ||||||
| 				requestHeaders: map[string][]string{ | 				requestHeaders: map[string][]string{ | ||||||
| 					"Cf-Cloudflared-Proxy-Src": {"non-blank-value"}, | 					"Cf-Cloudflared-Proxy-Src": {"non-blank-value"}, | ||||||
| 				}, | 				}, | ||||||
| 			}, | 			}, | ||||||
| 			want: want{ | 			want: want{ | ||||||
| 				message: []byte{}, | 				message: nil, | ||||||
| 				err:     true, | 				err:     true, | ||||||
| 			}, | 			}, | ||||||
| 		}, | 		}, | ||||||
|  | @ -661,6 +674,7 @@ func TestConnections(t *testing.T) { | ||||||
| 
 | 
 | ||||||
| 	for _, test := range tests { | 	for _, test := range tests { | ||||||
| 		t.Run(test.name, func(t *testing.T) { | 		t.Run(test.name, func(t *testing.T) { | ||||||
|  | 			replayer := &replayer{rw: &bytes.Buffer{}} | ||||||
| 			ctx, cancel := context.WithCancel(context.Background()) | 			ctx, cancel := context.WithCancel(context.Background()) | ||||||
| 			ln, err := net.Listen("tcp", "127.0.0.1:0") | 			ln, err := net.Listen("tcp", "127.0.0.1:0") | ||||||
| 			require.NoError(t, err) | 			require.NoError(t, err) | ||||||
|  | @ -681,15 +695,18 @@ func TestConnections(t *testing.T) { | ||||||
| 			require.NoError(t, err) | 			require.NoError(t, err) | ||||||
| 
 | 
 | ||||||
| 			req.Header = test.args.requestHeaders | 			req.Header = test.args.requestHeaders | ||||||
| 			respWriter := test.args.eyeballResponseWriter |  | ||||||
| 
 | 
 | ||||||
|  | 			var respWriter connection.ResponseWriter | ||||||
| 			if pipedReqBody, ok := test.args.eyeballRequestBody.(*pipedRequestBody); ok { | 			if pipedReqBody, ok := test.args.eyeballRequestBody.(*pipedRequestBody); ok { | ||||||
| 				respWriter = newTCPRespWriter(pipedReqBody.pipedConn) | 				respWriter = newTCPRespWriter(pipedReqBody.pipedConn) | ||||||
| 				go func() { | 				go func() { | ||||||
| 					resp := pipedReqBody.roundtrip(test.args.ingressServiceScheme + ln.Addr().String()) | 					resp := pipedReqBody.roundtrip(test.args.ingressServiceScheme + ln.Addr().String()) | ||||||
| 					replayer.Write(resp) | 					replayer.Write(resp) | ||||||
| 				}() | 				}() | ||||||
|  | 			} else { | ||||||
|  | 				respWriter = test.args.eyeballResponseWriter(replayer) | ||||||
| 			} | 			} | ||||||
|  | 
 | ||||||
| 			if test.args.connectionType == connection.TypeTCP { | 			if test.args.connectionType == connection.TypeTCP { | ||||||
| 				rwa := connection.NewHTTPResponseReadWriterAcker(respWriter, req) | 				rwa := connection.NewHTTPResponseReadWriterAcker(respWriter, req) | ||||||
| 				err = proxy.ProxyTCP(ctx, rwa, &connection.TCPRequest{Dest: dest}) | 				err = proxy.ProxyTCP(ctx, rwa, &connection.TCPRequest{Dest: dest}) | ||||||
|  |  | ||||||
		Loading…
	
		Reference in New Issue