stackforge_core/layer/dns/
builder.rs

1//! DNS packet builder with fluent API.
2//!
3//! Provides a builder pattern for constructing DNS packets,
4//! similar to Scapy's `DNS()` constructor.
5
6use std::collections::HashMap;
7
8use super::header;
9use super::query::DnsQuestion;
10use super::rr::DnsResourceRecord;
11use super::types;
12
13/// Builder for constructing DNS packets.
14#[derive(Debug, Clone)]
15pub struct DnsBuilder {
16    /// Transaction ID.
17    pub id: u16,
18    /// Query/Response flag.
19    pub qr: bool,
20    /// Operation code.
21    pub opcode: u8,
22    /// Authoritative Answer flag.
23    pub aa: bool,
24    /// Truncation flag.
25    pub tc: bool,
26    /// Recursion Desired flag.
27    pub rd: bool,
28    /// Recursion Available flag.
29    pub ra: bool,
30    /// Reserved (Z) flag.
31    pub z: bool,
32    /// Authenticated Data flag (DNSSEC).
33    pub ad: bool,
34    /// Checking Disabled flag (DNSSEC).
35    pub cd: bool,
36    /// Response code.
37    pub rcode: u8,
38    /// Question section.
39    pub questions: Vec<DnsQuestion>,
40    /// Answer section.
41    pub answers: Vec<DnsResourceRecord>,
42    /// Authority section.
43    pub authorities: Vec<DnsResourceRecord>,
44    /// Additional section.
45    pub additionals: Vec<DnsResourceRecord>,
46    /// Whether to use DNS name compression when building.
47    pub compress: bool,
48}
49
50impl DnsBuilder {
51    /// Create a new DNS builder with default values (standard query with RD=1).
52    pub fn new() -> Self {
53        Self {
54            id: 0,
55            qr: false,
56            opcode: types::opcode::QUERY,
57            aa: false,
58            tc: false,
59            rd: true,
60            ra: false,
61            z: false,
62            ad: false,
63            cd: false,
64            rcode: types::rcode::NOERROR,
65            questions: Vec::new(),
66            answers: Vec::new(),
67            authorities: Vec::new(),
68            additionals: Vec::new(),
69            compress: true,
70        }
71    }
72
73    /// Create a builder for a standard query.
74    pub fn query(qname: &str, qtype: u16) -> Self {
75        let mut b = Self::new();
76        if let Ok(q) = DnsQuestion::from_name(qname) {
77            let mut q = q;
78            q.qtype = qtype;
79            b.questions.push(q);
80        }
81        b
82    }
83
84    /// Create a builder for a standard response.
85    pub fn response() -> Self {
86        let mut b = Self::new();
87        b.qr = true;
88        b.ra = true;
89        b
90    }
91
92    // Fluent setters
93
94    pub fn id(mut self, id: u16) -> Self {
95        self.id = id;
96        self
97    }
98
99    pub fn qr(mut self, qr: bool) -> Self {
100        self.qr = qr;
101        self
102    }
103
104    pub fn opcode(mut self, opcode: u8) -> Self {
105        self.opcode = opcode;
106        self
107    }
108
109    pub fn aa(mut self, aa: bool) -> Self {
110        self.aa = aa;
111        self
112    }
113
114    pub fn tc(mut self, tc: bool) -> Self {
115        self.tc = tc;
116        self
117    }
118
119    pub fn rd(mut self, rd: bool) -> Self {
120        self.rd = rd;
121        self
122    }
123
124    pub fn ra(mut self, ra: bool) -> Self {
125        self.ra = ra;
126        self
127    }
128
129    pub fn z(mut self, z: bool) -> Self {
130        self.z = z;
131        self
132    }
133
134    pub fn ad(mut self, ad: bool) -> Self {
135        self.ad = ad;
136        self
137    }
138
139    pub fn cd(mut self, cd: bool) -> Self {
140        self.cd = cd;
141        self
142    }
143
144    pub fn rcode(mut self, rcode: u8) -> Self {
145        self.rcode = rcode;
146        self
147    }
148
149    pub fn compress(mut self, compress: bool) -> Self {
150        self.compress = compress;
151        self
152    }
153
154    /// Add a question to the question section.
155    pub fn question(mut self, q: DnsQuestion) -> Self {
156        self.questions.push(q);
157        self
158    }
159
160    /// Add a resource record to the answer section.
161    pub fn answer(mut self, rr: DnsResourceRecord) -> Self {
162        self.answers.push(rr);
163        self
164    }
165
166    /// Add a resource record to the authority section.
167    pub fn authority(mut self, rr: DnsResourceRecord) -> Self {
168        self.authorities.push(rr);
169        self
170    }
171
172    /// Add a resource record to the additional section.
173    pub fn additional(mut self, rr: DnsResourceRecord) -> Self {
174        self.additionals.push(rr);
175        self
176    }
177
178    /// Build the DNS packet bytes.
179    pub fn build(&self) -> Vec<u8> {
180        if self.compress {
181            self.build_compressed()
182        } else {
183            self.build_uncompressed()
184        }
185    }
186
187    /// Build without name compression.
188    fn build_uncompressed(&self) -> Vec<u8> {
189        let mut out = Vec::with_capacity(512);
190
191        // Header (12 bytes)
192        out.extend_from_slice(&self.id.to_be_bytes());
193        let flags = header::build_flags(
194            self.qr,
195            self.opcode,
196            self.aa,
197            self.tc,
198            self.rd,
199            self.ra,
200            self.z,
201            self.ad,
202            self.cd,
203            self.rcode,
204        );
205        out.extend_from_slice(&flags.to_be_bytes());
206        out.extend_from_slice(&(self.questions.len() as u16).to_be_bytes());
207        out.extend_from_slice(&(self.answers.len() as u16).to_be_bytes());
208        out.extend_from_slice(&(self.authorities.len() as u16).to_be_bytes());
209        out.extend_from_slice(&(self.additionals.len() as u16).to_be_bytes());
210
211        // Questions
212        for q in &self.questions {
213            out.extend_from_slice(&q.build());
214        }
215
216        // Answers
217        for rr in &self.answers {
218            out.extend_from_slice(&rr.build());
219        }
220
221        // Authorities
222        for rr in &self.authorities {
223            out.extend_from_slice(&rr.build());
224        }
225
226        // Additionals
227        for rr in &self.additionals {
228            out.extend_from_slice(&rr.build());
229        }
230
231        out
232    }
233
234    /// Build with name compression.
235    fn build_compressed(&self) -> Vec<u8> {
236        let mut out = Vec::with_capacity(512);
237        let mut compression_map: HashMap<String, u16> = HashMap::new();
238
239        // Header (12 bytes)
240        out.extend_from_slice(&self.id.to_be_bytes());
241        let flags = header::build_flags(
242            self.qr,
243            self.opcode,
244            self.aa,
245            self.tc,
246            self.rd,
247            self.ra,
248            self.z,
249            self.ad,
250            self.cd,
251            self.rcode,
252        );
253        out.extend_from_slice(&flags.to_be_bytes());
254        out.extend_from_slice(&(self.questions.len() as u16).to_be_bytes());
255        out.extend_from_slice(&(self.answers.len() as u16).to_be_bytes());
256        out.extend_from_slice(&(self.authorities.len() as u16).to_be_bytes());
257        out.extend_from_slice(&(self.additionals.len() as u16).to_be_bytes());
258
259        // Questions
260        for q in &self.questions {
261            let encoded = q.build_compressed(out.len(), &mut compression_map);
262            out.extend_from_slice(&encoded);
263        }
264
265        // Answers
266        for rr in &self.answers {
267            let encoded = rr.build_compressed(out.len(), &mut compression_map);
268            out.extend_from_slice(&encoded);
269        }
270
271        // Authorities
272        for rr in &self.authorities {
273            let encoded = rr.build_compressed(out.len(), &mut compression_map);
274            out.extend_from_slice(&encoded);
275        }
276
277        // Additionals
278        for rr in &self.additionals {
279            let encoded = rr.build_compressed(out.len(), &mut compression_map);
280            out.extend_from_slice(&encoded);
281        }
282
283        out
284    }
285
286    /// Get the minimum header size.
287    pub fn header_size(&self) -> usize {
288        header::DNS_HEADER_LEN
289    }
290}
291
292impl Default for DnsBuilder {
293    fn default() -> Self {
294        Self::new()
295    }
296}
297
298#[cfg(test)]
299mod tests {
300    use super::super::types::rr_type;
301    use super::*;
302    use crate::layer::field_ext::DnsName;
303
304    #[test]
305    fn test_builder_default_query() {
306        let b = DnsBuilder::new();
307        assert!(!b.qr);
308        assert!(b.rd);
309        assert_eq!(b.opcode, 0);
310        assert_eq!(b.rcode, 0);
311        assert!(b.questions.is_empty());
312    }
313
314    #[test]
315    fn test_builder_query_shortcut() {
316        let b = DnsBuilder::query("example.com", rr_type::A);
317        assert_eq!(b.questions.len(), 1);
318        assert_eq!(b.questions[0].qtype, rr_type::A);
319        assert_eq!(b.questions[0].qname.labels, vec!["example", "com"]);
320    }
321
322    #[test]
323    fn test_builder_fluent_api() {
324        let b = DnsBuilder::new()
325            .id(0x1234)
326            .qr(true)
327            .aa(true)
328            .rd(true)
329            .ra(true)
330            .rcode(0);
331        assert_eq!(b.id, 0x1234);
332        assert!(b.qr);
333        assert!(b.aa);
334    }
335
336    #[test]
337    fn test_builder_build_simple_query() {
338        let b = DnsBuilder::query("example.com", rr_type::A).id(0x1234);
339        let packet = b.build();
340
341        // Verify header
342        assert_eq!(packet.len() >= 12, true);
343        assert_eq!(u16::from_be_bytes([packet[0], packet[1]]), 0x1234); // ID
344        let qdcount = u16::from_be_bytes([packet[4], packet[5]]);
345        assert_eq!(qdcount, 1);
346    }
347
348    #[test]
349    fn test_builder_build_uncompressed() {
350        let b = DnsBuilder::query("example.com", rr_type::A)
351            .id(0x5678)
352            .compress(false);
353        let packet = b.build();
354        assert!(packet.len() >= 12 + 4 + 13); // header + type/class + "example.com" encoded
355    }
356
357    #[test]
358    fn test_builder_compression_reduces_size() {
359        let q1 = DnsQuestion::from_name("www.example.com").unwrap();
360        let q2 = DnsQuestion::from_name("mail.example.com").unwrap();
361        let b = DnsBuilder::new()
362            .question(q1.clone())
363            .question(q2.clone())
364            .compress(true);
365        let compressed = b.build();
366
367        let b2 = DnsBuilder::new().question(q1).question(q2).compress(false);
368        let uncompressed = b2.build();
369
370        assert!(compressed.len() < uncompressed.len());
371    }
372
373    #[test]
374    fn test_builder_response() {
375        let b = DnsBuilder::response();
376        assert!(b.qr);
377        assert!(b.ra);
378    }
379}
stackforge_core/layer/dns/builder.rs

stackforge_core/layer/dns/
builder.rs
