azure-storage-blob-go/azblob/zc_credential_token.go

package azblob

import (
	"context"
	"errors"
	"sync/atomic"

	"runtime"
	"sync"
	"time"

	"github.com/Azure/azure-pipeline-go/pipeline"
)

// TokenRefresher represents a callback method that you write; this method is called periodically
// so you can refresh the token credential's value.
type TokenRefresher func(credential TokenCredential) time.Duration

// TokenCredential represents a token credential (which is also a pipeline.Factory).
type TokenCredential interface {
	Credential
	Token() string
	SetToken(newToken string)
}

// NewTokenCredential creates a token credential for use with role-based access control (RBAC) access to Azure Storage
// resources. You initialize the TokenCredential with an initial token value. If you pass a non-nil value for
// tokenRefresher, then the function you pass will be called immediately so it can refresh and change the
// TokenCredential's token value by calling SetToken. Your tokenRefresher function must return a time.Duration
// indicating how long the TokenCredential object should wait before calling your tokenRefresher function again.
// If your tokenRefresher callback fails to refresh the token, you can return a duration of 0 to stop your
// TokenCredential object from ever invoking tokenRefresher again. Also, oen way to deal with failing to refresh a
// token is to cancel a context.Context object used by requests that have the TokenCredential object in their pipeline.
func NewTokenCredential(initialToken string, tokenRefresher TokenRefresher) TokenCredential {
	tc := &tokenCredential{}
	tc.SetToken(initialToken) // We don't set it above to guarantee atomicity
	if tokenRefresher == nil {
		return tc // If no callback specified, return the simple tokenCredential
	}

	tcwr := &tokenCredentialWithRefresh{token: tc}
	tcwr.token.startRefresh(tokenRefresher)
	runtime.SetFinalizer(tcwr, func(deadTC *tokenCredentialWithRefresh) {
		deadTC.token.stopRefresh()
		deadTC.token = nil //  Sanity (not really required)
	})
	return tcwr
}

// tokenCredentialWithRefresh is a wrapper over a token credential.
// When this wrapper object gets GC'd, it stops the tokenCredential's timer
// which allows the tokenCredential object to also be GC'd.
type tokenCredentialWithRefresh struct {
	token *tokenCredential
}

// credentialMarker is a package-internal method that exists just to satisfy the Credential interface.
func (*tokenCredentialWithRefresh) credentialMarker() {}

// Token returns the current token value
func (f *tokenCredentialWithRefresh) Token() string { return f.token.Token() }

// SetToken changes the current token value
func (f *tokenCredentialWithRefresh) SetToken(token string) { f.token.SetToken(token) }

// New satisfies pipeline.Factory's New method creating a pipeline policy object.
func (f *tokenCredentialWithRefresh) New(next pipeline.Policy, po *pipeline.PolicyOptions) pipeline.Policy {
	return f.token.New(next, po)
}

///////////////////////////////////////////////////////////////////////////////

// tokenCredential is a pipeline.Factory is the credential's policy factory.
type tokenCredential struct {
	token atomic.Value

	// The members below are only used if the user specified a tokenRefresher callback function.
	timer          *time.Timer
	tokenRefresher TokenRefresher
	lock           sync.Mutex
	stopped        bool
}

// credentialMarker is a package-internal method that exists just to satisfy the Credential interface.
func (*tokenCredential) credentialMarker() {}

// Token returns the current token value
func (f *tokenCredential) Token() string { return f.token.Load().(string) }

// SetToken changes the current token value
func (f *tokenCredential) SetToken(token string) { f.token.Store(token) }

// startRefresh calls refresh which immediately calls tokenRefresher
// and then starts a timer to call tokenRefresher in the future.
func (f *tokenCredential) startRefresh(tokenRefresher TokenRefresher) {
	f.tokenRefresher = tokenRefresher
	f.stopped = false // In case user calls StartRefresh, StopRefresh, & then StartRefresh again
	f.refresh()
}

// refresh calls the user's tokenRefresher so they can refresh the token (by
// calling SetToken) and then starts another time (based on the returned duration)
// in order to refresh the token again in the future.
func (f *tokenCredential) refresh() {
	d := f.tokenRefresher(f) // Invoke the user's refresh callback outside of the lock
	if d > 0 {               // If duration is 0 or negative, refresher wants to not be called again
		f.lock.Lock()
		if !f.stopped {
			f.timer = time.AfterFunc(d, f.refresh)
		}
		f.lock.Unlock()
	}
}

// stopRefresh stops any pending timer and sets stopped field to true to prevent
// any new timer from starting.
// NOTE: Stopping the timer allows the GC to destroy the tokenCredential object.
func (f *tokenCredential) stopRefresh() {
	f.lock.Lock()
	f.stopped = true
	if f.timer != nil {
		f.timer.Stop()
	}
	f.lock.Unlock()
}

// New satisfies pipeline.Factory's New method creating a pipeline policy object.
func (f *tokenCredential) New(next pipeline.Policy, po *pipeline.PolicyOptions) pipeline.Policy {
	return pipeline.PolicyFunc(func(ctx context.Context, request pipeline.Request) (pipeline.Response, error) {
		if request.URL.Scheme != "https" {
			// HTTPS must be used, otherwise the tokens are at the risk of being exposed
			return nil, errors.New("token credentials require a URL using the https protocol scheme")
		}
		request.Header[headerAuthorization] = []string{"Bearer " + f.Token()}
		return next.Do(ctx, request)
	})
}