zond_engine/core/models/
target.rs1use crate::core::models::ip::set::IpSet;
14use crate::core::models::port::{PortSet, Protocol};
15use std::{net::IpAddr, sync::Arc};
16use thiserror::Error;
17
18#[derive(Error, Debug, PartialEq, Eq)]
20pub enum TargetError {
21 #[error("Target collection is in a dirty state; call `canonicalize()` before concurrent reads")]
22 UncanonicalizedState,
23
24 #[error("Target calculation resulted in an integer overflow")]
25 CapacityOverflow,
26}
27
28#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
30pub struct Target {
31 pub ip: IpAddr,
32 pub port: u16,
33 pub protocol: Protocol,
34}
35
36#[derive(Debug, Clone, Default)]
41pub struct TargetSet {
42 ips: IpSet,
44 ports: PortSet,
46 is_canonicalized: bool,
48}
49
50impl TargetSet {
51 pub fn new(ips: IpSet, ports: PortSet) -> Self {
53 Self {
54 ips,
55 ports,
56 is_canonicalized: false,
57 }
58 }
59
60 pub fn ips(&self) -> &IpSet {
62 &self.ips
63 }
64
65 pub fn ports(&self) -> &PortSet {
67 &self.ports
68 }
69
70 pub fn ip_count(&mut self) -> u128 {
73 self.ips.len()
74 }
75
76 pub fn port_count(&self) -> usize {
78 self.ports.len()
79 }
80
81 pub fn canonicalize(&mut self) {
83 if !self.is_canonicalized {
84 self.ips.canonicalize();
85 self.is_canonicalized = true;
86 }
87 }
88
89 pub fn total_targets(&mut self) -> Result<u128, TargetError> {
93 self.canonicalize();
94
95 let port_len = self.ports.len() as u128;
96 self.ips
97 .len()
98 .checked_mul(port_len)
99 .ok_or(TargetError::CapacityOverflow)
100 }
101
102 pub fn iter(&mut self) -> impl Iterator<Item = Target> + Send + '_ {
107 self.canonicalize();
108
109 let ports_arc: Arc<[(u16, Protocol)]> = self.ports.to_vec().into();
110
111 self.ips.iter().flat_map(move |ip| {
112 let local_ports = Arc::clone(&ports_arc);
113 (0..local_ports.len()).map(move |i| Target {
114 ip,
115 port: local_ports[i].0,
116 protocol: local_ports[i].1,
117 })
118 })
119 }
120
121 pub fn total_targets_canonical(&self) -> Result<u128, TargetError> {
126 if !self.is_canonicalized {
127 return Err(TargetError::UncanonicalizedState);
128 }
129
130 let port_len = self.ports.len() as u128;
131 self.ips
132 .len_canonical()
133 .checked_mul(port_len)
134 .ok_or(TargetError::CapacityOverflow)
135 }
136
137 pub fn is_empty(&self) -> bool {
139 self.ips.is_empty() || self.ports.is_empty()
140 }
141}
142
143#[derive(Debug, Clone, Default)]
145pub struct TargetMap {
146 pub units: Vec<TargetSet>,
147}
148
149impl TargetMap {
150 pub fn new() -> Self {
152 Self::default()
153 }
154
155 pub fn add_unit(&mut self, unit: TargetSet) {
157 self.units.push(unit);
158 }
159
160 pub fn canonicalize(&mut self) {
162 for unit in &mut self.units {
163 unit.canonicalize();
164 }
165 }
166
167 pub fn gross_targets(&mut self) -> Result<u128, TargetError> {
170 let mut total: u128 = 0;
171 for unit in &mut self.units {
172 let unit_total = unit.total_targets()?;
173 total = total
174 .checked_add(unit_total)
175 .ok_or(TargetError::CapacityOverflow)?;
176 }
177 Ok(total)
178 }
179
180 pub fn gross_ips(&mut self) -> Result<u128, TargetError> {
183 let mut total: u128 = 0;
184 for unit in &mut self.units {
185 total = total
186 .checked_add(unit.ip_count())
187 .ok_or(TargetError::CapacityOverflow)?;
188 }
189 Ok(total)
190 }
191
192 pub fn is_empty(&self) -> bool {
194 self.units.is_empty() || self.units.iter().all(|u| u.is_empty())
195 }
196
197 pub fn iter(&mut self) -> impl Iterator<Item = Target> + Send + '_ {
199 self.units.iter_mut().flat_map(|unit| unit.iter())
200 }
201}
202
203#[cfg(test)]
213mod tests {
214 use super::*;
215
216 fn mock_ip_set(input: &str) -> IpSet {
218 input.parse().expect("Valid IP input")
219 }
220
221 fn mock_port_set(input: &str) -> PortSet {
222 input.parse().expect("Valid Port input")
223 }
224
225 #[test]
226 fn target_set_lazy_math() {
227 let mut ts = TargetSet::new(mock_ip_set("192.168.1.0/24"), mock_port_set("80, 443"));
228 assert_eq!(ts.total_targets().unwrap(), 256 * 2);
229 }
230
231 #[test]
232 fn thread_safe_reads_require_canonicalization() {
233 let ts = TargetSet::new(mock_ip_set("192.168.1.0/24"), mock_port_set("80"));
234
235 let err = ts.total_targets_canonical().unwrap_err();
237 assert_eq!(err, TargetError::UncanonicalizedState);
238 }
239
240 #[test]
241 fn thread_safe_reads_succeed_when_prepared() {
242 let mut ts = TargetSet::new(mock_ip_set("192.168.1.0/24"), mock_port_set("80"));
243 ts.canonicalize();
244
245 assert_eq!(ts.total_targets_canonical().unwrap(), 256);
246 }
247
248 #[test]
249 fn target_map_aggregation() {
250 let mut map = TargetMap::new();
251 map.add_unit(TargetSet::new(
252 mock_ip_set("10.0.0.1-10.0.0.5"),
253 mock_port_set("80,443"),
254 ));
255 assert_eq!(map.gross_targets().unwrap(), 10);
256 }
257}