Use mirrored private base images in source builds

lib/builds/dockerfile.go

@@ -7,11 +7,31 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+	"regexp"
 	"strings"
 
 	"github.com/google/go-containerregistry/pkg/name"
 )
 
+var fromImagePattern = regexp.MustCompile(`(?i)^(\s*FROM\s+(?:--\S+\s+)*)(\S+)`)
+
+// rewriteMirroredFROMs routes cached non-Docker-Hub bases through the host's
+// registry so the guest only needs its scoped local-registry token.
+func rewriteMirroredFROMs(content string, replacements map[string]string) string {
+	lines := strings.Split(content, "\n")
+	for i, line := range lines {
+		match := fromImagePattern.FindStringSubmatchIndex(line)
+		if match == nil {
+			continue
+		}
+		start, end := match[4], match[5]
+		if replacement, ok := replacements[normalizeImageRef(line[start:end])]; ok {
+			lines[i] = line[:start] + replacement + line[end:]
+		}
+	}
+	return strings.Join(lines, "\n")
+}
+
 // ParseDockerfileFROMs extracts and deduplicates base image references from
 // Dockerfile content. It reuses the same parsing logic as the builder agent's
 // rewriteDockerfileFROMs: split lines, find FROM, skip flags/comments/scratch,
@@ -125,6 +145,48 @@ func normalizeImageRef(ref string) string {
 	return repo + suffix
 }
 
+// writeSourceWithDockerfile prepares a build input without modifying the stored
+// source archive. Older builder images prefer the archive's Dockerfile to config.
+func writeSourceWithDockerfile(source io.Reader, destination io.Writer, dockerfile string) error {
+	reader, err := gzip.NewReader(source)
+	if err != nil {
+		return err
+	}
+	defer reader.Close()
+	compressed := gzip.NewWriter(destination)
+	defer compressed.Close()
+	writer := tar.NewWriter(compressed)
+	archive := tar.NewReader(reader)
+	for {
+		header, err := archive.Next()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return err
+		}
+		if header.Name == "Dockerfile" || header.Name == "./Dockerfile" {
+			continue
+		}
+		if err := writer.WriteHeader(header); err != nil {
+			return err
+		}
+		if _, err := io.Copy(writer, archive); err != nil {
+			return err
+		}
+	}
+	if err := writer.WriteHeader(&tar.Header{Name: "Dockerfile", Mode: 0644, Size: int64(len(dockerfile))}); err != nil {
+		return err
+	}
+	if _, err := io.WriteString(writer, dockerfile); err != nil {
+		return err
+	}
+	if err := writer.Close(); err != nil {
+		return err
+	}
+	return compressed.Close()
+}
+
 // ExtractDockerfileFromTarball reads just the Dockerfile entry from a .tar.gz
 // archive and returns its content as a string. It looks for entries named
 // "Dockerfile" or "./Dockerfile" at the root of the archive.

lib/builds/dockerfile_test.go

@@ -3,6 +3,7 @@ package builds
 import (
 	"archive/tar"
 	"compress/gzip"
+	"io"
 	"os"
 	"path/filepath"
 	"testing"
@@ -198,6 +199,69 @@ func TestExtractDockerfileFromTarball_DotSlashPrefix(t *testing.T) {
 	}
 }
 
+func TestWriteSourceWithDockerfile(t *testing.T) {
+	for _, originalName := range []string{"Dockerfile", "./Dockerfile", "other.txt"} {
+		t.Run(originalName, func(t *testing.T) {
+			dir := t.TempDir()
+			original := filepath.Join(dir, "source.tar.gz")
+			createTarball(t, original, map[string]string{originalName: "FROM old:latest", "app.txt": "application"})
+			source, err := os.Open(original)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer source.Close()
+			prepared, err := os.Create(filepath.Join(dir, "prepared.tar.gz"))
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer prepared.Close()
+			const replacement = "FROM mirror.example/runtime@sha256:abc\n"
+			if err := writeSourceWithDockerfile(source, prepared, replacement); err != nil {
+				t.Fatal(err)
+			}
+			got, err := ExtractDockerfileFromTarball(prepared.Name())
+			if err != nil || got != replacement {
+				t.Fatalf("replacement: %q, %v", got, err)
+			}
+			if originalName != "other.txt" {
+				got, err = ExtractDockerfileFromTarball(original)
+				if err != nil || got != "FROM old:latest" {
+					t.Fatalf("stored source changed: %q, %v", got, err)
+				}
+			}
+			if _, err := prepared.Seek(0, 0); err != nil {
+				t.Fatal(err)
+			}
+			gz, err := gzip.NewReader(prepared)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer gz.Close()
+			tr := tar.NewReader(gz)
+			foundApp := false
+			for {
+				h, err := tr.Next()
+				if err == io.EOF {
+					break
+				}
+				if err != nil {
+					t.Fatal(err)
+				}
+				if h.Name == "app.txt" {
+					content, err := io.ReadAll(tr)
+					if err != nil || string(content) != "application" {
+						t.Fatalf("app content: %q, %v", content, err)
+					}
+					foundApp = true
+				}
+			}
+			if !foundApp {
+				t.Fatal("lost app file")
+			}
+		})
+	}
+}
+
 // createTarball creates a .tar.gz file with the given files (name -> content).
 func createTarball(t *testing.T, path string, files map[string]string) {
 	t.Helper()

lib/builds/manager.go

@@ -6,6 +6,7 @@ import (
 	_ "embed"
 	"encoding/json"
 	"fmt"
+	"io"
 	"log/slog"
 	"net"
 	"os"
@@ -556,8 +557,8 @@ func (m *manager) CreateBuild(ctx context.Context, req CreateBuildRequest, sourc
 		refs := ParseDockerfileFROMs(dockerfileContent)
 		seen := make(map[string]bool)
 		for _, ref := range refs {
-			repo := ref
-			if idx := strings.LastIndex(repo, ":"); idx > 0 {
+			repo := strings.SplitN(ref, "@", 2)[0]
+			if idx := strings.LastIndex(repo, ":"); idx > strings.LastIndex(repo, "/") {
 				repo = repo[:idx]
 			}
 			if !seen[repo] {
@@ -688,13 +689,12 @@ func (m *manager) runBuild(ctx context.Context, id string, req CreateBuildReques
 	// Mirror base images to the local registry before launching the VM.
 	// BuildKit is configured with our registry as a mirror for docker.io,
 	// so pre-cached images will be served locally without pulling from Docker Hub.
-	if err := m.mirrorBaseImagesForBuild(buildCtx, id, req); err != nil {
-		m.logger.Warn("failed to mirror base images", "id", id, "error", err)
+	var result *BuildResult
+	err := m.mirrorBaseImagesForBuild(buildCtx, id, req)
+	if err == nil {
+		result, err = m.executeBuild(buildCtx, id, req, policy)
 	}
 
-	// Run the build in a builder VM
-	result, err := m.executeBuild(buildCtx, id, req, policy)
-
 	// executeBuild's defers release the builder acquisition when it returns,
 	// so same-builder builds may start while this build waits for image
 	// conversion and re-tagging below.
@@ -813,6 +813,26 @@ func (m *manager) executeBuild(ctx context.Context, id string, req CreateBuildRe
 	}
 	defer sourceFile.Close()
 
+	config, err := readBuildConfig(m.paths, id)
+	if err != nil {
+		return nil, err
+	}
+	if config.Dockerfile != "" && config.Dockerfile != req.Dockerfile {
+		prepared, err := os.CreateTemp(m.paths.BuildSourceDir(id), "mirrored-*.tar.gz")
+		if err != nil {
+			return nil, err
+		}
+		defer os.Remove(prepared.Name())
+		defer prepared.Close()
+		if err := writeSourceWithDockerfile(sourceFile, prepared, config.Dockerfile); err != nil {
+			return nil, fmt.Errorf("prepare mirrored Dockerfile: %w", err)
+		}
+		if _, err := prepared.Seek(0, io.SeekStart); err != nil {
+			return nil, err
+		}
+		sourceFile = prepared
+	}
+
 	// Create volume with source (using the volume manager's archive import)
 	_, err = m.volumeManager.CreateVolumeFromArchive(ctx, volumes.CreateVolumeFromArchiveRequest{
 		Id:     &sourceVolID,
@@ -1558,8 +1578,8 @@ func (m *manager) refreshBuildToken(buildID string, req *CreateBuildRequest) err
 		refs := ParseDockerfileFROMs(dockerfileContent)
 		seen := make(map[string]bool)
 		for _, ref := range refs {
-			repo := ref
-			if idx := strings.LastIndex(repo, ":"); idx > 0 {
+			repo := strings.SplitN(ref, "@", 2)[0]
+			if idx := strings.LastIndex(repo, ":"); idx > strings.LastIndex(repo, "/") {
 				repo = repo[:idx]
 			}
 			if !seen[repo] {

lib/builds/mirror.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/google/go-containerregistry/pkg/authn"
+	"github.com/google/go-containerregistry/pkg/name"
 	"github.com/kernel/hypeman/lib/images"
 )
 
@@ -72,18 +73,37 @@ func (m *manager) mirrorBaseImagesForBuild(ctx context.Context, id string, req C
 	// our token handler extracts the JWT from the username field only.
 	authConfig := &authn.AuthConfig{Username: registryToken, Password: "x"}
 
+	replacements := make(map[string]string)
 	for _, ref := range refs {
+		parsed, err := name.ParseReference(ref)
+		if err != nil {
+			m.logger.Warn("failed to mirror base image", "id", id, "image", ref, "error", err)
+			continue
+		}
+		useLocalRegistry := parsed.Context().RegistryStr() != "index.docker.io"
 		result, err := images.MirrorBaseImage(ctx, m.config.RegistryURL, images.MirrorRequest{
-			SourceImage: ref,
+			SourceImage:   ref,
+			PreserveIndex: useLocalRegistry,
 		}, authConfig)
 		if err != nil {
 			m.logger.Warn("failed to mirror base image",
 				"id", id, "image", ref, "error", err)
 			continue
 		}
+		if useLocalRegistry {
+			replacements[ref] = stripRegistryScheme(m.config.RegistryURL) + "/" + parsed.Context().Name() + "@" + result.Digest
+		}
 		m.logger.Info("mirrored base image",
 			"id", id, "image", ref, "local_ref", result.LocalRef, "digest", result.Digest)
 	}
 
+	if len(replacements) > 0 {
+		config, err := readBuildConfig(m.paths, id)
+		if err != nil {
+			return err
+		}
+		config.Dockerfile = rewriteMirroredFROMs(dockerfileContent, replacements)
+		return writeBuildConfig(m.paths, id, config)
+	}
 	return nil
 }