From df990bf49a07e5d2ce853aecd3b4fd5512145b8c Mon Sep 17 00:00:00 2001
From: Itxaka <itxaka.garcia@spectrocloud.com>
Date: Tue, 20 Jun 2023 08:43:43 +0200
Subject: [PATCH] Improve image extractor (#32)

---
 bundles/bundles.go |  4 +--
 utils/image.go     | 61 +++++++++++++++++++++++++++++++++++++---------
 2 files changed, 52 insertions(+), 13 deletions(-)

diff --git a/bundles/bundles.go b/bundles/bundles.go
index 62b65b2..7a7db51 100644
--- a/bundles/bundles.go
+++ b/bundles/bundles.go
@@ -159,7 +159,7 @@ func (e OCIImageExtractor) Install(config *BundleConfig) error {
 			return fmt.Errorf("could not create destination path %s: %s", config.RootPath, err)
 		}
 	}
-	return utils.ExtractOCIImage(config.Target, config.RootPath, utils.GetCurrentPlatform(), config.LocalFile)
+	return utils.ExtractOCIImage(config.Target, config.RootPath, utils.GetCurrentPlatform())
 }
 
 // OCIImageRunner will extract an OCI image and then run its run.sh
@@ -172,7 +172,7 @@ func (e OCIImageRunner) Install(config *BundleConfig) error {
 	}
 	defer os.RemoveAll(tempDir)
 
-	err = utils.ExtractOCIImage(config.Target, tempDir, utils.GetCurrentPlatform(), config.LocalFile)
+	err = utils.ExtractOCIImage(config.Target, tempDir, utils.GetCurrentPlatform())
 	if err != nil {
 		return err
 	}
diff --git a/utils/image.go b/utils/image.go
index 27510b1..229f534 100644
--- a/utils/image.go
+++ b/utils/image.go
@@ -2,21 +2,48 @@ package utils
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"github.com/containerd/containerd/archive"
 	"github.com/google/go-containerregistry/pkg/authn"
+	"github.com/google/go-containerregistry/pkg/logs"
 	"github.com/google/go-containerregistry/pkg/name"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/daemon"
 	"github.com/google/go-containerregistry/pkg/v1/mutate"
 	"github.com/google/go-containerregistry/pkg/v1/remote"
+	"github.com/google/go-containerregistry/pkg/v1/remote/transport"
+	"io"
 	"net/http"
 	"runtime"
+	"strings"
+	"syscall"
+	"time"
 )
 
-// ExtractOCIImage will extract a given targetImage into a given targetDestination and pull from the local repo if set.
-func ExtractOCIImage(targetImage, targetDestination, targetPlatform string, isLocal bool) error {
+var defaultRetryBackoff = remote.Backoff{
+	Duration: 1.0 * time.Second,
+	Factor:   3.0,
+	Jitter:   0.1,
+	Steps:    3,
+}
+
+var defaultRetryPredicate = func(err error) bool {
+	if err == nil {
+		return false
+	}
+
+	if errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, io.EOF) || errors.Is(err, syscall.EPIPE) || errors.Is(err, syscall.ECONNRESET) || strings.Contains(err.Error(), "connection refused") {
+		logs.Warn.Printf("retrying %v", err)
+		return true
+	}
+	return false
+}
+
+// ExtractOCIImage will extract a given targetImage into a given targetDestination
+func ExtractOCIImage(targetImage, targetDestination, targetPlatform string) error {
 	var platform *v1.Platform
+	var img v1.Image
 	var err error
 
 	if targetPlatform != "" {
@@ -36,7 +63,7 @@ func ExtractOCIImage(targetImage, targetDestination, targetPlatform string, isLo
 		return err
 	}
 
-	img, err := getimage(ref, *platform, isLocal)
+	img, err = getimage(ref, *platform)
 	if err != nil {
 		return err
 	}
@@ -51,14 +78,26 @@ func ExtractOCIImage(targetImage, targetDestination, targetPlatform string, isLo
 }
 
 // image returns the proper image to pull with transport and auth
-func getimage(ref name.Reference, platform v1.Platform, local bool) (v1.Image, error) {
-	if local {
-		return daemon.Image(ref)
+// tries local daemon first and then fallbacks into remote
+func getimage(ref name.Reference, platform v1.Platform) (v1.Image, error) {
+	var image v1.Image
+	var err error
+	tr := transport.NewRetry(http.DefaultTransport,
+		transport.WithRetryBackoff(defaultRetryBackoff),
+		transport.WithRetryPredicate(defaultRetryPredicate),
+	)
+
+	image, err = daemon.Image(ref)
+	fmt.Println("lo")
+
+	if err != nil {
+		fmt.Println("re")
+		image, err = remote.Image(ref,
+			remote.WithTransport(tr),
+			remote.WithPlatform(platform),
+			remote.WithAuthFromKeychain(authn.DefaultKeychain),
+		)
 	}
 
-	return remote.Image(ref,
-		remote.WithTransport(http.DefaultTransport),
-		remote.WithPlatform(platform),
-		remote.WithAuthFromKeychain(authn.DefaultKeychain),
-	)
+	return image, err
 }