Add MSK IAM auth support

Author: DJAndries

Cargo.lock

@@ -32,6 +32,9 @@ tokio-util = "0.7"
 reqwest = { version = "0.11", features = ["json"] }
 aws-config = { version = "1.6", features = ["behavior-version-latest"] }
 aws-sdk-s3 = { version = "1.85", features = ["behavior-version-latest"] }
+aws-sigv4 = "1.3"
+aws-credential-types = "1.2"
+url = "2.5"
 prometheus-client = "0.22"
 sentry = "0.36"
 jemallocator = "0.5"

Cargo.toml

@@ -25,6 +25,8 @@ DATABASE_PASSWORD="$(aws rds generate-db-auth-token \
 
 export DATABASE_URL="postgres://$DATABASE_USER:$DATABASE_PASSWORD@$DATABASE_HOST"
 
+unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN
+
 constellation-processors $@
 
 get_auth_token

misc/rds_iam_bootstrap

@@ -10,8 +10,7 @@ use crate::epoch::EpochConfig;
 use crate::models::{DBConnectionType, DBPool, DBStorageConnections, PgStoreError};
 use crate::profiler::{Profiler, ProfilerStat};
 use crate::record_stream::{
-  get_data_channel_topic_from_env, KafkaRecordStream, KafkaRecordStreamConfig, RecordStream,
-  RecordStreamArc, RecordStreamError,
+  get_data_channel_topic_from_env, KafkaRecordStreamConfig, KafkaRecordStreamFactory, RecordStream, RecordStreamArc, RecordStreamError
 };
 use crate::star::AppSTARError;
 use crate::util::parse_env_var;
@@ -48,14 +47,15 @@ pub enum AggregatorError {
 }
 
 fn create_output_stream(
+  rec_stream_factory: &KafkaRecordStreamFactory,
   output_measurements_to_stdout: bool,
   channel_name: &str,
 ) -> Result<Option<RecordStreamArc>, AggregatorError> {
   let topic = get_data_channel_topic_from_env(true, channel_name);
   Ok(if output_measurements_to_stdout {
     None
   } else {
-    let out_stream = Arc::new(KafkaRecordStream::new(KafkaRecordStreamConfig {
+    let out_stream = Arc::new(rec_stream_factory.create_record_stream(KafkaRecordStreamConfig {
       enable_producer: true,
       enable_consumer: false,
       topic,
@@ -96,12 +96,13 @@ pub async fn start_aggregation(
 
   info!("Starting aggregation...");
 
-  let out_stream = create_output_stream(output_measurements_to_stdout, channel_name)?;
+  let rec_stream_factory = KafkaRecordStreamFactory::new();
+  let out_stream = create_output_stream(&rec_stream_factory, output_measurements_to_stdout, channel_name)?;
 
   let mut in_streams: Vec<RecordStreamArc> = Vec::new();
   let in_stream_topic = get_data_channel_topic_from_env(false, channel_name);
   for _ in 0..CONSUMER_COUNT {
-    in_streams.push(Arc::new(KafkaRecordStream::new(KafkaRecordStreamConfig {
+    in_streams.push(Arc::new(rec_stream_factory.create_record_stream(KafkaRecordStreamConfig {
       enable_producer: false,
       enable_consumer: true,
       topic: in_stream_topic.clone(),
@@ -208,7 +209,7 @@ pub async fn start_aggregation(
   // Delete pending/recovered messages from DB.
   info!("Checking/processing expired epochs");
   let profiler = Arc::new(Profiler::default());
-  let out_stream = create_output_stream(output_measurements_to_stdout, channel_name)?;
+  let out_stream = create_output_stream(&rec_stream_factory, output_measurements_to_stdout, channel_name)?;
   let db_conn = Arc::new(Mutex::new(db_pool.get().await?));
   process_expired_epochs(db_conn.clone(), &epoch_config, out_stream, profiler.clone()).await?;
   info!("Profiler summary:\n{}", profiler.summary().await);

src/aggregator/mod.rs

@@ -1,7 +1,7 @@
 use crate::lake::{DataLake, DataLakeError};
 use crate::prometheus::DataLakeMetrics;
 use crate::record_stream::{
-  DynRecordStream, KafkaRecordStream, KafkaRecordStreamConfig, RecordStream, RecordStreamError,
+  DynRecordStream, KafkaRecordStreamConfig, KafkaRecordStreamFactory, RecordStream, RecordStreamError
 };
 use crate::util::parse_env_var;
 use derive_more::{Display, Error, From};
@@ -53,7 +53,8 @@ pub async fn start_lakesink(
 ) -> Result<(), LakeSinkError> {
   let batch_size = parse_env_var::<usize>(BATCH_SIZE_ENV_KEY, BATCH_SIZE_DEFAULT);
 
-  let rec_stream = KafkaRecordStream::new(KafkaRecordStreamConfig {
+  let rec_stream_factory = KafkaRecordStreamFactory::new();
+  let rec_stream = rec_stream_factory.create_record_stream(KafkaRecordStreamConfig {
     enable_producer: false,
     enable_consumer: true,
     topic: stream_topic,

src/lakesink.rs

@@ -4,6 +4,7 @@ mod epoch;
 mod lake;
 mod lakesink;
 mod models;
+mod msk_iam;
 mod profiler;
 mod prometheus;
 mod record_stream;

src/main.rs

@@ -0,0 +1,138 @@
+use aws_config::SdkConfig;
+use aws_credential_types::provider::ProvideCredentials;
+use aws_sigv4::{http_request::{self, SignableBody, SignableRequest, SigningSettings}, sign::v4};
+use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
+use futures::executor;
+use time::{macros::format_description, Duration, OffsetDateTime};
+use std::{error::Error, time::SystemTime};
+use url::Url;
+
+const ACTION_TYPE: &str = "Action";
+const ACTION_NAME: &str = "kafka-cluster:Connect";
+const SIGNING_NAME: &str = "kafka-cluster";
+const USER_AGENT_KEY: &str = "User-Agent";
+const DATE_QUERY_KEY: &str = "X-Amz-Date";
+const EXPIRES_QUERY_KEY: &str = "X-Amz-Expires";
+const DEFAULT_EXPIRY_SECONDS: u64 = 900;
+
+const APP_NAME: &str = "constellation-processors";
+const USER_AGENT_VERSION: &str = "1.0";
+
+#[derive(Clone)]
+pub struct TokenInfo {
+    pub token: String,
+    pub expiration_time: OffsetDateTime,
+}
+pub struct MSKIAMAuthManager {
+    token_info: Option<TokenInfo>,
+}
+
+impl MSKIAMAuthManager {
+    pub fn new() -> Self {
+        Self {
+            token_info: None,
+        }
+    }
+    
+    pub fn get_auth_token(&mut self) -> Result<TokenInfo, Box<dyn Error>> {
+        if let Some(token_info) = &self.token_info {
+            if token_info.expiration_time > OffsetDateTime::now_utc() {
+                return Ok(token_info.clone());
+            }
+        }
+
+        let token_info = executor::block_on(generate_auth_token_async())?;
+        self.token_info = Some(token_info.clone());
+        Ok(token_info)
+    }
+}
+
+async fn generate_auth_token_async() -> Result<TokenInfo, Box<dyn Error>> {
+    let config = aws_config::from_env()
+        .load()
+        .await;
+    
+    let mut url = build_request_url(&config)?;
+    
+    sign_request_url(&mut url, &config).await?;
+    
+    let expiration_time = get_expiration_time(&url)?;
+    
+    add_user_agent(&mut url);
+    
+    let encoded = URL_SAFE_NO_PAD.encode(url.as_str().as_bytes());
+    
+    Ok(TokenInfo {
+        token: encoded,
+        expiration_time,
+    })
+}
+
+fn build_request_url(config: &SdkConfig) -> Result<Url, Box<dyn Error>> {
+    let endpoint_url = format!("https://kafka.{}.amazonaws.com/", config.region().ok_or_else(|| "AWS region is not set")?.to_string());
+    let mut url = Url::parse(&endpoint_url)?;
+    
+    {
+        let mut query_pairs = url.query_pairs_mut();
+        query_pairs.append_pair(ACTION_TYPE, ACTION_NAME);
+        query_pairs.append_pair(EXPIRES_QUERY_KEY, &DEFAULT_EXPIRY_SECONDS.to_string());
+    }
+
+    Ok(url)
+}
+
+async fn sign_request_url(
+    url: &mut Url,
+    config: &SdkConfig,
+) -> Result<(), Box<dyn Error>> {
+    let credentials_provider = config.credentials_provider().ok_or_else(|| "AWS credentials provider is not set")?;
+    let credentials = credentials_provider.provide_credentials().await?;
+
+    let signable_request = SignableRequest::new("GET", url.as_str(), std::iter::empty(), SignableBody::Bytes(&[]))?;
+    
+    let identity = credentials.into();
+    let region = config.region().ok_or_else(|| "AWS region is not set")?.to_string();
+    let signing_params = v4::SigningParams::builder()
+        .identity(&identity)
+        .region(&region)
+        .name(SIGNING_NAME)
+        .time(SystemTime::now())
+        .settings(SigningSettings::default())
+        .build()?
+        .into();
+    
+    let signing_output = http_request::sign(
+        signable_request,
+        &signing_params,
+    )?;
+
+    for (key, value) in signing_output.output().params() {
+        url.query_pairs_mut().append_pair(key, value);
+    }
+    
+    Ok(())
+}
+
+fn get_expiration_time(url: &Url) -> Result<OffsetDateTime, Box<dyn Error>> {
+    let date_str = url.query_pairs()
+        .find_map(|(k, v)| if k == DATE_QUERY_KEY { Some(v.to_string()) } else { None })
+        .ok_or_else(|| "failed to find AWS signed date parameter")?;
+    
+    let date_format_description = format_description!("[year][month][day]T[hour][minute][second]Z");
+    let date = OffsetDateTime::parse(&date_str, date_format_description)?;
+    
+    let expiry_duration_seconds = url.query_pairs()
+        .find_map(|(k, v)| if k == EXPIRES_QUERY_KEY { Some(v.to_string()) } else { None })
+        .ok_or_else(|| "failed to find X-Amz-Expires parameter")?
+        .parse::<i64>()?;
+    
+    let expiry_duration = Duration::seconds(expiry_duration_seconds);
+    let expiry_time = date + expiry_duration;
+    
+    Ok(expiry_time)
+}
+
+fn add_user_agent(url: &mut Url) {
+    let user_agent = format!("{}/{}/{}", APP_NAME, USER_AGENT_VERSION, USER_AGENT_VERSION);
+    url.query_pairs_mut().append_pair(USER_AGENT_KEY, &user_agent);
+}

src/msk_iam.rs

@@ -2,7 +2,7 @@ use async_trait::async_trait;
 use derive_more::{Display, Error, From};
 use futures::future::try_join_all;
 use rand::{seq::SliceRandom, thread_rng};
-use rdkafka::client::ClientContext;
+use rdkafka::client::{ClientContext, OAuthToken};
 use rdkafka::config::ClientConfig;
 use rdkafka::consumer::{
   stream_consumer::StreamConsumer, CommitMode, Consumer, ConsumerContext, Rebalance,
@@ -14,20 +14,23 @@ use rdkafka::types::RDKafkaErrorCode;
 use rdkafka::TopicPartitionList;
 use std::collections::HashMap;
 use std::env;
-use std::sync::Arc;
+use std::error::Error as StdError;
+use std::sync::{Arc, Mutex as StdMutex};
 use std::time::Duration;
 use tokio::sync::mpsc::{error::SendError, unbounded_channel, UnboundedSender};
 use tokio::sync::{Mutex, RwLock};
 use tokio::task::{JoinError, JoinHandle};
 use tokio::time::sleep;
 
 use crate::channel::{get_data_channel_map_from_env, get_data_channel_value_from_env};
+use crate::msk_iam::MSKIAMAuthManager;
 use crate::util::parse_env_var;
 
 const KAFKA_ENC_TOPICS_ENV_KEY: &str = "KAFKA_ENCRYPTED_TOPICS";
 const KAFKA_OUT_TOPICS_ENV_KEY: &str = "KAFKA_OUTPUT_TOPICS";
 const DEFAULT_ENC_KAFKA_TOPICS: &str = "typical=p3a-star-enc";
 const DEFAULT_OUT_KAFKA_TOPICS: &str = "typical=p3a-star-out";
+const KAFKA_IAM_BROKERS_ENV_KEY: &str = "KAFKA_IAM_BROKERS";
 const KAFKA_BROKERS_ENV_KEY: &str = "KAFKA_BROKERS";
 const KAFKA_ENABLE_PLAINTEXT_ENV_KEY: &str = "KAFKA_ENABLE_PLAINTEXT";
 const KAFKA_PRODUCER_QUEUE_TASK_COUNT_ENV_KEY: &str = "KAFKA_PRODUCE_QUEUE_TASK_COUNT";
@@ -54,9 +57,24 @@ pub enum RecordStreamError {
   Join(JoinError),
 }
 
-struct KafkaContext;
+#[derive(Clone)]
+struct KafkaContext {
+  msk_iam_auth_manager: Arc<StdMutex<MSKIAMAuthManager>>,
+}
+
 
-impl ClientContext for KafkaContext {}
+impl ClientContext for KafkaContext {
+  const ENABLE_REFRESH_OAUTH_TOKEN: bool = true;
+
+  fn generate_oauth_token(&self, _oauthbearer_config: Option<&str>) -> Result<OAuthToken, Box<dyn StdError>> {
+    let token_info = self.msk_iam_auth_manager.lock().unwrap().get_auth_token()?;
+    Ok(OAuthToken {
+      token: token_info.token,
+      lifetime_ms: (token_info.expiration_time.unix_timestamp_nanos() / 1_000_000) as i64,
+      principal_name: String::new(),
+    })
+  }
+}
 
 impl ConsumerContext for KafkaContext {
   fn pre_rebalance(&self, rebalance: &Rebalance) {
@@ -115,6 +133,25 @@ pub struct KafkaRecordStreamConfig {
   pub use_output_group_id: bool,
 }
 
+pub struct KafkaRecordStreamFactory {
+  msk_iam_auth_manager: Arc<StdMutex<MSKIAMAuthManager>>,
+}
+
+impl KafkaRecordStreamFactory {
+  pub fn new() -> Self {
+    Self {
+      msk_iam_auth_manager: Arc::new(StdMutex::new(MSKIAMAuthManager::new())),
+    }
+  }
+
+  pub fn create_record_stream(&self, stream_config: KafkaRecordStreamConfig) -> KafkaRecordStream {
+    let context = KafkaContext {
+      msk_iam_auth_manager: self.msk_iam_auth_manager.clone(),
+    };
+    KafkaRecordStream::new(stream_config, context)
+  }
+}
+
 pub struct KafkaRecordStream {
   producer: Option<Arc<FutureProducer<KafkaContext>>>,
   consumer: Option<StreamConsumer<KafkaContext>>,
@@ -150,7 +187,7 @@ pub fn get_data_channel_topic_from_env(use_output_topic: bool, channel_name: &st
 }
 
 impl KafkaRecordStream {
-  pub fn new(stream_config: KafkaRecordStreamConfig) -> Self {
+  fn new(stream_config: KafkaRecordStreamConfig, context: KafkaContext) -> Self {
     let group_id = match stream_config.use_output_group_id {
       true => "star-agg-dec",
       false => "star-agg-enc",
@@ -163,7 +200,6 @@ impl KafkaRecordStream {
       producer_queues: RwLock::new(Vec::new()),
     };
     if stream_config.enable_producer {
-      let context = KafkaContext;
       let mut config = Self::new_client_config();
       let mut config_ref = &mut config;
       if stream_config.use_output_group_id {
@@ -175,13 +211,12 @@ impl KafkaRecordStream {
           .set("transaction.timeout.ms", "3600000")
           .set("request.timeout.ms", "900000")
           .set("socket.timeout.ms", "300000")
-          .create_with_context(context)
+          .create_with_context(context.clone())
           .unwrap(),
       ));
       info!("Producing to topic: {}", stream_config.topic);
     }
     if stream_config.enable_consumer {
-      let context = KafkaContext;
       let mut config = Self::new_client_config();
       result.consumer = Some(
         config
@@ -210,6 +245,13 @@ impl KafkaRecordStream {
   }
 
   fn new_client_config() -> ClientConfig {
+    if let Some(brokers) = env::var(KAFKA_IAM_BROKERS_ENV_KEY).ok() {
+      let mut result = ClientConfig::new();
+      result.set("bootstrap.servers", brokers);
+      result.set("security.protocol", "SASL_SSL");
+      result.set("sasl.mechanism", "OAUTHBEARER");
+      return result;
+    }
     let brokers = env::var(KAFKA_BROKERS_ENV_KEY)
       .unwrap_or_else(|_| panic!("{} env var must be defined", KAFKA_BROKERS_ENV_KEY));
     let mut result = ClientConfig::new();