stackforge_core/layer/icmp/
builder.rs

1//! ICMP packet builder.
2//!
3//! Provides a fluent API for constructing ICMP packets with type-specific fields
4//! and automatic checksum calculation.
5//!
6//! # Example
7//!
8//! ```rust
9//! use stackforge_core::layer::icmp::IcmpBuilder;
10//!
11//! // Build an echo request
12//! let packet = IcmpBuilder::echo_request(0x1234, 1)
13//!     .payload(b"ping data")
14//!     .build();
15//!
16//! // Build a destination unreachable message
17//! let packet = IcmpBuilder::dest_unreach(3) // Port unreachable
18//!     .build();
19//! ```
20
21use std::net::Ipv4Addr;
22
23use super::checksum::icmp_checksum;
24use super::types::types;
25use super::{ICMP_MIN_HEADER_LEN, offsets};
26use crate::layer::field::FieldError;
27
28/// Builder for ICMP packets.
29///
30/// Due to ICMP's type-specific fields, this builder provides factory methods
31/// for common ICMP message types rather than a generic constructor.
32#[derive(Debug, Clone)]
33pub struct IcmpBuilder {
34    // Base fields (all ICMP messages)
35    icmp_type: u8,
36    code: u8,
37    checksum: Option<u16>,
38
39    // Type-specific data (4 bytes after checksum)
40    // Layout depends on ICMP type:
41    // - Echo: [id: u16, seq: u16]
42    // - Redirect: [gateway: 4 bytes IP]
43    // - Dest Unreach (code 4): [unused: u16, mtu: u16]
44    // - Param Problem: [ptr: u8, unused: 3 bytes]
45    // - Timestamp: handled separately (has additional fields)
46    type_specific: [u8; 4],
47
48    // Additional data for timestamp messages (12 bytes)
49    // ts_ori, ts_rx, ts_tx (3 x u32)
50    timestamp_data: Option<[u8; 12]>,
51
52    // Payload data
53    payload: Vec<u8>,
54
55    // Build options
56    auto_checksum: bool,
57}
58
59impl Default for IcmpBuilder {
60    fn default() -> Self {
61        Self {
62            icmp_type: types::ECHO_REQUEST,
63            code: 0,
64            checksum: None,
65            type_specific: [0; 4],
66            timestamp_data: None,
67            payload: Vec::new(),
68            auto_checksum: true,
69        }
70    }
71}
72
73impl IcmpBuilder {
74    /// Create a new ICMP builder with default values (echo request).
75    #[must_use]
76    pub fn new() -> Self {
77        Self::default()
78    }
79
80    // ========== Factory Methods for Common Types ==========
81
82    /// Create an echo request (ping) packet.
83    ///
84    /// # Arguments
85    /// * `id` - Identifier
86    /// * `seq` - Sequence number
87    #[must_use]
88    pub fn echo_request(id: u16, seq: u16) -> Self {
89        let mut builder = Self::new();
90        builder.icmp_type = types::ECHO_REQUEST;
91        builder.code = 0;
92        builder.set_id_seq(id, seq);
93        builder
94    }
95
96    /// Create an echo reply (pong) packet.
97    ///
98    /// # Arguments
99    /// * `id` - Identifier (should match request)
100    /// * `seq` - Sequence number (should match request)
101    #[must_use]
102    pub fn echo_reply(id: u16, seq: u16) -> Self {
103        let mut builder = Self::new();
104        builder.icmp_type = types::ECHO_REPLY;
105        builder.code = 0;
106        builder.set_id_seq(id, seq);
107        builder
108    }
109
110    /// Create a destination unreachable message.
111    ///
112    /// # Arguments
113    /// * `code` - Specific unreachable code (0-15)
114    ///   - 0: Network unreachable
115    ///   - 1: Host unreachable
116    ///   - 2: Protocol unreachable
117    ///   - 3: Port unreachable
118    ///   - 4: Fragmentation needed (use `dest_unreach_need_frag` for this)
119    #[must_use]
120    pub fn dest_unreach(code: u8) -> Self {
121        let mut builder = Self::new();
122        builder.icmp_type = types::DEST_UNREACH;
123        builder.code = code;
124        builder.type_specific = [0; 4]; // Unused for most codes
125        builder
126    }
127
128    /// Create a destination unreachable - fragmentation needed message.
129    ///
130    /// # Arguments
131    /// * `mtu` - Next-hop MTU value
132    #[must_use]
133    pub fn dest_unreach_need_frag(mtu: u16) -> Self {
134        let mut builder = Self::new();
135        builder.icmp_type = types::DEST_UNREACH;
136        builder.code = 4; // Fragmentation needed
137        builder.type_specific[0] = 0; // Unused byte
138        builder.type_specific[1] = 0; // Unused byte
139        builder.type_specific[2..4].copy_from_slice(&mtu.to_be_bytes());
140        builder
141    }
142
143    /// Create a redirect message.
144    ///
145    /// # Arguments
146    /// * `code` - Redirect code (0-3)
147    ///   - 0: Redirect for network
148    ///   - 1: Redirect for host
149    ///   - 2: Redirect for TOS and network
150    ///   - 3: Redirect for TOS and host
151    /// * `gateway` - Gateway IP address to redirect to
152    #[must_use]
153    pub fn redirect(code: u8, gateway: Ipv4Addr) -> Self {
154        let mut builder = Self::new();
155        builder.icmp_type = types::REDIRECT;
156        builder.code = code;
157        builder.type_specific.copy_from_slice(&gateway.octets());
158        builder
159    }
160
161    /// Create a time exceeded message.
162    ///
163    /// # Arguments
164    /// * `code` - Time exceeded code
165    ///   - 0: TTL exceeded in transit
166    ///   - 1: Fragment reassembly time exceeded
167    #[must_use]
168    pub fn time_exceeded(code: u8) -> Self {
169        let mut builder = Self::new();
170        builder.icmp_type = types::TIME_EXCEEDED;
171        builder.code = code;
172        builder.type_specific = [0; 4]; // Unused
173        builder
174    }
175
176    /// Create a parameter problem message.
177    ///
178    /// # Arguments
179    /// * `ptr` - Pointer to the problematic byte in the original packet
180    #[must_use]
181    pub fn param_problem(ptr: u8) -> Self {
182        let mut builder = Self::new();
183        builder.icmp_type = types::PARAM_PROBLEM;
184        builder.code = 0;
185        builder.type_specific[0] = ptr;
186        builder.type_specific[1] = 0; // Unused
187        builder.type_specific[2] = 0; // Length
188        builder.type_specific[3] = 0; // Unused
189        builder
190    }
191
192    /// Create a source quench message (deprecated).
193    #[must_use]
194    pub fn source_quench() -> Self {
195        let mut builder = Self::new();
196        builder.icmp_type = types::SOURCE_QUENCH;
197        builder.code = 0;
198        builder.type_specific = [0; 4]; // Unused
199        builder
200    }
201
202    /// Create a timestamp request message.
203    ///
204    /// # Arguments
205    /// * `id` - Identifier
206    /// * `seq` - Sequence number
207    /// * `ts_ori` - Originate timestamp (milliseconds since midnight UT)
208    /// * `ts_rx` - Receive timestamp (0 for request)
209    /// * `ts_tx` - Transmit timestamp (0 for request)
210    #[must_use]
211    pub fn timestamp_request(id: u16, seq: u16, ts_ori: u32, ts_rx: u32, ts_tx: u32) -> Self {
212        let mut builder = Self::new();
213        builder.icmp_type = types::TIMESTAMP;
214        builder.code = 0;
215        builder.set_id_seq(id, seq);
216
217        let mut ts_data = [0u8; 12];
218        ts_data[0..4].copy_from_slice(&ts_ori.to_be_bytes());
219        ts_data[4..8].copy_from_slice(&ts_rx.to_be_bytes());
220        ts_data[8..12].copy_from_slice(&ts_tx.to_be_bytes());
221        builder.timestamp_data = Some(ts_data);
222
223        builder
224    }
225
226    /// Create a timestamp reply message.
227    ///
228    /// # Arguments
229    /// * `id` - Identifier (should match request)
230    /// * `seq` - Sequence number (should match request)
231    /// * `ts_ori` - Originate timestamp from request
232    /// * `ts_rx` - Receive timestamp (when request was received)
233    /// * `ts_tx` - Transmit timestamp (when reply is sent)
234    #[must_use]
235    pub fn timestamp_reply(id: u16, seq: u16, ts_ori: u32, ts_rx: u32, ts_tx: u32) -> Self {
236        let mut builder = Self::new();
237        builder.icmp_type = types::TIMESTAMP_REPLY;
238        builder.code = 0;
239        builder.set_id_seq(id, seq);
240
241        let mut ts_data = [0u8; 12];
242        ts_data[0..4].copy_from_slice(&ts_ori.to_be_bytes());
243        ts_data[4..8].copy_from_slice(&ts_rx.to_be_bytes());
244        ts_data[8..12].copy_from_slice(&ts_tx.to_be_bytes());
245        builder.timestamp_data = Some(ts_data);
246
247        builder
248    }
249
250    /// Create an address mask request message.
251    ///
252    /// # Arguments
253    /// * `id` - Identifier
254    /// * `seq` - Sequence number
255    #[must_use]
256    pub fn address_mask_request(id: u16, seq: u16) -> Self {
257        let mut builder = Self::new();
258        builder.icmp_type = types::ADDRESS_MASK_REQUEST;
259        builder.code = 0;
260        builder.set_id_seq(id, seq);
261        builder
262    }
263
264    /// Create an address mask reply message.
265    ///
266    /// # Arguments
267    /// * `id` - Identifier (should match request)
268    /// * `seq` - Sequence number (should match request)
269    /// * `mask` - Address mask
270    #[must_use]
271    pub fn address_mask_reply(id: u16, seq: u16, mask: Ipv4Addr) -> Self {
272        let mut builder = Self::new();
273        builder.icmp_type = types::ADDRESS_MASK_REPLY;
274        builder.code = 0;
275        builder.set_id_seq(id, seq);
276        // For address mask, the mask goes in the payload area
277        builder.payload = mask.octets().to_vec();
278        builder
279    }
280
281    // ========== Helper Methods ==========
282
283    /// Set ID and sequence number (for echo, timestamp, etc.)
284    fn set_id_seq(&mut self, id: u16, seq: u16) {
285        self.type_specific[0..2].copy_from_slice(&id.to_be_bytes());
286        self.type_specific[2..4].copy_from_slice(&seq.to_be_bytes());
287    }
288
289    // ========== Field Setters ==========
290
291    /// Set the ICMP type manually (use factory methods instead when possible).
292    #[must_use]
293    pub fn icmp_type(mut self, t: u8) -> Self {
294        self.icmp_type = t;
295        self
296    }
297
298    /// Set the ICMP code manually.
299    #[must_use]
300    pub fn code(mut self, c: u8) -> Self {
301        self.code = c;
302        self
303    }
304
305    /// Set the checksum manually.
306    ///
307    /// If not set, the checksum will be calculated automatically.
308    #[must_use]
309    pub fn checksum(mut self, csum: u16) -> Self {
310        self.checksum = Some(csum);
311        self.auto_checksum = false;
312        self
313    }
314
315    /// Alias for checksum (Scapy compatibility).
316    #[must_use]
317    pub fn chksum(self, csum: u16) -> Self {
318        self.checksum(csum)
319    }
320
321    /// Enable automatic checksum calculation (default).
322    #[must_use]
323    pub fn enable_auto_checksum(mut self) -> Self {
324        self.auto_checksum = true;
325        self.checksum = None;
326        self
327    }
328
329    /// Disable automatic checksum calculation.
330    #[must_use]
331    pub fn disable_auto_checksum(mut self) -> Self {
332        self.auto_checksum = false;
333        self
334    }
335
336    /// Set the payload data.
337    pub fn payload<T: Into<Vec<u8>>>(mut self, data: T) -> Self {
338        self.payload = data.into();
339        self
340    }
341
342    /// Append to the payload data.
343    pub fn append_payload<T: AsRef<[u8]>>(mut self, data: T) -> Self {
344        self.payload.extend_from_slice(data.as_ref());
345        self
346    }
347
348    // ========== Size Calculation ==========
349
350    /// Get the total packet size (header + optional timestamp + payload).
351    #[must_use]
352    pub fn packet_size(&self) -> usize {
353        let mut size = ICMP_MIN_HEADER_LEN; // Base 8 bytes
354
355        // Timestamp messages have 12 additional bytes
356        if self.timestamp_data.is_some() {
357            size += 12;
358        }
359
360        size + self.payload.len()
361    }
362
363    /// Get the header size (8 bytes for most, 20 for timestamp).
364    #[must_use]
365    pub fn header_size(&self) -> usize {
366        if self.timestamp_data.is_some() {
367            20 // 8 base + 12 timestamp data
368        } else {
369            ICMP_MIN_HEADER_LEN
370        }
371    }
372
373    // ========== Build Methods ==========
374
375    /// Build the ICMP packet into a new buffer.
376    #[must_use]
377    pub fn build(&self) -> Vec<u8> {
378        let total_size = self.packet_size();
379        let mut buf = vec![0u8; total_size];
380        self.build_into(&mut buf)
381            .expect("buffer is correctly sized");
382        buf
383    }
384
385    /// Build the ICMP packet into an existing buffer.
386    pub fn build_into(&self, buf: &mut [u8]) -> Result<usize, FieldError> {
387        let total_size = self.packet_size();
388
389        if buf.len() < total_size {
390            return Err(FieldError::BufferTooShort {
391                offset: 0,
392                need: total_size,
393                have: buf.len(),
394            });
395        }
396
397        // Type
398        buf[offsets::TYPE] = self.icmp_type;
399
400        // Code
401        buf[offsets::CODE] = self.code;
402
403        // Checksum (initially 0, calculated later if auto_checksum is enabled)
404        buf[offsets::CHECKSUM..offsets::CHECKSUM + 2].copy_from_slice(&[0, 0]);
405
406        // Type-specific data (4 bytes)
407        buf[4..8].copy_from_slice(&self.type_specific);
408
409        let mut offset = 8;
410
411        // Timestamp data if present (12 bytes)
412        if let Some(ts_data) = &self.timestamp_data {
413            buf[offset..offset + 12].copy_from_slice(ts_data);
414            offset += 12;
415        }
416
417        // Payload
418        if !self.payload.is_empty() {
419            buf[offset..offset + self.payload.len()].copy_from_slice(&self.payload);
420        }
421
422        // Calculate checksum if enabled
423        if self.auto_checksum {
424            let csum = icmp_checksum(&buf[..total_size]);
425            buf[offsets::CHECKSUM..offsets::CHECKSUM + 2].copy_from_slice(&csum.to_be_bytes());
426        } else if let Some(csum) = self.checksum {
427            buf[offsets::CHECKSUM..offsets::CHECKSUM + 2].copy_from_slice(&csum.to_be_bytes());
428        }
429
430        Ok(total_size)
431    }
432
433    /// Build just the ICMP header (without payload).
434    #[must_use]
435    pub fn build_header(&self) -> Vec<u8> {
436        let header_size = self.header_size();
437        let mut buf = vec![0u8; header_size];
438
439        // Create a copy without payload for header-only build
440        let builder = Self {
441            payload: Vec::new(),
442            ..self.clone()
443        };
444        builder
445            .build_into(&mut buf)
446            .expect("buffer is correctly sized");
447
448        buf
449    }
450}
451
452#[cfg(test)]
453mod tests {
454    use super::*;
455
456    #[test]
457    fn test_echo_request() {
458        let packet = IcmpBuilder::echo_request(0x1234, 5)
459            .payload(b"Hello")
460            .build();
461
462        assert_eq!(packet[0], types::ECHO_REQUEST); // type
463        assert_eq!(packet[1], 0); // code
464        // bytes 2-3 are checksum
465        assert_eq!(u16::from_be_bytes([packet[4], packet[5]]), 0x1234); // id
466        assert_eq!(u16::from_be_bytes([packet[6], packet[7]]), 5); // seq
467        assert_eq!(&packet[8..], b"Hello"); // payload
468    }
469
470    #[test]
471    fn test_echo_reply() {
472        let packet = IcmpBuilder::echo_reply(0x5678, 10).build();
473
474        assert_eq!(packet[0], types::ECHO_REPLY);
475        assert_eq!(u16::from_be_bytes([packet[4], packet[5]]), 0x5678);
476        assert_eq!(u16::from_be_bytes([packet[6], packet[7]]), 10);
477    }
478
479    #[test]
480    fn test_dest_unreach() {
481        let packet = IcmpBuilder::dest_unreach(3).build(); // Port unreachable
482
483        assert_eq!(packet[0], types::DEST_UNREACH);
484        assert_eq!(packet[1], 3); // code
485    }
486
487    #[test]
488    fn test_dest_unreach_need_frag() {
489        let packet = IcmpBuilder::dest_unreach_need_frag(1500).build();
490
491        assert_eq!(packet[0], types::DEST_UNREACH);
492        assert_eq!(packet[1], 4); // fragmentation needed
493        assert_eq!(u16::from_be_bytes([packet[6], packet[7]]), 1500); // MTU
494    }
495
496    #[test]
497    fn test_redirect() {
498        let gateway = Ipv4Addr::new(192, 168, 1, 1);
499        let packet = IcmpBuilder::redirect(1, gateway).build();
500
501        assert_eq!(packet[0], types::REDIRECT);
502        assert_eq!(packet[1], 1); // code: redirect host
503        assert_eq!(&packet[4..8], &[192, 168, 1, 1]); // gateway IP
504    }
505
506    #[test]
507    fn test_time_exceeded() {
508        let packet = IcmpBuilder::time_exceeded(0).build(); // TTL exceeded
509
510        assert_eq!(packet[0], types::TIME_EXCEEDED);
511        assert_eq!(packet[1], 0);
512    }
513
514    #[test]
515    fn test_param_problem() {
516        let packet = IcmpBuilder::param_problem(20).build();
517
518        assert_eq!(packet[0], types::PARAM_PROBLEM);
519        assert_eq!(packet[4], 20); // pointer
520    }
521
522    #[test]
523    fn test_timestamp_request() {
524        let packet = IcmpBuilder::timestamp_request(0x1234, 1, 1000, 0, 0).build();
525
526        assert_eq!(packet[0], types::TIMESTAMP);
527        assert_eq!(packet.len(), 20); // 8 base + 12 timestamp data
528        assert_eq!(u16::from_be_bytes([packet[4], packet[5]]), 0x1234); // id
529        assert_eq!(u16::from_be_bytes([packet[6], packet[7]]), 1); // seq
530        assert_eq!(
531            u32::from_be_bytes([packet[8], packet[9], packet[10], packet[11]]),
532            1000
533        ); // ts_ori
534    }
535
536    #[test]
537    fn test_checksum_calculation() {
538        let packet = IcmpBuilder::echo_request(1, 1).payload(b"test").build();
539
540        // Checksum should be non-zero
541        let checksum = u16::from_be_bytes([packet[2], packet[3]]);
542        assert_ne!(checksum, 0);
543    }
544
545    #[test]
546    fn test_manual_checksum() {
547        let packet = IcmpBuilder::echo_request(1, 1).checksum(0xABCD).build();
548
549        assert_eq!(u16::from_be_bytes([packet[2], packet[3]]), 0xABCD);
550    }
551
552    #[test]
553    fn test_build_header_only() {
554        let header = IcmpBuilder::echo_request(1, 1)
555            .payload(b"this should not be included")
556            .build_header();
557
558        assert_eq!(header.len(), 8); // Header only
559    }
560
561    #[test]
562    fn test_timestamp_header_size() {
563        let header = IcmpBuilder::timestamp_request(1, 1, 1000, 2000, 3000).build_header();
564
565        assert_eq!(header.len(), 20); // 8 + 12 for timestamps
566    }
567}
stackforge_core/layer/icmp/builder.rs

stackforge_core/layer/icmp/
builder.rs
