dns_parser_revived/
builder.rs

1use byteorder::{BigEndian, ByteOrder, WriteBytesExt};
2
3use crate::{Header, Opcode, QueryClass, QueryType, ResponseCode};
4
5/// Allows to build a DNS packet
6///
7/// Both query and answer packets may be built with this interface, although,
8/// much of functionality is not implemented yet.
9#[derive(Debug)]
10pub struct Builder {
11    buf: Vec<u8>,
12}
13
14impl Builder {
15    /// Creates a new query
16    ///
17    /// Initially all sections are empty. You're expected to fill
18    /// the questions section with `add_question`
19    pub fn new_query(id: u16, recursion: bool) -> Builder {
20        let mut buf = Vec::with_capacity(512);
21        let head = Header {
22            id,
23            query: true,
24            opcode: Opcode::StandardQuery,
25            authoritative: false,
26            truncated: false,
27            recursion_desired: recursion,
28            recursion_available: false,
29            authenticated_data: false,
30            checking_disabled: false,
31            response_code: ResponseCode::NoError,
32            questions: 0,
33            answers: 0,
34            nameservers: 0,
35            additional: 0,
36        };
37        buf.extend([0u8; 12].iter());
38        head.write(&mut buf[..12]);
39        Builder { buf }
40    }
41    /// Adds a question to the packet
42    ///
43    /// # Panics
44    ///
45    /// * Answers, nameservers or additional section has already been written
46    /// * There are already 65535 questions in the buffer.
47    /// * When name is invalid
48    pub fn add_question(
49        &mut self,
50        qname: &str,
51        prefer_unicast: bool,
52        qtype: QueryType,
53        qclass: QueryClass,
54    ) -> &mut Builder {
55        if &self.buf[6..12] != b"\x00\x00\x00\x00\x00\x00" {
56            panic!("Too late to add a question");
57        }
58        self.write_name(qname);
59        self.buf.write_u16::<BigEndian>(qtype as u16).unwrap();
60        let prefer_unicast: u16 = if prefer_unicast { 0x8000 } else { 0x0000 };
61        self.buf
62            .write_u16::<BigEndian>(qclass as u16 | prefer_unicast)
63            .unwrap();
64        let oldq = BigEndian::read_u16(&self.buf[4..6]);
65        if oldq == 65535 {
66            panic!("Too many questions");
67        }
68        BigEndian::write_u16(&mut self.buf[4..6], oldq + 1);
69        self
70    }
71    fn write_name(&mut self, name: &str) {
72        for part in name.split('.') {
73            assert!(part.len() < 63);
74            let ln = part.len() as u8;
75            self.buf.push(ln);
76            self.buf.extend(part.as_bytes());
77        }
78        self.buf.push(0);
79    }
80    /// Returns the final packet
81    ///
82    /// When packet is not truncated method returns `Ok(packet)`. If
83    /// packet is truncated the method returns `Err(packet)`. In both
84    /// cases the packet is fully valid.
85    ///
86    /// In the server implementation you may use
87    /// `x.build().unwrap_or_else(|x| x)`.
88    ///
89    /// In the client implementation it's probably unwise to send truncated
90    /// packet, as it doesn't make sense. Even panicking may be more
91    /// appropriate.
92    // TODO(tailhook) does the truncation make sense for TCP, and how
93    // to treat it for EDNS0?
94    pub fn build(mut self) -> Result<Vec<u8>, Vec<u8>> {
95        // TODO(tailhook) optimize labels
96        if self.buf.len() > 512 {
97            Header::set_truncated(&mut self.buf[..12]);
98            Err(self.buf)
99        } else {
100            Ok(self.buf)
101        }
102    }
103}
104
105#[cfg(test)]
106mod test {
107    use super::Builder;
108    use crate::QueryClass as QC;
109    use crate::QueryType as QT;
110
111    #[test]
112    fn build_query() {
113        let mut bld = Builder::new_query(1573, true);
114        bld.add_question("example.com", false, QT::A, QC::IN);
115        let result = b"\x06%\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\
116                      \x07example\x03com\x00\x00\x01\x00\x01";
117        assert_eq!(&bld.build().unwrap()[..], &result[..]);
118    }
119
120    #[test]
121    fn build_unicast_query() {
122        let mut bld = Builder::new_query(1573, true);
123        bld.add_question("example.com", true, QT::A, QC::IN);
124        let result = b"\x06%\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\
125                      \x07example\x03com\x00\x00\x01\x80\x01";
126        assert_eq!(&bld.build().unwrap()[..], &result[..]);
127    }
128
129    #[test]
130    fn build_srv_query() {
131        let mut bld = Builder::new_query(23513, true);
132        bld.add_question("_xmpp-server._tcp.gmail.com", false, QT::SRV, QC::IN);
133        let result = b"[\xd9\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\
134            \x0c_xmpp-server\x04_tcp\x05gmail\x03com\x00\x00!\x00\x01";
135        assert_eq!(&bld.build().unwrap()[..], &result[..]);
136    }
137}
dns_parser_revived/builder.rs

dns_parser_revived/
builder.rs
