1use byteorder::{ByteOrder, BigEndian, WriteBytesExt};
2
3use {Opcode, ResponseCode, Header, QueryType, QueryClass};
4
5#[derive(Debug)]
10pub struct Builder {
11 buf: Vec<u8>,
12}
13
14impl Builder {
15 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 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 }
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 pub fn build(mut self) -> Result<Vec<u8>,Vec<u8>> {
90 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)]
101mod test {
102 use QueryType as QT;
103 use QueryClass as QC;
104 use super::Builder;
105
106 #[test]
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]
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]
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}