move probe generation to client

Author: matthieugouel

src/agent/handler.rs

@@ -1,9 +1,6 @@
 use anyhow::Result;
-use caracat::models::Probe;
-use ipnet::IpNet;
 use log::{info, trace};
-use rand::seq::SliceRandom;
-use rand::thread_rng;
+
 use rdkafka::consumer::{CommitMode, Consumer};
 use rdkafka::message::Headers;
 use rdkafka::Message;
@@ -16,65 +13,9 @@ use crate::agent::receiver::ReceiveLoop;
 use crate::agent::sender::send;
 use crate::auth::{KafkaAuth, SaslAuth};
 use crate::config::AppConfig;
-use crate::target::{decode_target, Target};
+use crate::probe::decode_probes;
 use crate::utils::test_id;
 
-fn generate_probes(target: &Target) -> Result<Vec<Probe>> {
-    // TODO: We should pass an iterator instead of a vector.
-    let mut probes = vec![];
-
-    // First start by dividing the prefix into /24s or /64s, if necessary.
-    let subnets = match target.prefix {
-        IpNet::V4(_) => {
-            let prefix_len = target.prefix.prefix_len();
-            let target_len = if prefix_len > 24 { prefix_len } else { 24 };
-            target.prefix.subnets(target_len)
-        }
-        IpNet::V6(_) => {
-            let prefix_len = target.prefix.prefix_len();
-            let target_len = if prefix_len > 64 { prefix_len } else { 64 };
-            target.prefix.subnets(target_len)
-        }
-    }?;
-
-    // Iterate over the subnets and generate the probes.
-    for subnet in subnets {
-        // Right now the probe generation is simplistic, we just iterate over the hosts.
-        // If we need more flows than hosts, we will we explicitely fail.
-        // TODO: implement mapper-like generator such as the ones in diamond-miner.
-        // https://github.com/dioptra-io/diamond-miner/blob/main/diamond_miner/mappers.py
-        let mut prefix_hosts = subnet.hosts();
-        if target.n_flows > prefix_hosts.count().try_into()? {
-            return Err(anyhow::anyhow!("Not enough hosts in the prefix"));
-        }
-
-        for _ in 0..target.n_flows {
-            let dst_addr = prefix_hosts.next().unwrap();
-
-            // Randomize the probes order within a flow.
-            // In YARRP we randomize the probes over the entire probing space.
-            // This is of course a very big simplication, but it's not that silly.
-            // The rational is to avoid results errors due to path changes.
-            // So, for now, probes belonging to the same traceroute flow will be sent close in time.
-            // TODO: is this shuffle fast?
-            let mut ttls: Vec<u8> = (target.min_ttl..target.max_ttl).collect();
-            ttls.shuffle(&mut thread_rng());
-
-            for i in ttls {
-                probes.push(Probe {
-                    dst_addr,
-                    src_port: 24000,
-                    dst_port: 33434,
-                    ttl: i,
-                    protocol: target.protocol.clone(),
-                });
-            }
-        }
-    }
-
-    Ok(probes)
-}
-
 pub async fn handle(config: &AppConfig) -> Result<()> {
     trace!("Agent handler");
     trace!("{:?}", config);
@@ -123,7 +64,7 @@ pub async fn handle(config: &AppConfig) -> Result<()> {
                 info!("Kafka error: {}", e);
             }
             Ok(m) => {
-                let target = match m.payload_view::<str>() {
+                let probes = match m.payload_view::<str>() {
                     None => "",
                     Some(Ok(s)) => s,
                     Some(Err(e)) => {
@@ -134,15 +75,6 @@ pub async fn handle(config: &AppConfig) -> Result<()> {
                     }
                 };
 
-                info!(
-                    "key: '{:?}', payload: '{}', topic: {}, partition: {}, offset: {}, timestamp: {:?}",
-                    m.key(),
-                    target,
-                    m.topic(),
-                    m.partition(),
-                    m.offset(),
-                    m.timestamp()
-                );
                 if let Some(headers) = m.headers() {
                     for header in headers.iter() {
                         info!("  Header {:#?}: {:?}", header.key, header.value);
@@ -168,9 +100,8 @@ pub async fn handle(config: &AppConfig) -> Result<()> {
                     continue;
                 }
 
-                // Probe Generation
-                let target = decode_target(target)?;
-                let probes_to_send = generate_probes(&target)?;
+                // Decode probes
+                let probes_to_send = decode_probes(probes)?;
 
                 // Probing
                 let config_clone = config.clone();

src/client/handler.rs

@@ -4,6 +4,7 @@ use log::trace;
 use crate::auth::{KafkaAuth, SaslAuth};
 use crate::client::producer::produce;
 use crate::config::AppConfig;
+use crate::probe::generate_probes;
 use crate::target::decode_target;
 
 pub async fn handle(config: &AppConfig, agents: &str, target: &str) -> Result<()> {
@@ -25,10 +26,14 @@ pub async fn handle(config: &AppConfig, agents: &str, target: &str) -> Result<()
         }
     };
 
+    // Split agents
     let agents = agents.split(',').collect::<Vec<&str>>();
-    decode_target(target)?;
 
-    produce(config, auth, agents, target).await;
+    // Probe Generation
+    let target = decode_target(target)?;
+    let probes = generate_probes(&target)?;
+
+    produce(config, auth, agents, probes).await;
 
     Ok(())
 }

src/client/producer.rs

@@ -1,3 +1,4 @@
+use caracat::models::Probe;
 use log::info;
 use rdkafka::config::ClientConfig;
 use rdkafka::message::{Header, OwnedHeaders};
@@ -6,11 +7,12 @@ use std::time::Duration;
 
 use crate::auth::KafkaAuth;
 use crate::config::AppConfig;
+use crate::probe::encode_probe;
 
 // TODO
 // - check target format
 
-pub async fn produce(config: &AppConfig, auth: KafkaAuth, agents: Vec<&str>, target: &str) {
+pub async fn produce(config: &AppConfig, auth: KafkaAuth, agents: Vec<&str>, probes: Vec<Probe>) {
     let producer: &FutureProducer = match auth {
         KafkaAuth::PlainText => &ClientConfig::new()
             .set("bootstrap.servers", config.kafka.brokers.clone())
@@ -29,7 +31,7 @@ pub async fn produce(config: &AppConfig, auth: KafkaAuth, agents: Vec<&str>, tar
     };
 
     let topic = config.kafka.in_topics.split(',').collect::<Vec<&str>>()[0];
-    info!("Producing to topic: {}", topic);
+    info!("Producing {} probes to {}", probes.len(), topic);
 
     // Construct headers
     let mut headers = OwnedHeaders::new();
@@ -40,15 +42,35 @@ pub async fn produce(config: &AppConfig, auth: KafkaAuth, agents: Vec<&str>, tar
         });
     }
 
-    let delivery_status = producer
-        .send(
-            FutureRecord::to(topic)
-                .payload(target)
-                .key(&format!("")) // TODO Client ID
-                .headers(headers),
-            Duration::from_secs(0),
-        )
-        .await;
-
-    info!("{:?}", delivery_status);
+    // Bucket probes into Kafka messages
+    let mut messages = Vec::new();
+    let mut current_message = String::new();
+    for probe in probes {
+        // Format probe
+        let probe_str = encode_probe(&probe);
+        // Max message size is 1048576 bytes
+        if current_message.len() + probe_str.len() > 1048576 {
+            messages.push(current_message);
+            current_message = String::new();
+        }
+        current_message.push_str(&probe_str);
+    }
+    messages.push(current_message);
+
+    info!("Sending {} messages", messages.len());
+
+    // Send to Kafka
+    for message in messages {
+        let delivery_status = producer
+            .send(
+                FutureRecord::to(topic)
+                    .payload(&format!("{}", message))
+                    .key(&format!("")) // TODO Client ID
+                    .headers(headers.clone()),
+                Duration::from_secs(0),
+            )
+            .await;
+
+        info!("{:?}", delivery_status);
+    }
 }

src/main.rs

@@ -2,6 +2,7 @@ mod agent;
 mod auth;
 mod client;
 mod config;
+mod probe;
 mod target;
 mod utils;
 

src/probe.rs

@@ -0,0 +1,112 @@
+use anyhow::Result;
+use caracat::models::Probe;
+use ipnet::IpNet;
+use rand::seq::SliceRandom;
+use rand::thread_rng;
+
+use crate::target::Target;
+
+pub fn generate_probes(target: &Target) -> Result<Vec<Probe>> {
+    // TODO: We should pass an iterator instead of a vector.
+    let mut probes = vec![];
+
+    // First start by dividing the prefix into /24s or /64s, if necessary.
+    let subnets = match target.prefix {
+        IpNet::V4(_) => {
+            let prefix_len = target.prefix.prefix_len();
+            let target_len = if prefix_len > 24 { prefix_len } else { 24 };
+            target.prefix.subnets(target_len)
+        }
+        IpNet::V6(_) => {
+            let prefix_len = target.prefix.prefix_len();
+            let target_len = if prefix_len > 64 { prefix_len } else { 64 };
+            target.prefix.subnets(target_len)
+        }
+    }?;
+
+    // Iterate over the subnets and generate the probes.
+    for subnet in subnets {
+        // Right now the probe generation is simplistic, we just iterate over the hosts.
+        // If we need more flows than hosts, we will we explicitely fail.
+        // TODO: implement mapper-like generator such as the ones in diamond-miner.
+        // https://github.com/dioptra-io/diamond-miner/blob/main/diamond_miner/mappers.py
+        let mut prefix_hosts = subnet.hosts();
+        if target.n_flows > prefix_hosts.count().try_into()? {
+            return Err(anyhow::anyhow!("Not enough hosts in the prefix"));
+        }
+
+        for _ in 0..target.n_flows {
+            let dst_addr = prefix_hosts.next().unwrap();
+
+            // Randomize the probes order within a flow.
+            // In YARRP we randomize the probes over the entire probing space.
+            // This is of course a very big simplication, but it's not that silly.
+            // The rational is to avoid results errors due to path changes.
+            // So, for now, probes belonging to the same traceroute flow will be sent close in time.
+            // TODO: is this shuffle fast?
+            let mut ttls: Vec<u8> = (target.min_ttl..target.max_ttl).collect();
+            ttls.shuffle(&mut thread_rng());
+
+            for i in ttls {
+                probes.push(Probe {
+                    dst_addr,
+                    src_port: 24000,
+                    dst_port: 33434,
+                    ttl: i,
+                    protocol: target.protocol.clone(),
+                });
+            }
+        }
+    }
+
+    Ok(probes)
+}
+
+pub fn encode_protocol(protocol: caracat::models::L4) -> String {
+    match protocol {
+        caracat::models::L4::ICMP => "icmp".to_string(),
+        caracat::models::L4::ICMPv6 => "icmpv6".to_string(),
+        caracat::models::L4::UDP => "udp".to_string(),
+    }
+}
+
+pub fn decode_protocol(protocol: &str) -> Result<caracat::models::L4> {
+    match protocol {
+        "icmp" => Ok(caracat::models::L4::ICMP),
+        "icmpv6" => Ok(caracat::models::L4::ICMPv6),
+        "udp" => Ok(caracat::models::L4::UDP),
+        _ => Err(anyhow::anyhow!("Invalid protocol: {}", protocol)),
+    }
+}
+
+pub fn encode_probe(probe: &Probe) -> String {
+    format!(
+        "{},{},{},{},{}\n",
+        probe.dst_addr,
+        probe.src_port,
+        probe.dst_port,
+        probe.ttl,
+        encode_protocol(probe.protocol),
+    )
+}
+
+pub fn decode_probes(probes: &str) -> Result<Vec<Probe>> {
+    let mut decoded_probes = vec![];
+
+    for probe in probes.lines() {
+        let fields: Vec<&str> = probe.split(',').collect();
+        if fields.len() != 5 {
+            return Err(anyhow::anyhow!("Invalid probe format: {}", probe));
+        }
+
+        decoded_probes.push(Probe {
+            dst_addr: fields[0].parse()?,
+            src_port: fields[1].parse()?,
+            dst_port: fields[2].parse()?,
+            ttl: fields[3].parse()?,
+            protocol: decode_protocol(fields[4])?,
+        });
+    }
+
+    Ok(decoded_probes)
+}