Add optional metadata for cache.ExportReplace (#395)

apps/cache/cache.go

@@ -29,18 +29,26 @@ type Serializer interface {
 	Unmarshaler
 }
 
+// ExportHints are suggestions for storing data.
+type ExportHints struct {
+	// PartitionKey is a suggested key for partitioning the cache
+	PartitionKey string
+}
+
+// ReplaceHints are suggestions for loading data.
+type ReplaceHints struct {
+	// PartitionKey is a suggested key for partitioning the cache
+	PartitionKey string
+}
+
 // ExportReplace exports and replaces in-memory cache data. It doesn't support nil Context or
 // define the outcome of passing one. A Context without a timeout must receive a default timeout
 // specified by the implementor. Retries must be implemented inside the implementation.
 type ExportReplace interface {
-	// Replace replaces the cache with what is in external storage.
-	// key is the suggested key which can be used for partitioning the cache.
-	// Implementors should honor Context cancellations and return a context.Canceled or
-	// context.DeadlineExceeded in those cases.
-	Replace(ctx context.Context, cache Unmarshaler, key string) error
-	// Export writes the binary representation of the cache (cache.Marshal()) to
-	// external storage. This is considered opaque.
-	// key is the suggested key which can be used for partitioning the cache.
-	// Context cancellations should be honored as in Replace.
-	Export(ctx context.Context, cache Marshaler, key string) error
+	// Replace replaces the cache with what is in external storage. Implementors should honor
+	// Context cancellations and return context.Canceled or context.DeadlineExceeded in those cases.
+	Replace(ctx context.Context, cache Unmarshaler, hints ReplaceHints) error
+	// Export writes the binary representation of the cache (cache.Marshal()) to external storage.
+	// This is considered opaque. Context cancellations should be honored as in Replace.
+	Export(ctx context.Context, cache Marshaler, hints ExportHints) error
 }

apps/confidential/confidential_test.go

@@ -633,20 +633,18 @@ func TestTokenProviderOptions(t *testing.T) {
 // testCache is a simple in-memory cache.ExportReplace implementation
 type testCache map[string][]byte
 
-func (c testCache) Export(ctx context.Context, m cache.Marshaler, key string) error {
-	v, err := m.Marshal()
-	if err == nil {
-		c[key] = v
+func (c testCache) Export(ctx context.Context, m cache.Marshaler, h cache.ExportHints) error {
+	if v, err := m.Marshal(); err == nil {
+		c[h.PartitionKey] = v
 	}
-	return err
+	return nil
 }
 
-func (c testCache) Replace(ctx context.Context, u cache.Unmarshaler, key string) error {
-	var err error
-	if v, has := c[key]; has {
-		err = u.Unmarshal(v)
+func (c testCache) Replace(ctx context.Context, u cache.Unmarshaler, h cache.ReplaceHints) error {
+	if v, has := c[h.PartitionKey]; has {
+		_ = u.Unmarshal(v)
 	}
-	return err
+	return nil
 }
 
 func TestWithCache(t *testing.T) {

apps/internal/base/base.go

@@ -45,10 +45,10 @@ type partitionedManager interface {
 
 type noopCacheAccessor struct{}
 
-func (n noopCacheAccessor) Replace(ctx context.Context, cache cache.Unmarshaler, key string) error {
+func (n noopCacheAccessor) Replace(ctx context.Context, u cache.Unmarshaler, h cache.ReplaceHints) error {
 	return nil
 }
-func (n noopCacheAccessor) Export(ctx context.Context, cache cache.Marshaler, key string) error {
+func (n noopCacheAccessor) Export(ctx context.Context, m cache.Marshaler, h cache.ExportHints) error {
 	return nil
 }
 
@@ -300,7 +300,7 @@ func (b Client) AcquireTokenSilent(ctx context.Context, silent AcquireTokenSilen
 	if authParams.AuthorizationType == authority.ATOnBehalfOf {
 		if s, ok := b.pmanager.(cache.Serializer); ok {
 			suggestedCacheKey := authParams.CacheKey(silent.IsAppCache)
-			err = b.cacheAccessor.Replace(ctx, s, suggestedCacheKey)
+			err = b.cacheAccessor.Replace(ctx, s, cache.ReplaceHints{PartitionKey: suggestedCacheKey})
 			if err != nil {
 				return ar, err
 			}
@@ -315,7 +315,7 @@ func (b Client) AcquireTokenSilent(ctx context.Context, silent AcquireTokenSilen
 	} else {
 		if s, ok := b.manager.(cache.Serializer); ok {
 			suggestedCacheKey := authParams.CacheKey(silent.IsAppCache)
-			err = b.cacheAccessor.Replace(ctx, s, suggestedCacheKey)
+			err = b.cacheAccessor.Replace(ctx, s, cache.ReplaceHints{PartitionKey: suggestedCacheKey})
 			if err != nil {
 				return ar, err
 			}
@@ -426,7 +426,7 @@ func (b Client) AuthResultFromToken(ctx context.Context, authParams authority.Au
 	if authParams.AuthorizationType == authority.ATOnBehalfOf {
 		if s, ok := b.pmanager.(cache.Serializer); ok {
 			suggestedCacheKey := token.CacheKey(authParams)
-			err = b.cacheAccessor.Replace(ctx, s, suggestedCacheKey)
+			err = b.cacheAccessor.Replace(ctx, s, cache.ReplaceHints{PartitionKey: suggestedCacheKey})
 			if err != nil {
 				return ar, err
 			}
@@ -441,7 +441,7 @@ func (b Client) AuthResultFromToken(ctx context.Context, authParams authority.Au
 	} else {
 		if s, ok := b.manager.(cache.Serializer); ok {
 			suggestedCacheKey := token.CacheKey(authParams)
-			err = b.cacheAccessor.Replace(ctx, s, suggestedCacheKey)
+			err = b.cacheAccessor.Replace(ctx, s, cache.ReplaceHints{PartitionKey: suggestedCacheKey})
 			if err != nil {
 				return ar, err
 			}
@@ -461,7 +461,7 @@ func (b Client) AuthResultFromToken(ctx context.Context, authParams authority.Au
 func (b Client) AllAccounts(ctx context.Context) (accts []shared.Account, err error) {
 	if s, ok := b.manager.(cache.Serializer); ok {
 		suggestedCacheKey := b.AuthParams.CacheKey(false)
-		err = b.cacheAccessor.Replace(ctx, s, suggestedCacheKey)
+		err = b.cacheAccessor.Replace(ctx, s, cache.ReplaceHints{PartitionKey: suggestedCacheKey})
 		if err != nil {
 			return accts, err
 		}
@@ -480,7 +480,7 @@ func (b Client) Account(ctx context.Context, homeAccountID string) (acct shared.
 	authParams.HomeAccountID = homeAccountID
 	if s, ok := b.manager.(cache.Serializer); ok {
 		suggestedCacheKey := b.AuthParams.CacheKey(false)
-		err = b.cacheAccessor.Replace(ctx, s, suggestedCacheKey)
+		err = b.cacheAccessor.Replace(ctx, s, cache.ReplaceHints{PartitionKey: suggestedCacheKey})
 		if err != nil {
 			return acct, err
 		}
@@ -496,7 +496,7 @@ func (b Client) Account(ctx context.Context, homeAccountID string) (acct shared.
 func (b Client) RemoveAccount(ctx context.Context, account shared.Account) (err error) {
 	if s, ok := b.manager.(cache.Serializer); ok {
 		suggestedCacheKey := b.AuthParams.CacheKey(false)
-		err = b.cacheAccessor.Replace(ctx, s, suggestedCacheKey)
+		err = b.cacheAccessor.Replace(ctx, s, cache.ReplaceHints{PartitionKey: suggestedCacheKey})
 		if err != nil {
 			return err
 		}
@@ -515,5 +515,5 @@ func (b Client) export(ctx context.Context, marshal cache.Marshaler, key string,
 	if err != nil {
 		return err
 	}
-	return b.cacheAccessor.Export(ctx, marshal, key)
+	return b.cacheAccessor.Export(ctx, marshal, cache.ExportHints{PartitionKey: key})
 }

apps/internal/base/base_test.go

@@ -200,12 +200,12 @@ type failCache struct {
 	exportErr, replaceErr error
 }
 
-func (c *failCache) Export(context.Context, cache.Marshaler, string) error {
+func (c *failCache) Export(context.Context, cache.Marshaler, cache.ExportHints) error {
 	c.exported = true
 	return c.exportErr
 }
 
-func (c failCache) Replace(context.Context, cache.Unmarshaler, string) error {
+func (c failCache) Replace(context.Context, cache.Unmarshaler, cache.ReplaceHints) error {
 	return c.replaceErr
 }
 

apps/public/public_test.go

@@ -363,27 +363,25 @@ func TestWithInstanceDiscovery(t *testing.T) {
 }
 
 // testCache is a simple in-memory cache.ExportReplace implementation
-type testCache struct {
-	store map[string][]byte
-}
+type testCache map[string][]byte
 
-func (c *testCache) Export(ctx context.Context, m cache.Marshaler, key string) error {
+func (c testCache) Export(ctx context.Context, m cache.Marshaler, h cache.ExportHints) error {
 	v, err := m.Marshal()
 	if err == nil {
-		c.store[key] = v
+		c[h.PartitionKey] = v
 	}
 	return err
 }
 
-func (c *testCache) Replace(ctx context.Context, u cache.Unmarshaler, key string) error {
-	if v, has := c.store[key]; has {
+func (c testCache) Replace(ctx context.Context, u cache.Unmarshaler, h cache.ReplaceHints) error {
+	if v, has := c[h.PartitionKey]; has {
 		return u.Unmarshal(v)
 	}
 	return nil
 }
 
 func TestWithCache(t *testing.T) {
-	cache := testCache{make(map[string][]byte)}
+	cache := make(testCache)
 	accessToken, refreshToken := "*", "rt"
 	clientInfo := base64.RawStdEncoding.EncodeToString([]byte(`{"uid":"uid","utid":"utid"}`))
 	lmo := "login.microsoftonline.com"

apps/tests/devapps/sample_cache_accessor.go

@@ -15,15 +15,15 @@ type TokenCache struct {
 	file string
 }
 
-func (t *TokenCache) Replace(ctx context.Context, cache cache.Unmarshaler, key string) error {
+func (t *TokenCache) Replace(ctx context.Context, cache cache.Unmarshaler, hints cache.ReplaceHints) error {
 	data, err := os.ReadFile(t.file)
 	if err != nil {
 		log.Println(err)
 	}
 	return cache.Unmarshal(data)
 }
 
-func (t *TokenCache) Export(ctx context.Context, cache cache.Marshaler, key string) error {
+func (t *TokenCache) Export(ctx context.Context, cache cache.Marshaler, hints cache.ExportHints) error {
 	data, err := cache.Marshal()
 	if err != nil {
 		log.Println(err)