client: allow grpc dial option passthrough

This allows consumers of buildkit's client API to manually specify grpc
options, such as underlying transport details, or creating interceptors
for requests. All custom options are appended *after* the internal
options, to allow for more specific overrides if desired.

Additionally, to prevent accidentally overwriding the internal
interceptors, we switch the client to using
`WithChain{Unary,Stream}Interceptor` instead of the singular form, which
will overwrite if multiple are specified.

Signed-off-by: Justin Chadwell <[email protected]>

Author: jedevc

client/client.go

@@ -11,7 +11,6 @@ import (
 
 	contentapi "github.com/containerd/containerd/api/services/content/v1"
 	"github.com/containerd/containerd/defaults"
-	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
 	controlapi "github.com/moby/buildkit/api/services/control"
 	"github.com/moby/buildkit/client/connhelper"
 	"github.com/moby/buildkit/session"
@@ -54,6 +53,7 @@ func New(ctx context.Context, address string, opts ...ClientOpt) (*Client, error
 	var tracerProvider trace.TracerProvider
 	var tracerDelegate TracerDelegate
 	var sessionDialer func(context.Context, string, map[string][]string) (net.Conn, error)
+	var customDialOptions []grpc.DialOption
 
 	for _, o := range opts {
 		if _, ok := o.(*withFailFast); ok {
@@ -82,6 +82,9 @@ func New(ctx context.Context, address string, opts ...ClientOpt) (*Client, error
 		if sd, ok := o.(*withSessionDialer); ok {
 			sessionDialer = sd.dialer
 		}
+		if opt, ok := o.(grpc.DialOption); ok {
+			customDialOptions = append(customDialOptions, opt)
+		}
 	}
 
 	if !customTracer {
@@ -131,17 +134,9 @@ func New(ctx context.Context, address string, opts ...ClientOpt) (*Client, error
 	unary = append(unary, grpcerrors.UnaryClientInterceptor)
 	stream = append(stream, grpcerrors.StreamClientInterceptor)
 
-	if len(unary) == 1 {
-		gopts = append(gopts, grpc.WithUnaryInterceptor(unary[0]))
-	} else if len(unary) > 1 {
-		gopts = append(gopts, grpc.WithUnaryInterceptor(grpc_middleware.ChainUnaryClient(unary...)))
-	}
-
-	if len(stream) == 1 {
-		gopts = append(gopts, grpc.WithStreamInterceptor(stream[0]))
-	} else if len(stream) > 1 {
-		gopts = append(gopts, grpc.WithStreamInterceptor(grpc_middleware.ChainStreamClient(stream...)))
-	}
+	gopts = append(gopts, grpc.WithChainUnaryInterceptor(unary...))
+	gopts = append(gopts, grpc.WithChainStreamInterceptor(stream...))
+	gopts = append(gopts, customDialOptions...)
 
 	conn, err := grpc.DialContext(ctx, address, gopts...)
 	if err != nil {

client/client_test.go

@@ -65,6 +65,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"golang.org/x/crypto/ssh/agent"
 	"golang.org/x/sync/errgroup"
+	"google.golang.org/grpc"
 )
 
 func init() {
@@ -195,6 +196,7 @@ func TestIntegration(t *testing.T) {
 		testMountStubsTimestamp,
 		testSourcePolicy,
 		testLLBMountPerformance,
+		testClientCustomGRPCOpts,
 	)
 }
 
@@ -9020,3 +9022,22 @@ func testLLBMountPerformance(t *testing.T, sb integration.Sandbox) {
 	_, err = c.Solve(timeoutCtx, def, SolveOpt{}, nil)
 	require.NoError(t, err)
 }
+
+func testClientCustomGRPCOpts(t *testing.T, sb integration.Sandbox) {
+	var interceptedMethods []string
+	intercept := func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+		interceptedMethods = append(interceptedMethods, method)
+		return invoker(ctx, method, req, reply, cc, opts...)
+	}
+	c, err := New(sb.Context(), sb.Address(), grpc.WithChainUnaryInterceptor(intercept))
+	require.NoError(t, err)
+	defer c.Close()
+
+	st := llb.Image("busybox:latest")
+	def, err := st.Marshal(sb.Context())
+	require.NoError(t, err)
+	_, err = c.Solve(sb.Context(), def, SolveOpt{}, nil)
+	require.NoError(t, err)
+
+	require.Contains(t, interceptedMethods, "/moby.buildkit.v1.Control/Solve")
+}