Enable caching oci authentication tokens

If implemented, pkg/oci/auth token will be cached with an expiration time in seconds. Signed-off-by: Soule BA <[email protected]>
fluxcd · May 3, 2024 · f79844e · f79844e
1 parent 1e86776
commit f79844e
Show file tree

Hide file tree

Showing 12 changed files with 338 additions and 149 deletions.
diff --git a/cache/cache_test.go b/cache/cache_test.go
@@ -18,7 +18,7 @@ package cache
 
 import (
 	"fmt"
-	"math/rand"
+	"math/rand/v2"
 	"sync"
 	"testing"
 	"time"
@@ -187,7 +187,7 @@ func TestCache_Concurrent(t *testing.T) {
 
 	// simulate concurrent read and write
 	for i := 0; i < concurrency; i++ {
-		key := rand.Intn(keysNum)
+		key := rand.IntN(keysNum)
 		wg.Add(2)
 		go func() {
 			defer wg.Done()

diff --git a/cache/lru_test.go b/cache/lru_test.go
@@ -17,7 +17,7 @@ limitations under the License.
 package cache
 
 import (
-	"math/rand"
+	"math/rand/v2"
 	"sync"
 	"testing"
 
@@ -375,7 +375,7 @@ func TestLRU_Concurrent(t *testing.T) {
 
 	// simulate concurrent read and write
 	for i := 0; i < concurrency; i++ {
-		key := rand.Intn(keysNum)
+		key := rand.IntN(keysNum)
 		wg.Add(2)
 		go func() {
 			defer wg.Done()

diff --git a/oci/auth/aws/auth.go b/oci/auth/aws/auth.go
@@ -24,6 +24,7 @@ import (
 	"regexp"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/config"
@@ -78,7 +79,7 @@ func (c *Client) WithConfig(cfg *aws.Config) {
 // be the case if it's running in EKS, and may need additional setup
 // otherwise (visit https://aws.github.io/aws-sdk-go-v2/docs/configuring-sdk/
 // as a starting point).
-func (c *Client) getLoginAuth(ctx context.Context, awsEcrRegion string) (authn.AuthConfig, error) {
+func (c *Client) getLoginAuth(ctx context.Context, awsEcrRegion string) (authn.AuthConfig, time.Duration, error) {
 	// No caching of tokens is attempted; the quota for getting an
 	// auth token is high enough that getting a token every time you
 	// scan an image is viable for O(500) images per region. See
@@ -94,7 +95,7 @@ func (c *Client) getLoginAuth(ctx context.Context, awsEcrRegion string) (authn.A
 		cfg, err = config.LoadDefaultConfig(ctx, config.WithRegion(awsEcrRegion))
 		if err != nil {
 			c.mu.Unlock()
-			return authConfig, fmt.Errorf("failed to load default configuration: %w", err)
+			return authConfig, 0, fmt.Errorf("failed to load default configuration: %w", err)
 		}
 		c.config = &cfg
 	}
@@ -105,65 +106,65 @@ func (c *Client) getLoginAuth(ctx context.Context, awsEcrRegion string) (authn.A
 	// pass nil input.
 	ecrToken, err := ecrService.GetAuthorizationToken(ctx, nil)
 	if err != nil {
-		return authConfig, err
+		return authConfig, 0, err
 	}
 
 	// Validate the authorization data.
 	if len(ecrToken.AuthorizationData) == 0 {
-		return authConfig, errors.New("no authorization data")
+		return authConfig, 0, errors.New("no authorization data")
 	}
 	if ecrToken.AuthorizationData[0].AuthorizationToken == nil {
-		return authConfig, fmt.Errorf("no authorization token")
+		return authConfig, 0, fmt.Errorf("no authorization token")
 	}
 	token, err := base64.StdEncoding.DecodeString(*ecrToken.AuthorizationData[0].AuthorizationToken)
 	if err != nil {
-		return authConfig, err
+		return authConfig, 0, err
 	}
 
 	tokenSplit := strings.Split(string(token), ":")
 	// Validate the tokens.
 	if len(tokenSplit) != 2 {
-		return authConfig, fmt.Errorf("invalid authorization token, expected the token to have two parts separated by ':', got %d parts", len(tokenSplit))
+		return authConfig, 0, fmt.Errorf("invalid authorization token, expected the token to have two parts separated by ':', got %d parts", len(tokenSplit))
 	}
 	authConfig = authn.AuthConfig{
 		Username: tokenSplit[0],
 		Password: tokenSplit[1],
 	}
-	return authConfig, nil
+	return authConfig, ecrToken.AuthorizationData[0].ExpiresAt.Sub(time.Now()), nil
 }
 
 // Login attempts to get the authentication material for ECR.
-func (c *Client) Login(ctx context.Context, autoLogin bool, image string) (authn.Authenticator, error) {
+func (c *Client) Login(ctx context.Context, autoLogin bool, image string) (authn.Authenticator, time.Duration, error) {
 	if autoLogin {
 		log.FromContext(ctx).Info("logging in to AWS ECR for " + image)
 		_, awsEcrRegion, ok := ParseRegistry(image)
 		if !ok {
-			return nil, errors.New("failed to parse AWS ECR image, invalid ECR image")
+			return nil, 0, errors.New("failed to parse AWS ECR image, invalid ECR image")
 		}
 
-		authConfig, err := c.getLoginAuth(ctx, awsEcrRegion)
+		authConfig, expiration, err := c.getLoginAuth(ctx, awsEcrRegion)
 		if err != nil {
-			return nil, err
+			return nil, 0, err
 		}
 
 		auth := authn.FromConfig(authConfig)
-		return auth, nil
+		return auth, expiration, nil
 	}
-	return nil, fmt.Errorf("ECR authentication failed: %w", oci.ErrUnconfiguredProvider)
+	return nil, 0, fmt.Errorf("ECR authentication failed: %w", oci.ErrUnconfiguredProvider)
 }
 
 // OIDCLogin attempts to get the authentication material for ECR.
-func (c *Client) OIDCLogin(ctx context.Context, registryURL string) (authn.Authenticator, error) {
+func (c *Client) OIDCLogin(ctx context.Context, registryURL string) (authn.Authenticator, time.Duration, error) {
 	_, awsEcrRegion, ok := ParseRegistry(registryURL)
 	if !ok {
-		return nil, errors.New("failed to parse AWS ECR image, invalid ECR image")
+		return nil, 0, errors.New("failed to parse AWS ECR image, invalid ECR image")
 	}
 
-	authConfig, err := c.getLoginAuth(ctx, awsEcrRegion)
+	authConfig, expiration, err := c.getLoginAuth(ctx, awsEcrRegion)
 	if err != nil {
-		return nil, err
+		return nil, 0, err
 	}
 
 	auth := authn.FromConfig(authConfig)
-	return auth, nil
+	return auth, expiration, nil
 }
diff --git a/oci/auth/aws/auth_test.go b/oci/auth/aws/auth_test.go
@@ -183,7 +183,7 @@ func TestGetLoginAuth(t *testing.T) {
 			cfg.Credentials = credentials.NewStaticCredentialsProvider("x", "y", "z")
 			ec.WithConfig(cfg)
 
-			a, err := ec.getLoginAuth(context.TODO(), "us-east-1")
+			a, _, err := ec.getLoginAuth(context.TODO(), "us-east-1")
 			g.Expect(err != nil).To(Equal(tt.wantErr))
 			if tt.statusCode == http.StatusOK {
 				g.Expect(a).To(Equal(tt.wantAuthConfig))
@@ -247,11 +247,11 @@ func TestLogin(t *testing.T) {
 			cfg.Credentials = credentials.NewStaticCredentialsProvider("x", "y", "z")
 			ecrClient.WithConfig(cfg)
 
-			_, err := ecrClient.Login(context.TODO(), tt.autoLogin, tt.image)
+			_, _, err := ecrClient.Login(context.TODO(), tt.autoLogin, tt.image)
 			g.Expect(err != nil).To(Equal(tt.wantErr))
 
 			if tt.testOIDC {
-				_, err = ecrClient.OIDCLogin(context.TODO(), tt.image)
+				_, _, err = ecrClient.OIDCLogin(context.TODO(), tt.image)
 				g.Expect(err != nil).To(Equal(tt.wantErr))
 			}
 		})

diff --git a/oci/auth/azure/auth.go b/oci/auth/azure/auth.go
@@ -20,6 +20,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	_ "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
@@ -33,6 +34,11 @@ import (
 	"github.com/fluxcd/pkg/oci"
 )
 
+// Default cache expiration time in seconds for ACR refresh token.
+// TODO @souleb: This is copied from https://github.com/Azure/msi-acrpull/blob/0ca921a7740e561c7204d9c3b3b55c4e0b9bd7b9/pkg/authorizer/token_retriever.go#L21C2-L21C39
+// as it not provided by the Azure SDK. See with the Azure SDK team to see if there is a better way to get this value.
+const defaultCacheExpirationInSeconds = 600
+
 // Client is an Azure ACR client which can log into the registry and return
 // authorization information.
 type Client struct {
@@ -60,7 +66,7 @@ func (c *Client) WithScheme(scheme string) *Client {
 // getLoginAuth returns authentication for ACR. The details needed for authentication
 // are gotten from environment variable so there is no need to mount a host path.
 // The endpoint is the registry server and will be queried for OAuth authorization token.
-func (c *Client) getLoginAuth(ctx context.Context, registryURL string) (authn.AuthConfig, error) {
+func (c *Client) getLoginAuth(ctx context.Context, registryURL string) (authn.AuthConfig, time.Duration, error) {
 	var authConfig authn.AuthConfig
 
 	// Use default credentials if no token credential is provided.
@@ -69,7 +75,7 @@ func (c *Client) getLoginAuth(ctx context.Context, registryURL string) (authn.Au
 	if c.credential == nil {
 		cred, err := azidentity.NewDefaultAzureCredential(nil)
 		if err != nil {
-			return authConfig, err
+			return authConfig, 0, err
 		}
 		c.credential = cred
 	}
@@ -80,22 +86,22 @@ func (c *Client) getLoginAuth(ctx context.Context, registryURL string) (authn.Au
 		Scopes: []string{configurationEnvironment.Services[cloud.ResourceManager].Endpoint + "/" + ".default"},
 	})
 	if err != nil {
-		return authConfig, err
+		return authConfig, 0, err
 	}
 
 	// Obtain ACR access token using exchanger.
 	ex := newExchanger(registryURL)
 	accessToken, err := ex.ExchangeACRAccessToken(string(armToken.Token))
 	if err != nil {
-		return authConfig, fmt.Errorf("error exchanging token: %w", err)
+		return authConfig, 0, fmt.Errorf("error exchanging token: %w", err)
 	}
 
 	return authn.AuthConfig{
 		// This is the acr username used by Azure
 		// See documentation: https://docs.microsoft.com/en-us/azure/container-registry/container-registry-authentication?tabs=azure-cli#az-acr-login-with---expose-token
 		Username: "00000000-0000-0000-0000-000000000000",
 		Password: accessToken,
-	}, nil
+	}, defaultCacheExpirationInSeconds * time.Second, nil
 }
 
 // getCloudConfiguration returns the cloud configuration based on the registry URL.
@@ -124,35 +130,35 @@ func ValidHost(host string) bool {
 
 // Login attempts to get the authentication material for ACR. The caller can
 // ensure that the passed image is a valid ACR image using ValidHost().
-func (c *Client) Login(ctx context.Context, autoLogin bool, image string, ref name.Reference) (authn.Authenticator, error) {
+func (c *Client) Login(ctx context.Context, autoLogin bool, image string, ref name.Reference) (authn.Authenticator, time.Duration, error) {
 	if autoLogin {
 		log.FromContext(ctx).Info("logging in to Azure ACR for " + image)
 		// get registry host from image
 		strArr := strings.SplitN(image, "/", 2)
 		endpoint := fmt.Sprintf("%s://%s", c.scheme, strArr[0])
-		authConfig, err := c.getLoginAuth(ctx, endpoint)
+		authConfig, expiration, err := c.getLoginAuth(ctx, endpoint)
 		if err != nil {
 			log.FromContext(ctx).Info("error logging into ACR " + err.Error())
-			return nil, err
+			return nil, 0, err
 		}
 
 		auth := authn.FromConfig(authConfig)
-		return auth, nil
+		return auth, expiration, nil
 	}
-	return nil, fmt.Errorf("ACR authentication failed: %w", oci.ErrUnconfiguredProvider)
+	return nil, 0, fmt.Errorf("ACR authentication failed: %w", oci.ErrUnconfiguredProvider)
 }
 
 // OIDCLogin attempts to get an Authenticator for the provided ACR registry URL endpoint.
 //
 // If you want to construct an Authenticator based on an image reference,
 // you may want to use Login instead.
-func (c *Client) OIDCLogin(ctx context.Context, registryUrl string) (authn.Authenticator, error) {
-	authConfig, err := c.getLoginAuth(ctx, registryUrl)
+func (c *Client) OIDCLogin(ctx context.Context, registryUrl string) (authn.Authenticator, time.Duration, error) {
+	authConfig, expiration, err := c.getLoginAuth(ctx, registryUrl)
 	if err != nil {
 		log.FromContext(ctx).Info("error logging into ACR " + err.Error())
-		return nil, err
+		return nil, 0, err
 	}
 
 	auth := authn.FromConfig(authConfig)
-	return auth, nil
+	return auth, expiration, nil
 }
diff --git a/oci/auth/azure/auth_test.go b/oci/auth/azure/auth_test.go
@@ -84,7 +84,7 @@ func TestGetAzureLoginAuth(t *testing.T) {
 				WithTokenCredential(tt.tokenCredential).
 				WithScheme("http")
 
-			auth, err := c.getLoginAuth(context.TODO(), srv.URL)
+			auth, _, err := c.getLoginAuth(context.TODO(), srv.URL)
 			g.Expect(err != nil).To(Equal(tt.wantErr))
 			if tt.statusCode == http.StatusOK {
 				g.Expect(auth).To(Equal(tt.wantAuthConfig))
@@ -187,11 +187,11 @@ func TestLogin(t *testing.T) {
 				WithTokenCredential(&FakeTokenCredential{Token: "foo"}).
 				WithScheme("http")
 
-			_, err = ac.Login(context.TODO(), tt.autoLogin, u.Host, ref)
+			_, _, err = ac.Login(context.TODO(), tt.autoLogin, u.Host, ref)
 			g.Expect(err != nil).To(Equal(tt.wantErr))
 
 			if tt.testOIDC {
-				_, err = ac.OIDCLogin(context.TODO(), srv.URL)
+				_, _, err = ac.OIDCLogin(context.TODO(), srv.URL)
 				g.Expect(err != nil).To(Equal(tt.wantErr))
 			}
 		})

diff --git a/oci/auth/gcp/auth.go b/oci/auth/gcp/auth.go
@@ -23,6 +23,7 @@ import (
 	"io"
 	"net/http"
 	"strings"
+	"time"
 
 	"github.com/google/go-containerregistry/pkg/authn"
 	"github.com/google/go-containerregistry/pkg/name"
@@ -66,66 +67,67 @@ func (c *Client) WithTokenURL(url string) *Client {
 // on GCP. This assumes that the pod has right to pull the image which would be
 // the case if it is hosted on GCP. It works with both service account and
 // workload identity enabled clusters.
-func (c *Client) getLoginAuth(ctx context.Context) (authn.AuthConfig, error) {
+func (c *Client) getLoginAuth(ctx context.Context) (authn.AuthConfig, time.Duration, error) {
 	var authConfig authn.AuthConfig
 
 	request, err := http.NewRequestWithContext(ctx, http.MethodGet, c.tokenURL, nil)
 	if err != nil {
-		return authConfig, err
+		return authConfig, 0, err
 	}
 
 	request.Header.Add("Metadata-Flavor", "Google")
 
 	client := &http.Client{}
 	response, err := client.Do(request)
 	if err != nil {
-		return authConfig, err
+		return authConfig, 0, err
 	}
 	defer response.Body.Close()
 	defer io.Copy(io.Discard, response.Body)
 
 	if response.StatusCode != http.StatusOK {
-		return authConfig, fmt.Errorf("unexpected status from metadata service: %s", response.Status)
+		return authConfig, 0, fmt.Errorf("unexpected status from metadata service: %s", response.Status)
 	}
 
 	var accessToken gceToken
 	decoder := json.NewDecoder(response.Body)
 	if err := decoder.Decode(&accessToken); err != nil {
-		return authConfig, err
+		return authConfig, 0, err
 	}
 
 	authConfig = authn.AuthConfig{
 		Username: "oauth2accesstoken",
 		Password: accessToken.AccessToken,
 	}
-	return authConfig, nil
+
+	return authConfig, time.Duration(accessToken.ExpiresIn) * time.Second, nil
 }
 
 // Login attempts to get the authentication material for GCR. The caller can
 // ensure that the passed image is a valid GCR image using ValidHost().
-func (c *Client) Login(ctx context.Context, autoLogin bool, image string, ref name.Reference) (authn.Authenticator, error) {
+func (c *Client) Login(ctx context.Context, autoLogin bool, image string, ref name.Reference) (authn.Authenticator, time.Duration, error) {
 	if autoLogin {
 		log.FromContext(ctx).Info("logging in to GCP GCR for " + image)
-		authConfig, err := c.getLoginAuth(ctx)
+		authConfig, expiration, err := c.getLoginAuth(ctx)
 		if err != nil {
 			log.FromContext(ctx).Info("error logging into GCP " + err.Error())
-			return nil, err
+			return nil, 0, err
 		}
 
 		auth := authn.FromConfig(authConfig)
-		return auth, nil
+		return auth, expiration, nil
 	}
-	return nil, fmt.Errorf("GCR authentication failed: %w", oci.ErrUnconfiguredProvider)
+	return nil, 0, fmt.Errorf("GCR authentication failed: %w", oci.ErrUnconfiguredProvider)
 }
 
 // OIDCLogin attempts to get the authentication material for GCR from the token url set in the client.
-func (c *Client) OIDCLogin(ctx context.Context) (authn.Authenticator, error) {
-	authConfig, err := c.getLoginAuth(ctx)
+func (c *Client) OIDCLogin(ctx context.Context) (authn.Authenticator, time.Duration, error) {
+	authConfig, expiration, err := c.getLoginAuth(ctx)
 	if err != nil {
 		log.FromContext(ctx).Info("error logging into GCP " + err.Error())
-		return nil, err
+		return nil, 0, err
 	}
 
 	auth := authn.FromConfig(authConfig)
-	return auth, nil
+	return auth, expiration, nil
 }