implement prober selection

Author: matthieugouel

Cargo.lock

@@ -15,6 +15,7 @@ env_logger = "0.11.5"
 hyperloglog = "1.0.2"
 ipnet = "2.10.1"
 log = "0.4.22"
+nanoid = "0.4.0"
 pcap = "2.2.0"
 rand = "0.8.5"
 rdkafka = { version = "0.37.0", features = ["sasl", "ssl"] }

Cargo.toml

@@ -1,8 +1,7 @@
 <p align="center">
-  <img src="logo/logo.png" height="256" width="256" alt="saimiris project logo" />
+  <img src="logo/logo.png" height="256" width="256" alt="Project Logo" />
 </p>
 
-
 # Saimiris
 
 > [!WARNING]

README.md

@@ -4,7 +4,7 @@ use crate::auth::{KafkaAuth, SaslAuth};
 use crate::client::producer::produce;
 use crate::config::AppConfig;
 
-pub async fn handle(config: &AppConfig, target: &str) -> Result<()> {
+pub async fn handle(config: &AppConfig, probers: &str, target: &str) -> Result<()> {
     // Configure Kafka authentication
     let auth = match config.kafka.auth_protocol.as_str() {
         "PLAINTEXT" => KafkaAuth::PlainText,
@@ -20,7 +20,9 @@ pub async fn handle(config: &AppConfig, target: &str) -> Result<()> {
         }
     };
 
-    produce(config, auth, target).await;
+    let probers = probers.split(',').collect::<Vec<&str>>();
+
+    produce(config, auth, probers, target).await;
 
     Ok(())
 }

src/client/handler.rs

@@ -11,7 +11,7 @@ use crate::config::AppConfig;
 // - check target format
 // - Put as header the IDs of the prober to use
 
-pub async fn produce(config: &AppConfig, auth: KafkaAuth, target: &str) {
+pub async fn produce(config: &AppConfig, auth: KafkaAuth, probers: Vec<&str>, target: &str) {
     let producer: &FutureProducer = match auth {
         KafkaAuth::PlainText => &ClientConfig::new()
             .set("bootstrap.servers", config.kafka.brokers.clone())
@@ -32,16 +32,21 @@ pub async fn produce(config: &AppConfig, auth: KafkaAuth, target: &str) {
     let topic = config.kafka.in_topics.split(',').collect::<Vec<&str>>()[0];
     info!("Producing to topic: {}", topic);
 
+    // Construct headers
+    let mut headers = OwnedHeaders::new();
+    for prober in probers {
+        headers = headers.insert(Header {
+            key: prober,
+            value: Some(prober),
+        });
+    }
+
     let delivery_status = producer
         .send(
             FutureRecord::to(topic)
                 .payload(target)
-                .key(&format!("Key")) // TODO
-                .headers(OwnedHeaders::new().insert(Header {
-                    // TODO
-                    key: "header_key",
-                    value: Some("header_value"),
-                })),
+                .key(&format!("")) // TODO Client ID
+                .headers(headers),
             Duration::from_secs(0),
         )
         .await;

src/client/producer.rs

@@ -1,10 +1,11 @@
+use caracat::rate_limiter::RateLimitingMethod;
+use config::Config;
 use std::{
     net::{Ipv4Addr, Ipv6Addr},
     time::Duration,
 };
 
-use caracat::rate_limiter::RateLimitingMethod;
-use config::Config;
+use crate::utils::generate_id;
 
 #[derive(Debug, Clone)]
 pub struct CaracatConfig {
@@ -68,7 +69,7 @@ pub struct KafkaConfig {
 #[derive(Debug, Clone)]
 pub struct ProberConfig {
     /// Prober identifier.
-    pub prober_id: u16,
+    pub prober_id: String,
 }
 
 #[derive(Debug, Clone)]
@@ -146,7 +147,9 @@ pub fn prober_config(config: Config) -> AppConfig {
 
         // Prober configuration
         prober: ProberConfig {
-            prober_id: config.get_int("prober.prober_id").unwrap_or(0) as u16,
+            prober_id: config
+                .get_string("prober.id")
+                .unwrap_or(generate_id(None, None)),
         },
     }
 }

src/config.rs

@@ -2,12 +2,14 @@ mod auth;
 mod client;
 mod config;
 mod prober;
+mod utils;
 
 use anyhow::Result;
 use chrono::Local;
 use clap::{Args, Parser, Subcommand};
 use clap_verbosity_flag::{InfoLevel, Verbosity};
 use env_logger::Builder;
+use log::error;
 use std::io::Write;
 
 use crate::config::{load_config, prober_config};
@@ -35,8 +37,12 @@ enum Command {
         #[arg(short, long)]
         config: String,
 
-        /// Target (eg., 2606:4700:4700::1111/128,1,32,1)
+        /// Probers IDs (comma separated)
         #[arg(index = 1)]
+        probers: String,
+
+        /// Target (eg., 2606:4700:4700::1111/128,1,32,1)
+        #[arg(index = 2)]
         target: String,
     },
 }
@@ -75,16 +81,20 @@ async fn main() -> Result<()> {
 
             match prober::handle(&prober_config).await {
                 Ok(_) => (),
-                Err(e) => log::error!("Error: {}", e),
+                Err(e) => error!("Error: {}", e),
             }
         }
-        Command::Client { config, target } => {
+        Command::Client {
+            config,
+            probers,
+            target,
+        } => {
             let app_config = load_config(&config);
             let prober_config = prober_config(app_config);
 
-            match client::handle(&prober_config, &target).await {
+            match client::handle(&prober_config, &probers, &target).await {
                 Ok(_) => (),
-                Err(e) => log::error!("Error: {}", e),
+                Err(e) => error!("Error: {}", e),
             }
         }
     }

src/main.rs

@@ -15,6 +15,7 @@ use crate::config::AppConfig;
 use crate::prober::consumer::init_consumer;
 use crate::prober::prober::probe;
 use crate::prober::producer::produce;
+use crate::utils::test_id;
 
 struct Target {
     prefix: IpNet,
@@ -95,6 +96,13 @@ fn generate_probes(target: &Target) -> Result<Vec<Probe>> {
 }
 
 pub async fn handle(config: &AppConfig) -> Result<()> {
+    // Test input ID
+    if !test_id(Some(config.prober.prober_id.clone()), None, None) {
+        return Err(anyhow::anyhow!("Invalid prober ID"));
+    }
+
+    info!("Prober ID: {}", config.prober.prober_id);
+
     // Configure Kafka authentication
     let out_auth = match config.kafka.auth_protocol.as_str() {
         "PLAINTEXT" => KafkaAuth::PlainText,
@@ -144,6 +152,22 @@ pub async fn handle(config: &AppConfig) -> Result<()> {
                     }
                 }
 
+                // Filter out the messages that are not intended for the prober
+                // by looking at the prober ID in the headers
+                let mut is_intended = false;
+                if let Some(headers) = m.headers() {
+                    for header in headers.iter() {
+                        if header.key == &config.prober.prober_id {
+                            is_intended = true;
+                            break;
+                        }
+                    }
+                }
+                if !is_intended {
+                    info!("Target not intended for this prober");
+                    continue;
+                }
+
                 // Probe Generation
                 let target = decode_payload(target)?;
                 let probes_to_send = generate_probes(&target)?;

src/prober/handler.rs

@@ -24,7 +24,7 @@ fn format_mpls_labels(mpls_labels: &Vec<MPLSLabel>) -> String {
         + "]"
 }
 
-fn format_reply(prober_id: u16, reply: &Reply) -> String {
+fn format_reply(prober_id: String, reply: &Reply) -> String {
     let mut output = vec![];
     output.push(format!("{}", reply.capture_timestamp.as_millis()));
     output.push(format!("{}", prober_id));
@@ -75,7 +75,7 @@ pub async fn produce(config: &AppConfig, auth: KafkaAuth, results: Arc<Mutex<Vec
                 FutureRecord::to(config.kafka.out_topic.as_str())
                     .payload(&format!(
                         "{}",
-                        format_reply(config.prober.prober_id, result)
+                        format_reply(config.prober.prober_id.clone(), result)
                     ))
                     .key(&format!("Key")) // TODO
                     .headers(OwnedHeaders::new().insert(Header {

src/prober/producer.rs

@@ -0,0 +1,21 @@
+use nanoid::nanoid;
+
+const DEFAULT_ID_ALPHABET: &str = "abcdefghijklmnopqrstuvwxyz0123456789";
+const DEFAULT_ID_SIZE: usize = 10;
+
+pub fn generate_id(alphabet: Option<&str>, size: Option<usize>) -> String {
+    let alphabet = alphabet.unwrap_or(DEFAULT_ID_ALPHABET);
+    let size = size.unwrap_or(DEFAULT_ID_SIZE);
+
+    nanoid!(size, &alphabet.chars().collect::<Vec<char>>())
+}
+
+pub fn test_id(id: Option<String>, alphabet: Option<&str>, size: Option<usize>) -> bool {
+    let alphabet = alphabet.unwrap_or(DEFAULT_ID_ALPHABET);
+    let size = size.unwrap_or(DEFAULT_ID_SIZE);
+
+    match id {
+        Some(id) => id.len() == size && id.chars().all(|c| alphabet.contains(c)),
+        None => false,
+    }
+}