Repository: deislabs/ratify (go)

Issue:
Azure workload identity fails to refresh token
### What happened in your environment?

follow doc https://github.com/deislabs/ratify/blob/main/docs/examples/ratify-verify-azure-cmd.md to create AKS cluster with workload identity.
Everything works as expected when related crds are created. Both oras and certificateStore can access the acr and the azure key vault.
However, after 3 hours, the oras will report 401 error to access the acr.
```
time="2023-06-03T12:15:35Z" level=warning msg="failed to resolve the subject descriptor from store oras with error HEAD \"https://xinhl0ceacr.azurecr.io/v2/net-monitor/manifests/v1\": GET \"https://xinhl0ceacr.azurecr.io/oauth2/token?scope=repository%3Anet-monitor%3Apull&service=xinhl0ceacr.azurecr.io\": response status code 401: unauthorized: authentication required, visit https://aka.ms/acr/authorization for more information.\n"
time="2023-06-03T12:15:35Z" level=debug msg="verification: execution time for image xinhl0ceacr.azurecr.io/net-monitor:v1: 50ms"
```
below is the log enabled debug level
[ratify.log](https://github.com/deislabs/ratify/files/11642651/ratify.log)

It seems the ratify pod fails to refresh the workload identity token. There is no log related to "uccessfully refreshed AAD token" or "could not refresh AAD token"
After restarts the ratify pod, everything will recover.

### What did you expect to happen?

ratify pod automatically refresh the workload identity token

### What version of Kubernetes are you running?

1.26

### What version of Ratify are you running?

v1.0.0-rc.4

### Anything else you would like to add?

~~I also tried to add more logs about expire date. The token expire date stored is around 1d + 2h. However, the workload identity token refresh is by default 3600 (1h) and i think the old token is in fact deprecated in 3h~~ might have no relationship please ignore
```
time="2023-06-03T15:01:48Z" level=info msg="oras aad token expriesOn is 2023-06-04 14:59:51.979514425 +0000 UTC m=+86400.413043987"
```
### Are you willing to submit PRs to contribute to this bug fix?

- [ ] Yes, I am willing to implement it.

Pull request:
fix: Azure workload identity fails to refresh token
# Description

## What this PR does / why we need it:
Two change are included in this PR:
- extract errorResponse so we could get the actual http status code to evict auth Cache
- set auth expiration to earlier time between aadToken and acrToken

## Which issue(s) this PR fixes:
Fixes #863 

## Type of change

Please delete options that are not relevant.

- [X ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
- [ ] Helm Chart Change (any edit/addition/update that is necessary for changes merged to the `main` branch)
- [ ] This change requires a documentation update

# How Has This Been Tested?

- performed long running test > 3 hrs to validate fetch from registry is successful using work load identity. Fetch at time 0 , and Fetched again after 3 hrs, validated cache refresh was successful
- Validate the following scenario:
1. identity had no access, registry pull failed
2. Grant access to identity after step1
3. Identity pull from registry successful

- below scenario pending discussion :
1. validated identity successfully pulled from registry
2. revoke access after step 1
3.  however the we still have a cache of the valid refresh token , the identity will have access till the refresh token expires in 3 hrs.

Patch:
diff --git a/pkg/common/oras/authprovider/azure/azureidentity.go b/pkg/common/oras/authprovider/azure/azureidentity.go
index 6ae0773..8cb2fa0 100644
--- a/pkg/common/oras/authprovider/azure/azureidentity.go
+++ b/pkg/common/oras/authprovider/azure/azureidentity.go
@@ -148,11 +148,13 @@ func (d *azureManagedIdentityAuthProvider) Provide(ctx context.Context, artifact
 		return provider.AuthConfig{}, fmt.Errorf("failed to get refresh token for container registry by azure managed identity token - %w", err)
 	}
 
+	expiresOn := getACRExpiryIfEarlier(d.identityToken.ExpiresOn)
+
 	authConfig := provider.AuthConfig{
 		Username:  dockerTokenLoginUsernameGUID,
 		Password:  *rt.RefreshToken,
 		Provider:  d,
-		ExpiresOn: d.identityToken.ExpiresOn,
+		ExpiresOn: expiresOn,
 	}
 
 	return authConfig, nil
diff --git a/pkg/common/oras/authprovider/azure/azureworkloadidentity.go b/pkg/common/oras/authprovider/azure/azureworkloadidentity.go
index 1acb409..a15f996 100644
--- a/pkg/common/oras/authprovider/azure/azureworkloadidentity.go
+++ b/pkg/common/oras/authprovider/azure/azureworkloadidentity.go
@@ -139,12 +139,24 @@ func (d *azureWIAuthProvider) Provide(ctx context.Context, artifact string) (pro
 	}
 	metrics.ReportACRExchangeDuration(ctx, time.Since(startTime).Milliseconds(), artifactHostName)
 
+	refreshTokenExpiry := getACRExpiryIfEarlier(d.aadToken.ExpiresOn)
 	authConfig := provider.AuthConfig{
 		Username:  dockerTokenLoginUsernameGUID,
 		Password:  *rt.RefreshToken,
 		Provider:  d,
-		ExpiresOn: d.aadToken.ExpiresOn,
+		ExpiresOn: refreshTokenExpiry,
 	}
 
 	return authConfig, nil
 }
+
+// Compare addExpiry with default ACR refresh token expiry
+func getACRExpiryIfEarlier(aadExpiry time.Time) time.Time {
+	// set default refresh token expiry to default ACR expiry - 5 minutes
+	acrExpiration := time.Now().Add(defaultACRExpiryDuration - 5*time.Minute)
+
+	if acrExpiration.Before(aadExpiry) {
+		return acrExpiration
+	}
+	return aadExpiry
+}
diff --git a/pkg/common/oras/authprovider/azure/const.go b/pkg/common/oras/authprovider/azure/const.go
index 2faadc1..53dcaee 100644
--- a/pkg/common/oras/authprovider/azure/const.go
+++ b/pkg/common/oras/authprovider/azure/const.go
@@ -15,7 +15,10 @@ limitations under the License.
 
 package azure
 
+import "time"
+
 const (
-	dockerTokenLoginUsernameGUID = "00000000-0000-0000-0000-000000000000"
-	AADResource                  = "https://containerregistry.azure.net/.default"
+	dockerTokenLoginUsernameGUID               = "00000000-0000-0000-0000-000000000000"
+	AADResource                                = "https://containerregistry.azure.net/.default"
+	defaultACRExpiryDuration     time.Duration = 3 * time.Hour
 )
diff --git a/pkg/referrerstore/oras/cosign.go b/pkg/referrerstore/oras/cosign.go
index 277b246..f3cfef3 100644
--- a/pkg/referrerstore/oras/cosign.go
+++ b/pkg/referrerstore/oras/cosign.go
@@ -19,7 +19,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"net/http"
 	"strings"
 
 	"github.com/deislabs/ratify/pkg/common"
@@ -28,7 +27,6 @@ import (
 
 	"oras.land/oras-go/v2/errdef"
 	"oras.land/oras-go/v2/registry"
-	"oras.land/oras-go/v2/registry/remote/errcode"
 )
 
 const CosignArtifactType = "application/vnd.dev.cosign.artifact.sig.v1+json"
@@ -48,11 +46,7 @@ func getCosignReferences(ctx context.Context, subjectReference common.Reference,
 		if errors.Is(err, errdef.ErrNotFound) {
 			return nil, nil
 		}
-		var ec errcode.Error
-		if errors.As(err, &ec) && (ec.Code == fmt.Sprint(http.StatusForbidden) || ec.Code == fmt.Sprint(http.StatusUnauthorized)) {
-			store.evictAuthCache(subjectReference.Original, err)
-			return nil, err
-		}
+		evictOnError(err, store, subjectReference.Original)
 		return nil, err
 	}
 
diff --git a/pkg/referrerstore/oras/oras.go b/pkg/referrerstore/oras/oras.go
index 7313950..3082056 100644
--- a/pkg/referrerstore/oras/oras.go
+++ b/pkg/referrerstore/oras/oras.go
@@ -216,10 +216,7 @@ func (store *orasStore) ListReferrers(ctx context.Context, subjectReference comm
 		resolvedSubjectDesc = subjectDesc
 	} else {
 		if resolvedSubjectDesc, err = store.GetSubjectDescriptor(ctx, subjectReference); err != nil {
-			var ec errcode.Error
-			if errors.As(err, &ec) && (ec.Code == fmt.Sprint(http.StatusForbidden) || ec.Code == fmt.Sprint(http.StatusUnauthorized)) {
-				store.evictAuthCache(subjectReference.Original, err)
-			}
+			evictOnError(err, store, subjectReference.Original)
 			return referrerstore.ListReferrersResult{}, err
 		}
 	}
@@ -231,10 +228,7 @@ func (store *orasStore) ListReferrers(ctx context.Context, subjectReference comm
 		referrerDescriptors = append(referrerDescriptors, referrers...)
 		return nil
 	}); err != nil && !errors.Is(err, errdef.ErrNotFound) {
-		var ec errcode.Error
-		if errors.As(err, &ec) && (ec.Code == fmt.Sprint(http.StatusForbidden) || ec.Code == fmt.Sprint(http.StatusUnauthorized)) {
-			store.evictAuthCache(subjectReference.Original, err)
-		}
+		evictOnError(err, store, subjectReference.Original)
 		return referrerstore.ListReferrersResult{}, err
 	}
 	// add the repository client to the auth cache if all repository operations successful
@@ -288,10 +282,7 @@ func (store *orasStore) GetBlobContent(ctx context.Context, subjectReference com
 		// fetch blob content from remote repository
 		blobDesc, rc, err := repository.Blobs().FetchReference(ctx, ref)
 		if err != nil {
-			var ec errcode.Error
-			if errors.As(err, &ec) && (ec.Code == fmt.Sprint(http.StatusForbidden) || ec.Code == fmt.Sprint(http.StatusUnauthorized)) {
-				store.evictAuthCache(subjectReference.Original, err)
-			}
+			evictOnError(err, store, subjectReference.Original)
 			return nil, err
 		}
 
@@ -326,10 +317,7 @@ func (store *orasStore) GetReferenceManifest(ctx context.Context, subjectReferen
 		// fetch manifest content from repository
 		manifestReader, err := repository.Fetch(ctx, referenceDesc.Descriptor)
 		if err != nil {
-			var ec errcode.Error
-			if errors.As(err, &ec) && (ec.Code == fmt.Sprint(http.StatusForbidden) || ec.Code == fmt.Sprint(http.StatusUnauthorized)) {
-				store.evictAuthCache(subjectReference.Original, err)
-			}
+			evictOnError(err, store, subjectReference.Original)
 			return ocispecs.ReferenceManifest{}, err
 		}
 
@@ -389,12 +377,10 @@ func (store *orasStore) GetSubjectDescriptor(ctx context.Context, subjectReferen
 
 	desc, err = repository.Resolve(ctx, subjectReference.Original)
 	if err != nil {
-		var ec errcode.Error
-		if errors.As(err, &ec) && (ec.Code == fmt.Sprint(http.StatusForbidden) || ec.Code == fmt.Sprint(http.StatusUnauthorized)) {
-			store.evictAuthCache(subjectReference.Original, err)
-		}
+		evictOnError(err, store, subjectReference.Original)
 		return nil, err
 	}
+
 	// add the subject descriptor to cache
 	store.subjectDescriptorCache.Store(desc.Digest, desc)
 	// add the repository client to the auth cache if all repository operations successful
@@ -403,6 +389,14 @@ func (store *orasStore) GetSubjectDescriptor(ctx context.Context, subjectReferen
 	return &ocispecs.SubjectDescriptor{Descriptor: desc}, nil
 }
 
+// evict from cache on non retry-able errors including 401 and 403
+func evictOnError(err error, store *orasStore, subjectReference string) {
+	var ec *errcode.ErrorResponse
+
+	if errors.As(err, &ec) && (ec.StatusCode == http.StatusForbidden || ec.StatusCode == http.StatusUnauthorized) {
+		store.authCache.Delete(subjectReference)
+	}
+}
 func createDefaultRepository(ctx context.Context, store *orasStore, targetRef common.Reference) (registry.Repository, time.Time, error) {
 	if store.authProvider == nil || !store.authProvider.Enabled(ctx) {
 		return nil, time.Now(), fmt.Errorf("auth provider not properly enabled")
@@ -481,7 +475,3 @@ func (store *orasStore) addAuthCache(ref string, repository registry.Repository,
 		expiresOn: expiry,
 	})
 }
-
-func (store *orasStore) evictAuthCache(ref string, err error) {
-	store.authCache.Delete(ref)
-}


Test patch:
diff --git a/pkg/common/oras/authprovider/azure/azureworkloadidentity_test.go b/pkg/common/oras/authprovider/azure/azureworkloadidentity_test.go
index 97c21f5..08544da 100644
--- a/pkg/common/oras/authprovider/azure/azureworkloadidentity_test.go
+++ b/pkg/common/oras/authprovider/azure/azureworkloadidentity_test.go
@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"os"
 	"testing"
+	"time"
 
 	"github.com/AzureAD/microsoft-authentication-library-for-go/apps/confidential"
 	"github.com/deislabs/ratify/pkg/common/oras/authprovider"
@@ -57,6 +58,20 @@ func TestAzureWIEnabled_ExpectedResults(t *testing.T) {
 	}
 }
 
+func TestGetEarliestExpiration(t *testing.T) {
+	var aadExpiry = time.Now().Add(12 * time.Hour)
+
+	if getACRExpiryIfEarlier(aadExpiry) == aadExpiry {
+		t.Fatal("expected acr token expiry time")
+	}
+
+	aadExpiry = time.Now().Add(12 * time.Minute)
+
+	if getACRExpiryIfEarlier(aadExpiry) != aadExpiry {
+		t.Fatal("expected aad token expiry time")
+	}
+}
+
 // Verifies that tenant id, client id, token file path, and authority host
 // environment variables are properly set
 func TestAzureWIValidation_EnvironmentVariables_ExpectedResults(t *testing.T) {
diff --git a/pkg/referrerstore/oras/oras_test.go b/pkg/referrerstore/oras/oras_test.go
index 07f6d2d..11e08ff 100644
--- a/pkg/referrerstore/oras/oras_test.go
+++ b/pkg/referrerstore/oras/oras_test.go
@@ -476,6 +476,56 @@ func TestORASGetBlobContent_NotCachedDesc(t *testing.T) {
 	}
 }
 
+func Test_EvictOnError(t *testing.T) {
+	// store setup
+	conf := config.StorePluginConfig{
+		"name":    "oras",
+		"useHttp": true,
+	}
+
+	store, err := createBaseStore("1.0.0", conf)
+
+	if err != nil {
+		t.Fatalf("failed to create oras store: %v", err)
+	}
+
+	testcases := []struct {
+		Method           string
+		StatusCode       int
+		subjectReference string
+	}{
+		{
+			Method:     "GET",
+			StatusCode: 401,
+		},
+		{
+			Method:     "GET",
+			StatusCode: 403,
+		},
+	}
+
+	for _, testcase := range testcases {
+		subjectReference := "testSubjectRef"
+		store.authCache.Store(subjectReference, "hello")
+		_, ok := store.authCache.Load(subjectReference)
+		if !ok {
+			t.Fatalf("failed to add entry to auth cache")
+		}
+
+		mockErrResponse := errcode.ErrorResponse{
+			Method:     testcase.Method,
+			StatusCode: testcase.StatusCode,
+		}
+		evictOnError(&mockErrResponse, store, subjectReference)
+
+		// validate the entry should no longer exist
+		_, ok = store.authCache.Load(subjectReference)
+		if ok {
+			t.Fatalf("Auth cache entry should have been evicted")
+		}
+	}
+}
+
 // Test_ORASRetryClient tests that the retry client retries on 429 for specified number of retries
 func Test_ORASRetryClient(t *testing.T) {
 	ctx := context.Background()