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