wireforge_packet/
lib.rs

1//! WireForge Packet — unified packet enum and protocol dispatch.
2//!
3//! This crate provides a single `Packet` enum that can represent any
4//! supported protocol layer, plus helper functions to automatically
5//! dispatch from raw bytes to the correct protocol parser.
6
7use wireforge_core::arp::ArpPacket;
8use wireforge_core::ether::EthernetPacket;
9use wireforge_core::icmp::{IcmpPacket, Icmpv6Packet};
10use wireforge_core::ipv4::Ipv4Packet;
11use wireforge_core::ipv6::Ipv6Packet;
12use wireforge_core::tcp::TcpPacket;
13use wireforge_core::udp::UdpPacket;
14use wireforge_core::types::{EtherType, IpProtocol};
15
16/// Unified packet enum covering all supported L2–L4 protocols.
17#[non_exhaustive]
18#[derive(Debug, Clone)]
19pub enum Packet<'a> {
20    Ethernet(EthernetPacket<'a>),
21    Arp(ArpPacket<'a>),
22    Ipv4(Ipv4Packet<'a>),
23    Ipv6(Ipv6Packet<'a>),
24    Tcp(TcpPacket<'a>),
25    Udp(UdpPacket<'a>),
26    Icmp(IcmpPacket<'a>),
27    Icmpv6(Icmpv6Packet<'a>),
28}
29
30// From impls for easy conversion
31impl<'a> From<EthernetPacket<'a>> for Packet<'a> {
32    fn from(p: EthernetPacket<'a>) -> Self { Packet::Ethernet(p) }
33}
34
35impl<'a> From<ArpPacket<'a>> for Packet<'a> {
36    fn from(p: ArpPacket<'a>) -> Self { Packet::Arp(p) }
37}
38
39impl<'a> From<Ipv4Packet<'a>> for Packet<'a> {
40    fn from(p: Ipv4Packet<'a>) -> Self { Packet::Ipv4(p) }
41}
42
43impl<'a> From<Ipv6Packet<'a>> for Packet<'a> {
44    fn from(p: Ipv6Packet<'a>) -> Self { Packet::Ipv6(p) }
45}
46
47impl<'a> From<TcpPacket<'a>> for Packet<'a> {
48    fn from(p: TcpPacket<'a>) -> Self { Packet::Tcp(p) }
49}
50
51impl<'a> From<UdpPacket<'a>> for Packet<'a> {
52    fn from(p: UdpPacket<'a>) -> Self { Packet::Udp(p) }
53}
54
55impl<'a> From<IcmpPacket<'a>> for Packet<'a> {
56    fn from(p: IcmpPacket<'a>) -> Self { Packet::Icmp(p) }
57}
58
59impl<'a> From<Icmpv6Packet<'a>> for Packet<'a> {
60    fn from(p: Icmpv6Packet<'a>) -> Self { Packet::Icmpv6(p) }
61}
62
63/// Parse an Ethernet frame and dispatch the encapsulated protocol.
64///
65/// Returns the parsed `EthernetPacket` and, if the EtherType matches a
66/// known protocol, the inner `Packet` as well.
67pub fn parse_ethernet_frame(buf: &[u8]) -> Option<(EthernetPacket<'_>, Option<Packet<'_>>)> {
68    let eth = EthernetPacket::new(buf)?;
69    let inner = match eth.ethertype() {
70        EtherType::Arp => ArpPacket::new(eth.payload()).map(Packet::Arp),
71        EtherType::Ipv4 => Ipv4Packet::new(eth.payload()).map(Packet::Ipv4),
72        EtherType::Ipv6 => Ipv6Packet::new(eth.payload()).map(Packet::Ipv6),
73        _ => None,
74    };
75    Some((eth, inner))
76}
77
78/// Dispatch the payload of an IPv4 packet to the correct transport-layer parser.
79pub fn parse_ipv4_payload<'a>(ip: &Ipv4Packet<'a>) -> Option<Packet<'a>> {
80    match ip.protocol() {
81        IpProtocol::Icmp => IcmpPacket::new(ip.payload()).map(Packet::Icmp),
82        IpProtocol::Tcp => TcpPacket::new(ip.payload()).map(Packet::Tcp),
83        IpProtocol::Udp => UdpPacket::new(ip.payload()).map(Packet::Udp),
84        _ => None,
85    }
86}
87
88/// Dispatch the payload of an IPv6 packet (skipping extension headers).
89pub fn parse_ipv6_payload<'a>(ip: &Ipv6Packet<'a>) -> Option<Packet<'a>> {
90    match ip.final_protocol() {
91        IpProtocol::Icmpv6 => Icmpv6Packet::new(ip.payload()).map(Packet::Icmpv6),
92        IpProtocol::Tcp => TcpPacket::new(ip.payload()).map(Packet::Tcp),
93        IpProtocol::Udp => UdpPacket::new(ip.payload()).map(Packet::Udp),
94        _ => None,
95    }
96}
97
98// ===========================================================================
99// Visitor pattern — add operations on Packet without modifying the enum
100// ===========================================================================
101
102/// Visitor for type-safe dispatch across all [`Packet`] variants.
103///
104/// Implement this trait to add new operations over the packet hierarchy
105/// without modifying the `Packet` enum itself (Open-Closed Principle).
106/// Each `visit_*` method receives a reference to a concrete packet type.
107pub trait PacketVisitor {
108    type Output;
109
110    fn visit_ethernet(&mut self, pkt: &EthernetPacket<'_>) -> Self::Output;
111    fn visit_arp(&mut self, pkt: &ArpPacket<'_>) -> Self::Output;
112    fn visit_ipv4(&mut self, pkt: &Ipv4Packet<'_>) -> Self::Output;
113    fn visit_ipv6(&mut self, pkt: &Ipv6Packet<'_>) -> Self::Output;
114    fn visit_tcp(&mut self, pkt: &TcpPacket<'_>) -> Self::Output;
115    fn visit_udp(&mut self, pkt: &UdpPacket<'_>) -> Self::Output;
116    fn visit_icmp(&mut self, pkt: &IcmpPacket<'_>) -> Self::Output;
117    fn visit_icmpv6(&mut self, pkt: &Icmpv6Packet<'_>) -> Self::Output;
118}
119
120impl<'a> Packet<'a> {
121    /// Accept a visitor — the double-dispatch entry point.
122    ///
123    /// This is the Visitor pattern's core mechanism: the `Packet` dispatches
124    /// to the correct `visit_*` method based on its variant at runtime.
125    pub fn accept<V: PacketVisitor>(&self, visitor: &mut V) -> V::Output {
126        match self {
127            Packet::Ethernet(p) => visitor.visit_ethernet(p),
128            Packet::Arp(p) => visitor.visit_arp(p),
129            Packet::Ipv4(p) => visitor.visit_ipv4(p),
130            Packet::Ipv6(p) => visitor.visit_ipv6(p),
131            Packet::Tcp(p) => visitor.visit_tcp(p),
132            Packet::Udp(p) => visitor.visit_udp(p),
133            Packet::Icmp(p) => visitor.visit_icmp(p),
134            Packet::Icmpv6(p) => visitor.visit_icmpv6(p),
135        }
136    }
137
138    // =======================================================================
139    // Composite pattern — recursive protocol nesting
140    // =======================================================================
141
142    /// Return the encapsulated inner protocol, if any.
143    ///
144    /// This is the Composite pattern: each protocol layer can contain
145    /// another, forming a recursive tree. Walk the full stack with:
146    ///
147    /// ```ignore
148    /// let mut current = Some(packet);
149    /// while let Some(p) = current {
150    ///     println!("{}", p.protocol_name());
151    ///     current = p.inner();
152    /// }
153    /// ```
154    pub fn inner(&self) -> Option<Packet<'a>> {
155        match self {
156            Packet::Ethernet(eth) => {
157                parse_ethernet_frame(eth.as_bytes()).and_then(|(_, inner)| inner)
158            }
159            Packet::Ipv4(ip) => parse_ipv4_payload(ip),
160            Packet::Ipv6(ip) => parse_ipv6_payload(ip),
161            // Leaf protocols — no further nesting
162            Packet::Arp(_) | Packet::Tcp(_) | Packet::Udp(_)
163            | Packet::Icmp(_) | Packet::Icmpv6(_) => None,
164        }
165    }
166
167    /// Human-readable protocol name for this packet layer.
168    pub fn protocol_name(&self) -> &'static str {
169        match self {
170            Packet::Ethernet(_) => "Ethernet",
171            Packet::Arp(_) => "ARP",
172            Packet::Ipv4(_) => "IPv4",
173            Packet::Ipv6(_) => "IPv6",
174            Packet::Tcp(_) => "TCP",
175            Packet::Udp(_) => "UDP",
176            Packet::Icmp(_) => "ICMP",
177            Packet::Icmpv6(_) => "ICMPv6",
178        }
179    }
180}
181
182// ===========================================================================
183// Concrete Visitors (built-in examples of the Visitor pattern)
184// ===========================================================================
185
186/// Visitor that collects the protocol stack as a Vec of human-readable names
187/// (e.g. `["Ethernet", "IPv4", "TCP"]`).
188pub struct ProtocolStackVisitor {
189    pub stack: Vec<&'static str>,
190}
191
192impl ProtocolStackVisitor {
193    pub fn new() -> Self { Self { stack: Vec::new() } }
194}
195
196impl Default for ProtocolStackVisitor {
197    fn default() -> Self { Self::new() }
198}
199
200impl PacketVisitor for ProtocolStackVisitor {
201    type Output = ();
202
203    fn visit_ethernet(&mut self, _pkt: &EthernetPacket<'_>) {
204        self.stack.push("Ethernet");
205    }
206    fn visit_arp(&mut self, _pkt: &ArpPacket<'_>) {
207        self.stack.push("ARP");
208    }
209    fn visit_ipv4(&mut self, _pkt: &Ipv4Packet<'_>) {
210        self.stack.push("IPv4");
211    }
212    fn visit_ipv6(&mut self, _pkt: &Ipv6Packet<'_>) {
213        self.stack.push("IPv6");
214    }
215    fn visit_tcp(&mut self, _pkt: &TcpPacket<'_>) {
216        self.stack.push("TCP");
217    }
218    fn visit_udp(&mut self, _pkt: &UdpPacket<'_>) {
219        self.stack.push("UDP");
220    }
221    fn visit_icmp(&mut self, _pkt: &IcmpPacket<'_>) {
222        self.stack.push("ICMP");
223    }
224    fn visit_icmpv6(&mut self, _pkt: &Icmpv6Packet<'_>) {
225        self.stack.push("ICMPv6");
226    }
227}
228
229/// Visitor that computes the total byte size of a packet and its payload.
230pub struct PacketSizeVisitor {
231    pub size: usize,
232}
233
234impl PacketSizeVisitor {
235    pub fn new() -> Self { Self { size: 0 } }
236}
237
238impl Default for PacketSizeVisitor {
239    fn default() -> Self { Self::new() }
240}
241
242impl PacketVisitor for PacketSizeVisitor {
243    type Output = ();
244
245    fn visit_ethernet(&mut self, pkt: &EthernetPacket<'_>) {
246        self.size = pkt.as_bytes().len();
247    }
248    fn visit_arp(&mut self, pkt: &ArpPacket<'_>) {
249        self.size = pkt.as_bytes().len();
250    }
251    fn visit_ipv4(&mut self, pkt: &Ipv4Packet<'_>) {
252        self.size = pkt.as_bytes().len();
253    }
254    fn visit_ipv6(&mut self, pkt: &Ipv6Packet<'_>) {
255        self.size = pkt.as_bytes().len();
256    }
257    fn visit_tcp(&mut self, pkt: &TcpPacket<'_>) {
258        self.size = pkt.as_bytes().len();
259    }
260    fn visit_udp(&mut self, pkt: &UdpPacket<'_>) {
261        self.size = pkt.as_bytes().len();
262    }
263    fn visit_icmp(&mut self, pkt: &IcmpPacket<'_>) {
264        self.size = pkt.as_bytes().len();
265    }
266    fn visit_icmpv6(&mut self, pkt: &Icmpv6Packet<'_>) {
267        self.size = pkt.as_bytes().len();
268    }
269}
270
271/// Walk the full protocol stack recursively, collecting names at each layer.
272///
273/// This combines the **Composite** pattern (recursive `inner()` traversal)
274/// with the **Visitor** pattern (type-safe dispatch at each level).
275pub fn collect_protocol_stack(pkt: &Packet<'_>) -> Vec<&'static str> {
276    let mut stack = Vec::new();
277    let mut current = Some(pkt.clone());
278    while let Some(p) = current {
279        let mut visitor = ProtocolStackVisitor::new();
280        p.accept(&mut visitor);
281        stack.extend(visitor.stack);
282        current = p.inner();
283    }
284    stack
285}
286
287#[cfg(test)]
288mod tests {
289    use super::*;
290    use wireforge_core::ipv4::Ipv4PacketBuilder;
291    use wireforge_core::udp::UdpPacketBuilder;
292
293    #[test]
294    fn dispatch_ethernet_to_ipv4_to_tcp() {
295        let ip_bytes = Ipv4PacketBuilder::new()
296            .source(std::net::Ipv4Addr::new(10, 0, 0, 1))
297            .destination(std::net::Ipv4Addr::new(10, 0, 0, 2))
298            .protocol(IpProtocol::Tcp)
299            .ttl(64)
300            .build();
301
302        let mut frame = vec![0xFFu8; 12]; // dst + src mac
303        frame.extend_from_slice(&[0x08, 0x00]); // EtherType::Ipv4
304        frame.extend_from_slice(&ip_bytes);
305
306        let (eth, inner) = parse_ethernet_frame(&frame).unwrap();
307        assert_eq!(eth.ethertype(), EtherType::Ipv4);
308        assert!(matches!(inner, Some(Packet::Ipv4(_))));
309    }
310
311    #[test]
312    fn dispatch_ethernet_to_arp() {
313        // Build a valid ARP request (28 bytes minimum)
314        let mut arp_bytes = vec![0x00u8, 0x01]; // hw_type=Ethernet
315        arp_bytes.extend_from_slice(&[0x08, 0x00]); // proto_type=IPv4
316        arp_bytes.push(6);  // hw_addr_len
317        arp_bytes.push(4);  // proto_addr_len
318        arp_bytes.extend_from_slice(&[0x00, 0x01]); // operation=Request
319        arp_bytes.extend_from_slice(&[0x00u8; 20]); // addresses (6+4+6+4)
320
321        let mut frame = vec![0xFFu8; 12];
322        frame.extend_from_slice(&[0x08, 0x06]); // EtherType::Arp
323        frame.extend_from_slice(&arp_bytes);
324
325        let (eth, inner) = parse_ethernet_frame(&frame).unwrap();
326        assert_eq!(eth.ethertype(), EtherType::Arp);
327        assert!(matches!(inner, Some(Packet::Arp(_))));
328    }
329
330    #[test]
331    fn dispatch_ethernet_to_ipv6() {
332        // Build a minimal IPv6 header (40 bytes)
333        let mut frame = vec![0xFFu8; 12];
334        frame.extend_from_slice(&[0x86, 0xDD]); // EtherType::Ipv6
335        let mut ip6 = vec![0x60u8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x3A, 0x40]; // TC+flow=0, payload=0, next=ICMPv6, hop=64
336        ip6.extend_from_slice(&[0u8; 32]); // src + dst = ::
337        frame.extend_from_slice(&ip6);
338
339        let (eth, inner) = parse_ethernet_frame(&frame).unwrap();
340        assert_eq!(eth.ethertype(), EtherType::Ipv6);
341        assert!(matches!(inner, Some(Packet::Ipv6(_))));
342    }
343
344    #[test]
345    fn dispatch_ipv4_to_udp() {
346        let udp_bytes = UdpPacketBuilder::new()
347            .source_port(12345)
348            .destination_port(53)
349            .payload(&[0x01])
350            .build(std::net::Ipv4Addr::new(10, 0, 0, 1), std::net::Ipv4Addr::new(10, 0, 0, 2));
351
352        let ip_bytes = Ipv4PacketBuilder::new()
353            .source(std::net::Ipv4Addr::new(10, 0, 0, 1))
354            .destination(std::net::Ipv4Addr::new(10, 0, 0, 2))
355            .protocol(IpProtocol::Udp)
356            .ttl(64)
357            .payload(&udp_bytes).unwrap()
358            .build();
359
360        let ip = Ipv4Packet::new(&ip_bytes).unwrap();
361        let inner = parse_ipv4_payload(&ip);
362        assert!(matches!(inner, Some(Packet::Udp(_))));
363    }
364
365    #[test]
366    fn dispatch_ipv4_to_icmp() {
367        use wireforge_core::icmp::IcmpPacketBuilder;
368        let icmp_bytes = IcmpPacketBuilder::echo_request()
369            .identifier(1).sequence_number(1)
370            .payload(b"ping")
371            .build();
372
373        let ip_bytes = Ipv4PacketBuilder::new()
374            .source(std::net::Ipv4Addr::new(10, 0, 0, 1))
375            .destination(std::net::Ipv4Addr::new(10, 0, 0, 2))
376            .protocol(IpProtocol::Icmp)
377            .ttl(64)
378            .payload(&icmp_bytes).unwrap()
379            .build();
380
381        let ip = Ipv4Packet::new(&ip_bytes).unwrap();
382        let inner = parse_ipv4_payload(&ip);
383        assert!(matches!(inner, Some(Packet::Icmp(_))));
384    }
385
386    // -- Visitor pattern tests --
387
388    #[test]
389    fn visitor_protocol_name() {
390        use wireforge_core::arp::ArpPacketBuilder;
391
392        let bytes = ArpPacketBuilder::new()
393            .sender_hw_addr(&[0xAA; 6]).unwrap()
394            .sender_proto_addr(&[10, 0, 0, 1]).unwrap()
395            .target_hw_addr(&[0xBB; 6]).unwrap()
396            .target_proto_addr(&[10, 0, 0, 2]).unwrap()
397            .build();
398
399        let pkt = Packet::Arp(ArpPacket::new(&bytes).unwrap());
400
401        let mut visitor = ProtocolStackVisitor::new();
402        pkt.accept(&mut visitor);
403        assert_eq!(visitor.stack, vec!["ARP"]);
404    }
405
406    #[test]
407    fn visitor_packet_size() {
408        let ip_bytes = Ipv4PacketBuilder::new()
409            .source(std::net::Ipv4Addr::new(10, 0, 0, 1))
410            .destination(std::net::Ipv4Addr::new(10, 0, 0, 2))
411            .protocol(IpProtocol::Tcp)
412            .ttl(64)
413            .build();
414
415        let ip = Ipv4Packet::new(&ip_bytes).unwrap();
416        let pkt = Packet::Ipv4(ip.clone());
417
418        let mut visitor = PacketSizeVisitor::new();
419        pkt.accept(&mut visitor);
420        assert_eq!(visitor.size, ip_bytes.len());
421    }
422
423    #[test]
424    fn visitor_dispatches_all_variants() {
425        // Build one of each packet type and verify accept() dispatches correctly.
426        let ip_bytes = Ipv4PacketBuilder::new()
427            .source(std::net::Ipv4Addr::new(10, 0, 0, 1))
428            .destination(std::net::Ipv4Addr::new(10, 0, 0, 2))
429            .protocol(IpProtocol::Tcp)
430            .build();
431
432        let ip = Ipv4Packet::new(&ip_bytes).unwrap();
433        let pkt = Packet::Ipv4(ip);
434
435        // Use accept() to get the protocol name
436        let name = pkt.protocol_name();
437        assert_eq!(name, "IPv4");
438
439        let mut psv = ProtocolStackVisitor::new();
440        pkt.accept(&mut psv);
441        assert_eq!(psv.stack, vec!["IPv4"]);
442    }
443
444    // -- Composite pattern tests --
445
446    #[test]
447    fn composite_ethernet_to_ipv4_to_tcp() {
448        use wireforge_core::tcp::TcpPacketBuilder;
449
450        let tcp_bytes = TcpPacketBuilder::new()
451            .source_port(12345)
452            .destination_port(80)
453            .sequence(1)
454            .syn(true)
455            .window(65535)
456            .build(
457                std::net::Ipv4Addr::new(10, 0, 0, 1),
458                std::net::Ipv4Addr::new(10, 0, 0, 2),
459            );
460
461        let ip_bytes = Ipv4PacketBuilder::new()
462            .source(std::net::Ipv4Addr::new(10, 0, 0, 1))
463            .destination(std::net::Ipv4Addr::new(10, 0, 0, 2))
464            .protocol(IpProtocol::Tcp)
465            .payload(&tcp_bytes).unwrap()
466            .build();
467
468        let mut frame = vec![0xFFu8; 12]; // dst + src mac
469        frame.extend_from_slice(&[0x08, 0x00]); // EtherType::Ipv4
470        frame.extend_from_slice(&ip_bytes);
471
472        let (eth, _inner) = parse_ethernet_frame(&frame).unwrap();
473        let ethernet_pkt = Packet::Ethernet(eth);
474
475        // Composite: walk down the stack
476        assert_eq!(ethernet_pkt.protocol_name(), "Ethernet");
477        let l3 = ethernet_pkt.inner().unwrap();
478        assert_eq!(l3.protocol_name(), "IPv4");
479        let l4 = l3.inner().unwrap();
480        assert_eq!(l4.protocol_name(), "TCP");
481        // TCP is a leaf — no further nesting
482        assert!(l4.inner().is_none());
483    }
484
485    #[test]
486    fn collect_stack_ethernet_ipv4_tcp() {
487        use wireforge_core::tcp::TcpPacketBuilder;
488
489        let tcp_bytes = TcpPacketBuilder::new()
490            .source_port(443)
491            .destination_port(54321)
492            .sequence(100)
493            .syn(true)
494            .window(8192)
495            .build(
496                std::net::Ipv4Addr::new(10, 0, 0, 1),
497                std::net::Ipv4Addr::new(10, 0, 0, 2),
498            );
499
500        let ip_bytes = Ipv4PacketBuilder::new()
501            .source(std::net::Ipv4Addr::new(10, 0, 0, 1))
502            .destination(std::net::Ipv4Addr::new(10, 0, 0, 2))
503            .protocol(IpProtocol::Tcp)
504            .payload(&tcp_bytes).unwrap()
505            .build();
506
507        let mut frame = vec![0xFFu8; 12];
508        frame.extend_from_slice(&[0x08, 0x00]);
509        frame.extend_from_slice(&ip_bytes);
510
511        let (eth, _) = parse_ethernet_frame(&frame).unwrap();
512        let pkt = Packet::Ethernet(eth);
513
514        let stack = collect_protocol_stack(&pkt);
515        assert_eq!(stack, vec!["Ethernet", "IPv4", "TCP"]);
516    }
517
518    #[test]
519    fn arp_is_leaf() {
520        use wireforge_core::arp::ArpPacketBuilder;
521
522        let bytes = ArpPacketBuilder::new()
523            .sender_hw_addr(&[0xAA; 6]).unwrap()
524            .sender_proto_addr(&[10, 0, 0, 1]).unwrap()
525            .target_hw_addr(&[0xBB; 6]).unwrap()
526            .target_proto_addr(&[10, 0, 0, 2]).unwrap()
527            .build();
528
529        let pkt = Packet::Arp(ArpPacket::new(&bytes).unwrap());
530        assert_eq!(pkt.protocol_name(), "ARP");
531        assert!(pkt.inner().is_none()); // ARP has no inner protocol
532    }
533}
wireforge_packet/lib.rs

wireforge_packet/
lib.rs
