net_agent/core/
poller.rs

1use std::{error::Error, sync::{atomic::{AtomicBool, Ordering}, Arc}};
2
3use net_file::translator::packet::Packet;
4use pcap::{Active, Capture};
5
6pub struct Poller<H> {
7    capture: Capture<Active>,
8    packet_cnt: Option<u64>,
9    handler: H,
10    running: Arc<AtomicBool>,
11}
12
13pub trait Handler {
14    fn decode(&self, packet: Packet) -> Result<(), Box<dyn Error + Send + Sync>>;
15}
16
17impl<H: Handler> Poller<H> {
18    fn new(capture: Capture<Active>, packet_cnt: Option<u64>, handler: H, running: Arc<AtomicBool>) -> Self {
19        Poller {
20            capture,
21            packet_cnt,
22            handler,
23            running,
24        }
25    }
26
27    pub fn builder() -> PollerBuilder<H> {
28        PollerBuilder::default()
29    }
30
31    pub fn poll(&mut self) -> u64 {
32        let mut cnt = 0_u64;
33
34        while self.running.load(Ordering::SeqCst) && cnt < u64::MAX && ((self.packet_cnt.is_some() && cnt != *self.packet_cnt.as_ref().unwrap()) || self.packet_cnt.is_none()) {
35            let packet = self.capture.next_packet();
36            let packet = match packet {
37                Ok(packet) => packet,
38                Err(err) => {
39                    log::error!("Something went wrong during capturing packets: {}", err);
40                    return cnt;
41                }
42            };
43            match self.handler.decode(Packet::from(packet)) {
44                Ok(_) => (),
45                Err(err) => {
46                    log::error!("{err}");
47                    return cnt;
48                }
49            };
50            cnt += 1;
51        }
52        cnt
53    }
54}
55
56
57pub struct PollerBuilder<H: Handler> {
58    capture: Option<Capture<Active>>,
59    packet_cnt: Option<u64>,
60    handler: Option<H>,
61    running: Option<Arc<AtomicBool>>,
62}
63
64impl<H: Handler> Default for PollerBuilder<H> {
65    fn default() -> Self {
66        Self {
67            capture: None,
68            packet_cnt: None,
69            handler: None,
70            running: None,
71        }
72    }
73}
74
75
76impl<H: Handler> PollerBuilder<H> {
77    pub fn with_packet_cnt(mut self, packet_cnt: u64) -> Self {
78        self.packet_cnt = Some(packet_cnt);
79        self
80    }
81
82    pub fn with_handler(mut self, handler: H) -> Self {
83        self.handler = Some(handler);
84        self
85    }
86
87    pub fn with_running(mut self, running: Arc<AtomicBool>) -> Self {
88        self.running = Some(running);
89        self
90    }
91
92    pub fn with_capture(mut self, capture: Capture<Active>) -> Self {
93        self.capture = Some(capture);
94        self
95    }
96
97    pub fn build(self) -> Poller<H> {
98        Poller::new(
99            self.capture.unwrap(),
100            self.packet_cnt,
101            self.handler.unwrap(),
102            self.running.unwrap(),
103        )
104    }
105}
net_agent/core/poller.rs

net_agent/core/
poller.rs
