libmdns/dns_parser/
builder.rs

1use std::marker::PhantomData;
2
3use byteorder::{BigEndian, ByteOrder, WriteBytesExt};
4
5use super::{Header, Name, Opcode, QueryClass, QueryType, RRData, ResponseCode};
6
7pub enum Questions {}
8pub enum Answers {}
9#[allow(dead_code)]
10pub enum Nameservers {}
11pub enum Additional {}
12
13pub trait MoveTo<T> {}
14impl<T> MoveTo<T> for T {}
15
16impl MoveTo<Answers> for Questions {}
17
18impl MoveTo<Nameservers> for Questions {}
19impl MoveTo<Nameservers> for Answers {}
20
21impl MoveTo<Additional> for Questions {}
22impl MoveTo<Additional> for Answers {}
23impl MoveTo<Additional> for Nameservers {}
24
25/// Allows to build a DNS packet
26///
27/// Both query and answer packets may be built with this interface, although,
28/// much of functionality is not implemented yet.
29pub struct Builder<S> {
30    buf: Vec<u8>,
31    max_size: Option<usize>,
32    _state: PhantomData<S>,
33}
34
35impl Builder<Questions> {
36    /// Creates a new query
37    ///
38    /// Initially all sections are empty. You're expected to fill
39    /// the questions section with `add_question`
40    #[allow(dead_code)]
41    pub fn new_query(id: u16, recursion: bool) -> Builder<Questions> {
42        let mut buf = Vec::with_capacity(512);
43        let head = Header {
44            id,
45            query: true,
46            opcode: Opcode::StandardQuery,
47            authoritative: false,
48            truncated: false,
49            recursion_desired: recursion,
50            recursion_available: false,
51            response_code: ResponseCode::NoError,
52            questions: 0,
53            answers: 0,
54            nameservers: 0,
55            additional: 0,
56        };
57        buf.extend([0u8; 12].iter());
58        head.write(&mut buf[..12]);
59        Builder {
60            buf,
61            max_size: Some(512),
62            _state: PhantomData,
63        }
64    }
65
66    pub fn new_response(id: u16, recursion: bool, authoritative: bool) -> Builder<Questions> {
67        let mut buf = Vec::with_capacity(512);
68        let head = Header {
69            id,
70            query: false,
71            opcode: Opcode::StandardQuery,
72            authoritative,
73            truncated: false,
74            recursion_desired: recursion,
75            recursion_available: false,
76            response_code: ResponseCode::NoError,
77            questions: 0,
78            answers: 0,
79            nameservers: 0,
80            additional: 0,
81        };
82        buf.extend([0u8; 12].iter());
83        head.write(&mut buf[..12]);
84        Builder {
85            buf,
86            max_size: Some(512),
87            _state: PhantomData,
88        }
89    }
90}
91
92impl<T> Builder<T> {
93    fn write_rr(&mut self, name: &Name<'_>, cls: QueryClass, ttl: u32, data: &RRData<'_>) {
94        name.write_to(&mut self.buf).unwrap();
95        self.buf.write_u16::<BigEndian>(data.typ() as u16).unwrap();
96        self.buf.write_u16::<BigEndian>(cls as u16).unwrap();
97        self.buf.write_u32::<BigEndian>(ttl).unwrap();
98
99        let size_offset = self.buf.len();
100        self.buf.write_u16::<BigEndian>(0).unwrap();
101
102        let data_offset = self.buf.len();
103        data.write_to(&mut self.buf).unwrap();
104        let data_size = self.buf.len() - data_offset;
105
106        assert!(
107            (data_offset <= 65535),
108            "{} is too long to write to a RR record",
109            data_offset
110        );
111        #[allow(clippy::cast_possible_truncation)]
112        BigEndian::write_u16(
113            &mut self.buf[size_offset..size_offset + 2],
114            data_size as u16,
115        );
116    }
117
118    /// Returns the final packet
119    ///
120    /// When packet is not truncated method returns `Ok(packet)`. If
121    /// packet is truncated the method returns `Err(packet)`. In both
122    /// cases the packet is fully valid.
123    ///
124    /// In the server implementation you may use
125    /// `x.build().unwrap_or_else(|x| x)`.
126    ///
127    /// In the client implementation it's probably unwise to send truncated
128    /// packet, as it doesn't make sense. Even panicking may be more
129    /// appropriate.
130    // TODO(tailhook) does the truncation make sense for TCP, and how
131    // to treat it for EDNS0?
132    pub fn build(mut self) -> Result<Vec<u8>, Vec<u8>> {
133        // TODO(tailhook) optimize labels
134        match self.max_size {
135            Some(max_size) if self.buf.len() > max_size => {
136                Header::set_truncated(&mut self.buf[..12]);
137                Err(self.buf)
138            }
139            _ => Ok(self.buf),
140        }
141    }
142
143    pub fn move_to<U>(self) -> Builder<U>
144    where
145        T: MoveTo<U>,
146    {
147        Builder {
148            buf: self.buf,
149            max_size: self.max_size,
150            _state: PhantomData,
151        }
152    }
153
154    pub fn set_max_size(&mut self, max_size: Option<usize>) {
155        self.max_size = max_size;
156    }
157
158    pub fn is_empty(&self) -> bool {
159        Header::question_count(&self.buf) == 0
160            && Header::answer_count(&self.buf) == 0
161            && Header::nameserver_count(&self.buf) == 0
162            && Header::additional_count(&self.buf) == 0
163    }
164}
165
166impl<T: MoveTo<Questions>> Builder<T> {
167    /// Adds a question to the packet
168    ///
169    /// # Panics
170    ///
171    /// * There are already 65535 questions in the buffer.
172    #[allow(dead_code)]
173    pub fn add_question(
174        self,
175        qname: &Name<'_>,
176        qtype: QueryType,
177        qclass: QueryClass,
178    ) -> Builder<Questions> {
179        let mut builder = self.move_to::<Questions>();
180
181        qname.write_to(&mut builder.buf).unwrap();
182        builder.buf.write_u16::<BigEndian>(qtype as u16).unwrap();
183        builder.buf.write_u16::<BigEndian>(qclass as u16).unwrap();
184        Header::inc_questions(&mut builder.buf).expect("Too many questions");
185        builder
186    }
187}
188
189impl<T: MoveTo<Answers>> Builder<T> {
190    pub fn add_answer(
191        self,
192        name: &Name<'_>,
193        cls: QueryClass,
194        ttl: u32,
195        data: &RRData<'_>,
196    ) -> Builder<Answers> {
197        let mut builder = self.move_to::<Answers>();
198
199        builder.write_rr(name, cls, ttl, data);
200        Header::inc_answers(&mut builder.buf).expect("Too many answers");
201
202        builder
203    }
204}
205
206impl<T: MoveTo<Nameservers>> Builder<T> {
207    #[allow(dead_code)]
208    pub fn add_nameserver(
209        self,
210        name: &Name<'_>,
211        cls: QueryClass,
212        ttl: u32,
213        data: &RRData<'_>,
214    ) -> Builder<Nameservers> {
215        let mut builder = self.move_to::<Nameservers>();
216
217        builder.write_rr(name, cls, ttl, data);
218        Header::inc_nameservers(&mut builder.buf).expect("Too many nameservers");
219
220        builder
221    }
222}
223
224impl Builder<Additional> {
225    #[allow(dead_code)]
226    pub fn add_additional(
227        self,
228        name: &Name<'_>,
229        cls: QueryClass,
230        ttl: u32,
231        data: &RRData<'_>,
232    ) -> Builder<Additional> {
233        let mut builder = self.move_to::<Additional>();
234
235        builder.write_rr(name, cls, ttl, data);
236        Header::inc_nameservers(&mut builder.buf).expect("Too many additional answers");
237
238        builder
239    }
240}
241
242#[cfg(test)]
243mod test {
244    use super::Builder;
245    use super::Name;
246    use super::QueryClass as QC;
247    use super::QueryType as QT;
248
249    #[test]
250    fn build_query() {
251        let mut bld = Builder::new_query(1573, true);
252        let name = Name::from_str("example.com");
253        bld = bld.add_question(&name, QT::A, QC::IN);
254        let result = b"\x06%\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\
255                      \x07example\x03com\x00\x00\x01\x00\x01";
256        assert_eq!(&bld.build().unwrap()[..], &result[..]);
257    }
258
259    #[test]
260    fn build_srv_query() {
261        let mut bld = Builder::new_query(23513, true);
262        let name = Name::from_str("_xmpp-server._tcp.gmail.com");
263        bld = bld.add_question(&name, QT::SRV, QC::IN);
264        let result = b"[\xd9\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\
265            \x0c_xmpp-server\x04_tcp\x05gmail\x03com\x00\x00!\x00\x01";
266        assert_eq!(&bld.build().unwrap()[..], &result[..]);
267    }
268}
libmdns/dns_parser/builder.rs

libmdns/dns_parser/
builder.rs
