refactor: remove cached layers in clean

Signed-off-by: nikpivkin <[email protected]>
aquasecurity · Dec 27, 2024 · bba405b · bba405b
1 parent a6ce2ed
commit bba405b
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 7 deletions.
diff --git a/pkg/fanal/artifact/image/image.go b/pkg/fanal/artifact/image/image.go
@@ -41,7 +41,7 @@ type Artifact struct {
 
 	artifactOption artifact.Option
 
-	cacheDir string
+	layerCacheDir string
 }
 
 type LayerInfo struct {
@@ -81,7 +81,7 @@ func NewArtifact(img types.Image, c cache.ArtifactCache, opt artifact.Option) (a
 		handlerManager: handlerManager,
 
 		artifactOption: opt,
-		cacheDir:       cacheDir,
+		layerCacheDir:  cacheDir,
 	}, nil
 }
 
@@ -100,8 +100,10 @@ func (a Artifact) Inspect(ctx context.Context) (artifact.Reference, error) {
 	diffIDs := a.diffIDs(configFile)
 	a.logger.Debug("Detected diff ID", log.Any("diff_ids", diffIDs))
 
-	defer os.RemoveAll(a.cacheDir)
 	if err := a.checkImageSize(ctx, diffIDs); err != nil {
+		if err := os.RemoveAll(a.layerCacheDir); err != nil {
+			log.Error("Failed to remove layer cache", log.Err(err))
+		}
 		return artifact.Reference{}, err
 	}
 
@@ -158,8 +160,8 @@ func (a Artifact) Inspect(ctx context.Context) (artifact.Reference, error) {
 	}, nil
 }
 
-func (Artifact) Clean(_ artifact.Reference) error {
-	return nil
+func (a Artifact) Clean(_ artifact.Reference) error {
+	return os.RemoveAll(a.layerCacheDir)
 }
 
 func (a Artifact) calcCacheKeys(imageID string, diffIDs []string) (string, []string, error) {
@@ -269,7 +271,7 @@ func (a Artifact) saveLayer(diffID string) (int64, error) {
 	}
 	defer rc.Close()
 
-	f, err := os.Create(filepath.Join(a.cacheDir, diffID))
+	f, err := os.Create(filepath.Join(a.layerCacheDir, diffID))
 	if err != nil {
 		return -1, xerrors.Errorf("failed to create a file: %w", err)
 	}
@@ -441,7 +443,7 @@ func (a Artifact) uncompressedLayer(diffID string) (string, io.ReadCloser, error
 		digest = d.String()
 	}
 
-	f, err := os.Open(filepath.Join(a.cacheDir, diffID))
+	f, err := os.Open(filepath.Join(a.layerCacheDir, diffID))
 	if err == nil {
 		return digest, f, nil
 	}

diff --git a/pkg/fanal/artifact/image/image_test.go b/pkg/fanal/artifact/image/image_test.go
@@ -2272,6 +2272,8 @@ func TestArtifact_Inspect(t *testing.T) {
 				assert.ErrorContains(t, err, tt.wantErr, tt.name)
 				return
 			}
+			defer a.Clean(got)
+
 			require.NoError(t, err, tt.name)
 			assert.Equal(t, tt.want, got)
 		})

diff --git a/pkg/fanal/artifact/image/remote_sbom_test.go b/pkg/fanal/artifact/image/remote_sbom_test.go
@@ -170,6 +170,8 @@ func TestArtifact_InspectRekorAttestation(t *testing.T) {
 				assert.ErrorContains(t, err, tt.wantErr)
 				return
 			}
+			defer a.Clean(got)
+
 			require.NoError(t, err, tt.name)
 			got.BOM = nil
 			assert.Equal(t, tt.want, got)
@@ -312,6 +314,7 @@ func TestArtifact_inspectOCIReferrerSBOM(t *testing.T) {
 				assert.ErrorContains(t, err, tt.wantErr)
 				return
 			}
+			defer a.Clean(got)
 
 			require.NoError(t, err, tt.name)
 			got.BOM = nil

diff --git a/pkg/fanal/test/integration/registry_test.go b/pkg/fanal/test/integration/registry_test.go
@@ -256,6 +256,7 @@ func analyze(ctx context.Context, imageRef string, opt types.ImageOptions) (*typ
 	if err != nil {
 		return nil, err
 	}
+	defer ar.Clean(imageInfo)
 
 	imageDetail, err := ap.ApplyLayers(imageInfo.ID, imageInfo.BlobIDs)
 	if err != nil {