mio_httpc/dns_parser/

name.rs

use std::fmt;
use std::fmt::Write;
use std::str::from_utf8;
// use std::ascii::AsciiExt;
use super::Error;
use byteorder::{BigEndian, ByteOrder};

/// The DNS name as stored in the original packet
///
/// This is contains just a reference to a slice that contains the data.
/// You may turn this into a string using `.to_string()`
#[derive(Debug, Clone, Copy)]
pub struct Name<'a> {
    labels: &'a [u8],
    /// This is the original buffer size. The compressed names in original
    /// are calculated in this buffer
    original: &'a [u8],
}

impl<'a> Name<'a> {
    pub fn scan(data: &'a [u8], original: &'a [u8]) -> Result<Name<'a>, Error> {
        let mut parse_data = data;
        let mut return_pos = None;
        let mut pos = 0;
        if parse_data.len() <= pos {
            return Err(Error::UnexpectedEOF);
        }
        // By setting the largest_pos to be the original len, a side effect
        // is that the pos variable can move forwards in the buffer once.
        let mut largest_pos = original.len();
        let mut byte = parse_data[pos];
        while byte != 0 {
            if parse_data.len() <= pos {
                return Err(Error::UnexpectedEOF);
            }
            if byte & 0b1100_0000 == 0b1100_0000 {
                if parse_data.len() < pos + 2 {
                    return Err(Error::UnexpectedEOF);
                }
                let off = (BigEndian::read_u16(&parse_data[pos..pos + 2]) & !0b1100_0000_0000_0000)
                    as usize;
                if off >= original.len() {
                    return Err(Error::UnexpectedEOF);
                }
                // Set value for return_pos which is the pos in the original
                // data buffer that should be used to return after validating
                // the offsetted labels.
                if let None = return_pos {
                    return_pos = Some(pos);
                }

                // Check then set largest_pos to ensure we never go backwards
                // in the buffer.
                if off >= largest_pos {
                    return Err(Error::BadPointer);
                }
                largest_pos = off;
                pos = 0;
                parse_data = &original[off..];
            } else if byte & 0b1100_0000 == 0 {
                let end = pos + byte as usize + 1;
                if parse_data.len() < end {
                    return Err(Error::UnexpectedEOF);
                }
                if !parse_data[pos + 1..end].is_ascii() {
                    return Err(Error::LabelIsNotAscii);
                }
                pos = end;
                if parse_data.len() <= pos {
                    return Err(Error::UnexpectedEOF);
                }
            } else {
                return Err(Error::UnknownLabelFormat);
            }
            byte = parse_data[pos];
        }
        if let Some(return_pos) = return_pos {
            return Ok(Name {
                labels: &data[..return_pos + 2],
                original: original,
            });
        } else {
            return Ok(Name {
                labels: &data[..pos + 1],
                original: original,
            });
        }
    }
    pub fn byte_len(&self) -> usize {
        self.labels.len()
    }
}

impl<'a> fmt::Display for Name<'a> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let data = self.labels;
        let original = self.original;
        let mut pos = 0;
        loop {
            let byte = data[pos];
            if byte == 0 {
                return Ok(());
            } else if byte & 0b1100_0000 == 0b1100_0000 {
                let off =
                    (BigEndian::read_u16(&data[pos..pos + 2]) & !0b1100_0000_0000_0000) as usize;
                if pos != 0 {
                    fmt.write_char('.')?;
                }
                return fmt::Display::fmt(&Name::scan(&original[off..], original).unwrap(), fmt);
            } else if byte & 0b1100_0000 == 0 {
                if pos != 0 {
                    fmt.write_char('.')?;
                }
                let end = pos + byte as usize + 1;
                fmt.write_str(from_utf8(&data[pos + 1..end]).unwrap())?;
                pos = end;
                continue;
            } else {
                unreachable!();
            }
        }
    }
}

#[cfg(test)]
mod test {
    use crate::dns_parser::{Error, Name};

    #[test]
    fn parse_badpointer_same_offset() {
        // A buffer where an offset points to itself,
        // which is a bad compression pointer.
        let same_offset = vec![192, 2, 192, 2];
        let is_match = matches!(
            Name::scan(&same_offset, &same_offset),
            Err(Error::BadPointer)
        );

        assert!(is_match);
    }

    #[test]
    fn parse_badpointer_forward_offset() {
        // A buffer where the offsets points back to each other which causes
        // infinite recursion if never checked, a bad compression pointer.
        let forwards_offset = vec![192, 2, 192, 4, 192, 2];
        let is_match = matches!(
            Name::scan(&forwards_offset, &forwards_offset),
            Err(Error::BadPointer)
        );

        assert!(is_match);
    }

    #[test]
    fn nested_names() {
        // A buffer where an offset points to itself, a bad compression pointer.
        let buf = b"\x02xx\x00\x02yy\xc0\x00\x02zz\xc0\x04";

        assert_eq!(Name::scan(&buf[..], buf).unwrap().to_string(), "xx");
        assert_eq!(Name::scan(&buf[..], buf).unwrap().labels, b"\x02xx\x00");
        assert_eq!(Name::scan(&buf[4..], buf).unwrap().to_string(), "yy.xx");
        assert_eq!(
            Name::scan(&buf[4..], buf).unwrap().labels,
            b"\x02yy\xc0\x00"
        );
        assert_eq!(Name::scan(&buf[9..], buf).unwrap().to_string(), "zz.yy.xx");
        assert_eq!(
            Name::scan(&buf[9..], buf).unwrap().labels,
            b"\x02zz\xc0\x04"
        );
    }
}