- add ManagedentityInfo

- add ServicePrincipalInfo
- unwrap ExecutionException
- auth with managedId s
- remove EntraIDTokenAuthConfig

Author: atakavci

core/src/main/java/redis/clients/authentication/core/TokenManager.java

@@ -8,6 +8,7 @@
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -17,28 +18,38 @@ public class TokenManager {
     private TokenManagerConfig tokenManagerConfig;
     private IdentityProvider identityProvider;
     private TokenListener listener;
-    private ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);
+    private ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
     private ExecutorService executor = Executors.newSingleThreadExecutor();
     private boolean stopped = false;
     private ScheduledFuture<?> scheduledTask;
     private int numberOfRetries = 0;
     private Exception lastException;
     private Logger logger = LoggerFactory.getLogger(getClass());
+    private Token currentToken = null;
+    private AtomicBoolean started = new AtomicBoolean(false);
 
     public TokenManager(IdentityProvider identityProvider, TokenManagerConfig tokenManagerConfig) {
         this.identityProvider = identityProvider;
         this.tokenManagerConfig = tokenManagerConfig;
     }
 
-    public void start(TokenListener listener, boolean blockForInitialToken)
-            throws InterruptedException, ExecutionException, TimeoutException {
+    public void start(TokenListener listener, boolean blockForInitialToken) {
 
+        if (!started.compareAndSet(false, true)) {
+            throw new AuthXException("Token manager already started!");
+        }
         this.listener = listener;
         ScheduledFuture<?> currentTask = scheduleNext(0);
         scheduledTask = currentTask;
         if (blockForInitialToken) {
-            while (currentTask.get() == null) {
-                currentTask = scheduledTask;
+            try {
+                while (currentTask.get() == null) {
+                    currentTask = scheduledTask;
+                }
+            } catch (RuntimeException e) {
+                throw e;
+            } catch (Exception e) {
+                throw new TokenRequestException(unwrap(e), lastException);
             }
         }
     }
@@ -66,7 +77,8 @@ protected Token renewToken() {
         try {
             Future<Token> requestResult = executor.submit(() -> requestToken());
             newToken = requestResult.get(tokenManagerConfig.getTokenRequestExecTimeoutInMs(),
-                    TimeUnit.MILLISECONDS);
+                TimeUnit.MILLISECONDS);
+            currentToken = newToken;
             long delay = calculateRenewalDelay(newToken.getExpiresAt(), newToken.getReceivedAt());
             scheduledTask = scheduleNext(delay);
             listener.onTokenRenewed(newToken);
@@ -76,7 +88,7 @@ protected Token renewToken() {
                 numberOfRetries++;
                 scheduledTask = scheduleNext(tokenManagerConfig.getRetryPolicy().getdelayInMs());
             } else {
-                TokenRequestException tre = new TokenRequestException(e, lastException);
+                TokenRequestException tre = new TokenRequestException(unwrap(e), lastException);
                 listener.onError(tre);
                 throw tre;
             }
@@ -95,6 +107,14 @@ protected Token requestToken() {
         }
     }
 
+    private Throwable unwrap(Exception e) {
+        return (e instanceof ExecutionException) ? e.getCause() : e;
+    }
+
+    public Token getCurrentToken() {
+        return currentToken;
+    }
+
     public long calculateRenewalDelay(long expireDate, long issueDate) {
         long ttlLowerRefresh = ttlForLowerRefresh(expireDate);
         long ttlRatioRefresh = ttlForRatioRefresh(expireDate, issueDate);

core/src/test/java/redis/clients/authentication/CoreAuthenticationUnitTests.java

@@ -28,6 +28,8 @@
 import redis.clients.authentication.core.TokenListener;
 import redis.clients.authentication.core.TokenManager;
 import redis.clients.authentication.core.TokenManagerConfig;
+import redis.clients.authentication.core.TokenRequestException;
+
 import static org.awaitility.Awaitility.await;
 import java.util.concurrent.TimeUnit;
 
@@ -157,10 +159,10 @@ public void testBlockForInitialToken() {
     TokenManager tokenManager = new TokenManager(identityProvider,
         new TokenManagerConfig(0.7F, 200, 2000, new TokenManagerConfig.RetryPolicy(5, 100)));
 
-    ExecutionException e = assertThrows(ExecutionException.class,
+    TokenRequestException e = assertThrows(TokenRequestException.class,
       () -> tokenManager.start(mock(TokenListener.class), true));
 
-    assertEquals("java.lang.RuntimeException: Test exception from identity provider!",
+    assertEquals("Test exception from identity provider!",
       e.getCause().getCause().getMessage());
   }
 

entraid/src/main/java/redis/clients/authentication/entraid/EntraIDIdentityProvider.java

@@ -1,54 +1,87 @@
 package redis.clients.authentication.entraid;
 
 import java.net.MalformedURLException;
-import java.security.PrivateKey;
-import java.security.cert.X509Certificate;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
+import java.util.function.Supplier;
 
 import com.microsoft.aad.msal4j.ClientCredentialFactory;
 import com.microsoft.aad.msal4j.ClientCredentialParameters;
 import com.microsoft.aad.msal4j.ConfidentialClientApplication;
 import com.microsoft.aad.msal4j.IAuthenticationResult;
 import com.microsoft.aad.msal4j.IClientCredential;
+import com.microsoft.aad.msal4j.ManagedIdentityApplication;
+import com.microsoft.aad.msal4j.ManagedIdentityParameters;
 
 import redis.clients.authentication.core.IdentityProvider;
 import redis.clients.authentication.core.Token;
 
 public final class EntraIDIdentityProvider implements IdentityProvider {
 
-    private ConfidentialClientApplication app;
-    private ClientCredentialParameters clientParams;
+    private Supplier<IAuthenticationResult> resultSupplier;
 
-    public EntraIDIdentityProvider(String clientId, String authority, String secret,
-            Set<String> scopes) {
-        IClientCredential credential = ClientCredentialFactory.createFromSecret(secret);
-        init(clientId, authority, credential, scopes);
-    }
-
-    public EntraIDIdentityProvider(String clientId, String authority, PrivateKey key, X509Certificate cert,
-            Set<String> scopes) {
-        IClientCredential credential = ClientCredentialFactory.createFromCertificate(key, cert);
-        init(clientId, authority, credential, scopes);
-    }
+    public EntraIDIdentityProvider(ServicePrincipalInfo servicePrincipalInfo, Set<String> scopes) {
+        IClientCredential credential = getClientCredential(servicePrincipalInfo);
+        ConfidentialClientApplication app;
 
-    protected void init(String clientId, String authority, IClientCredential credential, Set<String> scopes) {
         try {
-            app = ConfidentialClientApplication.builder(clientId, credential).authority(authority).build();
+            String authority = servicePrincipalInfo.getAuthority();
+            authority = authority == null ? ConfidentialClientApplication.DEFAULT_AUTHORITY
+                    : authority;
+            app = ConfidentialClientApplication
+                    .builder(servicePrincipalInfo.getClientId(), credential).authority(authority)
+                    .build();
         } catch (MalformedURLException e) {
             throw new RedisEntraIDException("Failed to init EntraID client!", e);
         }
-        clientParams = ClientCredentialParameters.builder(scopes).build();
+        ClientCredentialParameters params = ClientCredentialParameters.builder(scopes).build();
+
+        resultSupplier = () -> supplierForConfidentialApp(app, params);
+    }
+
+    public EntraIDIdentityProvider(ManagedIdentityInfo info, Set<String> scopes) {
+        ManagedIdentityApplication app = ManagedIdentityApplication.builder(info.getId()).build();
+
+        ManagedIdentityParameters params = ManagedIdentityParameters
+                .builder(scopes.iterator().next()).build();
+        resultSupplier = () -> supplierForManagedIdentityApp(app, params);
+    }
+
+    private IClientCredential getClientCredential(ServicePrincipalInfo servicePrincipalInfo) {
+        switch (servicePrincipalInfo.getAccessWith()) {
+        case WithSecret:
+            return ClientCredentialFactory.createFromSecret(servicePrincipalInfo.getSecret());
+        case WithCert:
+            return ClientCredentialFactory.createFromCertificate(servicePrincipalInfo.getKey(),
+                servicePrincipalInfo.getCert());
+        default:
+            throw new RedisEntraIDException("Invalid ServicePrincipalAccess type!");
+        }
     }
 
     @Override
     public Token requestToken() {
+        return new JWToken(resultSupplier.get().accessToken());
+    }
+
+    public IAuthenticationResult supplierForConfidentialApp(ConfidentialClientApplication app,
+            ClientCredentialParameters params) {
         try {
-            Future<IAuthenticationResult> tokenRequest = app.acquireToken(clientParams);
-            return new JWToken(tokenRequest.get().accessToken());
+            Future<IAuthenticationResult> tokenRequest = app.acquireToken(params);
+            return tokenRequest.get();
         } catch (InterruptedException | ExecutionException e) {
             throw new RedisEntraIDException("Failed to acquire token!", e);
         }
     }
+
+    public IAuthenticationResult supplierForManagedIdentityApp(ManagedIdentityApplication app,
+            ManagedIdentityParameters params) {
+        try {
+            Future<IAuthenticationResult> tokenRequest = app.acquireTokenForManagedIdentity(params);
+            return tokenRequest.get();
+        } catch (Exception e) {
+            throw new RedisEntraIDException("Failed to acquire token!", e);
+        }
+    }
 }

entraid/src/main/java/redis/clients/authentication/entraid/EntraIDIdentityProviderConfig.java

@@ -1,58 +1,31 @@
 package redis.clients.authentication.entraid;
 
-import java.security.PrivateKey;
-import java.security.cert.X509Certificate;
 import java.util.Set;
 
 import redis.clients.authentication.core.IdentityProvider;
 import redis.clients.authentication.core.IdentityProviderConfig;
 
 public final class EntraIDIdentityProviderConfig implements IdentityProviderConfig, AutoCloseable {
 
-    public enum EntraIDAccess {
-        WithSecret,
-        WithCert,
-    }
-
-    private String clientId;
-    private EntraIDAccess accessWith;
-    private String secret;
-    private PrivateKey key;
-    private X509Certificate cert;
-    private String authority;
+    private ServicePrincipalInfo servicePrincipalInfo;
     private Set<String> scopes;
+    private ManagedIdentityInfo managedIdentityInfo;
 
-    public EntraIDIdentityProviderConfig(String clientId,
-            EntraIDAccess accessWith,
-            String secret,
-            PrivateKey key,
-            X509Certificate cert,
-            String authority, Set<String> scopes) {
-        this.clientId = clientId;
-        this.accessWith = accessWith;
-        this.secret = secret;
-        this.key = key;
-        this.cert = cert;
-        this.authority = authority;
+    public EntraIDIdentityProviderConfig(ServicePrincipalInfo servicePrincipalInfo,
+            ManagedIdentityInfo info, Set<String> scopes) {
+        this.servicePrincipalInfo = servicePrincipalInfo;
         this.scopes = scopes;
+        this.managedIdentityInfo = info;
     }
 
     @Override
     public IdentityProvider getProvider() {
         IdentityProvider identityProvider = null;
-        switch (accessWith) {
-            case WithSecret:
-                identityProvider = new EntraIDIdentityProvider(clientId, authority,
-                        secret, scopes);
-                break;
-            case WithCert:
-                identityProvider = new EntraIDIdentityProvider(clientId, authority,
-                        key, cert, scopes);
-                break;
-            default:
-                throw new RedisEntraIDException("Access type and credentials must be set!");
+        if (managedIdentityInfo != null) {
+            identityProvider = new EntraIDIdentityProvider(managedIdentityInfo, scopes);
+        } else {
+            identityProvider = new EntraIDIdentityProvider(servicePrincipalInfo, scopes);
         }
-
         clear();
         return identityProvider;
     }
@@ -63,11 +36,8 @@ public void close() throws Exception {
     }
 
     private void clear() {
-        clientId = null;
-        secret = null;
-        key = null;
-        cert = null;
-        authority = null;
+        servicePrincipalInfo = null;
+        managedIdentityInfo = null;
         scopes = null;
     }
 }