public/grafana
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

AzureMonitor: Fix Azure token provider national clouds (#34615)

Browse Source 
* Fix AAD authority for sovereign clouds

* Update Azure SDK with scopes fix

* Credential initialization in cache
This commit is contained in:
Sergey Kostrukov
2021-05-25 08:19:08 +02:00
committed by GitHub
parent 9b518669dd
commit a337f70469
5 changed files with 260 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/api/pluginproxy/token_cache.go
+53 -12
View File
@@ -5,6 +5,7 @@ import (
"sort"
"strings"
"sync"
"sync/atomic"
"time"
)
@@ -15,6 +16,7 @@ type AccessToken struct {
type TokenCredential interface {
GetCacheKey() string
Init() error
GetAccessToken(ctx context.Context, scopes []string) (*AccessToken, error)
}
@@ -31,7 +33,10 @@ type tokenCacheImpl struct {
}
type credentialCacheEntry struct {
credential TokenCredential
cache sync.Map // of *scopesCacheEntry
credInit uint32
credMutex sync.Mutex
cache sync.Map // of *scopesCacheEntry
}
type scopesCacheEntry struct {
@@ -44,31 +49,67 @@ type scopesCacheEntry struct {
}
func (c *tokenCacheImpl) GetAccessToken(ctx context.Context, credential TokenCredential, scopes []string) (string, error) {
return c.getEntryFor(credential).getAccessToken(ctx, scopes)
}
func (c *tokenCacheImpl) getEntryFor(credential TokenCredential) *credentialCacheEntry {
var entry interface{}
var ok bool
credentialKey := credential.GetCacheKey()
scopesKey := getKeyForScopes(scopes)
key := credential.GetCacheKey()
if entry, ok = c.cache.Load(credentialKey); !ok {
entry, _ = c.cache.LoadOrStore(credentialKey, &credentialCacheEntry{
if entry, ok = c.cache.Load(key); !ok {
entry, _ = c.cache.LoadOrStore(key, &credentialCacheEntry{
credential: credential,
})
}
credentialEntry := entry.(*credentialCacheEntry)
return entry.(*credentialCacheEntry)
}
if entry, ok = credentialEntry.cache.Load(scopesKey); !ok {
entry, _ = credentialEntry.cache.LoadOrStore(scopesKey, &scopesCacheEntry{
credential: credentialEntry.credential,
func (c *credentialCacheEntry) getAccessToken(ctx context.Context, scopes []string) (string, error) {
err := c.ensureInitialized()
if err != nil {
return "", err
}
return c.getEntryFor(scopes).getAccessToken(ctx)
}
func (c *credentialCacheEntry) ensureInitialized() error {
if atomic.LoadUint32(&c.credInit) == 0 {
c.credMutex.Lock()
defer c.credMutex.Unlock()
if c.credInit == 0 {
// Initialize credential
err := c.credential.Init()
if err != nil {
return err
}
atomic.StoreUint32(&c.credInit, 1)
}
}
return nil
}
func (c *credentialCacheEntry) getEntryFor(scopes []string) *scopesCacheEntry {
var entry interface{}
var ok bool
key := getKeyForScopes(scopes)
if entry, ok = c.cache.Load(key); !ok {
entry, _ = c.cache.LoadOrStore(key, &scopesCacheEntry{
credential: c.credential,
scopes: scopes,
cond: sync.NewCond(&sync.Mutex{}),
})
}
scopesEntry := entry.(*scopesCacheEntry)
return scopesEntry.getAccessToken(ctx)
return entry.(*scopesCacheEntry)
}
func (c *scopesCacheEntry) getAccessToken(ctx context.Context) (string, error) {
pkg/api/pluginproxy/token_cache_test.go
+177 -6
View File
@@ -14,7 +14,9 @@ import (
type fakeCredential struct {
key string
initCalledTimes int
calledTimes int
initFunc func() error
getAccessTokenFunc func(ctx context.Context, scopes []string) (*AccessToken, error)
}
@@ -22,6 +24,19 @@ func (c *fakeCredential) GetCacheKey() string {
return c.key
}
func (c *fakeCredential) Reset() {
c.initCalledTimes = 0
c.calledTimes = 0
}
func (c *fakeCredential) Init() error {
c.initCalledTimes = c.initCalledTimes + 1
if c.initFunc != nil {
return c.initFunc()
}
return nil
}
func (c *fakeCredential) GetAccessToken(ctx context.Context, scopes []string) (*AccessToken, error) {
c.calledTimes = c.calledTimes + 1
if c.getAccessTokenFunc != nil {
@@ -103,12 +118,168 @@ func TestConcurrentTokenCache_GetAccessToken(t *testing.T) {
})
}
func TestCredentialCacheEntry_EnsureInitialized(t *testing.T) {
t.Run("when credential init returns error", func(t *testing.T) {
credential := &fakeCredential{
initFunc: func() error {
return errors.New("unable to initialize")
},
}
t.Run("should return error", func(t *testing.T) {
cacheEntry := &credentialCacheEntry{
credential: credential,
}
err := cacheEntry.ensureInitialized()
assert.Error(t, err)
})
t.Run("should call init again each time and return error", func(t *testing.T) {
credential.Reset()
cacheEntry := &credentialCacheEntry{
credential: credential,
}
var err error
err = cacheEntry.ensureInitialized()
assert.Error(t, err)
err = cacheEntry.ensureInitialized()
assert.Error(t, err)
err = cacheEntry.ensureInitialized()
assert.Error(t, err)
assert.Equal(t, 3, credential.initCalledTimes)
})
})
t.Run("when credential init returns error only once", func(t *testing.T) {
var times = 0
credential := &fakeCredential{
initFunc: func() error {
times = times + 1
if times == 1 {
return errors.New("unable to initialize")
}
return nil
},
}
t.Run("should call credential init again only while it returns error", func(t *testing.T) {
cacheEntry := &credentialCacheEntry{
credential: credential,
}
var err error
err = cacheEntry.ensureInitialized()
assert.Error(t, err)
err = cacheEntry.ensureInitialized()
assert.NoError(t, err)
err = cacheEntry.ensureInitialized()
assert.NoError(t, err)
assert.Equal(t, 2, credential.initCalledTimes)
})
})
t.Run("when credential init panics", func(t *testing.T) {
credential := &fakeCredential{
initFunc: func() error {
panic(errors.New("unable to initialize"))
},
}
t.Run("should call credential init again each time", func(t *testing.T) {
credential.Reset()
cacheEntry := &credentialCacheEntry{
credential: credential,
}
func() {
defer func() {
assert.NotNil(t, recover(), "credential expected to panic")
}()
_ = cacheEntry.ensureInitialized()
}()
func() {
defer func() {
assert.NotNil(t, recover(), "credential expected to panic")
}()
_ = cacheEntry.ensureInitialized()
}()
func() {
defer func() {
assert.NotNil(t, recover(), "credential expected to panic")
}()
_ = cacheEntry.ensureInitialized()
}()
assert.Equal(t, 3, credential.initCalledTimes)
})
})
t.Run("when credential init panics only once", func(t *testing.T) {
var times = 0
credential := &fakeCredential{
initFunc: func() error {
times = times + 1
if times == 1 {
panic(errors.New("unable to initialize"))
}
return nil
},
}
t.Run("should call credential init again only while it panics", func(t *testing.T) {
cacheEntry := &credentialCacheEntry{
credential: credential,
}
var err error
func() {
defer func() {
assert.NotNil(t, recover(), "credential expected to panic")
}()
_ = cacheEntry.ensureInitialized()
}()
func() {
defer func() {
assert.Nil(t, recover(), "credential not expected to panic")
}()
err = cacheEntry.ensureInitialized()
assert.NoError(t, err)
}()
func() {
defer func() {
assert.Nil(t, recover(), "credential not expected to panic")
}()
err = cacheEntry.ensureInitialized()
assert.NoError(t, err)
}()
assert.Equal(t, 2, credential.initCalledTimes)
})
})
}
func TestScopesCacheEntry_GetAccessToken(t *testing.T) {
ctx := context.Background()
scopes := []string{"Scope1"}
t.Run("when credential returns error", func(t *testing.T) {
t.Run("when credential getAccessToken returns error", func(t *testing.T) {
credential := &fakeCredential{
getAccessTokenFunc: func(ctx context.Context, scopes []string) (*AccessToken, error) {
invalidToken := &AccessToken{Token: "invalid_token", ExpiresOn: timeNow().Add(time.Hour)}
@@ -130,7 +301,7 @@ func TestScopesCacheEntry_GetAccessToken(t *testing.T) {
})
t.Run("should call credential again each time and return error", func(t *testing.T) {
credential.calledTimes = 0
credential.Reset()
cacheEntry := &scopesCacheEntry{
credential: credential,
@@ -152,7 +323,7 @@ func TestScopesCacheEntry_GetAccessToken(t *testing.T) {
})
})
t.Run("when credential returns error only once", func(t *testing.T) {
t.Run("when credential getAccessToken returns error only once", func(t *testing.T) {
var times = 0
credential := &fakeCredential{
getAccessTokenFunc: func(ctx context.Context, scopes []string) (*AccessToken, error) {
@@ -191,7 +362,7 @@ func TestScopesCacheEntry_GetAccessToken(t *testing.T) {
})
})
t.Run("when credential panics", func(t *testing.T) {
t.Run("when credential getAccessToken panics", func(t *testing.T) {
credential := &fakeCredential{
getAccessTokenFunc: func(ctx context.Context, scopes []string) (*AccessToken, error) {
panic(errors.New("unable to get access token"))
@@ -199,7 +370,7 @@ func TestScopesCacheEntry_GetAccessToken(t *testing.T) {
}
t.Run("should call credential again each time", func(t *testing.T) {
credential.calledTimes = 0
credential.Reset()
cacheEntry := &scopesCacheEntry{
credential: credential,
@@ -232,7 +403,7 @@ func TestScopesCacheEntry_GetAccessToken(t *testing.T) {
})
})
t.Run("when credential panics only once", func(t *testing.T) {
t.Run("when credential getAccessToken panics only once", func(t *testing.T) {
var times = 0
credential := &fakeCredential{
getAccessTokenFunc: func(ctx context.Context, scopes []string) (*AccessToken, error) {
pkg/api/pluginproxy/token_provider_azure.go
+18 -59
View File
@@ -3,11 +3,8 @@ package pluginproxy
import (
"context"
"crypto/sha256"
"errors"
"fmt"
"strings"
"sync"
"sync/atomic"
"github.com/Azure/azure-sdk-for-go/sdk/azcore"
"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
@@ -108,9 +105,8 @@ func (provider *azureAccessTokenProvider) resolveAuthorityHost(cloudName string)
}
type managedIdentityCredential struct {
clientId string
credLock sync.Mutex
credValue atomic.Value // of azcore.TokenCredential
clientId string
credential azcore.TokenCredential
}
func (c *managedIdentityCredential) GetCacheKey() string {
@@ -121,39 +117,17 @@ func (c *managedIdentityCredential) GetCacheKey() string {
return fmt.Sprintf("azure|msi|%s", clientId)
}
func (c *managedIdentityCredential) getCredential() (azcore.TokenCredential, error) {
credential := c.credValue.Load()
if credential == nil {
c.credLock.Lock()
defer c.credLock.Unlock()
var err error
credential, err = azidentity.NewManagedIdentityCredential(c.clientId, nil)
if err != nil {
return nil, err
}
c.credValue.Store(credential)
func (c *managedIdentityCredential) Init() error {
if credential, err := azidentity.NewManagedIdentityCredential(c.clientId, nil); err != nil {
return err
} else {
c.credential = credential
return nil
}
return credential.(azcore.TokenCredential), nil
}
func (c *managedIdentityCredential) GetAccessToken(ctx context.Context, scopes []string) (*AccessToken, error) {
credential, err := c.getCredential()
if err != nil {
return nil, err
}
// Implementation of ManagedIdentityCredential doesn't support scopes, converting to resource
if len(scopes) == 0 {
return nil, errors.New("scopes not provided")
}
resource := strings.TrimSuffix(scopes[0], "/.default")
scopes = []string{resource}
accessToken, err := credential.GetToken(ctx, azcore.TokenRequestOptions{Scopes: scopes})
accessToken, err := c.credential.GetToken(ctx, azcore.TokenRequestOptions{Scopes: scopes})
if err != nil {
return nil, err
}
@@ -166,40 +140,25 @@ type clientSecretCredential struct {
tenantId string
clientId string
clientSecret string
credLock sync.Mutex
credValue atomic.Value // of azcore.TokenCredential
credential azcore.TokenCredential
}
func (c *clientSecretCredential) GetCacheKey() string {
return fmt.Sprintf("azure|clientsecret|%s|%s|%s|%s", c.authority, c.tenantId, c.clientId, hashSecret(c.clientSecret))
}
func (c *clientSecretCredential) getCredential() (azcore.TokenCredential, error) {
credential := c.credValue.Load()
if credential == nil {
c.credLock.Lock()
defer c.credLock.Unlock()
var err error
credential, err = azidentity.NewClientSecretCredential(c.tenantId, c.clientId, c.clientSecret, nil)
if err != nil {
return nil, err
}
c.credValue.Store(credential)
func (c *clientSecretCredential) Init() error {
options := &azidentity.ClientSecretCredentialOptions{AuthorityHost: c.authority}
if credential, err := azidentity.NewClientSecretCredential(c.tenantId, c.clientId, c.clientSecret, options); err != nil {
return err
} else {
c.credential = credential
return nil
}
return credential.(azcore.TokenCredential), nil
}
func (c *clientSecretCredential) GetAccessToken(ctx context.Context, scopes []string) (*AccessToken, error) {
credential, err := c.getCredential()
if err != nil {
return nil, err
}
accessToken, err := credential.GetToken(ctx, azcore.TokenRequestOptions{Scopes: scopes})
accessToken, err := c.credential.GetToken(ctx, azcore.TokenRequestOptions{Scopes: scopes})
if err != nil {
return nil, err
}