From 6081335ee527a9e11b5d3ae6fcc06d86ed92094b Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Tue, 26 Mar 2019 11:33:56 -0700
Subject: [PATCH] client: allow setting custom dialer

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 client/client.go | 32 +++++++++++++++++++++++---------
 1 file changed, 23 insertions(+), 9 deletions(-)

diff --git a/client/client.go b/client/client.go
index 0c66e9ff..6b824094 100644
--- a/client/client.go
+++ b/client/client.go
@@ -26,15 +26,8 @@ type ClientOpt interface{}
 
 // New returns a new buildkit client. Address can be empty for the system-default address.
 func New(ctx context.Context, address string, opts ...ClientOpt) (*Client, error) {
-	dialFn, err := resolveDialer(address)
-	if err != nil {
-		return nil, err
-	}
-	gopts := []grpc.DialOption{
-		// TODO(AkihiroSuda): use WithContextDialer (requires grpc 1.19)
-		// https://github.com/grpc/grpc-go/commit/40cb5618f475e7b9d61aa7920ae4b04ef9bbaf89
-		grpc.WithDialer(dialFn),
-	}
+	gopts := []grpc.DialOption{}
+	needDialer := true
 	needWithInsecure := true
 	for _, o := range opts {
 		if _, ok := o.(*withFailFast); ok {
@@ -53,6 +46,19 @@ func New(ctx context.Context, address string, opts ...ClientOpt) (*Client, error
 				grpc.WithUnaryInterceptor(otgrpc.OpenTracingClientInterceptor(wt.tracer, otgrpc.LogPayloads())),
 				grpc.WithStreamInterceptor(otgrpc.OpenTracingStreamClientInterceptor(wt.tracer)))
 		}
+		if wd, ok := o.(*withDialer); ok {
+			gopts = append(gopts, grpc.WithDialer(wd.dialer))
+			needDialer = false
+		}
+	}
+	if needDialer {
+		dialFn, err := resolveDialer(address)
+		if err != nil {
+			return nil, err
+		}
+		// TODO(AkihiroSuda): use WithContextDialer (requires grpc 1.19)
+		// https://github.com/grpc/grpc-go/commit/40cb5618f475e7b9d61aa7920ae4b04ef9bbaf89
+		gopts = append(gopts, grpc.WithDialer(dialFn))
 	}
 	if needWithInsecure {
 		gopts = append(gopts, grpc.WithInsecure())
@@ -84,6 +90,14 @@ func WithFailFast() ClientOpt {
 	return &withFailFast{}
 }
 
+type withDialer struct {
+	dialer func(string, time.Duration) (net.Conn, error)
+}
+
+func WithDialer(df func(string, time.Duration) (net.Conn, error)) ClientOpt {
+	return &withDialer{dialer: df}
+}
+
 type withCredentials struct {
 	ServerName string
 	CACert     string