Skip to main content

packet_strata/tracker/
tuple.rs

1use serde::{Deserialize, Serialize};
2use smallvec::SmallVec;
3use std::hash::{Hash, Hasher};
4use std::net::{Ipv4Addr, Ipv6Addr};
5
6use crate::packet::ether::EthAddr;
7use crate::packet::protocol::{EtherProto, IpProto};
8use crate::{
9    packet::{
10        header::{LinkLayer, NetworkLayer},
11        Packet,
12    },
13    tracker::vni::{VniId, VniLayer, VniMapper},
14};
15
16/// Common trait for all flow tuple types
17///
18/// This trait defines the interface for creating flow tuples from packets.
19/// (e.g., 5-tuple with VNI, 3-tuple, MAC-based, etc.)
20pub trait Tuple: Sized + Hash + Eq + Clone + Copy {
21    type Addr: Eq;
22
23    /// Create a new flow tuple from a packet
24    fn from_packet(pkt: &Packet<'_>, vni_mapper: &mut VniMapper) -> Option<Self>;
25
26    /// Flip the source and destination fields of the flow tuple
27    fn flip(&self) -> Self;
28
29    /// Hashes the tuple in a canonical (symmetric) way without creating a new instance.
30    /// Used by `Symmetric` wrapper to ensure `Hash(A->B) == Hash(B->A)`.
31    fn hash_canonical<H: Hasher>(&self, state: &mut H);
32
33    /// Checks equality in a canonical (symmetric) way without creating a new instance.
34    /// Used by `Symmetric` wrapper to ensure `Eq(A->B, B->A)`.
35    fn eq_canonical(&self, other: &Self) -> bool;
36
37    /// Checks if the tuple is symmetric (source equals destination).
38    ///
39    /// A tuple is considered symmetric when both the source address equals the destination
40    /// address and the source port equals the destination port. This is useful for
41    /// identifying self-referential connections.
42    #[inline]
43    fn is_symmetric(&self) -> bool {
44        self.source_port() == self.dest_port() && self.source() == self.dest()
45    }
46
47    /// Returns the source address of the flow tuple.
48    fn source(&self) -> Self::Addr;
49
50    /// Returns the destination address of the flow tuple.
51    fn dest(&self) -> Self::Addr;
52
53    /// Returns the source port of the flow tuple.
54    fn source_port(&self) -> u16;
55
56    /// Returns the destination port of the flow tuple.
57    fn dest_port(&self) -> u16;
58
59    /// Returns the IP protocol of the flow tuple.
60    fn protocol(&self) -> IpProto;
61
62    /// Returns the VNI (VXLAN Network Identifier) of the flow tuple.
63    fn vni(&self) -> VniId;
64}
65
66/// Helper function to extract VNI from packet tunnels
67///
68/// This function is shared between all Tuple implementations to avoid code duplication.
69/// Returns `VNI_NULL` if there are no tunnel layers, otherwise extracts and maps the VNI stack.
70#[inline]
71fn extract_vni(pkt: &Packet<'_>, vni_mapper: &mut VniMapper) -> Option<VniId> {
72    let network_tunnel_layers = pkt.tunnels();
73
74    if network_tunnel_layers.is_empty() {
75        return Some(VniId::default());
76    }
77
78    let vni_stack = network_tunnel_layers
79        .iter()
80        .map(TryInto::try_into)
81        .collect::<Result<SmallVec<[VniLayer; 4]>, _>>()
82        .ok()?;
83
84    Some(vni_mapper.get_or_create_vni_id(&vni_stack))
85}
86
87/// Helper function to extract transport layer ports
88///
89/// Returns (src_port, dst_port) or (0, 0) if no transport layer is present.
90#[inline]
91fn extract_ports(pkt: &Packet<'_>) -> (u16, u16) {
92    pkt.transport().map(|t| t.ports()).unwrap_or((0, 0))
93}
94
95#[derive(Debug, Copy, Clone)]
96#[repr(transparent)]
97pub struct Symmetric<T: Tuple>(pub T);
98
99impl<T: Tuple> PartialEq for Symmetric<T> {
100    #[inline]
101    fn eq(&self, other: &Self) -> bool {
102        self.0.eq_canonical(&other.0)
103    }
104}
105
106impl<T: Tuple> Eq for Symmetric<T> {}
107
108impl<T: Tuple> Hash for Symmetric<T> {
109    #[inline]
110    fn hash<H: Hasher>(&self, state: &mut H) {
111        self.0.hash_canonical(state);
112    }
113}
114
115impl<T: Tuple> From<T> for Symmetric<T> {
116    fn from(t: T) -> Self {
117        Self(t)
118    }
119}
120
121/// IPv4 5-tuple flow with VNI support
122///
123/// This tuple uniquely identifies a network flow using:
124/// - Source and destination IPv4 addresses
125/// - Source and destination ports
126/// - IP protocol number
127/// - Virtual Network Identifier (VNI) for tunnel-aware flow tracking
128#[derive(Hash, Eq, PartialEq, Debug, Copy, Clone, Serialize, Deserialize)]
129pub struct TupleV4 {
130    pub src_ip: Ipv4Addr,
131    pub dst_ip: Ipv4Addr,
132    pub src_port: u16,
133    pub dst_port: u16,
134    pub protocol: IpProto,
135    pub vni: VniId,
136}
137
138impl Default for TupleV4 {
139    fn default() -> Self {
140        Self {
141            src_ip: Ipv4Addr::UNSPECIFIED,
142            dst_ip: Ipv4Addr::UNSPECIFIED,
143            src_port: 0,
144            dst_port: 0,
145            protocol: IpProto::default(),
146            vni: VniId::default(),
147        }
148    }
149}
150
151impl TupleV4 {
152    /// Create a new IPv4 flow tuple from a packet
153    ///
154    /// Returns `None` if the packet does not contain an IPv4 header or if VNI extraction fails.
155    pub fn new(pkt: &Packet<'_>, vni_mapper: &mut VniMapper) -> Option<Self> {
156        let NetworkLayer::Ipv4(ipv4) = pkt.network()? else {
157            return None;
158        };
159
160        let src_ip = ipv4.header.src_ip();
161        let dst_ip = ipv4.header.dst_ip();
162        let protocol = ipv4.header.protocol();
163        let (src_port, dst_port) = extract_ports(pkt);
164        let vni = extract_vni(pkt, vni_mapper)?;
165
166        Some(Self {
167            src_ip,
168            dst_ip,
169            src_port,
170            dst_port,
171            protocol,
172            vni,
173        })
174    }
175}
176
177impl Tuple for TupleV4 {
178    type Addr = Ipv4Addr;
179
180    #[inline]
181    fn source(&self) -> Self::Addr {
182        self.src_ip
183    }
184
185    #[inline]
186    fn dest(&self) -> Self::Addr {
187        self.dst_ip
188    }
189
190    #[inline]
191    fn source_port(&self) -> u16 {
192        self.src_port
193    }
194
195    #[inline]
196    fn dest_port(&self) -> u16 {
197        self.dst_port
198    }
199
200    #[inline]
201    fn protocol(&self) -> IpProto {
202        self.protocol
203    }
204
205    #[inline]
206    fn vni(&self) -> VniId {
207        self.vni
208    }
209
210    #[inline]
211    fn from_packet(pkt: &Packet<'_>, vni_mapper: &mut VniMapper) -> Option<Self> {
212        Self::new(pkt, vni_mapper)
213    }
214
215    #[inline]
216    fn flip(&self) -> Self {
217        Self {
218            src_ip: self.dst_ip,
219            dst_ip: self.src_ip,
220            src_port: self.dst_port,
221            dst_port: self.src_port,
222            protocol: self.protocol,
223            vni: self.vni,
224        }
225    }
226
227    #[inline]
228    fn hash_canonical<H: Hasher>(&self, state: &mut H) {
229        // Hash invariant fields
230        self.protocol.hash(state);
231        self.vni.hash(state);
232
233        // Hash variant fields in sorted order (src < dst)
234        // This avoids creating a new struct or flipping
235        if (self.src_ip, self.src_port) <= (self.dst_ip, self.dst_port) {
236            self.src_ip.hash(state);
237            self.src_port.hash(state);
238            self.dst_ip.hash(state);
239            self.dst_port.hash(state);
240        } else {
241            self.dst_ip.hash(state);
242            self.dst_port.hash(state);
243            self.src_ip.hash(state);
244            self.src_port.hash(state);
245        }
246    }
247
248    #[inline]
249    fn eq_canonical(&self, other: &Self) -> bool {
250        if self.protocol != other.protocol || self.vni != other.vni {
251            return false;
252        }
253
254        // Check direct equality OR crossed equality
255        // This is much cheaper than constructing a new struct
256        (self.src_ip == other.src_ip
257            && self.dst_ip == other.dst_ip
258            && self.src_port == other.src_port
259            && self.dst_port == other.dst_port)
260            || (self.src_ip == other.dst_ip
261                && self.dst_ip == other.src_ip
262                && self.src_port == other.dst_port
263                && self.dst_port == other.src_port)
264    }
265}
266
267/// IPv6 5-tuple flow tuple with VNI support
268///
269/// This tuple uniquely identifies a network flow using:
270/// - Source and destination IPv6 addresses
271/// - Source and destination ports
272/// - IP protocol number (next header)
273/// - Virtual Network Identifier (VNI) for tunnel-aware flow tracking
274#[derive(Hash, Eq, PartialEq, Debug, Copy, Clone, Serialize, Deserialize)]
275pub struct TupleV6 {
276    pub src_ip: Ipv6Addr,
277    pub dst_ip: Ipv6Addr,
278    pub src_port: u16,
279    pub dst_port: u16,
280    pub protocol: IpProto,
281    pub vni: VniId,
282}
283
284impl Default for TupleV6 {
285    fn default() -> Self {
286        Self {
287            src_ip: Ipv6Addr::UNSPECIFIED,
288            dst_ip: Ipv6Addr::UNSPECIFIED,
289            src_port: 0,
290            dst_port: 0,
291            protocol: IpProto::default(),
292            vni: VniId::default(),
293        }
294    }
295}
296
297impl TupleV6 {
298    /// Create a new IPv6 flow tuple from a packet
299    ///
300    /// Returns `None` if the packet does not contain an IPv6 header or if VNI extraction fails.
301    pub fn new(pkt: &Packet<'_>, vni_mapper: &mut VniMapper) -> Option<Self> {
302        let NetworkLayer::Ipv6(ipv6) = pkt.network()? else {
303            return None;
304        };
305
306        let src_ip = ipv6.header.src_ip();
307        let dst_ip = ipv6.header.dst_ip();
308        let protocol = ipv6.header.next_header();
309        let (src_port, dst_port) = extract_ports(pkt);
310        let vni = extract_vni(pkt, vni_mapper)?;
311
312        Some(Self {
313            src_ip,
314            dst_ip,
315            src_port,
316            dst_port,
317            protocol,
318            vni,
319        })
320    }
321}
322
323impl Tuple for TupleV6 {
324    type Addr = Ipv6Addr;
325
326    #[inline]
327    fn source(&self) -> Self::Addr {
328        self.src_ip
329    }
330
331    #[inline]
332    fn dest(&self) -> Self::Addr {
333        self.dst_ip
334    }
335
336    #[inline]
337    fn source_port(&self) -> u16 {
338        self.src_port
339    }
340
341    #[inline]
342    fn dest_port(&self) -> u16 {
343        self.dst_port
344    }
345
346    #[inline]
347    fn protocol(&self) -> IpProto {
348        self.protocol
349    }
350
351    #[inline]
352    fn vni(&self) -> VniId {
353        self.vni
354    }
355
356    #[inline]
357    fn from_packet(pkt: &Packet<'_>, vni_mapper: &mut VniMapper) -> Option<Self> {
358        Self::new(pkt, vni_mapper)
359    }
360
361    #[inline]
362    fn flip(&self) -> Self {
363        Self {
364            src_ip: self.dst_ip,
365            dst_ip: self.src_ip,
366            src_port: self.dst_port,
367            dst_port: self.src_port,
368            protocol: self.protocol,
369            vni: self.vni,
370        }
371    }
372
373    #[inline]
374    fn hash_canonical<H: Hasher>(&self, state: &mut H) {
375        self.protocol.hash(state);
376        self.vni.hash(state);
377
378        if (self.src_ip, self.src_port) <= (self.dst_ip, self.dst_port) {
379            self.src_ip.hash(state);
380            self.src_port.hash(state);
381            self.dst_ip.hash(state);
382            self.dst_port.hash(state);
383        } else {
384            self.dst_ip.hash(state);
385            self.dst_port.hash(state);
386            self.src_ip.hash(state);
387            self.src_port.hash(state);
388        }
389    }
390
391    #[inline]
392    fn eq_canonical(&self, other: &Self) -> bool {
393        if self.protocol != other.protocol || self.vni != other.vni {
394            return false;
395        }
396
397        (self.src_ip == other.src_ip
398            && self.dst_ip == other.dst_ip
399            && self.src_port == other.src_port
400            && self.dst_port == other.dst_port)
401            || (self.src_ip == other.dst_ip
402                && self.dst_ip == other.src_ip
403                && self.src_port == other.dst_port
404                && self.dst_port == other.src_port)
405    }
406}
407
408#[derive(Hash, Eq, PartialEq, Debug, Copy, Clone, Serialize, Deserialize)]
409pub struct TupleEth {
410    pub src: EthAddr,
411    pub dst: EthAddr,
412    pub protocol: EtherProto,
413}
414
415impl Default for TupleEth {
416    fn default() -> Self {
417        Self {
418            src: EthAddr::default(),
419            dst: EthAddr::default(),
420            protocol: EtherProto::default(),
421        }
422    }
423}
424
425impl TupleEth {
426    /// Create a new Ethernet flow tuple from a packet
427    ///
428    /// Returns `None` if the packet is not Ethernet.
429    pub fn new(pkt: &Packet<'_>) -> Option<Self> {
430        let LinkLayer::Ethernet(eth) = pkt.link() else {
431            return None;
432        };
433
434        Some(Self {
435            src: *eth.source(),
436            dst: *eth.dest(),
437            protocol: eth.inner_type(),
438        })
439    }
440}
441
442impl Tuple for TupleEth {
443    type Addr = EthAddr;
444
445    #[inline]
446    fn source(&self) -> Self::Addr {
447        self.src
448    }
449
450    #[inline]
451    fn dest(&self) -> Self::Addr {
452        self.dst
453    }
454
455    #[inline]
456    fn source_port(&self) -> u16 {
457        0
458    }
459
460    #[inline]
461    fn dest_port(&self) -> u16 {
462        0
463    }
464
465    #[inline]
466    fn protocol(&self) -> IpProto {
467        IpProto::default()
468    }
469
470    #[inline]
471    fn vni(&self) -> VniId {
472        VniId::default()
473    }
474
475    #[inline]
476    fn from_packet(pkt: &Packet<'_>, _vni_mapper: &mut VniMapper) -> Option<Self> {
477        Self::new(pkt)
478    }
479
480    #[inline]
481    fn flip(&self) -> Self {
482        Self {
483            src: self.dst,
484            dst: self.src,
485            protocol: self.protocol,
486        }
487    }
488
489    #[inline]
490    fn hash_canonical<H: Hasher>(&self, state: &mut H) {
491        self.protocol.hash(state);
492
493        // EthAddr implements Ord so we can compare directly
494        if self.src <= self.dst {
495            self.src.hash(state);
496            self.dst.hash(state);
497        } else {
498            self.dst.hash(state);
499            self.src.hash(state);
500        }
501    }
502
503    #[inline]
504    fn eq_canonical(&self, other: &Self) -> bool {
505        if self.protocol != other.protocol {
506            return false;
507        }
508
509        (self.src == other.src && self.dst == other.dst)
510            || (self.src == other.dst && self.dst == other.src)
511    }
512}