1 use byteorder::{ByteOrder, BigEndian, WriteBytesExt};
2 
3 use {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.
9 #[derive(Debug)]
10 pub struct Builder {
11     buf: Vec<u8>,
12 }
13 
14 impl 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`
new_query(id: u16, recursion: bool) -> Builder19     pub fn new_query(id: u16, recursion: bool) -> Builder {
20         let mut buf = Vec::with_capacity(512);
21         let head = Header {
22             id: 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: 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
add_question(&mut self, qname: &str, prefer_unicast: bool, qtype: QueryType, qclass: QueryClass) -> &mut Builder48     pub fn add_question(&mut self, qname: &str, prefer_unicast: bool,
49         qtype: QueryType, qclass: QueryClass)
50         -> &mut Builder
51     {
52         if &self.buf[6..12] != b"\x00\x00\x00\x00\x00\x00" {
53             panic!("Too late to add a question");
54         }
55         self.write_name(qname);
56         self.buf.write_u16::<BigEndian>(qtype as u16).unwrap();
57         let prefer_unicast: u16 = if prefer_unicast { 0x8000 } else { 0x0000 };
58         self.buf.write_u16::<BigEndian>(qclass as u16 | prefer_unicast).unwrap();
59         let oldq = BigEndian::read_u16(&self.buf[4..6]);
60         if oldq == 65535 {
61             panic!("Too many questions");
62         }
63         BigEndian::write_u16(&mut self.buf[4..6], oldq+1);
64         self
65     }
write_name(&mut self, name: &str)66     fn write_name(&mut self, name: &str) {
67         for part in name.split('.') {
68             assert!(part.len() < 63);
69             let ln = part.len() as u8;
70             self.buf.push(ln);
71             self.buf.extend(part.as_bytes());
72         }
73         self.buf.push(0);
74     }
75     /// Returns the final packet
76     ///
77     /// When packet is not truncated method returns `Ok(packet)`. If
78     /// packet is truncated the method returns `Err(packet)`. In both
79     /// cases the packet is fully valid.
80     ///
81     /// In the server implementation you may use
82     /// `x.build().unwrap_or_else(|x| x)`.
83     ///
84     /// In the client implementation it's probably unwise to send truncated
85     /// packet, as it doesn't make sense. Even panicking may be more
86     /// appropriate.
87     // TODO(tailhook) does the truncation make sense for TCP, and how
88     // to treat it for EDNS0?
build(mut self) -> Result<Vec<u8>,Vec<u8>>89     pub fn build(mut self) -> Result<Vec<u8>,Vec<u8>> {
90         // TODO(tailhook) optimize labels
91         if self.buf.len() > 512 {
92             Header::set_truncated(&mut self.buf[..12]);
93             Err(self.buf)
94         } else {
95             Ok(self.buf)
96         }
97     }
98 }
99 
100 #[cfg(test)]
101 mod test {
102     use QueryType as QT;
103     use QueryClass as QC;
104     use super::Builder;
105 
106     #[test]
build_query()107     fn build_query() {
108         let mut bld = Builder::new_query(1573, true);
109         bld.add_question("example.com", false, QT::A, QC::IN);
110         let result = b"\x06%\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\
111                       \x07example\x03com\x00\x00\x01\x00\x01";
112         assert_eq!(&bld.build().unwrap()[..], &result[..]);
113     }
114 
115     #[test]
build_unicast_query()116     fn build_unicast_query() {
117         let mut bld = Builder::new_query(1573, true);
118         bld.add_question("example.com", true, QT::A, QC::IN);
119         let result = b"\x06%\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\
120                       \x07example\x03com\x00\x00\x01\x80\x01";
121         assert_eq!(&bld.build().unwrap()[..], &result[..]);
122     }
123 
124     #[test]
build_srv_query()125     fn build_srv_query() {
126         let mut bld = Builder::new_query(23513, true);
127         bld.add_question("_xmpp-server._tcp.gmail.com", false, QT::SRV, QC::IN);
128         let result = b"[\xd9\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\
129             \x0c_xmpp-server\x04_tcp\x05gmail\x03com\x00\x00!\x00\x01";
130         assert_eq!(&bld.build().unwrap()[..], &result[..]);
131     }
132 }
133