out_kafka: add AWS MSK IAM authentication support

Enable AWS MSK IAM authentication in the Kafka output plugin:
- Add AWS MSK IAM configuration options
- Integrate with OAuth callback mechanism
- Support automatic credential refresh
- Add TLS configuration for secure connections

Signed-off-by: Arbin <[email protected]>

Author: kalavt

plugins/out_kafka/kafka.c

@@ -678,19 +678,6 @@ static struct flb_config_map config_map[] = {
     "that key will be sent to Kafka."
    },
 
-#ifdef FLB_HAVE_AWS_MSK_IAM
-   {
-    FLB_CONFIG_MAP_STR, "aws_msk_iam_cluster_arn", NULL,
-    0, FLB_TRUE, offsetof(struct flb_out_kafka, aws_msk_iam_cluster_arn),
-    "ARN of the MSK cluster when using AWS IAM authentication"
-   },
-   {
-    FLB_CONFIG_MAP_BOOL, "aws_msk_iam", "false",
-    0, FLB_TRUE, offsetof(struct flb_out_kafka, aws_msk_iam),
-    "Enable AWS MSK IAM authentication"
-   },
-#endif
-
    /* EOF */
    {0}
 };

plugins/out_kafka/kafka_config.c

@@ -58,37 +58,33 @@ struct flb_out_kafka *flb_out_kafka_create(struct flb_output_instance *ins,
         return NULL;
     }
 
+    /* Retrieve SASL mechanism if configured */
+    tmp = flb_output_get_property("rdkafka.sasl.mechanism", ins);
+    if (tmp) {
+        ctx->sasl_mechanism = flb_sds_create(tmp);
+        flb_plg_info(ins, "SASL mechanism configured: %s", ctx->sasl_mechanism);
+        
 #ifdef FLB_HAVE_AWS_MSK_IAM
-    /*
-     * When MSK IAM auth is enabled, default the required
-     * security settings so users don't need to specify them.
-     */
-    if (ctx->aws_msk_iam && ctx->aws_msk_iam_cluster_arn) {
-        tmp = flb_output_get_property("rdkafka.security.protocol", ins);
-        if (!tmp) {
-            flb_output_set_property(ins, "rdkafka.security.protocol", "SASL_SSL");
-        }
-
-        tmp = flb_output_get_property("rdkafka.sasl.mechanism", ins);
-        if (!tmp) {
+        /* Check if using aws_msk_iam as SASL mechanism */
+        if (strcasecmp(tmp, "aws_msk_iam") == 0) {
+            /* Mark that user explicitly requested AWS MSK IAM */
+            ctx->aws_msk_iam = FLB_TRUE;
+            
+            /* Set SASL mechanism to OAUTHBEARER for librdkafka */
             flb_output_set_property(ins, "rdkafka.sasl.mechanism", "OAUTHBEARER");
+            flb_sds_destroy(ctx->sasl_mechanism);
             ctx->sasl_mechanism = flb_sds_create("OAUTHBEARER");
+            
+            /* Ensure security protocol is set */
+            tmp = flb_output_get_property("rdkafka.security.protocol", ins);
+            if (!tmp) {
+                flb_output_set_property(ins, "rdkafka.security.protocol", "SASL_SSL");
+            }
+            
+            flb_plg_info(ins, "AWS MSK IAM authentication enabled via rdkafka.sasl.mechanism");
         }
-        else {
-            ctx->sasl_mechanism = flb_sds_create(tmp);
-        }
-    }
-    else {
 #endif
-        /* Retrieve SASL mechanism if configured */
-        tmp = flb_output_get_property("rdkafka.sasl.mechanism", ins);
-        if (tmp) {
-            ctx->sasl_mechanism = flb_sds_create(tmp);
-        }
-
-#ifdef FLB_HAVE_AWS_MSK_IAM
     }
-#endif
 
     /* rdkafka config context */
     ctx->conf = flb_kafka_conf_create(&ctx->kafka, &ins->properties, 0);
@@ -210,18 +206,38 @@ struct flb_out_kafka *flb_out_kafka_create(struct flb_output_instance *ins,
     flb_kafka_opaque_set(ctx->opaque, ctx, NULL);
     rd_kafka_conf_set_opaque(ctx->conf, ctx->opaque);
 
+    /*
+     * Enable SASL queue for all OAUTHBEARER configurations.
+     * This allows librdkafka to handle OAuth token refresh in a background thread,
+     * which is essential for idle connections where rd_kafka_poll() is not called.
+     * This benefits all OAUTHBEARER methods: AWS IAM, OIDC, custom OAuth, etc.
+     */
+    if (ctx->sasl_mechanism && strcasecmp(ctx->sasl_mechanism, "OAUTHBEARER") == 0) {
+        rd_kafka_conf_enable_sasl_queue(ctx->conf, 1);
+        flb_plg_debug(ins, "SASL queue enabled for OAUTHBEARER mechanism");
+    }
+
 #ifdef FLB_HAVE_AWS_MSK_IAM
-    if (ctx->aws_msk_iam && ctx->aws_msk_iam_cluster_arn && ctx->sasl_mechanism &&
+    /* Only register MSK IAM if user explicitly requested it via rdkafka.sasl.mechanism=aws_msk_iam */
+    if (ctx->aws_msk_iam && ctx->sasl_mechanism && 
         strcasecmp(ctx->sasl_mechanism, "OAUTHBEARER") == 0) {
-
-        ctx->msk_iam = flb_aws_msk_iam_register_oauth_cb(config,
-                                                         ctx->conf,
-                                                         ctx->aws_msk_iam_cluster_arn,
-                                                         ctx->opaque);
-        if (!ctx->msk_iam) {
-            flb_plg_error(ctx->ins, "failed to setup MSK IAM authentication");
-        }
-        else {
+        /* Check if brokers are configured for MSK IAM */
+        if (ctx->kafka.brokers && 
+            (strstr(ctx->kafka.brokers, ".kafka.") || strstr(ctx->kafka.brokers, ".kafka-serverless.")) && 
+            strstr(ctx->kafka.brokers, ".amazonaws.com")) {
+
+            /* Register MSK IAM OAuth callback - pass brokers string directly */
+            flb_plg_info(ins, "registering AWS MSK IAM authentication OAuth callback");
+            ctx->msk_iam = flb_aws_msk_iam_register_oauth_cb(config,
+                                                             ctx->conf,
+                                                             ctx->opaque,
+                                                             ctx->kafka.brokers);
+            if (!ctx->msk_iam) {
+                flb_plg_error(ctx->ins, "failed to setup MSK IAM authentication OAuth callback");
+                flb_out_kafka_destroy(ctx);
+                return NULL;
+            }
+            
             res = rd_kafka_conf_set(ctx->conf, "sasl.oauthbearer.config",
                                     "principal=admin", errstr, sizeof(errstr));
             if (res != RD_KAFKA_CONF_OK) {
@@ -236,13 +252,38 @@ struct flb_out_kafka *flb_out_kafka_create(struct flb_output_instance *ins,
     /* Kafka Producer */
     ctx->kafka.rk = rd_kafka_new(RD_KAFKA_PRODUCER, ctx->conf,
                                  errstr, sizeof(errstr));
+    
     if (!ctx->kafka.rk) {
         flb_plg_error(ctx->ins, "failed to create producer: %s",
                       errstr);
+        /* rd_kafka_new() did NOT take ownership on failure; ctx->conf is
+         * still valid and will be destroyed by flb_out_kafka_destroy(). */
         flb_out_kafka_destroy(ctx);
         return NULL;
     }
 
+    /* rd_kafka_new() takes ownership of ctx->conf on success */
+    ctx->conf = NULL;
+
+    /*
+     * Enable SASL background callbacks for all OAUTHBEARER configurations.
+     * This ensures OAuth tokens are refreshed automatically even on idle connections.
+     * This benefits all OAUTHBEARER methods: AWS IAM, OIDC, custom OAuth, etc.
+     */
+    if (ctx->sasl_mechanism && strcasecmp(ctx->sasl_mechanism, "OAUTHBEARER") == 0) {
+        rd_kafka_error_t *error;
+        error = rd_kafka_sasl_background_callbacks_enable(ctx->kafka.rk);
+        if (error) {
+            flb_plg_warn(ctx->ins, "failed to enable SASL background callbacks: %s. "
+                         "OAuth tokens may not refresh on idle connections.",
+                         rd_kafka_error_string(error));
+            rd_kafka_error_destroy(error);
+        }
+        else {
+            flb_plg_info(ctx->ins, "OAUTHBEARER: SASL background callbacks enabled");
+        }
+    }
+
 #ifdef FLB_HAVE_AVRO_ENCODER
     /* Config AVRO */
     tmp = flb_output_get_property("schema_str", ins);
@@ -301,8 +342,13 @@ int flb_out_kafka_destroy(struct flb_out_kafka *ctx)
     flb_kafka_topic_destroy_all(ctx);
 
     if (ctx->kafka.rk) {
+        /* rd_kafka_destroy also destroys the conf that was passed to rd_kafka_new */
         rd_kafka_destroy(ctx->kafka.rk);
     }
+    else if (ctx->conf) {
+        /* If rd_kafka was never created, we need to destroy conf manually */
+        rd_kafka_conf_destroy(ctx->conf);
+    }
 
     if (ctx->opaque) {
         flb_kafka_opaque_destroy(ctx->opaque);

plugins/out_kafka/kafka_config.h

@@ -126,12 +126,10 @@ struct flb_out_kafka {
 #endif
 
 #ifdef FLB_HAVE_AWS_MSK_IAM
-    flb_sds_t aws_msk_iam_cluster_arn;
     struct flb_aws_msk_iam *msk_iam;
+    int aws_msk_iam;  /* Flag to indicate user explicitly requested AWS MSK IAM */
 #endif
 
-    int aws_msk_iam;
-
     struct flb_kafka_opaque *opaque;
 
     /* SASL mechanism configured in rdkafka.sasl.mechanism */