manycastr/
custom_module.rs

1use crate::custom_module::manycastr::{Configuration, MeasurementType, Origin, ProtocolType};
2use manycastr::{
3    Ack, Address, Empty, IPv6, address::Value::UnicastV4, address::Value::UnicastV6,
4    address::Value::V4, address::Value::V6,
5};
6use std::fmt;
7use std::fmt::Display;
8use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
9use std::str::FromStr;
10
11pub mod manycastr {
12    tonic::include_proto!("manycastr");
13}
14
15impl Ack {
16    /// A successful acknowledgement
17    pub fn ok() -> Self {
18        Ack {
19            is_success: true,
20            error_message: String::new(),
21        }
22    }
23}
24
25/// Write Address to string (e.g., 1.1.1.1)
26impl Display for Address {
27    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
28        match &self.value {
29            Some(V4(v4)) => {
30                write!(f, "{}", Ipv4Addr::from(*v4))
31            }
32            Some(V6(_)) => {
33                let val: u128 = self.into();
34                write!(f, "{}", Ipv6Addr::from(val))
35            }
36            Some(UnicastV4(_)) => write!(f, "unicastv4"),
37            Some(UnicastV6(_)) => write!(f, "unicastv6"),
38            None => write!(f, "None"),
39        }
40    }
41}
42
43impl Address {
44    /// A unicast origin source address that each worker replaces
45    /// with its own local IPv4 unicast address at measurement start.
46    pub fn unicast_v4() -> Self {
47        Address {
48            value: Some(UnicastV4(Empty {})),
49        }
50    }
51
52    /// A unicast origin source address that each worker replaces
53    /// with its own local IPv6 unicast address at measurement start.
54    pub fn unicast_v6() -> Self {
55        Address {
56            value: Some(UnicastV6(Empty {})),
57        }
58    }
59
60    /// Returns the integer representation of the IP address as u128.
61    pub fn as_numeric(&self) -> u128 {
62        match &self.value {
63            Some(V4(v4)) => *v4 as u128,
64            Some(V6(_)) => self.into(),
65            _ => 0,
66        }
67    }
68
69    /// Whether this address is IPv6.
70    pub fn is_v6(&self) -> bool {
71        matches!(self.value, Some(V6(_)) | Some(UnicastV6(_)))
72    }
73
74    /// Whether this address is unicast.
75    pub fn is_unicast(&self) -> bool {
76        matches!(self.value, Some(UnicastV4(_)) | Some(UnicastV6(_)))
77    }
78
79    /// Get the prefix of the address (/24 for IPv4 and /48 for IPv6)
80    pub fn get_prefix(&self) -> u64 {
81        match &self.value {
82            // /24: Shift right by 8 bits
83            Some(V4(v4)) => (v4 >> 8) as u64,
84            // /48: high is 64 bits, we want top 48. Shift right by (64 - 48) = 16
85            Some(V6(v6)) => v6.high >> 16,
86            _ => 0,
87        }
88    }
89
90    /// The base address of this address' prefix (/24 for IPv4, /48 for IPv6)
91    pub fn prefix_base(&self) -> Address {
92        match &self.value {
93            Some(V4(v4)) => Address::from(v4 & !0xff),
94            Some(V6(_)) => {
95                let val: u128 = self.into();
96                Address::from(val & !((1u128 << 80) - 1))
97            }
98            _ => *self,
99        }
100    }
101
102    /// Convert Address to bytes (big-endian)
103    pub fn to_be_bytes(self) -> Vec<u8> {
104        match &self.value {
105            Some(V4(v4)) => v4.to_be_bytes().to_vec(),
106            Some(V6(_)) => {
107                let val: u128 = self.into();
108                val.to_be_bytes().to_vec()
109            }
110            _ => vec![],
111        }
112    }
113
114    /// Convert to 16-byte IPv4-mapped-IPv6 representation (big-endian).
115    /// IPv4 `x.x.x.x` becomes `::ffff:x.x.x.x` (RFC 4291 §2.5.5.2).
116    /// IPv6 addresses are stored as-is.
117    pub fn to_ipv6_mapped_bytes(self) -> [u8; 16] {
118        match self.value {
119            Some(V4(v4)) => {
120                let octets = v4.to_be_bytes();
121                [
122                    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xFF, 0xFF, octets[0], octets[1], octets[2],
123                    octets[3],
124                ]
125            }
126            Some(V6(v6)) => {
127                let mut bytes = [0u8; 16];
128                bytes[..8].copy_from_slice(&v6.high.to_be_bytes());
129                bytes[8..].copy_from_slice(&v6.low.to_be_bytes());
130                bytes
131            }
132            _ => [0u8; 16],
133        }
134    }
135}
136
137impl Origin {
138    /// Whether this origin probes from workers' local unicast addresses
139    pub fn is_unicast(&self) -> bool {
140        self.src.is_some_and(|s| s.is_unicast())
141    }
142
143    /// Whether this origin's source address is IPv6
144    pub fn is_v6(&self) -> bool {
145        self.src.is_some_and(|s| s.is_v6())
146    }
147}
148
149/// Parse an origin source address token: an anycast IP address, or one of the
150/// unicast keywords `unicastv4`/`unicastv6`.
151///
152/// # Panics
153/// * If the token is not a valid address or unicast keyword.
154pub fn parse_src_address(token: &str) -> Address {
155    if token.eq_ignore_ascii_case("unicastv4") {
156        Address::unicast_v4()
157    } else if token.eq_ignore_ascii_case("unicastv6") {
158        Address::unicast_v6()
159    } else if token.eq_ignore_ascii_case("unicast") {
160        panic!("'unicast' must specify an IP version: use 'unicastv4' or 'unicastv6'");
161    } else {
162        Address::from(token)
163    }
164}
165
166/// Whether any configuration probes from an anycast source address
167pub fn has_anycast_origin(configurations: &[Configuration]) -> bool {
168    configurations
169        .iter()
170        .filter_map(|c| c.origin.as_ref().and_then(|o| o.src.as_ref()))
171        .any(|src| !src.is_unicast())
172}
173
174/// Address -> u32 (panic if not V4)
175impl From<Address> for u32 {
176    fn from(addr: Address) -> Self {
177        match addr.value {
178            Some(V4(v4)) => v4,
179            _ => panic!("Attempted to convert non-IPv4 Address to u32"),
180        }
181    }
182}
183
184/// Address -> u128 (Panic if not V6)
185impl From<Address> for u128 {
186    fn from(addr: Address) -> Self {
187        match addr.value {
188            Some(V6(v6)) => (v6.high as u128) << 64 | v6.low as u128,
189            _ => panic!("Attempted to convert non-IPv6 Address to u128"),
190        }
191    }
192}
193
194/// convert bytes into Address
195impl From<&[u8]> for Address {
196    fn from(bytes: &[u8]) -> Self {
197        match bytes.len() {
198            4 => {
199                let array: [u8; 4] = bytes.try_into().unwrap();
200                Address::from(array)
201            }
202            16 => {
203                let array: [u8; 16] = bytes.try_into().unwrap();
204                Address::from(array)
205            }
206            _ => panic!("Invalid IP address length: {}", bytes.len()),
207        }
208    }
209}
210
211impl From<[u8; 4]> for Address {
212    fn from(bytes: [u8; 4]) -> Self {
213        Address {
214            value: Some(V4(u32::from_be_bytes(bytes))),
215        }
216    }
217}
218
219impl From<[u8; 16]> for Address {
220    fn from(bytes: [u8; 16]) -> Self {
221        Address {
222            value: Some(V6(IPv6 {
223                high: u64::from_be_bytes(bytes[0..8].try_into().unwrap()),
224                low: u64::from_be_bytes(bytes[8..16].try_into().unwrap()),
225            })),
226        }
227    }
228}
229
230impl From<u32> for Address {
231    fn from(bytes: u32) -> Self {
232        Address {
233            value: Some(V4(bytes)),
234        }
235    }
236}
237
238impl From<u128> for Address {
239    fn from(bytes: u128) -> Self {
240        Address {
241            value: Some(V6(IPv6 {
242                high: (bytes >> 64) as u64,
243                low: (bytes & 0xFFFFFFFFFFFFFFFF) as u64,
244            })),
245        }
246    }
247}
248
249impl From<&Address> for u32 {
250    fn from(addr: &Address) -> Self {
251        match &addr.value {
252            Some(V4(v4)) => *v4,
253            _ => panic!("Attempted to convert non-IPv4 &Address to u32"),
254        }
255    }
256}
257
258impl From<&Address> for u128 {
259    fn from(addr: &Address) -> Self {
260        match &addr.value {
261            Some(V6(v6)) => (v6.high as u128) << 64 | v6.low as u128,
262            _ => panic!("Attempted to convert non-IPv6 &Address to u128"),
263        }
264    }
265}
266
267/// Convert String into an Address (can be an IP number or a standard IP string)
268impl FromStr for Address {
269    type Err = String;
270
271    fn from_str(s: &str) -> Result<Self, Self::Err> {
272        // Standard IP string format (e.g., "1.1.1.1" or "2001::1")
273        if let Ok(ip) = s.parse::<IpAddr>() {
274            return Ok(Address::from(ip));
275        }
276
277        // IP number
278        if let Ok(ip_number) = s.parse::<u128>() {
279            return Ok(Address::from(ip_number));
280        }
281
282        Err(format!("Invalid IP address or IP number: {s}"))
283    }
284}
285
286impl From<&str> for Address {
287    fn from(s: &str) -> Self {
288        s.parse().unwrap_or_else(|e| panic!("{}", e))
289    }
290}
291
292impl From<String> for Address {
293    fn from(s: String) -> Self {
294        Address::from(s.as_str())
295    }
296}
297
298impl From<&String> for Address {
299    fn from(s: &String) -> Self {
300        Address::from(s.as_str())
301    }
302}
303
304/// Convert IpAddr to Address (used for converting local unicast addresses)
305impl From<IpAddr> for Address {
306    fn from(ip: IpAddr) -> Self {
307        match ip {
308            IpAddr::V4(v4) => Address::from(u32::from(v4)),
309            IpAddr::V6(v6) => Address::from(v6.octets()),
310        }
311    }
312}
313
314impl From<SocketAddr> for Address {
315    fn from(addr: SocketAddr) -> Self {
316        Address::from(addr.ip())
317    }
318}
319
320impl From<&Address> for IpAddr {
321    fn from(addr: &Address) -> Self {
322        match addr.value {
323            Some(V4(v4_u32)) => IpAddr::V4(Ipv4Addr::from(v4_u32)),
324            Some(V6(v6_msg)) => {
325                let combined = ((v6_msg.high as u128) << 64) | (v6_msg.low as u128);
326                IpAddr::V6(Ipv6Addr::from(combined))
327            }
328            _ => IpAddr::V4(Ipv4Addr::UNSPECIFIED), // None and Unicast
329        }
330    }
331}
332
333pub trait Separated {
334    fn with_separator(&self) -> String;
335}
336
337fn format_number(number: usize) -> String {
338    let number_str = number.to_string();
339    let chunks: Vec<&str> = number_str
340        .as_bytes()
341        .rchunks(3)
342        .rev()
343        .map(std::str::from_utf8)
344        .collect::<Result<Vec<&str>, _>>()
345        .expect("Unable to format number");
346
347    chunks.join(",")
348}
349
350/// Print integer types with a thousand separator (e.g., 1000 -> 1,000)
351macro_rules! impl_separated {
352    ($($t:ty),*) => {
353        $(
354            impl Separated for $t {
355                fn with_separator(&self) -> String {
356                    format_number(*self as usize)
357                }
358            }
359        )*
360    };
361}
362
363impl_separated!(u32, usize, u64, i32);
364
365/// Pretty print measurement types
366impl Display for MeasurementType {
367    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
368        let s = match self {
369            Self::Laces => "LACeS",
370            Self::Catchment => "Catchment Mapping",
371            Self::AnycastLatency => "Latency",
372            Self::AnycastTraceroute => "Anycast Traceroute",
373            Self::Tracemap => "Tracemap",
374            Self::Feed => "Live Feed",
375            Self::FeedTrace => "Live Feed Traceroute",
376        };
377        write!(f, "{}", s)
378    }
379}
380
381impl MeasurementType {
382    pub fn as_str(&self) -> &'static str {
383        match self {
384            Self::Laces => "laces",
385            Self::Catchment => "catchment",
386            Self::AnycastLatency => "latency",
387            Self::AnycastTraceroute => "anycast-traceroute",
388            Self::Tracemap => "tracemap",
389            Self::Feed => "feed",
390            Self::FeedTrace => "feed-trace",
391        }
392    }
393
394    pub fn from_str(s: &str) -> Option<Self> {
395        match s.to_lowercase().as_str() {
396            "laces" => Some(Self::Laces),
397            "catchment" => Some(Self::Catchment),
398            "latency" => Some(Self::AnycastLatency),
399            "anycast-traceroute" => Some(Self::AnycastTraceroute),
400            "tracemap" => Some(Self::Tracemap),
401            "feed" => Some(Self::Feed),
402            "feed-trace" => Some(Self::FeedTrace),
403            _ => None,
404        }
405    }
406
407    /// Whether this is a live (feed-based) measurement type (targets streamed over stdin).
408    pub fn is_feed(&self) -> bool {
409        matches!(self, Self::Feed | Self::FeedTrace)
410    }
411}
412
413impl Display for ProtocolType {
414    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
415        let s = match self {
416            ProtocolType::Icmp => "ICMP",
417            ProtocolType::ADns => "DNS (A)",
418            ProtocolType::Tcp => "TCP",
419            ProtocolType::ChaosDns => "DNS (CHAOS)",
420        };
421        write!(f, "{}", s)
422    }
423}
424
425impl ProtocolType {
426    pub fn as_str(&self) -> &'static str {
427        match self {
428            ProtocolType::Icmp => "icmp",
429            ProtocolType::ADns => "dns",
430            ProtocolType::Tcp => "tcp",
431            ProtocolType::ChaosDns => "chaos",
432        }
433    }
434
435    pub fn from_str(s: &str) -> Option<Self> {
436        match s.to_lowercase().as_str() {
437            "icmp" => Some(Self::Icmp),
438            "dns" => Some(Self::ADns),
439            "tcp" => Some(Self::Tcp),
440            "chaos" => Some(Self::ChaosDns),
441            _ => None,
442        }
443    }
444}
manycastr/custom_module.rs

manycastr/
custom_module.rs
