From bb0cee2ffe1346638b0d45bda5a80fc95c3dc2c7 Mon Sep 17 00:00:00 2001 From: Martin Hoffmann Date: Sun, 22 Nov 2015 14:49:17 +0100 Subject: [PATCH] Add message assembling. --- src/egress.rs | 670 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 670 insertions(+) create mode 100644 src/egress.rs diff --git a/src/egress.rs b/src/egress.rs new file mode 100644 index 00000000..11f85ba5 --- /dev/null +++ b/src/egress.rs @@ -0,0 +1,670 @@ +//! Construction of DNS messages for sending +//! + +use std::collections::HashMap; +use std::io; +use std::mem; +use std::net; +use std::ptr; +use super::header::Header; +use super::name; +use super::record::{RecordData, RecordDataAssembly}; + +//------------ Assembly ----------------------------------------------------- + +/// A type to assemble an outgoing DNS message. +/// +/// When assembling messages you sadly need to know what is going on later +/// or come back to parts written earlier. The Assembly type does exactly +/// that. +/// +/// Use this type directly if you already know what to put into the message +/// before starting to create it. Otherwise, use `Message` which allows you +/// put things together step by step. +/// +#[derive(Debug)] +pub struct Assembly { + // The message's data. + inner: Vec, + + // If we are writing to a stream, the message actually starts with two + // octets for the length, but references for name pointers and such + // are relative to after that. To make the math easier, we simply keep + // the index of the real start of the message in here. + origin: usize, + + // Maximum size of the resulting package (and thus, `self.inner`). + // + // This is limited to 65535 due either to the maximum UDP datagram + // size or the 16 bit length marker in a stream. We still use usize + // here because with origin we can overrun the 16 bit boundary. + maxlen: usize, + + // If we overrun our maxlen while writing, we set this flag. The + // higher level methods know what to do in this case. + tc: bool, + + // We need a copy of the header because we may have to modify the + // tc bit when we are done. + header: Header, + + // If we do compression, then we will store labels in this here map. + // If we don't do compression, there is no map. + compress: Option>, +} + +// Management +// +impl Assembly { + + /// Creates a new assembly. + /// + /// The maximum package size is given in `maxlen`. If more data is + /// pushed to the assembly, it will be truncated at a record boundary + /// and the TC bit will be set in the header. + /// + /// If the assembly is for a stream, set `for_stream` to true. In this + /// case the first two octets of the resulting data will be the size + /// of the message in network order. + /// + /// Set `compress` to false, if you want domain names never to be + /// compressed. + /// + pub fn new(maxlen: u16, for_stream: bool, compress: bool) -> Assembly { + let map = match compress { + true => Some(HashMap::new()), + false => None + }; + if for_stream { + let mut res = Assembly { inner: Vec::new(), + maxlen: (maxlen as usize) + 2, + origin: 2, tc: false, + header: Header::new(), + compress: map }; + res.push_u16(0); + res + } + else { + Assembly { inner: Vec::new(), maxlen: maxlen as usize, origin: 0, + tc: false, header: Header::new(), compress: map } + } + } +} + +// Low level writing -- use this in record data implementations. +// +impl Assembly { + + /// Pushes a u8 to the end of the assembly. + /// + pub fn push_u8(&mut self, data: u8) { + if self.keep_pushing(1) { + self.inner.push(data) + } + } + + /// Pushes a u16 to the end of the assembly. + /// + pub fn push_u16(&mut self, data: u16) { + if self.keep_pushing(2) { + let data = data.to_be(); + let bytes: [u8; 2] = unsafe { mem::transmute(data) }; + self.inner.extend(&bytes); + } + } + + /// Pushes a u32 to the end of the assembly. + /// + pub fn push_u32(&mut self, data: u32) { + if self.keep_pushing(4) { + let data = data.to_be(); + let bytes: [u8; 4] = unsafe { mem::transmute(data) }; + self.inner.extend(&bytes); + } + } + + /// Pushes `data` to the end of the assembly. + /// + pub fn push_bytes(&mut self, data: &[u8]) { + if self.keep_pushing(data.len()) { + self.inner.extend(data); + } + } + + /// Pushes a domain name to the end of the assembly. + /// + /// The domain name will not be compressed, even if compression is + /// enabled for the assembly. This is the default because RFC 3597 + /// limits name compression to those record types intially defined in + /// RFC 1035. So, if you implement a `RecordDataAssembly`, you likely + /// want to use uncompressed name pushage. + /// + pub fn push_name(&mut self, name: &name::DomainName) { + assert!(!name.is_empty()); + + if self.compress.is_none() { + self._push_name_simple(name); + } + else { + self._push_name_uncompressed(name); + } + } + + /// Pushes a domain name to the end of the assembly employing compression. + /// + /// This is similar to `push_name()` except that name compression will + /// be used if it has been enabled for the assembly. + /// + pub fn push_name_compressed(&mut self, name: &name::DomainName) { + assert!(!name.is_empty()); + + if self.compress.is_none() { + self._push_name_simple(name); + } + else { + self._push_name_compressed(name); + } + } + + // Simply pushes `name`'s bytes. Luckily, our internal domain name + // encoding is actually the wire format. Smart, huh?! + fn _push_name_simple(&mut self, name: &name::DomainName) { + if self.keep_pushing(name.as_bytes().len()) { + self.push_bytes(name.as_bytes()); + } + } + + // Pushes `name` storing the label positions for later referencing + // if compression is enabled. + fn _push_name_uncompressed(&mut self, name: &name::DomainName) { + let mut name = name; + while !name.is_empty() { + name = self._push_label(name); + } + } + + // Pushes `name`, using a reference for the longest known tail. + fn _push_name_compressed(&mut self, name: &name::DomainName) { + let mut name = name; + while !name.is_empty() { + let pos; + { + pos = match self.compress { + None => None, + Some(ref map) => { + match map.get(name) { + None => None, + Some(&pos) => Some(pos) + } + } + } + } + match pos { + None => name = self._push_label(name), + Some(pos) => { + if self.keep_pushing(2) { + self.push_u16(pos | 0xC000); + } + return + } + } + } + } + + // Pushes the first label in `name`, store a reference to `name` for + // later use in compression. + fn _push_label<'a>(&mut self, name: &'a name::DomainName) + -> &'a name::DomainName + { + let pos = self.pos(); + match self.compress { + None => { }, + Some(ref mut map) => { + map.insert(name.to_owned(), (pos - self.origin) as u16); + } + } + let (head, tail) = name.split_first().unwrap(); + assert!(head.len() < 64); + if self.keep_pushing(head.as_bytes().len() + 1) { + self.push_u8(head.len() as u8); + self.push_bytes(head.as_bytes()); + } + tail + } + + // Returns whether it is fine to push `len` octets. + // + fn keep_pushing(&mut self, len: usize) -> bool { + if self.inner.len() + len > self.maxlen { + self.tc = true; + false + } + else { true } + } + + /// Returns the current position of the assembly. + /// + /// You must use the returned position for later use in the replace + /// methods only. Do not assign any other meaning to it. + /// + pub fn pos(&self) -> usize { + self.inner.len() + } + + /// Replaces the u8 at position `pos` with `data`. + /// + pub fn replace_u8(&mut self, pos: usize, data: u8) { + assert!(pos < self.inner.len()); + self.inner[pos] = data; + } + + /// Replaces the u16 beginning at position `pos` with `data`. + /// + pub fn replace_u16(&mut self, pos: usize, data: u16) { + let data = data.to_be(); + assert!(pos + 1 < self.inner.len()); + unsafe { + let src: [u8; 2] = mem::transmute(data); + ptr::copy_nonoverlapping(&src.as_ptr(), + &mut self.inner[pos .. pos + 2].as_ptr(), + 2); + } + } + + /// Replaces the u32 beginning at position `pos` with `data`. + /// + pub fn replace_u32(&mut self, pos: usize, data: u32) { + let data = data.to_be(); + assert!(pos + 1 < self.inner.len()); + unsafe { + let src: [u8; 4] = mem::transmute(data); + ptr::copy_nonoverlapping(&src.as_ptr(), + &mut self.inner[pos .. pos + 4].as_ptr(), + 4); + } + } +} + +// High-level writing +// +impl Assembly { + /// Pushes the message header to the assembly. + /// + /// A call to this method must be the first thing on a new assembly. + /// Later there must not be any more calls to it on threat of a panic. + /// + /// The arguments are the actual header, followed by the number of + /// elements in each seection. You must follow up this call with the + /// correct number of calls to `push_question()` and `push_rr()` or + /// `push_raw_rr()`. Finally, you must finish with calling `finish()`. + /// + pub fn push_header(&mut self, header: &Header, qdcount: u16, ancount: u16, + nscount: u16, arcount: u16) + { + assert!(self.pos() == self.origin); + self.header = header.clone(); + self.header.set_tc(false); + let header_bytes = self.header.as_u32(); + self.push_u32(header_bytes); + self.push_u16(qdcount); + self.push_u16(ancount); + self.push_u16(nscount); + self.push_u16(arcount); + } + + /// Finishes the assembly. + /// + /// This must be the last thing you do with the assembly. To make this + /// clear, the method moves the inner vec back to the caller. This + /// returned vec can be sent over the wire as is. + /// + pub fn finish(mut self) -> Vec { + if self.origin == 2 { + let len = (self.inner.len() - self.origin) as u16; + self.replace_u16(0, len); + } + if self.tc { + self.header.set_tc(true); + let pos = self.origin; + let header_bytes = self.header.as_u32(); + self.replace_u32(pos, header_bytes); + } + self.inner + } + + /// Pushes a question to the assembly. + /// + /// If doing so would exceed the maximum length, nothing will be + /// written, the TC bit will be set in the header later, and false + /// will be returned now. + /// + /// If all goes well, true will be returned. + /// + pub fn push_question(&mut self, qname: &name::DomainName, qtype: u16, + qclass: u16) -> bool + { + if self.tc { return false } + let pos = self.pos(); + self.push_name_compressed(qname); + self.push_u16(qtype); + self.push_u16(qclass); + if self.tc { + self.inner.truncate(pos); + } + !self.tc + } + + /// Pushes a raw resource record to the assembly. + /// + /// For our purposes, a raw resource record is one for which the type + /// is explicitely given and record data is already given in its wire + /// format. + /// + /// Returns true if the entire record could be pushed to the assembly. + /// If this would exceed the maximum length, nothing is pushed at all + /// and false is returned. Upon finishing the assembly, the TC bit will + /// be set in the header. + /// + pub fn push_raw_rr(&mut self, name: &name::DomainName, rtype: u16, + rclass: u16, ttl: u32, rdata: &[u8]) -> bool + { + assert!(rdata.len() <= ::std::u16::MAX as usize); + if self.tc { return false } + let pos = self.pos(); + self.push_name_compressed(name); + self.push_u16(rtype); + self.push_u16(rclass); + self.push_u32(ttl); + self.push_u16(rdata.len() as u16); + self.push_bytes(rdata); + if self.tc { + self.inner.truncate(pos); + } + !self.tc + } + + /// Pushes a resource record to the assembly. + /// + /// The record data is being assembled by calling `assembly()` on the + /// trait object `rdata`. If pushing + /// the record would exceed the maximum length, it is removed again and + /// the TC bit will later be set in the header when the assembly is + /// being finished. The method returns whether this did not happen. + /// + pub fn push_rr(&mut self, name: &name::DomainName, rtype: u16, + rclass: u16, ttl: u32, + rdata: &RecordDataAssembly) -> bool + { + if self.tc { return false } + let pos = self.pos(); + self.push_name_compressed(name); + self.push_u16(rtype); + self.push_u16(rclass); + self.push_u32(ttl); + let rdlen_pos = self.pos(); + self.push_u16(0); + rdata.assemble(self); + if self.tc { + self.inner.truncate(pos); + } + else { + let rdlen = self.pos() - rdlen_pos; + self.replace_u16(rdlen_pos, rdlen as u16); + } + !self.tc + } + +} + +//------------ Message ------------------------------------------------------ + +/// An outgoing DNS message. +/// +/// The message consists of four sections that are vecs over questions or +/// resource records, keeping their own data. +/// +/// Once you have all data collected, you can have the actual wire data +/// assembled or even sent out directly. +/// +#[derive(Debug)] +pub struct Message { + header: Header, + question: QuestionSection, + answer: RecordSection, + authority: RecordSection, + additional: RecordSection +} + +impl Message { + /// Creates a new, empty message. + /// + pub fn new() -> Message { + Message { header: Header::new(), question: QuestionSection::new(), + answer: RecordSection::new(), + authority: RecordSection::new(), + additional: RecordSection::new() } + } + + /// Gives access to the message header. + /// + pub fn header(&mut self) -> &mut Header { + &mut self.header + } + + /// Gives access to the question section. + /// + /// The question section contains the question being asked in a query. + /// + pub fn question(&mut self) -> &mut QuestionSection { + &mut self.question + } + + /// Gives access to the zone section of an Update query. + /// + /// In an update query, the question section is actually the zone + /// section specifing the zone to be updated. + /// + pub fn zone(&mut self) -> &mut QuestionSection { + &mut self.question + } + + /// Gives access to the answer section. + /// + /// The answer sections contains those resource records that answer + /// the question. + /// + /// In an Update query, this + /// + pub fn answer(&mut self) -> &mut RecordSection { + &mut self.answer + } + + /// Gives access to the prerequiste section of an Update query. + /// + /// The prerequisite section contains resource records or resource + /// record sets that must or must not preexist. + /// + pub fn prerequisite(&mut self) -> &mut RecordSection { + &mut self.answer + } + + /// Gives access to the authority section. + /// + /// The authority section contains resource records identifying the + /// authoritative name servers responsible for answering the question. + /// + pub fn authority(&mut self) -> &mut RecordSection { + &mut self.authority + } + + /// Gives access to the update section of an Update query. + /// + /// The update section contains resource records or resource record + /// sets to be added or delete. + /// + pub fn update(&mut self) -> &mut RecordSection { + &mut self.authority + } + + /// Gives access to the additional section. + /// + /// The additional section contains resource records that help making + /// sense of the answer. What exactly ought to be in here depends on + /// the record type being asked for. + /// + pub fn additional(&mut self) -> &mut RecordSection { + &mut self.additional + } + + //--- Finish it off ... + + /// Assembles the message. + /// + /// Returns a vector containing the wire data of the message and whether + /// the wire message had to be truncated because it would have been + /// longer than `maxlen`. + /// + /// If you plan on writing the data into a stream, such as a + /// `TcpStream`, set `for_stream` to true. In this case, the resulting + /// data will start with the required length indication. If you wish + /// domain names to be compressed where applicable, set `compress` to + /// true. + pub fn assemble(&self, maxlen: u16, for_stream: bool, compress: bool) + -> (Vec, bool) + { + let mut asm = Assembly::new(maxlen, for_stream, compress); + asm.push_header(&self.header, self.question.len(), + self.answer.len(), self.authority.len(), + self.additional.len()); + let complete = self.question.assemble(&mut asm) + && self.answer.assemble(&mut asm) + && self.authority.assemble(&mut asm) + && self.additional.assemble(&mut asm); + (asm.finish(), complete) + } + + pub fn to_udp(&self, sock: &net::UdpSocket, + addr: A, maxlen: u16, compress: bool) + -> io::Result<(usize, bool)> + { + let (vec, complete) = self.assemble(maxlen, false, compress); + sock.send_to(&vec, addr).map(|x| (x, complete)) + } + + pub fn to_writer(&self, w: &mut W, compress: bool) + -> io::Result<(usize, bool)> + { + let (vec, complete) = self.assemble(::std::u16::MAX, true, compress); + w.write(&vec).map(|x| (x, complete)) + } +} + + +//------------ QuestionSection ---------------------------------------------- + +#[derive(Debug)] +struct Question { + qname: name::DomainNameBuf, + qtype: u16, + qclass: u16, +} + + +/// Collects the questions for a message. +/// +#[derive(Debug)] +pub struct QuestionSection { + inner: Vec, +} + +impl QuestionSection { + fn new() -> QuestionSection { + QuestionSection { inner: Vec::new() } + } + + /// Adds a questions. + /// + pub fn push(&mut self, qname: name::DomainNameBuf, qtype: u16, + qclass: u16) + { + assert!(self.inner.len() < ::std::u16::MAX as usize); + self.inner.push(Question { qname: qname, qtype: qtype, + qclass: qclass }) + } + + /// Returns the length of the question section. + /// + pub fn len(&self) -> u16 { + self.inner.len() as u16 + } + + fn assemble(&self, asm: &mut Assembly) -> bool { + for q in self.inner.iter() { + if !asm.push_question(&q.qname, q.qtype, q.qclass) { + return false + } + } + true + } +} + + +//------------ RecordSection ------------------------------------------------ + +#[derive(Debug)] +struct Record { + name: name::DomainNameBuf, + rtype: u16, + rclass: u16, + ttl: u32, + rdata: Box +} + + +/// Collects the resource records of one of the three record sections. +/// +#[derive(Debug)] +pub struct RecordSection { + inner: Vec, +} + +impl RecordSection { + fn new() -> RecordSection { + RecordSection { inner: Vec::new() } + } + + pub fn push(&mut self, name: name::DomainNameBuf, + rclass: u16, ttl: u32, + data: R) + { + self.inner.push(Record { name: name, rtype: R::rtype(), + rclass: rclass, ttl: ttl, + rdata: Box::new(data) }) + } + + /// Returns the number of records in the section. + /// + pub fn len(&self) -> u16 { + self.inner.len() as u16 + } + + fn assemble(&self, asm: &mut Assembly) -> bool { + for rr in self.inner.iter() { + if !asm.push_rr(&rr.name, rr.rtype, rr.rclass, rr.ttl, &*rr.rdata) + { + return false + } + } + true + } +} + +//------------ Tests -------------------------------------------------------- + +#[cfg(test)] +mod tests { + + #[test] + fn test_it() { + } +}