use idpResponse instead of plain raw token

Author: ndyakov

azure_default_identity_provider.go

@@ -2,7 +2,7 @@ package entraid
 
 import (
 	"context"
-	"time"
+	"fmt"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
@@ -32,16 +32,16 @@ func NewDefaultAzureIdentityProvider(opts DefaultAzureIdentityProviderOptions) (
 
 // RequestToken requests a token from the Azure Default Identity provider.
 // It returns the token, the expiration time, and an error if any.
-func (a *DefaultAzureIdentityProvider) RequestToken() (string, time.Time, error) {
+func (a *DefaultAzureIdentityProvider) RequestToken() (IdentityProviderResponse, error) {
 	cred, err := azidentity.NewDefaultAzureCredential(a.options)
 	if err != nil {
-		return "", time.Time{}, err
+		return nil, fmt.Errorf("failed to create default azure credential: %w", err)
 	}
 
 	token, err := cred.GetToken(context.TODO(), policy.TokenRequestOptions{Scopes: a.scopes})
 	if err != nil {
-		return "", time.Time{}, err
+		return nil, fmt.Errorf("failed to get token: %w", err)
 	}
 
-	return token.Token, token.ExpiresOn.UTC(), nil
+	return newIDPResponse(typeAccessToken, &token)
 }

confidential_identity_provider.go

@@ -5,7 +5,6 @@ import (
 	"crypto"
 	"crypto/x509"
 	"fmt"
-	"time"
 
 	confidential "github.com/AzureAD/microsoft-authentication-library-for-go/apps/confidential"
 )
@@ -117,17 +116,16 @@ func NewConfidentialIdentityProvider(opts ConfidentialIdentityProviderOptions) (
 }
 
 // RequestToken requests a token from the identity provider.
-// It returns the token, the expiration time, and an error if any.
-// The token is used to authenticate the identity when requesting a token.
-func (c *ConfidentialIdentityProvider) RequestToken() (string, time.Time, error) {
+// It returns the identity provider response, including the auth result.
+func (c *ConfidentialIdentityProvider) RequestToken() (IdentityProviderResponse, error) {
 	if c.client == nil {
-		return "", time.Time{}, fmt.Errorf("client is not initialized")
+		return nil, fmt.Errorf("client is not initialized")
 	}
 
 	result, err := c.client.AcquireTokenByCredential(context.TODO(), c.scopes)
 	if err != nil {
-		return "", time.Time{}, fmt.Errorf("failed to acquire token: %w", err)
+		return nil, fmt.Errorf("failed to acquire token: %w", err)
 	}
 
-	return result.AccessToken, result.ExpiresOn.UTC(), nil
+	return newIDPResponse(typeAuthResult, &result)
 }

identity_provider.go

@@ -1,12 +1,76 @@
 package entraid
 
-import "time"
+import (
+	"fmt"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	"github.com/AzureAD/microsoft-authentication-library-for-go/apps/public"
+)
+
+const (
+	// typeAuthResult is the type of the auth result.
+	typeAuthResult = "AuthResult"
+	// typeAccessToken is the type of the access token.
+	typeAccessToken = "AccessToken"
+)
+
+// IdentityProviderResponse is an interface that defines the methods for an identity provider authentication result.
+// It is used to get the type of the authentication result, the authentication result itself (can be AuthResult or AccessToken),
+type IdentityProviderResponse interface {
+	// Type returns the type of the auth result
+	Type() string
+	AuthResult() *public.AuthResult
+	AccessToken() *azcore.AccessToken
+}
 
 // IdentityProvider is an interface that defines the methods for an identity provider.
 // It is used to request a token for authentication.
 // The identity provider is responsible for providing the raw authentication token.
 type IdentityProvider interface {
 	// RequestToken requests a token from the identity provider.
 	// It returns the token, the expiration time, and an error if any.
-	RequestToken() (string, time.Time, error)
+	RequestToken() (IdentityProviderResponse, error)
+}
+
+type authResult struct {
+	resultType  string
+	authResult  *public.AuthResult
+	accessToken *azcore.AccessToken
+}
+
+func (a *authResult) Type() string {
+	return a.resultType
+}
+
+func (a *authResult) AuthResult() *public.AuthResult {
+	return a.authResult
+}
+
+func (a *authResult) AccessToken() *azcore.AccessToken {
+	return a.accessToken
+}
+
+// newAuthResult creates a new auth result based on the type provided.
+// It returns an IdentityProviderResponse interface.
+func newIDPResponse(t string, result interface{}) (IdentityProviderResponse, error) {
+	r := &authResult{resultType: t}
+
+	switch t {
+	case typeAuthResult:
+		if typed, ok := result.(*public.AuthResult); !ok {
+			return nil, fmt.Errorf("expected AuthResult, got %T", result)
+		} else {
+			r.authResult = typed
+		}
+	case typeAccessToken:
+		if typed, ok := result.(*azcore.AccessToken); !ok {
+			return nil, fmt.Errorf("expected AccessToken, got %T", result)
+		} else {
+			r.accessToken = typed
+		}
+	default:
+		return nil, fmt.Errorf("unknown type: %s", t)
+	}
+
+	return r, nil
 }

managed_identity_provider.go

@@ -4,11 +4,12 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"time"
 
 	mi "github.com/AzureAD/microsoft-authentication-library-for-go/apps/managedidentity"
 )
 
+// ManagedIdentityProviderOptions represents the options for the managed identity provider.
+// It is used to configure the identity provider when requesting a token.
 type ManagedIdentityProviderOptions struct {
 	// UserAssignedClientID is the client ID of the user assigned identity.
 	// This is used to identify the identity when requesting a token.
@@ -21,6 +22,7 @@ type ManagedIdentityProviderOptions struct {
 	Scopes []string
 }
 
+// ManagedIdentityProvider represents a managed identity provider.
 type ManagedIdentityProvider struct {
 	// userAssignedClientID is the client ID of the user assigned identity.
 	// This is used to identify the identity when requesting a token.
@@ -38,6 +40,8 @@ type ManagedIdentityProvider struct {
 	client *mi.Client
 }
 
+// NewManagedIdentityProvider creates a new managed identity provider for Azure with managed identity.
+// It is used to configure the identity provider when requesting a token.
 func NewManagedIdentityProvider(opts ManagedIdentityProviderOptions) (*ManagedIdentityProvider, error) {
 	var client mi.Client
 	var err error
@@ -73,9 +77,11 @@ func NewManagedIdentityProvider(opts ManagedIdentityProviderOptions) (*ManagedId
 	}, nil
 }
 
-func (m *ManagedIdentityProvider) RequestToken() (string, time.Time, error) {
+// RequestToken requests a token from the managed identity provider.
+// It returns IdentityProviderResponse, which contains the Acc and the expiration time.
+func (m *ManagedIdentityProvider) RequestToken() (IdentityProviderResponse, error) {
 	if m.client == nil {
-		return "", time.Time{}, errors.New("managed identity client is not initialized")
+		return nil, errors.New("managed identity client is not initialized")
 	}
 
 	// default resource is RedisResource == "https://redis.azure.com"
@@ -87,11 +93,10 @@ func (m *ManagedIdentityProvider) RequestToken() (string, time.Time, error) {
 	}
 	// acquire token using the managed identity client
 	// the resource is the URL of the resource that the identity has access to
-	token, err := m.client.AcquireToken(context.TODO(), resource)
+	authResult, err := m.client.AcquireToken(context.TODO(), resource)
 	if err != nil {
-		return "", time.Time{}, fmt.Errorf("coudn't acquire token: %w", err)
+		return nil, fmt.Errorf("coudn't acquire token: %w", err)
 	}
 
-	// return the access token
-	return token.AccessToken, token.ExpiresOn.UTC(), nil
+	return newIDPResponse(typeAuthResult, &authResult)
 }

providers.go

@@ -6,6 +6,9 @@ import (
 	"github.com/redis/go-redis/v9/auth"
 )
 
+// CredentialsProviderOptions is a struct that holds the options for the credentials provider.
+// It is used to configure the credentials provider when requesting a token.
+// It is used to specify the client ID, TokenManagerOptions, and callback functions for re-authentication and retryable errors.
 type CredentialsProviderOptions struct {
 	// ClientID is the client ID of the identity.
 	// This is used to identify the identity when requesting a token.

token.go

@@ -20,8 +20,8 @@ type Token struct {
 	RawToken string `json:"raw_token"`
 }
 
-// TokenParserFunc is a function that parses the token and returns the username and password.
-type TokenParserFunc func(token string, expiresOn time.Time) (*Token, error)
+// IdentityProviderResponseParserFunc is a function that parses the token and returns the username and password.
+type IdentityProviderResponseParserFunc func(response IdentityProviderResponse) (*Token, error)
 
 // copyToken creates a copy of the token.
 func copyToken(token *Token) *Token {

token_manager.go

@@ -31,12 +31,13 @@ type TokenManagerOptions struct {
 	// default: 0 ms (no lower bound, refresh based on ExpirationRefreshRatio)
 	LowerRefreshBoundMs int
 
-	// TokenParser is a function that parses the raw token and returns a Token object.
-	// The function takes the raw token as a string and returns a Token object and an error.
+	// IdentityProviderResponseParser is a function that parses the IdentityProviderResponse.
+	// The function takes the response and based on its type returns the populated Token object.
 	// If this function is not provided, the default implementation will be used.
 	//
 	// required: true
-	TokenParser TokenParserFunc
+	// default: defaultIdentityProviderResponseParser
+	IdentityProviderResponseParser IdentityProviderResponseParserFunc
 
 	// RetryOptions is a struct that contains the options for retrying the token request.
 	// It contains the maximum number of attempts, initial delay, maximum delay, and backoff multiplier.
@@ -84,17 +85,28 @@ type TokenManager interface {
 	Close() error
 }
 
-// defaultTokenParser is a function that parses the raw token and returns Token object.
-var defaultTokenParser = func(rawToken string, expiresOn time.Time) (*Token, error) {
-	// Parse the token and return the username and password.
-	// In this example, we are just returning the raw token as the password.
-	// In a real implementation, you would parse the token and extract the username and password.
-	// For example, if the token is a JWT, you would decode the JWT and extract the claims.
-	// This is just a placeholder implementation.
-	// You should replace this with your own implementation.
-	if rawToken == "" {
-		return nil, fmt.Errorf("raw token is empty")
+// defaultIdentityProviderResponseParser is a function that parses the token and returns the username and password.
+var defaultIdentityProviderResponseParser = func(response IdentityProviderResponse) (*Token, error) {
+	var username, password, rawToken string
+	var expiresOn time.Time
+	if response == nil {
+		return nil, fmt.Errorf("response is nil")
 	}
+	switch response.Type() {
+	case typeAuthResult:
+		authResult := response.AuthResult()
+		if authResult == nil {
+			return nil, fmt.Errorf("auth result is nil")
+		}
+	case typeAccessToken:
+		accessToken := response.AccessToken()
+		if accessToken == nil {
+			return nil, fmt.Errorf("access token is nil")
+		}
+	default:
+		return nil, fmt.Errorf("unknown response type: %s", response.Type())
+	}
+
 	if expiresOn.IsZero() {
 		return nil, fmt.Errorf("expires on is zero")
 	}
@@ -143,8 +155,9 @@ func NewTokenManager(idp IdentityProvider, options TokenManagerOptions) (TokenMa
 type entraidTokenManager struct {
 	idp   IdentityProvider
 	token *Token
+
 	// TokenParser is a function that parses the token.
-	tokenParser TokenParserFunc
+	identityProviderResponseParser IdentityProviderResponseParserFunc
 
 	// retryOptions is a struct that contains the options for retrying the token request.
 	// It contains the maximum number of attempts, initial delay, maximum delay, and backoff multiplier.
@@ -175,12 +188,12 @@ func (e *entraidTokenManager) GetToken() (*Token, error) {
 		return copyToken(e.token), nil
 	}
 
-	rawToken, expiresOn, err := e.idp.RequestToken()
+	idpResult, err := e.idp.RequestToken()
 	if err != nil {
-		return nil, fmt.Errorf("failed to request token: %w", err)
+		return nil, fmt.Errorf("failed to request token from idp: %w", err)
 	}
 
-	token, err := e.tokenParser(rawToken, expiresOn)
+	token, err := e.identityProviderResponseParser(idpResult)
 	if err != nil {
 		return nil, fmt.Errorf("failed to parse token: %w", err)
 	}
@@ -299,7 +312,7 @@ var defaultRetryableFunc = func(err error) bool {
 	}
 
 	if ok := errors.As(err, netErr); ok {
-		return netErr.Timeout() || netErr.Temporary()
+		return netErr.Timeout()
 	}
 	return false
 }
@@ -328,19 +341,19 @@ func defaultRetryOptionsOr(retryOptions RetryOptions) RetryOptions {
 	return retryOptions
 }
 
-// defaultTokenParserOr returns the default token parser if the provided token parser is not set.
-// It sets the default token parser to the defaultTokenParser function.
+// defaultIdentityProviderResponseParserOr returns the default token parser if the provided token parser is not set.
+// It sets the default token parser to the defaultIdentityProviderResponseParser function.
 // The default token parser is used to parse the raw token and return a Token object.
-func defaultTokenParserOr(tokenParser TokenParserFunc) TokenParserFunc {
-	if tokenParser == nil {
-		return defaultTokenParser
+func defaultIdentityProviderResponseParserOr(idpResponseParser IdentityProviderResponseParserFunc) IdentityProviderResponseParserFunc {
+	if idpResponseParser == nil {
+		return defaultIdentityProviderResponseParser
 	}
-	return tokenParser
+	return idpResponseParser
 }
 
 func defaultTokenManagerOptionsOr(options TokenManagerOptions) TokenManagerOptions {
 	options.RetryOptions = defaultRetryOptionsOr(options.RetryOptions)
-	options.TokenParser = defaultTokenParserOr(options.TokenParser)
+	options.IdentityProviderResponseParser = defaultIdentityProviderResponseParserOr(options.IdentityProviderResponseParser)
 	if options.ExpirationRefreshRatio <= 0 || options.ExpirationRefreshRatio > 1 {
 		options.ExpirationRefreshRatio = 0.7
 	}