frontend: fix gwclient.Client wrappers losing CurrentFrontend interface

Add internal/gwutil package with WithCurrentFrontend helper that
preserves CurrentFrontend and SpecLoader interfaces when wrapping a
gwclient.Client. This replaces ad-hoc manual CurrentFrontend() forwarding
methods that were fragile and easy to forget when adding new wrappers.

Update all wrapping sites (setClientOptOption, maybeSetDalecTargetKey,
WithDefaultPlatform, and test wrappers) to use gwutil.WithCurrentFrontend.

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

Author: cpuguy83

frontend/build.go

@@ -15,6 +15,7 @@ import (
 	ocispecs "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
 	"github.com/project-dalec/dalec"
+	"github.com/project-dalec/dalec/internal/gwutil"
 )
 
 type LoadConfig struct {
@@ -23,10 +24,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 +196,16 @@ func WithDefaultPlatform(platform ocispecs.Platform, build gwclient.BuildFunc) g
 		if client.BuildOpts().Opts["platform"] != "" {
 			return build(ctx, client)
 		}
-		client = &clientWithPlatform{
+		client = gwutil.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 +218,7 @@ func (c *clientWithPlatform) BuildOpts() gwclient.BuildOpts {
 }
 
 func GetCurrentFrontend(client gwclient.Client) (llb.State, error) {
-	f, err := client.(frontendClient).CurrentFrontend()
+	f, err := client.(gwutil.CurrentFrontend).CurrentFrontend()
 	if err != nil {
 		return llb.Scratch(), err
 	}

frontend/client.go

@@ -8,10 +8,10 @@ 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"
+	"github.com/project-dalec/dalec/internal/gwutil"
 )
 
 const (
@@ -39,15 +39,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 +52,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 gwutil.WithCurrentFrontend(client, &clientWithCustomOpts{
 		Client: client,
 		opts:   opts,
-	}
+	})
 }
 
 func maybeSetDalecTargetKey(client gwclient.Client, key string) gwclient.Client {
@@ -82,13 +71,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 = gwutil.WithCurrentFrontend(client, &clientWithCustomOpts{opts: opts, Client: client})
 	}
 	return setClientOptOption(client, map[string]string{keyTopLevelTarget: key, "build-arg:" + dalec.KeyDalecTarget: key})
 }

internal/gwutil/gwutil.go

@@ -0,0 +1,83 @@
+package gwutil
+
+import (
+	"context"
+
+	"github.com/moby/buildkit/client/llb"
+	gwclient "github.com/moby/buildkit/frontend/gateway/client"
+	"github.com/project-dalec/dalec"
+)
+
+// 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)
+}
+
+// SpecLoader is an optional interface that a [gwclient.Client] wrapper can
+// implement to provide a cached or pre-loaded dalec spec. This avoids
+// redundant loads when many callers need the spec from the same client.
+type SpecLoader interface {
+	LoadSpec(context.Context) (*dalec.Spec, 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.
+// Otherwise, if either the wrapper or the inner client implements [SpecLoader],
+// that interface is also preserved on the returned client (wrapper is preferred
+// over inner).
+func WithCurrentFrontend(inner gwclient.Client, wrapper gwclient.Client) gwclient.Client {
+	cf, ok := inner.(CurrentFrontend)
+	if !ok {
+		return wrapper
+	}
+
+	w := &clientWithCurrentFrontend{Client: wrapper, cf: cf}
+
+	// Preserve SpecLoader if present on either wrapper or inner.
+	// Prefer wrapper (it's the more specific layer).
+	switch {
+	case isSpecLoader(wrapper):
+		return &clientWithCurrentFrontendAndSpecLoader{clientWithCurrentFrontend: w, sl: wrapper.(SpecLoader)}
+	case isSpecLoader(inner):
+		return &clientWithCurrentFrontendAndSpecLoader{clientWithCurrentFrontend: w, sl: inner.(SpecLoader)}
+	}
+
+	return w
+}
+
+func isSpecLoader(c gwclient.Client) bool {
+	_, ok := c.(SpecLoader)
+	return ok
+}
+
+// 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()
+}
+
+// clientWithCurrentFrontendAndSpecLoader extends clientWithCurrentFrontend
+// with SpecLoader forwarding.
+type clientWithCurrentFrontendAndSpecLoader struct {
+	*clientWithCurrentFrontend
+	sl SpecLoader
+}
+
+var (
+	_ gwclient.Client = (*clientWithCurrentFrontendAndSpecLoader)(nil)
+	_ SpecLoader      = (*clientWithCurrentFrontendAndSpecLoader)(nil)
+)
+
+func (c *clientWithCurrentFrontendAndSpecLoader) LoadSpec(ctx context.Context) (*dalec.Spec, error) {
+	return c.sl.LoadSpec(ctx)
+}

test/testenv/buildx.go

@@ -25,6 +25,7 @@ import (
 	"github.com/moby/buildkit/solver/pb"
 	spb "github.com/moby/buildkit/sourcepolicy/pb"
 	pkgerrors "github.com/pkg/errors"
+	"github.com/project-dalec/dalec/internal/gwutil"
 	"github.com/project-dalec/dalec/sessionutil/socketprovider"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
@@ -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 = gwutil.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 gwutil.WithCurrentFrontend(c, &gwClientInputInject{
 		Client: c,
 		id:     id,
 		f:      f,
-	}
+	})
 }
 
 func (c *gwClientInputInject) Solve(ctx context.Context, req gwclient.SolveRequest) (*gwclient.Result, error) {