Merge pull request #1187 from seanconroy2021/bug960

Refactored the opts into a function & added a back off strategy
enterprise-contract · Dec 7, 2023 · ce1eae9 · ce1eae9
2 parents 8a0ab94 + cb28ad8
commit ce1eae9
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 27 deletions.
diff --git a/internal/evaluation_target/application_snapshot_image/application_snapshot_image.go b/internal/evaluation_target/application_snapshot_image/application_snapshot_image.go
@@ -24,6 +24,7 @@ import (
 	"fmt"
 	"os"
 	"path"
+	"time"
 
 	ecc "github.com/enterprise-contract/enterprise-contract-controller/api/v1alpha1"
 	"github.com/google/go-containerregistry/pkg/authn"
@@ -76,6 +77,22 @@ func (a ApplicationSnapshotImage) GetReference() name.Reference {
 	return a.reference
 }
 
+func createRemoteOptions(ctx context.Context) []remote.Option {
+	backoff := remote.Backoff{
+		Duration: 1.0 * time.Second,
+		Factor:   2.0,
+		Jitter:   0.1,
+		Steps:    3,
+	}
+
+	return []remote.Option{
+		imageRefTransport,
+		remote.WithContext(ctx),
+		remote.WithAuthFromKeychain(authn.DefaultKeychain),
+		remote.WithRetryBackoff(backoff),
+	}
+}
+
 // NewApplicationSnapshotImage returns an ApplicationSnapshotImage struct with reference, checkOpts, and evaluator ready to use.
 func NewApplicationSnapshotImage(ctx context.Context, component app.SnapshotComponent, p policy.Policy) (*ApplicationSnapshotImage, error) {
 	opts, err := p.CheckOpts()
@@ -141,11 +158,7 @@ func fetchPolicySources(s ecc.Source) ([]source.PolicySource, error) {
 
 // ValidateImageAccess executes the remote.Head method on the ApplicationSnapshotImage image ref
 func (a *ApplicationSnapshotImage) ValidateImageAccess(ctx context.Context) error {
-	opts := []remote.Option{
-		imageRefTransport,
-		remote.WithContext(ctx),
-		remote.WithAuthFromKeychain(authn.DefaultKeychain),
-	}
+	opts := createRemoteOptions(ctx)
 	resp, err := NewClient(ctx).Head(a.reference, opts...)
 	if err != nil {
 		return err
@@ -175,23 +188,14 @@ func (a *ApplicationSnapshotImage) SetImageURL(url string) error {
 }
 
 func (a *ApplicationSnapshotImage) FetchImageConfig(ctx context.Context) error {
-	opts := []remote.Option{
-		imageRefTransport,
-		remote.WithContext(ctx),
-		remote.WithAuthFromKeychain(authn.DefaultKeychain),
-	}
+	opts := createRemoteOptions(ctx)
 	var err error
 	a.configJSON, err = config.FetchImageConfig(ctx, a.reference, opts...)
 	return err
 }
 
 func (a *ApplicationSnapshotImage) FetchParentImageConfig(ctx context.Context) error {
-	opts := []remote.Option{
-		imageRefTransport,
-		remote.WithContext(ctx),
-		remote.WithAuthFromKeychain(authn.DefaultKeychain),
-	}
-
+	opts := createRemoteOptions(ctx)
 	var err error
 	a.parentRef, err = config.FetchParentImage(ctx, a.reference, opts...)
 	if err != nil {
@@ -202,11 +206,7 @@ func (a *ApplicationSnapshotImage) FetchParentImageConfig(ctx context.Context) e
 }
 
 func (a *ApplicationSnapshotImage) FetchImageFiles(ctx context.Context) error {
-	opts := []remote.Option{
-		imageRefTransport,
-		remote.WithContext(ctx),
-		remote.WithAuthFromKeychain(authn.DefaultKeychain),
-	}
+	opts := createRemoteOptions(ctx)
 	var err error
 	a.files, err = files.ImageFiles(ctx, a.reference, opts...)
 	return err

diff --git a/internal/evaluation_target/application_snapshot_image/application_snapshot_image_test.go b/internal/evaluation_target/application_snapshot_image/application_snapshot_image_test.go
@@ -74,11 +74,24 @@ func TestApplicationSnapshotImage_ValidateImageAccess(t *testing.T) {
 	}
 	ref, _ := name.ParseReference("registry/image:tag")
 	tests := []struct {
-		name    string
-		fields  fields
-		args    args
-		wantErr bool
+		name      string
+		fields    fields
+		args      args
+		wantErr   bool
+		wantRetry bool
 	}{
+		{
+			name: "Retries until timeout when unable to access image ref",
+			fields: fields{
+				reference:    ref,
+				checkOpts:    cosign.CheckOpts{},
+				attestations: nil,
+				Evaluator:    nil,
+			},
+			args:      args{ctx: context.Background()},
+			wantErr:   false,
+			wantRetry: true,
+		},
 		{
 			name: "Returns no error when able to access image ref",
 			fields: fields{
@@ -104,7 +117,9 @@ func TestApplicationSnapshotImage_ValidateImageAccess(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			if tt.wantErr {
+			if tt.wantRetry {
+				imageRefTransport = remote.WithTransport(&mocks.HttpTransportTimeoutFailure{})
+			} else if tt.wantErr {
 				imageRefTransport = remote.WithTransport(&mocks.HttpTransportMockFailure{})
 			} else {
 				imageRefTransport = remote.WithTransport(&mocks.HttpTransportMockSuccess{})

diff --git a/internal/mocks/mocks.go b/internal/mocks/mocks.go
@@ -16,12 +16,19 @@
 
 package mocks
 
-import "net/http"
+import (
+	"bytes"
+	"io"
+	"net/http"
+)
 
 type HttpTransportMockSuccess struct {
 }
 type HttpTransportMockFailure struct {
 }
+type HttpTransportTimeoutFailure struct {
+	CallCount int
+}
 
 func (h *HttpTransportMockSuccess) RoundTrip(_ *http.Request) (*http.Response, error) {
 	return &http.Response{
@@ -41,3 +48,20 @@ func (h *HttpTransportMockFailure) RoundTrip(_ *http.Request) (*http.Response, e
 		},
 	}, nil
 }
+
+func (h *HttpTransportTimeoutFailure) RoundTrip(req *http.Request) (*http.Response, error) {
+	h.CallCount++
+	if h.CallCount <= 2 {
+		return &http.Response{
+			StatusCode: http.StatusRequestTimeout,
+			Body:       io.NopCloser(bytes.NewBufferString(`{"error": "Gateway Timeout"}`)),
+		}, nil
+	}
+	return &http.Response{
+		StatusCode: 200,
+		Header: http.Header{
+			"Content-Type":          {"application/json"},
+			"Docker-Content-Digest": {"sha256:11db66166c3d16c8251134e538b794ec08dfbe5f11bcc8066c6fe50e3282d6ed"},
+		},
+	}, nil
+}