Skip to main content

stackforge_core/layer/ipv6/
builder.rs

1//! IPv6 packet builder.
2//!
3//! Provides a fluent API for constructing IPv6 packet headers.
4//!
5//! # Example
6//!
7//! ```rust
8//! use stackforge_core::layer::ipv6::Ipv6Builder;
9//! use std::net::Ipv6Addr;
10//!
11//! let bytes = Ipv6Builder::new()
12//!     .src(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 1))
13//!     .dst(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 2))
14//!     .hop_limit(64)
15//!     .next_header(58) // ICMPv6
16//!     .build();
17//!
18//! assert_eq!(bytes.len(), 40); // IPv6 header only (no payload)
19//! ```
20
21use std::net::Ipv6Addr;
22
23use super::IPV6_HEADER_LEN;
24
25/// Builder for IPv6 packet headers.
26#[derive(Debug, Clone)]
27pub struct Ipv6Builder {
28    /// Traffic Class (8 bits, DSCP+ECN)
29    traffic_class: u8,
30    /// Flow Label (20 bits)
31    flow_label: u32,
32    /// Next Header protocol number
33    next_header: u8,
34    /// Hop Limit (analogous to TTL in IPv4)
35    hop_limit: u8,
36    /// Source IPv6 address
37    src: Ipv6Addr,
38    /// Destination IPv6 address
39    dst: Ipv6Addr,
40    /// Payload bytes (appended after the 40-byte header)
41    payload: Vec<u8>,
42    /// If true, automatically compute payload_len from payload.len()
43    auto_length: bool,
44    /// Manual override for payload length (ignored if auto_length=true)
45    payload_len_override: Option<u16>,
46}
47
48impl Default for Ipv6Builder {
49    fn default() -> Self {
50        Self {
51            traffic_class: 0,
52            flow_label: 0,
53            next_header: 59, // No next header
54            hop_limit: 64,
55            src: Ipv6Addr::UNSPECIFIED,
56            dst: Ipv6Addr::UNSPECIFIED,
57            payload: Vec::new(),
58            auto_length: true,
59            payload_len_override: None,
60        }
61    }
62}
63
64impl Ipv6Builder {
65    /// Create a new IPv6 builder with default values.
66    pub fn new() -> Self {
67        Self::default()
68    }
69
70    // ========== Field Setters ==========
71
72    /// Set the source IPv6 address.
73    pub fn src(mut self, src: Ipv6Addr) -> Self {
74        self.src = src;
75        self
76    }
77
78    /// Set the destination IPv6 address.
79    pub fn dst(mut self, dst: Ipv6Addr) -> Self {
80        self.dst = dst;
81        self
82    }
83
84    /// Set the Hop Limit (equivalent to TTL in IPv4).
85    pub fn hop_limit(mut self, hlim: u8) -> Self {
86        self.hop_limit = hlim;
87        self
88    }
89
90    /// Set the Traffic Class field.
91    pub fn traffic_class(mut self, tc: u8) -> Self {
92        self.traffic_class = tc;
93        self
94    }
95
96    /// Set the Flow Label (20 bits; upper 12 bits are ignored).
97    pub fn flow_label(mut self, fl: u32) -> Self {
98        self.flow_label = fl & 0x000F_FFFF;
99        self
100    }
101
102    /// Set the Next Header field.
103    pub fn next_header(mut self, nh: u8) -> Self {
104        self.next_header = nh;
105        self
106    }
107
108    /// Set the payload bytes (appended after the 40-byte header).
109    pub fn payload<T: Into<Vec<u8>>>(mut self, data: T) -> Self {
110        self.payload = data.into();
111        self
112    }
113
114    /// Configure automatic payload length calculation (default: true).
115    pub fn auto_length(mut self, auto: bool) -> Self {
116        self.auto_length = auto;
117        self
118    }
119
120    /// Manually override the payload length field.
121    ///
122    /// Only used when auto_length is false.
123    pub fn payload_len(mut self, len: u16) -> Self {
124        self.payload_len_override = Some(len);
125        self.auto_length = false;
126        self
127    }
128
129    // ========== Size Calculation ==========
130
131    /// Get the header size (always 40 bytes for IPv6).
132    pub fn header_size(&self) -> usize {
133        IPV6_HEADER_LEN
134    }
135
136    /// Get the total packet size (header + payload).
137    pub fn packet_size(&self) -> usize {
138        IPV6_HEADER_LEN + self.payload.len()
139    }
140
141    // ========== Build Methods ==========
142
143    /// Build the IPv6 header + payload into a byte vector.
144    ///
145    /// # Byte layout:
146    /// - Byte 0: version (6) in high 4 bits | TC high nibble in low 4 bits
147    /// - Byte 1: TC low nibble in high 4 bits | FL bits 19-16 in low 4 bits
148    /// - Bytes 2-3: FL bits 15-0 (big-endian)
149    /// - Bytes 4-5: payload_len (big-endian)
150    /// - Byte 6: next_header
151    /// - Byte 7: hop_limit
152    /// - Bytes 8-23: source address (16 bytes)
153    /// - Bytes 24-39: destination address (16 bytes)
154    /// - Bytes 40+: payload
155    pub fn build(&self) -> Vec<u8> {
156        let total = self.packet_size();
157        let mut buf = vec![0u8; total];
158
159        // Byte 0: version (6) high nibble | TC high nibble low nibble
160        buf[0] = 0x60 | ((self.traffic_class >> 4) & 0x0F);
161
162        // Byte 1: TC low nibble high nibble | FL bits 19-16 low nibble
163        buf[1] = ((self.traffic_class & 0x0F) << 4) | (((self.flow_label >> 16) & 0x0F) as u8);
164
165        // Bytes 2-3: FL bits 15-0
166        buf[2] = ((self.flow_label >> 8) & 0xFF) as u8;
167        buf[3] = (self.flow_label & 0xFF) as u8;
168
169        // Bytes 4-5: payload length
170        let plen: u16 = if self.auto_length {
171            self.payload.len() as u16
172        } else {
173            self.payload_len_override
174                .unwrap_or(self.payload.len() as u16)
175        };
176        buf[4] = (plen >> 8) as u8;
177        buf[5] = (plen & 0xFF) as u8;
178
179        // Byte 6: next header
180        buf[6] = self.next_header;
181
182        // Byte 7: hop limit
183        buf[7] = self.hop_limit;
184
185        // Bytes 8-23: source address
186        buf[8..24].copy_from_slice(&self.src.octets());
187
188        // Bytes 24-39: destination address
189        buf[24..40].copy_from_slice(&self.dst.octets());
190
191        // Bytes 40+: payload
192        if !self.payload.is_empty() {
193            buf[40..40 + self.payload.len()].copy_from_slice(&self.payload);
194        }
195
196        buf
197    }
198}
199
200#[cfg(test)]
201mod tests {
202    use super::*;
203
204    #[test]
205    fn test_ipv6_builder_basic() {
206        let src = Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 1);
207        let dst = Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 2);
208
209        let bytes = Ipv6Builder::new()
210            .src(src)
211            .dst(dst)
212            .hop_limit(64)
213            .next_header(58)
214            .build();
215
216        assert_eq!(bytes.len(), 40);
217        // Version = 6
218        assert_eq!((bytes[0] >> 4) & 0x0F, 6);
219        // Next header = 58 (ICMPv6)
220        assert_eq!(bytes[6], 58);
221        // Hop limit = 64
222        assert_eq!(bytes[7], 64);
223        // Source address
224        let mut src_bytes = [0u8; 16];
225        src_bytes.copy_from_slice(&bytes[8..24]);
226        assert_eq!(Ipv6Addr::from(src_bytes), src);
227        // Destination address
228        let mut dst_bytes = [0u8; 16];
229        dst_bytes.copy_from_slice(&bytes[24..40]);
230        assert_eq!(Ipv6Addr::from(dst_bytes), dst);
231    }
232
233    #[test]
234    fn test_ipv6_builder_with_payload() {
235        let payload = b"hello world";
236        let bytes = Ipv6Builder::new()
237            .src(Ipv6Addr::LOCALHOST)
238            .dst(Ipv6Addr::LOCALHOST)
239            .next_header(59)
240            .payload(payload.as_ref())
241            .build();
242
243        assert_eq!(bytes.len(), 40 + payload.len());
244        // Payload length in header
245        let plen = u16::from_be_bytes([bytes[4], bytes[5]]);
246        assert_eq!(plen as usize, payload.len());
247        // Payload bytes at offset 40
248        assert_eq!(&bytes[40..], payload.as_ref());
249    }
250
251    #[test]
252    fn test_ipv6_traffic_class() {
253        let bytes = Ipv6Builder::new().traffic_class(0xAB).build();
254
255        // Byte 0: 0x60 | (0xAB >> 4 = 0x0A) = 0x6A
256        assert_eq!(bytes[0], 0x6A);
257        // Byte 1: (0xAB & 0x0F = 0x0B) << 4 = 0xB0
258        assert_eq!(bytes[1] & 0xF0, 0xB0);
259        // Reconstruct TC
260        let tc = ((bytes[0] & 0x0F) << 4) | ((bytes[1] >> 4) & 0x0F);
261        assert_eq!(tc, 0xAB);
262    }
263
264    #[test]
265    fn test_ipv6_flow_label() {
266        let bytes = Ipv6Builder::new().flow_label(0x12345).build();
267
268        // FL bits 19-16 in byte 1 low nibble
269        assert_eq!(bytes[1] & 0x0F, 0x01);
270        // FL bits 15-8 in byte 2
271        assert_eq!(bytes[2], 0x23);
272        // FL bits 7-0 in byte 3
273        assert_eq!(bytes[3], 0x45);
274        // Reconstruct FL
275        let fl = ((bytes[1] as u32 & 0x0F) << 16) | ((bytes[2] as u32) << 8) | (bytes[3] as u32);
276        assert_eq!(fl, 0x12345);
277    }
278
279    #[test]
280    fn test_ipv6_manual_payload_len() {
281        let bytes = Ipv6Builder::new()
282            .payload(vec![0u8; 20])
283            .payload_len(100) // Override with manual value
284            .build();
285
286        let plen = u16::from_be_bytes([bytes[4], bytes[5]]);
287        assert_eq!(plen, 100);
288    }
289
290    #[test]
291    fn test_ipv6_header_size() {
292        let builder = Ipv6Builder::new();
293        assert_eq!(builder.header_size(), 40);
294    }
295}