Fix memory usage regression

.github/workflows/ci.yml

@@ -0,0 +1,40 @@
+name: CI
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+env:
+  GO_VERSION: "1.22.x"
+
+jobs:
+  roothash-tests:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Setup Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: ${{ env.GO_VERSION }}
+
+      - name: Run dm-verity roothash tests
+        run: go test ./cmd/dmverity-vhd -run TestRootHash -count=1 -v
+
+  create-vhd-tests:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Setup Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: ${{ env.GO_VERSION }}
+
+      - name: Run dm-verity create VHD tests
+        run: go test ./cmd/dmverity-vhd -run TestCreateVHD -count=1 -v

cmd/dmverity-vhd/args.go

@@ -0,0 +1,46 @@
+package main
+
+import (
+	"errors"
+
+	"github.com/urfave/cli"
+)
+
+func getImageParsers(ctx *cli.Context) (
+	imageFetcher ImageFetcher,
+	imageParser ImageParser,
+	manifestParser ManifestParser,
+	err error,
+) {
+
+	// Get args
+	imageName := ctx.String(inputFlag)
+	username := ctx.String(usernameFlag)
+	password := ctx.String(passwordFlag)
+	tarballPath := ctx.GlobalString(tarballFlag)
+	useDocker := ctx.GlobalBool(dockerFlag)
+
+	// Validation
+	if useDocker && tarballPath != "" {
+		err = errors.New("cannot use both docker and tarball for image source")
+		return
+	}
+
+	if tarballPath != "" {
+		imageFetcher = func() (ImageSource, error) { return fetchImageTarball(tarballPath) }
+		imageParser = parseLocalImage
+	} else if useDocker {
+		imageFetcher = func() (ImageSource, error) { return fetchDockerImage(imageName) }
+		imageParser = parseLocalImage
+	} else {
+		imageFetcher = func() (ImageSource, error) { return fetchContainerRegistryImage(imageName, username, password) }
+		imageParser = parseContainerRegistryImage
+	}
+
+	manifestParser = combineManifestParsers([]ManifestParser{
+		parseOCIImage,
+		parseDockerImage,
+	})
+
+	return
+}

cmd/dmverity-vhd/containerregistry.go

@@ -0,0 +1,122 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"path"
+
+	"github.com/google/go-containerregistry/pkg/authn"
+	"github.com/google/go-containerregistry/pkg/name"
+	v1 "github.com/google/go-containerregistry/pkg/v1"
+	"github.com/google/go-containerregistry/pkg/v1/remote"
+	log "github.com/sirupsen/logrus"
+)
+
+func fetchContainerRegistryImage(
+	imageName string,
+	username string,
+	password string,
+) (
+	image v1.Image,
+	err error,
+) {
+	log.Tracef("fetchContainerRegistryImage called for image: %s", imageName)
+	TraceMemUsage()
+
+	ref, err := name.ParseReference(imageName)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse image reference %s: %w", imageName, err)
+	}
+
+	var remoteOpts []remote.Option
+	if username != "" && password != "" {
+
+		auth := authn.Basic{
+			Username: username,
+			Password: password,
+		}
+
+		authConf, err := auth.Authorization()
+		if err != nil {
+			return nil, fmt.Errorf("failed to set remote: %w", err)
+		}
+
+		log.Debug("using basic auth")
+		authOpt := remote.WithAuth(authn.FromConfig(*authConf))
+		remoteOpts = append(remoteOpts, authOpt)
+	}
+
+	image, err = remote.Image(ref, remoteOpts...)
+	if err != nil {
+		return nil, fmt.Errorf("unable to fetch image %q, make sure it exists: %w", imageName, err)
+	}
+
+	return
+}
+
+func parseContainerRegistryImage(imageSource ImageSource, onLayer LayerParser) (
+	layerPathToHash map[string]string,
+	manifestFiles map[string]any,
+	err error,
+) {
+	layerPathToHash = make(map[string]string)
+	manifestFiles = make(map[string]any)
+
+	image, ok := imageSource.(v1.Image)
+	if !ok {
+		return nil, nil, fmt.Errorf("container registry image parser expects v1.Image, got %T", imageSource)
+	}
+
+	manifest, err := image.Manifest()
+	if err != nil {
+		return nil, nil, fmt.Errorf("unable to fetch image manifest: %w", err)
+	}
+	manifestBytes, err := json.Marshal(manifest)
+	if err != nil {
+		return nil, nil, err
+	}
+	var manifestJson map[string]any
+	if err := json.Unmarshal(manifestBytes, &manifestJson); err != nil {
+		return nil, nil, err
+	}
+	manifestFiles["manifest.json"] = manifestJson
+
+	configName, err := image.ConfigName()
+	if err != nil {
+		return nil, nil, fmt.Errorf("unable to fetch image config name: %w", err)
+	}
+	configBytes, err := image.RawConfigFile()
+	if err != nil {
+		return nil, nil, fmt.Errorf("unable to fetch image config file: %w", err)
+	}
+	var configJson map[string]any
+	if err := json.Unmarshal(configBytes, &configJson); err != nil {
+		return nil, nil, fmt.Errorf("unable to decode image config file: %w", err)
+	}
+	manifestFiles[path.Join("blobs", "sha256", configName.Hex)] = configJson
+
+	// Read the layers
+	layers, err := image.Layers()
+	if err != nil {
+		return nil, nil, fmt.Errorf("unable to fetch image layers: %w", err)
+	}
+
+	for layerNumber, layer := range layers {
+		layerReader, err := layer.Uncompressed()
+		if err != nil {
+			return nil, nil, fmt.Errorf("failed to uncompress layer %d: %w", layerNumber, err)
+		}
+		defer layerReader.Close()
+		layerDigest, err := layer.Digest()
+		if err != nil {
+			return nil, nil, fmt.Errorf("failed to read layer digest %d: %w", layerNumber, err)
+		}
+		hash, err := onLayer(layerDigest.Hex, layerReader)
+		if err != nil {
+			return nil, nil, fmt.Errorf("failed to process layer %d: %w", layerNumber, err)
+		}
+		layerPathToHash[path.Join("blobs", layerDigest.Algorithm, layerDigest.Hex)] = hash
+	}
+
+	return
+}

cmd/dmverity-vhd/createvhd.go

@@ -0,0 +1,104 @@
+package main
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+
+	log "github.com/sirupsen/logrus"
+	"github.com/urfave/cli"
+)
+
+func parseCreateVhdArgs(ctx *cli.Context) (
+	imageName string,
+	outDir string,
+	verityHashDev bool,
+	verityData bool,
+	imageFetcher ImageFetcher,
+	imageParser ImageParser,
+	manifestParser ManifestParser,
+	err error,
+) {
+
+	imageName = ctx.String(inputFlag)
+	outDir = ctx.String(outputDirFlag)
+	verityHashDev = ctx.Bool(hashDeviceVhdFlag)
+	verityData = ctx.Bool(dataVhdFlag)
+
+	imageFetcher, imageParser, manifestParser, err = getImageParsers(ctx)
+	if err != nil {
+		return "", "", false, false, nil, nil, nil, err
+	}
+
+	return
+}
+
+func createVhd(
+	imageFetcher ImageFetcher,
+	imageParser ImageParser,
+	manifestParser ManifestParser,
+	imageName string,
+	outDir string,
+	verityHashDev bool,
+	verityData bool,
+) error {
+
+	// Ensure output directory exists
+	err := ensureDirExists(outDir)
+	if err != nil {
+		return err
+	}
+
+	if verityData {
+		return saveDirTarAsVhd(imageName, verityHashDev, outDir)
+	}
+
+	layerParser := func(layerID string, layerReader io.Reader) (string, error) {
+		return "", createVHDLayer(layerID, layerReader, verityHashDev, outDir)
+	}
+
+	log.Debug("creating layer VHDs with dm-verity")
+	image, err := imageFetcher()
+	if err != nil {
+		return err
+	}
+
+	_, manifestFiles, err := imageParser(image, layerParser)
+	if err != nil {
+		return err
+	}
+
+	layerDiffIds, layerDigests, err := manifestParser(manifestFiles)
+	if err != nil {
+		return err
+	}
+
+	// Move the VHDs to the output directory
+	// They can't immediately be in the output directory because they have
+	// temporary file names based on the layer id which isn't necessarily
+	// the layer digest
+	for layerNumber := 0; layerNumber < len(layerDigests); layerNumber++ {
+		layerDigest := layerDiffIds[layerNumber]
+		layerID := layerDigests[layerNumber]
+		sanitisedFileName := sanitiseVHDFilename(layerID)
+
+		suffixes := []string{".vhd"}
+
+		for _, srcSuffix := range suffixes {
+			src := filepath.Join(os.TempDir(), sanitisedFileName+srcSuffix)
+			if _, err := os.Stat(src); os.IsNotExist(err) {
+				return fmt.Errorf("layer VHD %s does not exist", src)
+			}
+
+			dst := filepath.Join(outDir, layerDigest+srcSuffix)
+			if err := moveFile(src, dst); err != nil {
+				return err
+			}
+
+			fmt.Fprintf(os.Stdout, "Layer VHD created at %s\n", dst)
+		}
+
+	}
+	return nil
+}