From 740eae8f2653bffe2891c5d79a66bb08245f6cb5 Mon Sep 17 00:00:00 2001
From: peusebiu <peusebiu@cisco.com>
Date: Thu, 29 Feb 2024 19:09:21 +0200
Subject: [PATCH] fix(sync): better cleaning sync's download dir (#2273)

added cleanup in the case of copy.Image() failures.

Signed-off-by: Petu Eusebiu <peusebiu@cisco.com>
---
 pkg/extensions/sync/destination.go | 21 +++++++++++++++++----
 pkg/extensions/sync/service.go     |  7 +++++++
 pkg/extensions/sync/sync.go        |  2 ++
 3 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/pkg/extensions/sync/destination.go b/pkg/extensions/sync/destination.go
index 4e4c57b1..3eaee8ab 100644
--- a/pkg/extensions/sync/destination.go
+++ b/pkg/extensions/sync/destination.go
@@ -186,6 +186,12 @@ func (registry *DestinationRegistry) CommitImage(imageReference types.ImageRefer
 	return nil
 }
 
+func (registry *DestinationRegistry) CleanupImage(imageReference types.ImageReference, repo, reference string) error {
+	tmpDir := getTempRootDirFromImageReference(imageReference, repo, reference)
+
+	return os.RemoveAll(tmpDir)
+}
+
 func (registry *DestinationRegistry) copyManifest(repo string, manifestContent []byte, reference string,
 	tempImageStore storageTypes.ImageStore,
 ) error {
@@ -275,17 +281,24 @@ func (registry *DestinationRegistry) copyBlob(repo string, blobDigest digest.Dig
 	return err
 }
 
+// use only with local imageReferences.
 func getImageStoreFromImageReference(imageReference types.ImageReference, repo, reference string,
 ) storageTypes.ImageStore {
-	var tempRootDir string
+	tmpRootDir := getTempRootDirFromImageReference(imageReference, repo, reference)
+
+	return getImageStore(tmpRootDir)
+}
+
+func getTempRootDirFromImageReference(imageReference types.ImageReference, repo, reference string) string {
+	var tmpRootDir string
 
 	if strings.HasSuffix(imageReference.StringWithinTransport(), reference) {
-		tempRootDir = strings.ReplaceAll(imageReference.StringWithinTransport(), fmt.Sprintf("%s:%s", repo, reference), "")
+		tmpRootDir = strings.ReplaceAll(imageReference.StringWithinTransport(), fmt.Sprintf("%s:%s", repo, reference), "")
 	} else {
-		tempRootDir = strings.ReplaceAll(imageReference.StringWithinTransport(), fmt.Sprintf("%s:", repo), "")
+		tmpRootDir = strings.ReplaceAll(imageReference.StringWithinTransport(), fmt.Sprintf("%s:", repo), "")
 	}
 
-	return getImageStore(tempRootDir)
+	return tmpRootDir
 }
 
 func getImageStore(rootDir string) storageTypes.ImageStore {
diff --git a/pkg/extensions/sync/service.go b/pkg/extensions/sync/service.go
index e0caedf6..237822f0 100644
--- a/pkg/extensions/sync/service.go
+++ b/pkg/extensions/sync/service.go
@@ -441,6 +441,13 @@ func (service *BaseService) syncTag(ctx context.Context, destinationRepo, remote
 
 		_, err = copy.Image(ctx, policyContext, localImageRef, remoteImageRef, &copyOptions)
 		if err != nil {
+			// cleanup in cases of copy.Image errors while copying.
+			if cErr := service.destination.CleanupImage(localImageRef, destinationRepo, tag); cErr != nil {
+				service.log.Error().Err(err).Str("errortype", common.TypeOf(err)).
+					Str("local image", fmt.Sprintf("%s:%s", destinationRepo, tag)).
+					Msg("couldn't cleanup temp local image")
+			}
+
 			service.log.Error().Err(err).Str("errortype", common.TypeOf(err)).
 				Str("remote image", remoteImageRef.DockerReference().String()).
 				Str("local image", fmt.Sprintf("%s:%s", destinationRepo, tag)).Msg("coulnd't sync image")
diff --git a/pkg/extensions/sync/sync.go b/pkg/extensions/sync/sync.go
index 4743ee34..cfd9eb85 100644
--- a/pkg/extensions/sync/sync.go
+++ b/pkg/extensions/sync/sync.go
@@ -75,6 +75,8 @@ type Destination interface {
 	CanSkipImage(repo, tag string, imageDigest digest.Digest) (bool, error)
 	// CommitImage moves a synced repo/ref from temporary oci layout to ImageStore
 	CommitImage(imageReference types.ImageReference, repo, tag string) error
+	// Removes image reference, used when copy.Image() errors out
+	CleanupImage(imageReference types.ImageReference, repo, reference string) error
 }
 
 type TaskGenerator struct {