fix test

grpc · Mar 8, 2017 · a9c135c · a9c135c
1 parent d3aa546
commit a9c135c
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/proxy/proxy.go b/proxy/proxy.go
@@ -163,8 +163,12 @@ func doHTTPConnectHandshake(ctx context.Context, conn net.Conn, addr string, hea
 // provided dialer, does handshake with the Handshaker and returns the connection.
 func NewDialer(pm Mapper, hs Handshaker, dialer func(string, time.Duration) (net.Conn, error)) func(string, time.Duration) (net.Conn, error) {
 	return func(addr string, d time.Duration) (conn net.Conn, err error) {
-		ctx, cancel := context.WithTimeout(context.Background(), d)
-		defer cancel()
+		ctx := context.Background()
+		var cancel context.CancelFunc
+		if d > 0 {
+			ctx, cancel = context.WithTimeout(context.Background(), d)
+			defer cancel()
+		}
 		var skipHandshake bool
 
 		newAddr, h, err := pm.MapAddress(ctx, addr)

diff --git a/transport/http2_client.go b/transport/http2_client.go
@@ -132,8 +132,12 @@ func dial(ctx context.Context, fn func(context.Context, string) (net.Conn, error
 		proxy.NewEnvironmentProxyMapper(),
 		proxy.NewHTTPConnectHandshaker(),
 		func(addr string, d time.Duration) (net.Conn, error) {
-			ctx, cancel := context.WithTimeout(context.Background(), d)
-			defer cancel()
+			ctx := context.Background()
+			var cancel context.CancelFunc
+			if d > 0 {
+				ctx, cancel = context.WithTimeout(context.Background(), d)
+				defer cancel()
+			}
 			return dialContext(ctx, "tcp", addr)
 		},
 	)