frontend: fix gwclient.Client wrappers losing CurrentFrontend interface

Signed-off-by: Brian Goff <cpuguy83@gmail.com>

Author: cpuguy83

frontend/build.go

@@ -23,10 +23,6 @@ type LoadConfig struct {
 
 type LoadOpt func(*LoadConfig)
 
-type frontendClient interface {
-	CurrentFrontend() (*llb.State, error)
-}
-
 func WithAllowArgs(args ...string) LoadOpt {
 	return func(cfg *LoadConfig) {
 		set := make(map[string]struct{}, len(args))
@@ -199,14 +195,16 @@ func WithDefaultPlatform(platform ocispecs.Platform, build gwclient.BuildFunc) g
 		if client.BuildOpts().Opts["platform"] != "" {
 			return build(ctx, client)
 		}
-		client = &clientWithPlatform{
+		client = withCurrentFrontend(client, &clientWithPlatform{
 			Client:   client,
 			platform: &platform,
-		}
+		})
 		return build(ctx, client)
 	}
 }
 
+var _ gwclient.Client = (*clientWithPlatform)(nil)
+
 type clientWithPlatform struct {
 	gwclient.Client
 	platform *ocispecs.Platform
@@ -219,7 +217,7 @@ func (c *clientWithPlatform) BuildOpts() gwclient.BuildOpts {
 }
 
 func GetCurrentFrontend(client gwclient.Client) (llb.State, error) {
-	f, err := client.(frontendClient).CurrentFrontend()
+	f, err := client.(currentFrontend).CurrentFrontend()
 	if err != nil {
 		return llb.Scratch(), err
 	}

frontend/client.go

@@ -8,7 +8,6 @@ import (
 
 	"github.com/containerd/platforms"
 	"github.com/goccy/go-yaml"
-	"github.com/moby/buildkit/client/llb"
 	"github.com/moby/buildkit/frontend/dockerui"
 	gwclient "github.com/moby/buildkit/frontend/gateway/client"
 	"github.com/project-dalec/dalec"
@@ -39,15 +38,8 @@ func (err *noSuchHandlerError) Error() string {
 	return fmt.Sprintf("no such handler for target %q: available targets: %s", err.Target, strings.Join(err.Available, ", "))
 }
 
-// CurrentFrontend is an interface typically implemented by a [gwclient.Client].
-// This is used to get the rootfs of the current frontend.
-type CurrentFrontend interface {
-	CurrentFrontend() (*llb.State, error)
-}
-
 var (
 	_ gwclient.Client = (*clientWithCustomOpts)(nil)
-	_ CurrentFrontend = (*clientWithCustomOpts)(nil)
 )
 
 type clientWithCustomOpts struct {
@@ -59,20 +51,16 @@ func (d *clientWithCustomOpts) BuildOpts() gwclient.BuildOpts {
 	return d.opts
 }
 
-func (d *clientWithCustomOpts) CurrentFrontend() (*llb.State, error) {
-	return d.Client.(CurrentFrontend).CurrentFrontend()
-}
-
-func setClientOptOption(client gwclient.Client, extraOpts map[string]string) *clientWithCustomOpts {
+func setClientOptOption(client gwclient.Client, extraOpts map[string]string) gwclient.Client {
 	opts := client.BuildOpts()
 
 	for key, value := range extraOpts {
 		opts.Opts[key] = value
 	}
-	return &clientWithCustomOpts{
+	return withCurrentFrontend(client, &clientWithCustomOpts{
 		Client: client,
 		opts:   opts,
-	}
+	})
 }
 
 func maybeSetDalecTargetKey(client gwclient.Client, key string) gwclient.Client {
@@ -82,13 +70,13 @@ func maybeSetDalecTargetKey(client gwclient.Client, key string) gwclient.Client
 		return client
 	}
 
-	// optimization to help prevent unnecessary grpc requests
+	// Optimization to help prevent unnecessary grpc requests.
 	// The gateway client will make a grpc request to get the build opts from the gateway.
 	// This just caches those opts locally.
 	// If the client is already a clientWithCustomOpts, then the opts are already cached.
 	if _, ok := client.(*clientWithCustomOpts); !ok {
 		// this forces the client to use our cached opts from above
-		client = &clientWithCustomOpts{opts: opts, Client: client}
+		client = withCurrentFrontend(client, &clientWithCustomOpts{opts: opts, Client: client})
 	}
 	return setClientOptOption(client, map[string]string{keyTopLevelTarget: key, "build-arg:" + dalec.KeyDalecTarget: key})
 }

frontend/gwutil.go

@@ -0,0 +1,38 @@
+package frontend
+
+import (
+	"github.com/moby/buildkit/client/llb"
+	gwclient "github.com/moby/buildkit/frontend/gateway/client"
+)
+
+// currentFrontend is an interface typically implemented by a [gwclient.Client].
+// This is used to get the rootfs of the current frontend.
+type currentFrontend interface {
+	CurrentFrontend() (*llb.State, error)
+}
+
+// withCurrentFrontend wraps a [gwclient.Client] to ensure that [currentFrontend]
+// is not lost when the client is wrapped.
+//
+// If inner does not implement [currentFrontend], wrapper is returned as-is.
+func withCurrentFrontend(inner gwclient.Client, wrapper gwclient.Client) gwclient.Client {
+	cf, ok := inner.(currentFrontend)
+	if !ok {
+		return wrapper
+	}
+
+	return &clientWithCurrentFrontend{Client: wrapper, cf: cf}
+}
+
+// clientWithCurrentFrontend wraps a gwclient.Client and forwards
+// CurrentFrontend calls to the inner client.
+type clientWithCurrentFrontend struct {
+	gwclient.Client
+	cf currentFrontend
+}
+
+var _ gwclient.Client = (*clientWithCurrentFrontend)(nil)
+
+func (c *clientWithCurrentFrontend) CurrentFrontend() (*llb.State, error) {
+	return c.cf.CurrentFrontend()
+}

test/testenv/buildx.go

@@ -18,6 +18,7 @@ import (
 
 	"github.com/cpuguy83/dockercfg"
 	"github.com/moby/buildkit/client"
+	"github.com/moby/buildkit/client/llb"
 	gwclient "github.com/moby/buildkit/frontend/gateway/client"
 	"github.com/moby/buildkit/session/auth"
 	"github.com/moby/buildkit/session/secrets/secretsprovider"
@@ -398,7 +399,7 @@ func (b *BuildxEnv) RunTest(ctx context.Context, t *testing.T, f TestFunc, opts
 	}
 
 	_, err = c.Build(ctx, so, "", func(ctx context.Context, gwc gwclient.Client) (*gwclient.Result, error) {
-		gwc = &clientForceDalecWithInput{gwc}
+		gwc = withCurrentFrontend(gwc, &clientForceDalecWithInput{gwc})
 
 		b.mu.Lock()
 		for id, f := range b.refs {
@@ -441,12 +442,12 @@ type gwClientInputInject struct {
 	f  gwclient.BuildFunc
 }
 
-func wrapWithInput(c gwclient.Client, id string, f gwclient.BuildFunc) *gwClientInputInject {
-	return &gwClientInputInject{
+func wrapWithInput(c gwclient.Client, id string, f gwclient.BuildFunc) gwclient.Client {
+	return withCurrentFrontend(c, &gwClientInputInject{
 		Client: c,
 		id:     id,
 		f:      f,
-	}
+	})
 }
 
 func (c *gwClientInputInject) Solve(ctx context.Context, req gwclient.SolveRequest) (*gwclient.Result, error) {
@@ -539,3 +540,28 @@ func (ap *authProvider) VerifyTokenAuthority(ctx context.Context, req *auth.Veri
 func (ap *authProvider) Register(server *grpc.Server) {
 	auth.RegisterAuthServer(server, ap)
 }
+
+// currentFrontend is the interface for obtaining the current frontend's rootfs.
+// This mirrors the unexported interface in the frontend package.
+type currentFrontend interface {
+	CurrentFrontend() (*llb.State, error)
+}
+
+// withCurrentFrontend wraps a gwclient.Client to preserve the currentFrontend
+// interface through client wrapping.
+func withCurrentFrontend(inner gwclient.Client, wrapper gwclient.Client) gwclient.Client {
+	cf, ok := inner.(currentFrontend)
+	if !ok {
+		return wrapper
+	}
+	return &clientWithCurrentFrontend{Client: wrapper, cf: cf}
+}
+
+type clientWithCurrentFrontend struct {
+	gwclient.Client
+	cf currentFrontend
+}
+
+func (c *clientWithCurrentFrontend) CurrentFrontend() (*llb.State, error) {
+	return c.cf.CurrentFrontend()
+}