Ensure we can access layer by diffID on a mutated image

Previously, calling LayerByDiffID() immediately after adding a layer would error with the layer not found. By ensuring that we call "compute" on the image before trying to access the diffID map, we ensure that the layer is found. Signed-off-by: Natalie Arellano <[email protected]>
google · Feb 5, 2024 · 220652e · 220652e
1 parent 8dadbe7
commit 220652e
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 0 deletions.
diff --git a/pkg/v1/mutate/image.go b/pkg/v1/mutate/image.go
@@ -277,6 +277,9 @@ func (i *image) LayerByDigest(h v1.Hash) (v1.Layer, error) {
 // LayerByDiffID is an analog to LayerByDigest, looking up by "diff id"
 // (the uncompressed hash).
 func (i *image) LayerByDiffID(h v1.Hash) (v1.Layer, error) {
+	if err := i.compute(); err != nil {
+		return nil, err
+	}
 	if layer, ok := i.diffIDMap[h]; ok {
 		return layer, nil
 	}

diff --git a/pkg/v1/mutate/mutate_test.go b/pkg/v1/mutate/mutate_test.go
@@ -28,6 +28,7 @@ import (
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
+
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/empty"
 	"github.com/google/go-containerregistry/pkg/v1/match"
@@ -200,6 +201,10 @@ func TestAppendLayers(t *testing.T) {
 		t.Fatalf("failed to append a layer: %v", err)
 	}
 
+	if !canGetLayerByDiffID(t, layer, result) {
+		t.Fatal("could not get layer by diffID before fetching manifest, config, or layers")
+	}
+
 	if manifestsAreEqual(t, source, result) {
 		t.Fatal("appending a layer did not mutate the manifest")
 	}
@@ -228,6 +233,12 @@ func TestAppendLayers(t *testing.T) {
 	}
 }
 
+func canGetLayerByDiffID(t *testing.T, layer v1.Layer, result v1.Image) bool {
+	diffID := getDiffID(t, layer)
+	_, err := result.LayerByDiffID(diffID)
+	return err == nil
+}
+
 func TestMutateConfig(t *testing.T) {
 	source := sourceImage(t)
 	cfg, err := source.ConfigFile()
@@ -700,6 +711,17 @@ func getLayers(t *testing.T, i v1.Image) []v1.Layer {
 	return l
 }
 
+func getDiffID(t *testing.T, l v1.Layer) v1.Hash {
+	t.Helper()
+
+	diffID, err := l.DiffID()
+	if err != nil {
+		t.Fatalf("Error fetching layer diffID: %v", err)
+	}
+
+	return diffID
+}
+
 func getConfigFile(t *testing.T, i v1.Image) *v1.ConfigFile {
 	t.Helper()