dns_parser_joe/
builder.rs

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