Added Azure Entra Auth support

Author: Thomas Newman

api/pom.xml

@@ -94,6 +94,12 @@
             <version>2.1.0</version>
         </dependency>
 
+        <dependency>
+            <groupId>com.azure</groupId>
+            <artifactId>azure-identity</artifactId>
+            <version>1.13.0</version>
+        </dependency>
+
         <dependency>
             <groupId>org.apache.avro</groupId>
             <artifactId>avro</artifactId>

api/src/main/java/io/kafbat/ui/sasl/azure/entra/AzureEntraLoginCallbackHandler.java

@@ -0,0 +1,101 @@
+package io.kafbat.ui.sasl.azure.entra;
+
+import static org.apache.kafka.clients.CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG;
+
+import com.azure.core.credential.AccessToken;
+import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRequestContext;
+import com.azure.identity.DefaultAzureCredentialBuilder;
+import java.net.URI;
+import java.util.List;
+import java.util.Map;
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.auth.login.AppConfigurationEntry;
+import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerTokenCallback;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class AzureEntraLoginCallbackHandler implements AuthenticateCallbackHandler {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(AzureEntraLoginCallbackHandler.class);
+
+  private static final String TOKEN_AUDIENCE_FORMAT = "%s://%s/.default";
+
+  static TokenCredential tokenCredential = new DefaultAzureCredentialBuilder().build();
+
+  private TokenRequestContext tokenRequestContext;
+
+  @Override
+  public void configure(
+      Map<String, ?> configs, String mechanism, List<AppConfigurationEntry> jaasConfigEntries) {
+    tokenRequestContext = buildTokenRequestContext(configs);
+  }
+
+  private TokenRequestContext buildTokenRequestContext(Map<String, ?> configs) {
+    URI uri = buildEventHubsServerUri(configs);
+    String tokenAudience = buildTokenAudience(uri);
+
+    TokenRequestContext request = new TokenRequestContext();
+    request.addScopes(tokenAudience);
+    return request;
+  }
+
+  private URI buildEventHubsServerUri(Map<String, ?> configs) {
+    final List<String> bootstrapServers = (List<String>) configs.get(BOOTSTRAP_SERVERS_CONFIG);
+
+    if (null == bootstrapServers) {
+      final String message = BOOTSTRAP_SERVERS_CONFIG + " is missing from the Kafka configuration.";
+      LOGGER.error(message);
+      throw new IllegalArgumentException(message);
+    }
+
+    if (bootstrapServers.size() != 1) {
+      final String message =
+          BOOTSTRAP_SERVERS_CONFIG
+              + " contains multiple bootstrap servers. Only a single bootstrap server is supported.";
+      LOGGER.error(message);
+      throw new IllegalArgumentException(message);
+    }
+
+    return URI.create("https://" + bootstrapServers.get(0));
+  }
+
+  private String buildTokenAudience(URI uri) {
+    return String.format(TOKEN_AUDIENCE_FORMAT, uri.getScheme(), uri.getHost());
+  }
+
+  @Override
+  public void handle(Callback[] callbacks) throws UnsupportedCallbackException {
+    for (Callback callback : callbacks) {
+      if (callback instanceof OAuthBearerTokenCallback oauthCallback) {
+        handleOAuthCallback(oauthCallback);
+      } else {
+        throw new UnsupportedCallbackException(callback);
+      }
+    }
+  }
+
+  private void handleOAuthCallback(OAuthBearerTokenCallback oauthCallback) {
+    try {
+      final AccessToken accessToken = tokenCredential.getTokenSync(tokenRequestContext);
+      final AzureEntraOAuthBearerTokenImpl oauthBearerToken = new AzureEntraOAuthBearerTokenImpl(accessToken);
+      oauthCallback.token(oauthBearerToken);
+    } catch (final RuntimeException e) {
+      final String message =
+          "Failed to acquire Azure token for Event Hub Authentication. "
+              + "Please ensure valid Azure credentials are configured.";
+      LOGGER.error(message, e);
+      oauthCallback.error("invalid_grant", message, null);
+    }
+  }
+
+  public void close() {
+    // NOOP
+  }
+
+  void setTokenCredential(final TokenCredential tokenCredential) {
+    AzureEntraLoginCallbackHandler.tokenCredential = tokenCredential;
+  }
+}

api/src/main/java/io/kafbat/ui/sasl/azure/entra/AzureEntraOAuthBearerTokenImpl.java

@@ -0,0 +1,64 @@
+package io.kafbat.ui.sasl.azure.entra;
+
+import com.azure.core.credential.AccessToken;
+import com.nimbusds.jwt.JWTClaimsSet;
+import com.nimbusds.jwt.JWTParser;
+import java.text.ParseException;
+import java.util.Arrays;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
+import org.apache.kafka.common.errors.SaslAuthenticationException;
+import org.apache.kafka.common.security.oauthbearer.OAuthBearerToken;
+
+public class AzureEntraOAuthBearerTokenImpl implements OAuthBearerToken {
+
+  private final AccessToken accessToken;
+
+  private final JWTClaimsSet claims;
+
+  public AzureEntraOAuthBearerTokenImpl(AccessToken accessToken) {
+    this.accessToken = accessToken;
+
+    try {
+      claims = JWTParser.parse(accessToken.getToken()).getJWTClaimsSet();
+    } catch (ParseException exception) {
+      throw new SaslAuthenticationException("Unable to parse the access token", exception);
+    }
+  }
+
+  @Override
+  public String value() {
+    return accessToken.getToken();
+  }
+
+  @Override
+  public Long startTimeMs() {
+    return claims.getIssueTime().getTime();
+  }
+
+  @Override
+  public long lifetimeMs() {
+    return claims.getExpirationTime().getTime();
+  }
+
+  @Override
+  public Set<String> scope() {
+    // Referring to
+    // https://docs.microsoft.com/azure/active-directory/develop/access-tokens#payload-claims, the
+    // scp
+    // claim is a String which is presented as a space separated list.
+    return Optional.ofNullable(claims.getClaim("scp"))
+        .map(s -> Arrays.stream(((String) s).split(" ")).collect(Collectors.toSet()))
+        .orElse(null);
+  }
+
+  @Override
+  public String principalName() {
+    return (String) claims.getClaim("upn");
+  }
+
+  public boolean isExpired() {
+    return accessToken.isExpired();
+  }
+}

frontend/src/lib/constants.ts

@@ -84,6 +84,7 @@ export const AUTH_OPTIONS = [
   { value: 'Delegation tokens', label: 'Delegation tokens' },
   { value: 'SASL/LDAP', label: 'SASL/LDAP' },
   { value: 'SASL/AWS IAM', label: 'SASL/AWS IAM' },
+  { value: 'SASL/Azure Entra', label: 'SASL/Azure Entra' },
   { value: 'mTLS', label: 'mTLS' },
 ];
 

frontend/src/widgets/ClusterConfigForm/schema.ts

@@ -121,6 +121,7 @@ const authPropsSchema = lazy((_, { parent }) => {
       return object({
         awsProfileName: string(),
       });
+    case 'SASL/Azure Entra':
     case 'mTLS':
     default:
       return mixed().optional();
@@ -157,6 +158,7 @@ const authSchema = lazy((value) => {
               'SASL/SCRAM-512',
               'SASL/LDAP',
               'SASL/AWS IAM',
+              'SASL/Azure Entra',
             ].includes(v);
           },
           then: (schema) => schema.required('required field'),

frontend/src/widgets/ClusterConfigForm/utils/getJaasConfig.ts

@@ -9,6 +9,7 @@ const JAAS_CONFIGS = {
     'org.apache.kafka.common.security.scram.ScramLoginModule',
   'SASL/LDAP': 'org.apache.kafka.common.security.plain.PlainLoginModule',
   'SASL/AWS IAM': 'software.amazon.msk.auth.iam.IAMLoginModule',
+  'SASL/Azure Entra': 'org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule',
 };
 
 type MethodName = keyof typeof JAAS_CONFIGS;

frontend/src/widgets/ClusterConfigForm/utils/transformFormDataToPayload.ts

@@ -215,6 +215,14 @@ export const transformFormDataToPayload = (data: ClusterConfigFormValues) => {
           }),
         };
         break;
+      case 'SASL/Azure Entra':
+        config.properties = {
+          'security.protocol': securityProtocol,
+          'sasl.mechanism': 'OAUTHBEARER',
+          'sasl.client.callback.handler.class': 'io.kafbat.ui.sasl.azure.entra.AzureEntraLoginCallbackHandler',
+          'sasl.jaas.config': getJaasConfig('SASL/Azure Entra', {}),
+        };
+        break;
       case 'mTLS':
         config.properties = {
           'security.protocol': 'SSL',