Don't accidentally DoS the netflow recipient, that isn't nice.

LibreQoE · Dec 4, 2024 · 4f13bc1 · 4f13bc1
1 parent 7cbd071
commit 4f13bc1
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 52 deletions.
diff --git a/src/rust/lqosd/src/throughput_tracker/flow_data/netflow5/mod.rs b/src/rust/lqosd/src/throughput_tracker/flow_data/netflow5/mod.rs
@@ -54,43 +54,41 @@ impl Netflow5 {
         target: &str,
         sequence: &AtomicU32,
     ) {
-        loop {
-            let num_records = (accumulator.len() * 2) as u16;
-            let sequence_number = sequence.load(std::sync::atomic::Ordering::Relaxed);
-            let header = Netflow5Header::new(sequence_number, num_records);
-            let header_bytes = unsafe {
-                std::slice::from_raw_parts(
-                    &header as *const _ as *const u8,
-                    std::mem::size_of::<Netflow5Header>(),
-                )
-            };
+        let num_records = (accumulator.len() * 2) as u16;
+        let sequence_number = sequence.load(std::sync::atomic::Ordering::Relaxed);
+        let header = Netflow5Header::new(sequence_number, num_records);
+        let header_bytes = unsafe {
+            std::slice::from_raw_parts(
+                &header as *const _ as *const u8,
+                std::mem::size_of::<Netflow5Header>(),
+            )
+        };
 
-            let mut buffer = Vec::with_capacity(
-                header_bytes.len() + (accumulator.len() * 2 * std::mem::size_of::<Netflow5Record>()),
-            );
+        let mut buffer = Vec::with_capacity(
+            header_bytes.len() + (accumulator.len() * 2 * std::mem::size_of::<Netflow5Record>()),
+        );
 
-            buffer.extend_from_slice(header_bytes);
-            for (key, (data, _)) in accumulator {
-                if let Ok((packet1, packet2)) = to_netflow_5(key, data) {
-                    let packet1_bytes = unsafe {
-                        std::slice::from_raw_parts(
-                            &packet1 as *const _ as *const u8,
-                            std::mem::size_of::<Netflow5Record>(),
-                        )
-                    };
-                    let packet2_bytes = unsafe {
-                        std::slice::from_raw_parts(
-                            &packet2 as *const _ as *const u8,
-                            std::mem::size_of::<Netflow5Record>(),
-                        )
-                    };
-                    buffer.extend_from_slice(packet1_bytes);
-                    buffer.extend_from_slice(packet2_bytes);
-                }
+        buffer.extend_from_slice(header_bytes);
+        for (key, (data, _)) in accumulator {
+            if let Ok((packet1, packet2)) = to_netflow_5(key, data) {
+                let packet1_bytes = unsafe {
+                    std::slice::from_raw_parts(
+                        &packet1 as *const _ as *const u8,
+                        std::mem::size_of::<Netflow5Record>(),
+                    )
+                };
+                let packet2_bytes = unsafe {
+                    std::slice::from_raw_parts(
+                        &packet2 as *const _ as *const u8,
+                        std::mem::size_of::<Netflow5Record>(),
+                    )
+                };
+                buffer.extend_from_slice(packet1_bytes);
+                buffer.extend_from_slice(packet2_bytes);
             }
-
-            socket.send_to(&buffer, target).unwrap();
-            sequence.fetch_add(num_records as u32, std::sync::atomic::Ordering::Relaxed);
         }
+
+        socket.send_to(&buffer, target).unwrap();
+        sequence.fetch_add(num_records as u32, std::sync::atomic::Ordering::Relaxed);
     }
 }
diff --git a/src/rust/lqosd/src/throughput_tracker/flow_data/netflow9/mod.rs b/src/rust/lqosd/src/throughput_tracker/flow_data/netflow9/mod.rs
@@ -49,26 +49,24 @@ impl Netflow9 {
         target: &str,
         sequence: &AtomicU32,
     ) {
-        loop {
-            let num_records = (accumulator.len() * 2) as u16 + 2; // +2 to include templates
-            let sequence_num = sequence.load(std::sync::atomic::Ordering::Relaxed);
-            let header = Netflow9Header::new(sequence_num, num_records);
-            let header_bytes = unsafe { std::slice::from_raw_parts(&header as *const _ as *const u8, std::mem::size_of::<Netflow9Header>()) };
-            let template1 = template_data_ipv4();
-            let template2 = template_data_ipv6();
-            let mut buffer = Vec::with_capacity(header_bytes.len() + template1.len() + template2.len() + (num_records as usize * 140));
-            buffer.extend_from_slice(header_bytes);
-            buffer.extend_from_slice(&template1);
-            buffer.extend_from_slice(&template2);
+        let num_records = (accumulator.len() * 2) as u16 + 2; // +2 to include templates
+        let sequence_num = sequence.load(std::sync::atomic::Ordering::Relaxed);
+        let header = Netflow9Header::new(sequence_num, num_records);
+        let header_bytes = unsafe { std::slice::from_raw_parts(&header as *const _ as *const u8, std::mem::size_of::<Netflow9Header>()) };
+        let template1 = template_data_ipv4();
+        let template2 = template_data_ipv6();
+        let mut buffer = Vec::with_capacity(header_bytes.len() + template1.len() + template2.len() + (num_records as usize * 140));
+        buffer.extend_from_slice(header_bytes);
+        buffer.extend_from_slice(&template1);
+        buffer.extend_from_slice(&template2);
 
-            for (key, (data, _)) in accumulator {
-                if let Ok((packet1, packet2)) = to_netflow_9(key, data) {
-                    buffer.extend_from_slice(&packet1);
-                    buffer.extend_from_slice(&packet2);
-                }
+        for (key, (data, _)) in accumulator {
+            if let Ok((packet1, packet2)) = to_netflow_9(key, data) {
+                buffer.extend_from_slice(&packet1);
+                buffer.extend_from_slice(&packet2);
             }
-            socket.send_to(&buffer, target).unwrap();
-            sequence.fetch_add(num_records as u32, std::sync::atomic::Ordering::Relaxed);
         }
+        socket.send_to(&buffer, target).unwrap();
+        sequence.fetch_add(num_records as u32, std::sync::atomic::Ordering::Relaxed);
     }
 }