client: allow setting custom dialer for session endpoint

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>

client/client.go

@@ -30,6 +30,7 @@ import (
 
 type Client struct {
 	conn          *grpc.ClientConn
+	sessionDialer func(ctx context.Context, proto string, meta map[string][]string) (net.Conn, error)
 }
 
 type ClientOpt interface{}
@@ -49,6 +50,7 @@ func New(ctx context.Context, address string, opts ...ClientOpt) (*Client, error
 	var customTracer bool // allows manually setting disabling tracing even if tracer in context
 	var tracerProvider trace.TracerProvider
 	var tracerDelegate TracerDelegate
+	var sessionDialer func(context.Context, string, map[string][]string) (net.Conn, error)
 
 	for _, o := range opts {
 		if _, ok := o.(*withFailFast); ok {
@@ -73,6 +75,9 @@ func New(ctx context.Context, address string, opts ...ClientOpt) (*Client, error
 		if wt, ok := o.(*withTracerDelegate); ok {
 			tracerDelegate = wt
 		}
+		if sd, ok := o.(*withSessionDialer); ok {
+			sessionDialer = sd.dialer
+		}
 	}
 
 	if !customTracer {
@@ -132,6 +137,7 @@ func New(ctx context.Context, address string, opts ...ClientOpt) (*Client, error
 
 	c := &Client{
 		conn:          conn,
+		sessionDialer: sessionDialer,
 	}
 
 	if tracerDelegate != nil {
@@ -244,6 +250,14 @@ type withTracerDelegate struct {
 	TracerDelegate
 }
 
+func WithSessionDialer(dialer func(context.Context, string, map[string][]string) (net.Conn, error)) ClientOpt {
+	return &withSessionDialer{dialer}
+}
+
+type withSessionDialer struct {
+	dialer func(context.Context, string, map[string][]string) (net.Conn, error)
+}
+
 func resolveDialer(address string) (func(context.Context, string) (net.Conn, error), error) {
 	ch, err := connhelper.GetConnectionHelper(address)
 	if err != nil {

client/solve.go

@@ -162,7 +162,11 @@ func (c *Client) solve(ctx context.Context, def *llb.Definition, runGateway runG
 		}
 
 		eg.Go(func() error {
-			return s.Run(statusContext, grpchijack.Dialer(c.controlClient()))
+			sd := c.sessionDialer
+			if sd == nil {
+				sd = grpchijack.Dialer(c.controlClient())
+			}
+			return s.Run(statusContext, sd)
 		})
 	}