feat: add SquashSubset

pkg/mutate/layer_selector.go

@@ -0,0 +1,80 @@
+// Copyright 2024 Sylabs Inc. All rights reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+package mutate
+
+import (
+	"errors"
+
+	v1 "github.com/google/go-containerregistry/pkg/v1"
+)
+
+// layerSelector is a list of selected layer indexs. Negative indexes are supported, for example
+// an index of -1 would select the last layer in the image. If the underlying slice is nil, all
+// layers are selected.
+type layerSelector []int
+
+// rangeLayerSelector returns a layerSelector that selects indicies from start up to end.
+func rangeLayerSelector(start, end int) layerSelector {
+	if start >= end {
+		return layerSelector([]int{})
+	}
+
+	var s layerSelector
+	if start < end {
+		for i := start; i < end; i++ {
+			s = append(s, i)
+		}
+	}
+	return s
+}
+
+var errLayerIndexOutOfRange = errors.New("layer index out of range")
+
+// layerSelected returns true if s indicates that layer i is selected in an image with n layers.
+func (s layerSelector) indexSelected(i, n int) (bool, error) {
+	if s == nil {
+		return true, nil
+	}
+
+	for _, index := range s {
+		if index < 0 {
+			index += n
+		}
+
+		if index < 0 || n <= index {
+			return false, errLayerIndexOutOfRange
+		}
+
+		if index == i {
+			return true, nil
+		}
+	}
+
+	return false, nil
+}
+
+// layersSelected returns the selected layers from im.
+func (s layerSelector) layersSelected(im v1.Image) ([]v1.Layer, error) {
+	ls, err := im.Layers()
+	if err != nil {
+		return nil, err
+	}
+
+	if s == nil {
+		return ls, nil
+	}
+
+	var selected []v1.Layer
+
+	for i, l := range ls {
+		if ok, err := s.indexSelected(i, len(ls)); err != nil {
+			return nil, err
+		} else if ok {
+			selected = append(selected, l)
+		}
+	}
+
+	return selected, nil
+}

pkg/mutate/mutate.go

@@ -28,15 +28,40 @@ func SetLayer(i int, l v1.Layer) Mutation {
 	}
 }
 
-// ReplaceLayers replaces all layers in the image with l. The layer is annotated with the specified
-// values.
+// ReplaceLayers replaces all layers in the image with l.
 func ReplaceLayers(l v1.Layer) Mutation {
 	return func(img *image) error {
 		img.overrides = []v1.Layer{l}
 		return nil
 	}
 }
 
+// replaceSelectedLayers replaces selected layers in the image with l.
+func replaceSelectedLayers(s layerSelector, l v1.Layer) Mutation {
+	return func(img *image) error {
+		var found bool
+		var overrides []v1.Layer
+
+		// Iterate over the current layers, replacing matching layers with rl.
+		for i, override := range img.overrides {
+			selected, err := s.indexSelected(i, len(img.overrides))
+			if err != nil {
+				return err
+			}
+
+			if !selected {
+				overrides = append(overrides, override)
+			} else if !found {
+				overrides = append(overrides, l)
+				found = true
+			}
+		}
+
+		img.overrides = overrides
+		return nil
+	}
+}
+
 // SetHistory replaces the history in an image with the specified entry.
 func SetHistory(history v1.History) Mutation {
 	return func(img *image) error {

pkg/mutate/squash.go

@@ -225,11 +225,11 @@ func (s *imageState) writeHardlinksFor(target string, root entry) (entry, error)
 	return root, nil
 }
 
-// squash writes a single, squashed TAR layer built from img to w.
-func squash(img v1.Image, w io.Writer) error {
-	ls, err := img.Layers()
+// squash writes a single, squashed TAR layer built from layers selected by s from img to w.
+func squash(img v1.Image, s layerSelector, w io.Writer) error {
+	ls, err := s.layersSelected(img)
 	if err != nil {
-		return fmt.Errorf("retrieving layers: %w", err)
+		return fmt.Errorf("selecting layers: %w", err)
 	}
 
 	tw := tar.NewWriter(w)
@@ -272,13 +272,14 @@ func squash(img v1.Image, w io.Writer) error {
 	return nil
 }
 
-// Squash replaces the layers in the base image with a single, squashed layer.
-func Squash(base v1.Image) (v1.Image, error) {
+// squashSelected replaces the layers selected by s in the base image with a single, squashed
+// layer.
+func squashSelected(base v1.Image, s layerSelector) (v1.Image, error) {
 	opener := func() (io.ReadCloser, error) {
 		pr, pw := io.Pipe()
 
 		go func() {
-			pw.CloseWithError(squash(base, pw))
+			pw.CloseWithError(squash(base, s, pw))
 		}()
 
 		return pr, nil
@@ -289,5 +290,16 @@ func Squash(base v1.Image) (v1.Image, error) {
 		return nil, err
 	}
 
-	return Apply(base, ReplaceLayers(l))
+	return Apply(base, replaceSelectedLayers(s, l))
+}
+
+// Squash replaces all layers in the base image with a single, squashed layer.
+func Squash(base v1.Image) (v1.Image, error) {
+	return squashSelected(base, nil)
+}
+
+// SquashSubset replaces the layers starting at start index and up to end index with a single,
+// squashed layer.
+func SquashSubset(base v1.Image, start, end int) (v1.Image, error) {
+	return squashSelected(base, rangeLayerSelector(start, end))
 }

pkg/mutate/squash_test.go

@@ -16,6 +16,7 @@ func TestSquash(t *testing.T) {
 	tests := []struct {
 		name string
 		base v1.Image
+		s    layerSelector
 	}{
 		{
 			name: "RootDirEntry",
@@ -69,12 +70,17 @@ func TestSquash(t *testing.T) {
 			name: "HardLinkDeleteXattr",
 			base: corpus.Image(t, "hard-link-delete-xattr"),
 		},
+		{
+			name: "LayerSelector",
+			base: corpus.Image(t, "hard-link-delete-4"),
+			s:    rangeLayerSelector(0, 2),
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			var b bytes.Buffer
 
-			if err := squash(tt.base, &b); err != nil {
+			if err := squash(tt.base, tt.s, &b); err != nil {
 				t.Fatal(err)
 			}