optionally read probes from a file with the client

Author: matthieugouel

src/agent/producer.rs

@@ -1,5 +1,5 @@
 use caracat::models::{MPLSLabel, Reply};
-use log::{debug, info, warn};
+use log::{debug, error, info, warn};
 use rdkafka::config::ClientConfig;
 use rdkafka::message::OwnedHeaders;
 use rdkafka::producer::{FutureProducer, FutureRecord};
@@ -108,6 +108,8 @@ pub async fn produce(config: &AppConfig, auth: KafkaAuth, rx: Receiver<Reply>) {
 
             let message = message.unwrap();
             let message_str = encode_reply(config.agent.id.clone(), &message);
+
+            // Max message size is 1048576 bytes (including headers)
             if final_message.len() + message_str.len() + 1 > config.kafka.message_max_bytes {
                 additional_message = Some(message);
                 break;
@@ -138,6 +140,11 @@ pub async fn produce(config: &AppConfig, auth: KafkaAuth, rx: Receiver<Reply>) {
             )
             .await;
 
-        info!("{:?}", delivery_status);
+        match delivery_status {
+            Ok(status) => info!("{:?}", status),
+            Err((error, _)) => {
+                error!("{}", error.to_string());
+            }
+        }
     }
 }

src/client/handler.rs

@@ -1,13 +1,24 @@
 use anyhow::Result;
+use caracat::models::Probe;
 use log::trace;
 use std::io::{stdin, BufRead};
+use std::path::PathBuf;
 
 use crate::auth::{KafkaAuth, SaslAuth};
 use crate::client::producer::produce;
 use crate::config::AppConfig;
 use crate::probe::decode_probe;
 
-pub async fn handle(config: &AppConfig, agents: &str) -> Result<()> {
+fn read_probes<R: BufRead>(buf_reader: R) -> Result<Vec<Probe>> {
+    let mut probes = Vec::new();
+    for line in buf_reader.lines() {
+        let probe = line?;
+        probes.push(decode_probe(&probe)?);
+    }
+    Ok(probes)
+}
+
+pub async fn handle(config: &AppConfig, agents: &str, probes_file: Option<PathBuf>) -> Result<()> {
     trace!("Client handler");
     trace!("{:?}", config);
 
@@ -26,16 +37,24 @@ pub async fn handle(config: &AppConfig, agents: &str) -> Result<()> {
         }
     };
 
-    // Get probes from stdin
-    let mut probes = Vec::new();
-    for line in stdin().lock().lines() {
-        let probe = line?;
-        probes.push(decode_probe(&probe)?);
-    }
+    // Read probes from file or stdin
+    let probes = match probes_file {
+        Some(probes_file) => {
+            let file = std::fs::File::open(probes_file)?;
+            let buf_reader = std::io::BufReader::new(file);
+            read_probes(buf_reader)?
+        }
+        None => {
+            let stdin = stdin();
+            let buf_reader = stdin.lock();
+            read_probes(buf_reader)?
+        }
+    };
 
-    // Split agents
+    // Split the agents
     let agents = agents.split(',').collect::<Vec<&str>>();
 
+    // Produce Kafka messages
     produce(config, auth, agents, probes).await;
 
     Ok(())

src/client/producer.rs

@@ -1,5 +1,5 @@
 use caracat::models::Probe;
-use log::{debug, info};
+use log::{debug, error, info};
 use rdkafka::config::ClientConfig;
 use rdkafka::message::{Header, OwnedHeaders};
 use rdkafka::producer::{FutureProducer, FutureRecord};
@@ -15,7 +15,7 @@ fn create_messages(probes: Vec<Probe>, message_max_bytes: usize) -> Vec<String>
     for probe in probes {
         // Format probe
         let probe_str = encode_probe(&probe);
-        // Max message size is 1048576 bytes
+        // Max message size is 1048576 bytes (including headers)
         if current_message.len() + probe_str.len() + 1 > message_max_bytes {
             // Remove the last newline character
             current_message.pop();
@@ -87,6 +87,12 @@ pub async fn produce(config: &AppConfig, auth: KafkaAuth, agents: Vec<&str>, pro
             )
             .await;
 
-        info!("{:?}", delivery_status);
+        match delivery_status {
+            Ok(status) => info!("{:?}", status),
+            Err((error, _)) => {
+                error!("{}", error.to_string());
+                error!("{}", config.kafka.message_max_bytes);
+            }
+        }
     }
 }

src/config.rs

@@ -184,8 +184,7 @@ pub fn app_config(config_path: &str) -> AppConfig {
             auth_sasl_mechanism: config
                 .get_string("kafka.auth_sasl_mechanism")
                 .unwrap_or("SCRAM-SHA-512".to_string()),
-            message_max_bytes: config.get_int("kafka.message_max_bytes").unwrap_or(1048576)
-                as usize,
+            message_max_bytes: config.get_int("kafka.message_max_bytes").unwrap_or(990000) as usize,
             in_topics: config
                 .get_string("kafka.in_topics")
                 .unwrap_or("saimiris-targets".to_string()),

src/main.rs

@@ -12,6 +12,7 @@ use clap_verbosity_flag::{InfoLevel, Verbosity};
 use env_logger::Builder;
 use log::error;
 use std::io::{stdin, IsTerminal, Write};
+use std::path::PathBuf;
 
 use crate::config::app_config;
 
@@ -41,6 +42,10 @@ enum Command {
         /// Agent IDs (comma separated)
         #[arg(index = 1)]
         agents: String,
+
+        /// Probes file (read stdin if not provided)
+        #[arg(short, long)]
+        probes_file: Option<PathBuf>,
     },
 }
 
@@ -79,14 +84,18 @@ async fn main() -> Result<()> {
                 Err(e) => error!("Error: {}", e),
             }
         }
-        Command::Client { config, agents } => {
+        Command::Client {
+            config,
+            agents,
+            probes_file,
+        } => {
             if stdin().is_terminal() {
                 App::command().print_help().unwrap();
                 ::std::process::exit(2);
             }
 
             let app_config = app_config(&config);
-            match client::handle(&app_config, &agents).await {
+            match client::handle(&app_config, &agents, probes_file).await {
                 Ok(_) => (),
                 Err(e) => error!("Error: {}", e),
             }

testbed/config/saimiris/saimiris.yml

@@ -8,5 +8,7 @@ kafka:
   out_enable: true
   out_topic: saimiris-results
 
+  message_max_bytes: 999995
+
 agent:
   id: wbmwwp9vna