Allow custom worker images

frontend/azlinux/azlinux3.go

@@ -16,8 +16,11 @@ const (
 	AzLinux3TargetKey     = "azlinux3"
 	tdnfCacheNameAzlinux3 = "azlinux3-tdnf-cache"
 
-	azlinux3Ref           = "azurelinuxpreview.azurecr.io/public/azurelinux/base/core:3.0"
-	azlinux3DistrolessRef = "azurelinuxpreview.azurecr.io/public/azurelinux/distroless/base:3.0"
+	// Azlinux3Ref is the image ref used for the base worker image
+	Azlinux3Ref = "azurelinuxpreview.azurecr.io/public/azurelinux/base/core:3.0"
+	// Azlinux3WorkerContextName is the build context name that can be used to lookup
+	Azlinux3WorkerContextName = "dalec-azlinux3-worker"
+	azlinux3DistrolessRef     = "azurelinuxpreview.azurecr.io/public/azurelinux/distroless/base:3.0"
 )
 
 func NewAzlinux3Handler() gwclient.BuildFunc {
@@ -26,11 +29,29 @@ func NewAzlinux3Handler() gwclient.BuildFunc {
 
 type azlinux3 struct{}
 
-func (w azlinux3) Base(resolver llb.ImageMetaResolver, opts ...llb.ConstraintsOpt) llb.State {
-	return llb.Image(azlinux3Ref, llb.WithMetaResolver(resolver), dalec.WithConstraints(opts...)).Run(
+func (w azlinux3) Base(sOpt dalec.SourceOpts, opts ...llb.ConstraintsOpt) (llb.State, error) {
+	base, err := sOpt.GetContext(Azlinux3Ref, dalec.WithConstraints(opts...))
+	if err != nil {
+		return llb.Scratch(), err
+	}
+
+	if base != nil {
+		return *base, nil
+	}
+
+	base, err = sOpt.GetContext(Azlinux3WorkerContextName, dalec.WithConstraints(opts...))
+	if err != nil {
+		return llb.Scratch(), nil
+	}
+	if base != nil {
+		return *base, nil
+	}
+
+	img := llb.Image(Azlinux3Ref, llb.WithMetaResolver(sOpt.Resolver), dalec.WithConstraints(opts...))
+	return img.Run(
 		w.Install([]string{"rpm-build", "mariner-rpm-macros", "build-essential", "ca-certificates"}, installWithConstraints(opts)),
 		dalec.WithConstraints(opts...),
-	).Root()
+	).Root(), nil
 }
 
 func (w azlinux3) Install(pkgs []string, opts ...installOpt) llb.RunOption {
@@ -53,6 +74,20 @@ func (azlinux3) DefaultImageConfig(ctx context.Context, resolver llb.ImageMetaRe
 	return &cfg, nil
 }
 
+func (azlinux3) WorkerImageConfig(ctx context.Context, resolver llb.ImageMetaResolver, platform *ocispecs.Platform) (*dalec.DockerImageSpec, error) {
+	_, _, dt, err := resolver.ResolveImageConfig(ctx, Azlinux3Ref, sourceresolver.Opt{Platform: platform})
+	if err != nil {
+		return nil, err
+	}
+
+	var cfg dalec.DockerImageSpec
+	if err := json.Unmarshal(dt, &cfg); err != nil {
+		return nil, err
+	}
+
+	return &cfg, nil
+}
+
 func (azlinux3) tdnfCacheMount(root string) llb.RunOption {
 	return llb.AddMount(filepath.Join(root, tdnfCacheDir), llb.Scratch(), llb.AsPersistentCacheDir(tdnfCacheNameAzlinux3, llb.CacheMountLocked))
 }

frontend/azlinux/handle_container.go

@@ -35,7 +35,7 @@ func handleContainer(w worker) gwclient.BuildFunc {
 				return nil, nil, err
 			}
 
-			st, err := specToContainerLLB(w, client, spec, targetKey, rpmDir, rpms, sOpt, pg)
+			st, err := specToContainerLLB(w, spec, targetKey, rpmDir, rpms, sOpt, pg)
 			if err != nil {
 				return nil, nil, err
 			}
@@ -62,12 +62,17 @@ func handleContainer(w worker) gwclient.BuildFunc {
 				return nil, nil, err
 			}
 
+			base, err := w.Base(sOpt, pg)
+			if err != nil {
+				return nil, nil, err
+			}
+
 			withTestDeps := func(in llb.State) llb.State {
 				deps := spec.GetTestDeps(targetKey)
 				if len(deps) == 0 {
 					return in
 				}
-				return w.Base(client, pg).Run(
+				return base.Run(
 					w.Install(spec.GetTestDeps(targetKey), atRoot("/tmp/rootfs")),
 					pg,
 					dalec.ProgressGroup("Install test dependencies"),
@@ -130,11 +135,14 @@ func readRPMs(ctx context.Context, client gwclient.Client, st llb.State) ([]stri
 	return out, nil
 }
 
-func specToContainerLLB(w worker, client gwclient.Client, spec *dalec.Spec, targetKey string, rpmDir llb.State, files []string, sOpt dalec.SourceOpts, opts ...llb.ConstraintsOpt) (llb.State, error) {
+func specToContainerLLB(w worker, spec *dalec.Spec, targetKey string, rpmDir llb.State, files []string, sOpt dalec.SourceOpts, opts ...llb.ConstraintsOpt) (llb.State, error) {
 	opts = append(opts, dalec.ProgressGroup("Install RPMs"))
 	const workPath = "/tmp/rootfs"
 
-	builderImg := w.Base(client, opts...)
+	builderImg, err := w.Base(sOpt, opts...)
+	if err != nil {
+		return llb.Scratch(), err
+	}
 
 	rootfs := llb.Scratch()
 	if ref := dalec.GetBaseOutputImage(spec, targetKey); ref != "" {

frontend/azlinux/handle_depsonly.go

@@ -15,7 +15,16 @@ func handleDepsOnly(w worker) gwclient.BuildFunc {
 	return func(ctx context.Context, client gwclient.Client) (*gwclient.Result, error) {
 		return frontend.BuildWithPlatform(ctx, client, func(ctx context.Context, client gwclient.Client, platform *ocispecs.Platform, spec *dalec.Spec, targetKey string) (gwclient.Reference, *dalec.DockerImageSpec, error) {
 			pg := dalec.ProgressGroup("Build mariner2 deps-only container: " + spec.Name)
-			baseImg := w.Base(client, pg)
+
+			sOpt, err := frontend.SourceOptFromClient(ctx, client)
+			if err != nil {
+				return nil, nil, err
+			}
+
+			baseImg, err := w.Base(sOpt, pg)
+			if err != nil {
+				return nil, nil, err
+			}
 			rpmDir := baseImg.Run(
 				dalec.ShArgs(`set -ex; dir="/tmp/rpms/RPMS/$(uname -m)"; mkdir -p "${dir}"; tdnf install -y --releasever=2.0 --downloadonly --alldeps --downloaddir "${dir}" `+strings.Join(spec.GetRuntimeDeps(targetKey), " ")),
 				pg,
@@ -27,11 +36,7 @@ func handleDepsOnly(w worker) gwclient.BuildFunc {
 				return nil, nil, err
 			}
 
-			sOpt, err := frontend.SourceOptFromClient(ctx, client)
-			if err != nil {
-				return nil, nil, err
-			}
-			st, err := specToContainerLLB(w, client, spec, targetKey, rpmDir, files, sOpt, pg)
+			st, err := specToContainerLLB(w, spec, targetKey, rpmDir, files, sOpt, pg)
 			if err != nil {
 				return nil, nil, err
 			}

frontend/azlinux/handle_rpm.go

@@ -65,7 +65,12 @@ func installBuildDeps(w worker, spec *dalec.Spec, targetKey string, opts ...llb.
 }
 
 func specToRpmLLB(ctx context.Context, w worker, client gwclient.Client, spec *dalec.Spec, sOpt dalec.SourceOpts, targetKey string, opts ...llb.ConstraintsOpt) (llb.State, error) {
-	base := w.Base(client, opts...).With(installBuildDeps(w, spec, targetKey, opts...))
+	base, err := w.Base(sOpt, opts...)
+	base = base.With(installBuildDeps(w, spec, targetKey, opts...))
+	if err != nil {
+		return llb.Scratch(), err
+	}
+
 	br, err := rpm.SpecToBuildrootLLB(base, spec, sOpt, targetKey, opts...)
 	if err != nil {
 		return llb.Scratch(), err

frontend/azlinux/handler.go

@@ -19,9 +19,10 @@ const (
 )
 
 type worker interface {
-	Base(resolver llb.ImageMetaResolver, opts ...llb.ConstraintsOpt) llb.State
+	Base(sOpt dalec.SourceOpts, opts ...llb.ConstraintsOpt) (llb.State, error)
 	Install(pkgs []string, opts ...installOpt) llb.RunOption
 	DefaultImageConfig(context.Context, llb.ImageMetaResolver, *ocispecs.Platform) (*dalec.DockerImageSpec, error)
+	WorkerImageConfig(context.Context, llb.ImageMetaResolver, *ocispecs.Platform) (*dalec.DockerImageSpec, error)
 }
 
 func newHandler(w worker) gwclient.BuildFunc {
@@ -44,18 +45,30 @@ func newHandler(w worker) gwclient.BuildFunc {
 		Description: "Builds a container image with only the runtime dependencies installed.",
 	})
 
+	mux.Add("worker", handleBaseImg(w), &targets.Target{
+		Name:        "worker",
+		Description: "Builds the base worker image responsible for building the rpm",
+	})
+
 	return mux.Handle
 }
 
 func handleDebug(w worker) gwclient.BuildFunc {
 	return func(ctx context.Context, client gwclient.Client) (*gwclient.Result, error) {
-		return rpm.HandleDebug(getSpecWorker(w))(ctx, client)
+		sOpt, err := frontend.SourceOptFromClient(ctx, client)
+		if err != nil {
+			return nil, err
+		}
+		return rpm.HandleDebug(getSpecWorker(w, sOpt))(ctx, client)
 	}
 }
 
-func getSpecWorker(w worker) rpm.WorkerFunc {
+func getSpecWorker(w worker, sOpt dalec.SourceOpts) rpm.WorkerFunc {
 	return func(resolver llb.ImageMetaResolver, spec *dalec.Spec, targetKey string, opts ...llb.ConstraintsOpt) (llb.State, error) {
-		st := w.Base(resolver, opts...)
+		st, err := w.Base(sOpt, opts...)
+		if err != nil {
+			return llb.Scratch(), err
+		}
 		if spec.HasGomods() {
 			deps := spec.GetBuildDeps(targetKey)
 			hasGolang := func(s string) bool {
@@ -70,3 +83,46 @@ func getSpecWorker(w worker) rpm.WorkerFunc {
 		return st, nil
 	}
 }
+
+func handleBaseImg(w worker) gwclient.BuildFunc {
+	return func(ctx context.Context, client gwclient.Client) (*gwclient.Result, error) {
+		return frontend.BuildWithPlatform(ctx, client, func(ctx context.Context, client gwclient.Client, platform *ocispecs.Platform, spec *dalec.Spec, targetKey string) (gwclient.Reference, *dalec.DockerImageSpec, error) {
+
+			sOpt, err := frontend.SourceOptFromClient(ctx, client)
+			if err != nil {
+				return nil, nil, err
+			}
+
+			st, err := w.Base(sOpt)
+			if err != nil {
+				return nil, nil, err
+			}
+
+			def, err := st.Marshal(ctx)
+			if err != nil {
+				return nil, nil, err
+			}
+
+			req := gwclient.SolveRequest{
+				Definition: def.ToPB(),
+			}
+
+			res, err := client.Solve(ctx, req)
+			if err != nil {
+				return nil, nil, err
+			}
+
+			ref, err := res.SingleRef()
+			if err != nil {
+				return nil, nil, err
+			}
+
+			cfg, err := w.DefaultImageConfig(ctx, client, platform)
+			if err != nil {
+				return nil, nil, err
+			}
+
+			return ref, cfg, nil
+		})
+	}
+}

frontend/azlinux/mariner2.go

@@ -16,8 +16,9 @@ const (
 	Mariner2TargetKey     = "mariner2"
 	tdnfCacheNameMariner2 = "mariner2-tdnf-cache"
 
-	mariner2Ref           = "mcr.microsoft.com/cbl-mariner/base/core:2.0"
-	mariner2DistrolessRef = "mcr.microsoft.com/cbl-mariner/distroless/base:2.0"
+	Mariner2Ref               = "mcr.microsoft.com/cbl-mariner/base/core:2.0"
+	Mariner2WorkerContextName = "dalec-mariner2-worker"
+	mariner2DistrolessRef     = "mcr.microsoft.com/cbl-mariner/distroless/base:2.0"
 )
 
 func NewMariner2Handler() gwclient.BuildFunc {
@@ -26,11 +27,28 @@ func NewMariner2Handler() gwclient.BuildFunc {
 
 type mariner2 struct{}
 
-func (w mariner2) Base(resolver llb.ImageMetaResolver, opts ...llb.ConstraintsOpt) llb.State {
-	return llb.Image(mariner2Ref, llb.WithMetaResolver(resolver), dalec.WithConstraints(opts...)).Run(
-		w.Install([]string{"rpm-build", "mariner-rpm-macros", "systemd-rpm-macros", "build-essential", "ca-certificates"}, installWithConstraints(opts)),
+func (w mariner2) Base(sOpt dalec.SourceOpts, opts ...llb.ConstraintsOpt) (llb.State, error) {
+	base, err := sOpt.GetContext(Mariner2Ref, dalec.WithConstraints(opts...))
+	if err != nil {
+		return llb.Scratch(), err
+	}
+
+	if base == nil {
+		base, err = sOpt.GetContext(Mariner2WorkerContextName, dalec.WithConstraints(opts...))
+		if err != nil {
+			return llb.Scratch(), nil
+		}
+	}
+
+	if base == nil {
+		st := llb.Image(Mariner2Ref, llb.WithMetaResolver(sOpt.Resolver), dalec.WithConstraints(opts...))
+		base = &st
+	}
+
+	return base.Run(
+		w.Install([]string{"rpm-build", "mariner-rpm-macros", "build-essential", "ca-certificates"}, installWithConstraints(opts)),
 		dalec.WithConstraints(opts...),
-	).Root()
+	).Root(), nil
 }
 
 func (w mariner2) Install(pkgs []string, opts ...installOpt) llb.RunOption {
@@ -55,6 +73,20 @@ func (mariner2) DefaultImageConfig(ctx context.Context, resolver llb.ImageMetaRe
 	return &cfg, nil
 }
 
+func (mariner2) WorkerImageConfig(ctx context.Context, resolver llb.ImageMetaResolver, platform *ocispecs.Platform) (*dalec.DockerImageSpec, error) {
+	_, _, dt, err := resolver.ResolveImageConfig(ctx, Mariner2Ref, sourceresolver.Opt{Platform: platform})
+	if err != nil {
+		return nil, err
+	}
+
+	var cfg dalec.DockerImageSpec
+	if err := json.Unmarshal(dt, &cfg); err != nil {
+		return nil, err
+	}
+
+	return &cfg, nil
+}
+
 func (mariner2) tdnfCacheMount(root string) llb.RunOption {
 	return llb.AddMount(filepath.Join(root, tdnfCacheDir), llb.Scratch(), llb.AsPersistentCacheDir(tdnfCacheNameMariner2, llb.CacheMountLocked))
 }

frontend/windows/handle_container.go

@@ -54,7 +54,10 @@ func handleContainer(ctx context.Context, client gwclient.Client) (*gwclient.Res
 		}
 
 		pg := dalec.ProgressGroup("Build windows container: " + spec.Name)
-		worker := workerImg(sOpt, pg)
+		worker, err := workerImg(sOpt, pg)
+		if err != nil {
+			return nil, nil, err
+		}
 
 		bin, err := buildBinaries(ctx, spec, worker, client, sOpt, targetKey)
 		if err != nil {

frontend/windows/handle_zip.go

@@ -18,10 +18,11 @@ import (
 )
 
 const (
-	workerImgRef    = "mcr.microsoft.com/mirror/docker/library/ubuntu:jammy"
-	outputDir       = "/tmp/output"
-	buildScriptName = "_build.sh"
-	aptCachePrefix  = "jammy-windowscross"
+	workerImgRef                  = "mcr.microsoft.com/mirror/docker/library/ubuntu:jammy"
+	WindowscrossWorkerContextName = "dalec-windowscross-worker"
+	outputDir                     = "/tmp/output"
+	buildScriptName               = "_build.sh"
+	aptCachePrefix                = "jammy-windowscross"
 )
 
 func handleZip(ctx context.Context, client gwclient.Client) (*gwclient.Result, error) {
@@ -32,7 +33,10 @@ func handleZip(ctx context.Context, client gwclient.Client) (*gwclient.Result, e
 		}
 
 		pg := dalec.ProgressGroup("Build windows container: " + spec.Name)
-		worker := workerImg(sOpt, pg)
+		worker, err := workerImg(sOpt, pg)
+		if err != nil {
+			return nil, nil, err
+		}
 
 		bin, err := buildBinaries(ctx, spec, worker, client, sOpt, targetKey)
 		if err != nil {
@@ -168,13 +172,29 @@ func generateInvocationScript(binaries []string) *strings.Builder {
 	return script
 }
 
-func workerImg(sOpt dalec.SourceOpts, opts ...llb.ConstraintsOpt) llb.State {
-	// TODO: support named context override... also this should possibly be its own image, maybe?
+func workerImg(sOpt dalec.SourceOpts, opts ...llb.ConstraintsOpt) (llb.State, error) {
+	base, err := sOpt.GetContext(workerImgRef, dalec.WithConstraints(opts...))
+	if err != nil {
+		return llb.Scratch(), err
+	}
+
+	if base != nil {
+		return *base, nil
+	}
+
+	base, err = sOpt.GetContext(WindowscrossWorkerContextName, dalec.WithConstraints(opts...))
+	if err != nil {
+		return llb.Scratch(), nil
+	}
+	if base != nil {
+		return *base, nil
+	}
+
 	return llb.Image(workerImgRef, llb.WithMetaResolver(sOpt.Resolver), dalec.WithConstraints(opts...)).
 		Run(
 			dalec.ShArgs("apt-get update && apt-get install -y build-essential binutils-mingw-w64 g++-mingw-w64-x86-64 gcc git make pkg-config quilt zip"),
 			dalec.WithMountedAptCache(aptCachePrefix),
-		).Root()
+		).Root(), nil
 }
 
 func createBuildScript(spec *dalec.Spec) llb.State {