wireforge_app/dns/
builder.rs

1//! DNS message builder (RFC 1035).
2
3use super::name::encode_name;
4use super::parser::DNS_HEADER_LEN;
5use super::types::{DnsClass, DnsRcode, DnsType};
6use wireforge_core::util::write_u16be;
7
8/// Builder for DNS messages.
9pub struct DnsPacketBuilder {
10    header: Vec<u8>,
11    questions: Vec<Vec<u8>>,
12    answers: Vec<Vec<u8>>,
13    authorities: Vec<Vec<u8>>,
14    additionals: Vec<Vec<u8>>,
15}
16
17impl Default for DnsPacketBuilder {
18    fn default() -> Self { Self::new() }
19}
20
21impl DnsPacketBuilder {
22    pub fn new() -> Self {
23        Self {
24            header: vec![0u8; DNS_HEADER_LEN],
25            questions: Vec::new(),
26            answers: Vec::new(),
27            authorities: Vec::new(),
28            additionals: Vec::new(),
29        }
30    }
31
32    /// Preset for a standard query (QR=0, RD=1).
33    pub fn query() -> Self {
34        let mut s = Self::new();
35        s.header[2] = 0x01; // RD=1
36        s
37    }
38
39    /// Preset for a response (QR=1, RA=1).
40    pub fn response() -> Self {
41        let mut s = Self::new();
42        s.header[2] = 0x80; // QR=1
43        s.header[3] = 0x80; // RA=1
44        s
45    }
46
47    pub fn transaction_id(mut self, id: u16) -> Self {
48        write_u16be(&mut self.header[..2], id);
49        self
50    }
51
52    pub fn authoritative(mut self, aa: bool) -> Self {
53        if aa { self.header[2] |= 0x04; } else { self.header[2] &= !0x04; }
54        self
55    }
56
57    pub fn truncated(mut self, tc: bool) -> Self {
58        if tc { self.header[2] |= 0x02; } else { self.header[2] &= !0x02; }
59        self
60    }
61
62    pub fn recursion_desired(mut self, rd: bool) -> Self {
63        if rd { self.header[2] |= 0x01; } else { self.header[2] &= !0x01; }
64        self
65    }
66
67    pub fn recursion_available(mut self, ra: bool) -> Self {
68        if ra { self.header[3] |= 0x80; } else { self.header[3] &= !0x80; }
69        self
70    }
71
72    pub fn rcode(mut self, rcode: DnsRcode) -> Self {
73        let v: u8 = match rcode {
74            DnsRcode::NoError => 0,
75            DnsRcode::FormatError => 1,
76            DnsRcode::ServerFailure => 2,
77            DnsRcode::NameError => 3,
78            DnsRcode::NotImplemented => 4,
79            DnsRcode::Refused => 5,
80            DnsRcode::Unknown(v) => v,
81        };
82        self.header[3] = (self.header[3] & 0xF0) | (v & 0x0F);
83        self
84    }
85
86    /// Add a question entry.
87    pub fn add_question(mut self, name: &str, qtype: DnsType, qclass: DnsClass) -> Self {
88        let mut entry = encode_name(name);
89        let qt: u16 = qtype.into();
90        let qc: u16 = qclass.into();
91        entry.extend_from_slice(&qt.to_be_bytes());
92        entry.extend_from_slice(&qc.to_be_bytes());
93        self.questions.push(entry);
94        self
95    }
96
97    /// Add an answer record (pre-encoded RDATA).
98    pub fn add_raw_answer(mut self, name: &str, qtype: DnsType, ttl: u32, rdata: &[u8]) -> Self {
99        let mut entry = encode_name(name);
100        let qt: u16 = qtype.into();
101        entry.extend_from_slice(&qt.to_be_bytes());
102        entry.extend_from_slice(&1u16.to_be_bytes()); // class IN
103        entry.extend_from_slice(&ttl.to_be_bytes());
104        entry.extend_from_slice(&(rdata.len() as u16).to_be_bytes());
105        entry.extend_from_slice(rdata);
106        self.answers.push(entry);
107        self
108    }
109
110    /// Add an authority record.
111    pub fn add_raw_authority(mut self, name: &str, qtype: DnsType, ttl: u32, rdata: &[u8]) -> Self {
112        let mut entry = encode_name(name);
113        let qt: u16 = qtype.into();
114        entry.extend_from_slice(&qt.to_be_bytes());
115        entry.extend_from_slice(&1u16.to_be_bytes());
116        entry.extend_from_slice(&ttl.to_be_bytes());
117        entry.extend_from_slice(&(rdata.len() as u16).to_be_bytes());
118        entry.extend_from_slice(rdata);
119        self.authorities.push(entry);
120        self
121    }
122
123    /// Add an additional record.
124    pub fn add_raw_additional(mut self, name: &str, qtype: DnsType, ttl: u32, rdata: &[u8]) -> Self {
125        let mut entry = encode_name(name);
126        let qt: u16 = qtype.into();
127        entry.extend_from_slice(&qt.to_be_bytes());
128        entry.extend_from_slice(&1u16.to_be_bytes());
129        entry.extend_from_slice(&ttl.to_be_bytes());
130        entry.extend_from_slice(&(rdata.len() as u16).to_be_bytes());
131        entry.extend_from_slice(rdata);
132        self.additionals.push(entry);
133        self
134    }
135
136    /// Build the complete DNS message.
137    pub fn build(mut self) -> Vec<u8> {
138        write_u16be(&mut self.header[4..6], self.questions.len() as u16);
139        write_u16be(&mut self.header[6..8], self.answers.len() as u16);
140        write_u16be(&mut self.header[8..10], self.authorities.len() as u16);
141        write_u16be(&mut self.header[10..12], self.additionals.len() as u16);
142
143        let mut packet = self.header;
144        for q in &self.questions {
145            packet.extend_from_slice(q);
146        }
147        for a in &self.answers {
148            packet.extend_from_slice(a);
149        }
150        for a in &self.authorities {
151            packet.extend_from_slice(a);
152        }
153        for a in &self.additionals {
154            packet.extend_from_slice(a);
155        }
156        packet
157    }
158}
wireforge_app/dns/builder.rs

wireforge_app/dns/
builder.rs
