From 77134a11bc367f256021c8af271862caedaf2a70 Mon Sep 17 00:00:00 2001
From: Natalie Arellano <narellano@vmware.com>
Date: Fri, 12 Apr 2024 10:32:44 -0400
Subject: [PATCH] Add custom error type for layer not found

This avoids the need for the lifecycle to do some brittle string checking

Signed-off-by: Natalie Arellano <narellano@vmware.com>
---
 cnb_image.go   | 20 ++++++++++++++++++++
 image.go       |  8 ++++++++
 local/local.go | 14 +++++++-------
 3 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/cnb_image.go b/cnb_image.go
index 54503c88..2881d446 100644
--- a/cnb_image.go
+++ b/cnb_image.go
@@ -82,6 +82,17 @@ func (i *CNBImageCore) GetAnnotateRefName() (string, error) {
 }
 
 func (i *CNBImageCore) GetLayer(diffID string) (io.ReadCloser, error) {
+	layerHash, err := v1.NewHash(diffID)
+	if err != nil {
+		return nil, err
+	}
+	configFile, err := i.ConfigFile()
+	if err != nil {
+		return nil, err
+	}
+	if !contains(configFile.RootFS.DiffIDs, layerHash) {
+		return nil, ErrLayerNotFound{DiffID: layerHash.String()}
+	}
 	hash, err := v1.NewHash(diffID)
 	if err != nil {
 		return nil, err
@@ -93,6 +104,15 @@ func (i *CNBImageCore) GetLayer(diffID string) (io.ReadCloser, error) {
 	return layer.Uncompressed()
 }
 
+func contains(diffIDs []v1.Hash, hash v1.Hash) bool {
+	for _, diffID := range diffIDs {
+		if diffID.String() == hash.String() {
+			return true
+		}
+	}
+	return false
+}
+
 // TBD Deprecated: History
 func (i *CNBImageCore) History() ([]v1.History, error) {
 	configFile, err := getConfigFile(i.Image)
diff --git a/image.go b/image.go
index 116bb5a0..f8009442 100644
--- a/image.go
+++ b/image.go
@@ -100,3 +100,11 @@ func (e SaveError) Error() string {
 	}
 	return fmt.Sprintf("failed to write image to the following tags: %s", strings.Join(errors, ","))
 }
+
+type ErrLayerNotFound struct {
+	DiffID string
+}
+
+func (e ErrLayerNotFound) Error() string {
+	return fmt.Sprintf("failed to find layer with diff ID %q", e.DiffID)
+}
diff --git a/local/local.go b/local/local.go
index 5c407f0f..77c7a032 100644
--- a/local/local.go
+++ b/local/local.go
@@ -53,6 +53,13 @@ func (i *Image) GetLayer(diffID string) (io.ReadCloser, error) {
 	if err != nil {
 		return nil, err
 	}
+	configFile, err := i.ConfigFile()
+	if err != nil {
+		return nil, err
+	}
+	if !contains(configFile.RootFS.DiffIDs, layerHash) {
+		return nil, imgutil.ErrLayerNotFound{DiffID: layerHash.String()}
+	}
 	layer, err := i.LayerByDiffID(layerHash)
 	if err == nil {
 		// this avoids downloading ALL the image layers from the daemon
@@ -62,13 +69,6 @@ func (i *Image) GetLayer(diffID string) (io.ReadCloser, error) {
 			return layer.Uncompressed()
 		}
 	}
-	configFile, err := i.ConfigFile()
-	if err != nil {
-		return nil, err
-	}
-	if !contains(configFile.RootFS.DiffIDs, layerHash) {
-		return nil, fmt.Errorf("image %q does not contain layer with diff ID %q", i.Name(), layerHash.String())
-	}
 	if err = i.ensureLayers(); err != nil {
 		return nil, err
 	}