#![deny(warnings, rust_2018_idioms)]

/// Implementation of `domain-name`s as defined by RFC1035 sections 3.1 and 3.3
use anyhow::{bail, Context, Result};
use bytes::{BufMut, BytesMut};
use std::collections::HashMap;

pub struct LabelOffsets {
    offsets: HashMap<Vec<u8>, usize>,
}

impl LabelOffsets {
    pub fn new() -> LabelOffsets {
        LabelOffsets {
            offsets: HashMap::new(),
        }
    }
}

pub fn write(
    name: &str,
    msg_buf: &mut BytesMut,
    ptr_offsets: &mut LabelOffsets,
    description: &str,
) -> Result<()> {
    // Encode domain string into length-delimited elements
    let name_buf: Vec<u8> = encode_uncompressed(name, description)?;

    // Keep the initial offset for storing ptr_offsets later
    let start_offset = msg_buf.len();

    // Scan the resulting name_buf by skipping across the label sizes we just wrote and search for matching substrings in ptr_offsets.
    // If any matching substring is found, then update name_buf with the associated pointer and write the result to buf.
    // Otherwise, proceed to the next loop where we record the offsets for the buffer labels since it wasn't "tainted" by any pointers of its own.
    let mut scan_offset: usize = 0;
    loop {
        if name_buf[scan_offset] == 0 {
            // Reached last label, of length zero
            break;
        }
        match ptr_offsets
            .offsets
            .get(&name_buf[scan_offset..name_buf.len()])
        {
            Some(ptr_offset) => {
                // Offset found for (remainder of) name, write pointer from here and exit
                if ptr_offset > &(0x3FFF as usize) {
                    // This shouldn't happen since we would be the ones who wrote the offset table, but it doesn't hurt to be careful
                    bail!(
                        "pointer offset within '{}' is {} which exceeds max {}",
                        name,
                        ptr_offset,
                        0x3FFF
                    );
                }
                msg_buf.reserve(scan_offset + 2);
                msg_buf.put_slice(&name_buf[0..scan_offset]);
                msg_buf.put_u8((ptr_offset >> 8) as u8 | 0xC0);
                msg_buf.put_u8((ptr_offset & 0xFF) as u8);
                return Ok(());
            }
            None => {
                // No offset found, skip ahead to the next label: skip past the size octet itself, plus the following label.
                // If we were being very conservative, we could check that we haven't passed the end of the string here,
                // but we just built this string ourselves so let's assume it's OK.
                scan_offset += 1 + name_buf[scan_offset] as usize;
            }
        }
    }

    // If no matching pointers were found, write the name to msg_buf and record its labels into ptr_offsets.
    // We avoid doing this as part of the above search because we want to avoid chaining pointers, which we ourselves disallow when reading.
    // In other words, if a string contains any pointers, then it cannot itself be registered in ptr_offsets.
    msg_buf.extend_from_slice(&name_buf);

    scan_offset = 0;
    loop {
        let label_len = name_buf[scan_offset] as usize;
        if label_len == 0 {
            // Reached last label of size zero, no more scanning left
            return Ok(());
        }
        // Record from here to the end of the string
        ptr_offsets.offsets.insert(
            name_buf[scan_offset..name_buf.len()].to_vec(),
            start_offset + scan_offset,
        );
        scan_offset += label_len + 1;
    }
}

pub fn write_nopointer(name: &str, msg_buf: &mut BytesMut, description: &str) -> Result<()> {
    // Encode domain string into length-delimited elements
    msg_buf.extend_from_slice(&encode_uncompressed(name, description)?);
    Ok(())
}

fn encode_uncompressed(name: &str, description: &str) -> Result<Vec<u8>> {
    // Name should be divided into dot-delimited 'labels' each starting with 1-byte len, with a final empty/NUL label
    // "www.n.geek.nz." => \3 www \1 n \4 geek \2 nz \0
    if !name.ends_with(".") {
        bail!("{} must end with a .: [{}]", description, name);
    }
    if name == "." {
        // Shortcut for common case (e.g. in EDNS labels)
        return Ok(vec![0]);
    }

    // Extra byte for leading size value
    let name_bytes = name.as_bytes();
    let mut name_buf: Vec<u8> = Vec::with_capacity(1 + name_bytes.len());
    let mut label_len: usize = 0;
    // Scan and search for labels. Copy the labels with their size prefixes into 'name_buf'.
    for name_bytes_off in 0..name_bytes.len() {
        let name_byte = name_bytes[name_bytes_off];
        if name_byte == b'.' {
            // Reached end of label. Write label size and copy label content
            if label_len == 0 {
                // Note: this can also be hit in the empty "." case, but we handle that above.
                bail!(
                    "{} contains empty label before the end of the string, at offset {}: {}",
                    description,
                    name_bytes_off,
                    name
                );
            }
            name_buf.push(label_len as u8);
            name_buf.extend_from_slice(&name_bytes[(name_bytes_off - label_len)..name_bytes_off]);
            label_len = 0;
        } else if label_len >= 63 {
            // Not at the end of the label, but already at the size limit.
            bail!(
                "{} contains label exceeding 63 byte size max, at offset {}: {}",
                description,
                name_bytes_off,
                name
            );
        } else {
            label_len += 1;
        }
    }

    if label_len == 0 {
        // Handle final \0 (effectively the 'label' following the trailing '.')
        name_buf.push(label_len as u8);
        Ok(name_buf)
    } else {
        // String didn't end with expected '.' (shouldn't happen, checked earlier!)
        bail!(
            "{} doesn't end with a zero-length label: {}",
            description,
            name
        )
    }
}

pub fn read(msg_buf: &[u8], original_offset: usize, description: &str) -> Result<(usize, String)> {
    let max_offset;
    if msg_buf.len() > 0 {
        max_offset = msg_buf.len() - 1;
    } else {
        max_offset = 0;
    }
    read_range(msg_buf, original_offset, max_offset, true, description)
}

pub fn read_noptr(
    msg_buf: &[u8],
    original_offset: usize,
    description: &str,
) -> Result<(usize, String)> {
    let max_offset;
    if msg_buf.len() > 0 {
        max_offset = msg_buf.len() - 1;
    } else {
        max_offset = 0;
    }
    read_range(msg_buf, original_offset, max_offset, false, description)
}

fn read_range(
    msg_buf: &[u8],
    original_offset: usize,
    max_offset: usize,
    enable_ptr: bool,
    description: &str,
) -> Result<(usize, String)> {
    let mut cur_offset = original_offset;
    let mut out = String::new();
    if max_offset >= msg_buf.len() {
        bail!(
            "max_offset={} equals or exceeds msg_buf.len={}",
            max_offset,
            msg_buf.len()
        );
    }
    loop {
        // Next byte should be size of label
        if cur_offset > max_offset {
            bail!(
                "{} expected label size at offset={} but max_offset={}",
                description,
                cur_offset,
                max_offset
            );
        }
        let label_size = msg_buf[cur_offset] as usize;
        cur_offset += 1;
        if label_size >= 192 {
            if !enable_ptr {
                // This field doesn't support pointers in its domain strings
                bail!(
                    "{} doesn't support pointers but a pointer was encountered at offset={}: {}",
                    description,
                    cur_offset,
                    out
                );
            }
            // This label is a pointer into msg_buf.
            // Read the next byte. 14 bits of this byte and the following byte should give us the pointer destination.
            if cur_offset > max_offset {
                bail!(
                    "{} expected 1 more byte for second half of pointer at offset={} but max_offset={}: {}",
                    description, cur_offset, max_offset, out,
                );
            }
            let label_size2 = msg_buf[cur_offset] as usize;
            cur_offset += 1;
            let pointer_offset = ((label_size - 192) << 8) | label_size2;
            // pointer loop: Check that the pointer is to a location preceding the area we already read in this round
            if pointer_offset >= original_offset {
                bail!(
                    "{} pointer at cur_offset={} to pointer_offset={} when original_offset={}: {}",
                    description,
                    cur_offset,
                    pointer_offset,
                    original_offset,
                    out
                );
            }
            // pointer loop: For the next round, ensure that we stay within data that we haven't already read in earlier rounds
            let next_max_offset: usize;
            if original_offset == 0 {
                // Shouldn't happen due to earlier "pointer_offset < original_offset" check: pointer_offset would have to be negative!
                next_max_offset = 0;
            } else {
                // Can read up to the byte BEFORE this round
                next_max_offset = original_offset - 1;
            }
            let (_consumed, pointer_labels) = read_range(
                &msg_buf,
                pointer_offset,
                next_max_offset,
                enable_ptr,
                description,
            )
            .with_context(|| {
                format!(
                    "{} failed to fetch content of pointer at offset {} for name '{}'",
                    description, cur_offset, out
                )
            })?;
            // RECURSE
            if !pointer_labels.is_empty() {
                if !out.is_empty() {
                    // String already has a prior label.
                    // Add '.' separator for the prior label before adding the next label(s) from the read() call.
                    out.push('.');
                }
                out.push_str(pointer_labels.as_str());
            }
            // Ignore amount 'consumed' at external pointer, just track what we consumed locally:
            return Ok((cur_offset - original_offset, out));
        } else if label_size > 63 {
            // Domain labels cannot exceed 63 bytes, this may indicate corrupt data or a parsing bug.
            // Note: I18n domains (RFC3492/punycode) continue to honor a 63 byte limit.
            bail!(
                "{} contains label with size {} at offset {}, exceeds max 63. bad data?: {:02X?}",
                description,
                label_size,
                cur_offset,
                msg_buf
            );
        } else if label_size == 0 {
            // Reached empty label for end of string: add trailing '.' then exit
            out.push('.');
            return Ok((cur_offset - original_offset, out));
        } else {
            // Fetch remaining bytes and include trailing '.'
            // Check that we are staying within the allowed range
            if max_offset < cur_offset + label_size as usize {
                bail!(
                    "{} max_offset={} is exceeded by requested cur_offset={} + label_size={}: {}",
                    description,
                    max_offset,
                    cur_offset,
                    label_size,
                    out
                );
            }
            if !out.is_empty() {
                // String already has a prior label.
                // Add '.' separator for the prior label before adding the next label.
                out.push('.');
            }
            let label = Vec::from(&msg_buf[cur_offset..cur_offset + label_size as usize]);
            if !valid_label(&label) {
                bail!("{} contains an invalid label: {:02X?}", description, label);
            }
            // unsafe: We have just checked that the label content consists only of ASCII characters ([a-zA-Z0-9-])
            unsafe {
                out.push_str(&String::from_utf8_unchecked(label));
            }
            cur_offset += label_size;
        }
    }
}

/// Per RFC1035 section 2.3.1: [a-zA-Z]([a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?
/// BUT in practice I'm seeing valid domains that start with a number, so maybe this has been loosened?
fn valid_label(s: &Vec<u8>) -> bool {
    if s.len() == 0 {
        return false;
    }

    // First char: [a-zA-Z0-9] (despite RFC1035 disallowing 0-9)
    match s[0] {
        65..=90 | 97..=122 | 48..=57 => {} // ascii: A-Z | a-z | 0-9
        _ => {
            return false;
        }
    }
    if s.len() == 1 {
        return true;
    }

    // Last char (for length > 1): [a-zA-Z0-9]
    match s[s.len() - 1] {
        65..=90 | 97..=122 | 48..=57 => {} // ascii: A-Z | a-z | 0-9
        _ => {
            return false;
        }
    }
    if s.len() == 2 {
        return true;
    }

    // Middle chars (for length > 2): [a-zA-Z0-9-]
    for i in 0..s.len() - 2 {
        match s[i] {
            65..=90 | 97..=122 | 48..=57 | 45 => {} // ascii: A-Z | a-z | 0-9 | -
            _ => {
                return false;
            }
        }
    }

    true
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn read_empty() {
        let buf = vec![0]; // .
        assert_eq!((buf.len(), String::from(".")), read(&buf, 0, "").unwrap());
    }

    #[test]
    fn invalid_values() {
        assert_eq!(false, valid_label(&Vec::<u8>::new()));
        assert_eq!(true, valid_label(&String::from("5nines").into_bytes()));
        assert_eq!(true, valid_label(&String::from("Fi5enines").into_bytes()));
        assert_eq!(
            false,
            valid_label(&String::from("n9ne-tO-fiv5-").into_bytes())
        );
        assert_eq!(
            true,
            valid_label(&String::from("n9ne-tO-fiv5").into_bytes())
        );
    }

    #[test]
    fn read_basic() {
        let buf = vec![
            1, 110, // n
            4, 103, 101, 101, 107, // geek
            2, 110, 122, // nz
            0,   // .
        ];

        // good offsets should succeed
        assert_eq!(
            (buf.len(), String::from("n.geek.nz.")),
            read(&buf, 0, "").unwrap()
        );
        assert_eq!((9, String::from("geek.nz.")), read(&buf, 2, "").unwrap());
        assert_eq!((4, String::from("nz.")), read(&buf, 7, "").unwrap());
        assert_eq!((1, String::from(".")), read(&buf, 10, "").unwrap());

        // bad offsets should fail
        read(&buf, 1, "").unwrap_err();
        read(&buf, 3, "").unwrap_err();
        read(&buf, 4, "").unwrap_err();
        read(&buf, 5, "").unwrap_err();
        read(&buf, 6, "").unwrap_err();
        read(&buf, 8, "").unwrap_err();
        read(&buf, 9, "").unwrap_err();
    }

    #[test]
    fn read_basic_ignore_rest() {
        let buf = vec![
            1, 110, // n
            4, 103, 101, 101, 107, // geek
            2, 110, 122, // nz
            0,   // .
            4, 110, 105, 99, 107, 192, 2, // nick->2
        ];

        // good offsets should succeed, and ignore anything past the '.'
        assert_eq!((11, String::from("n.geek.nz.")), read(&buf, 0, "").unwrap());
        assert_eq!((9, String::from("geek.nz.")), read(&buf, 2, "").unwrap());
        assert_eq!((4, String::from("nz.")), read(&buf, 7, "").unwrap());
        assert_eq!((1, String::from(".")), read(&buf, 10, "").unwrap());
    }

    #[test]
    fn read_pointer_small() {
        let buf = vec![
            1, 110, // n
            4, 103, 101, 101, 107, // geek
            2, 110, 122, // nz
            0,   // .
            4, 110, 105, 99, 107, 192, 2, // nick->2
        ];
        // only 'consumes' the 'nick' segment, ignores the post-pointer data
        assert_eq!(
            (7, String::from("nick.geek.nz.")),
            read(&buf, 11, "").unwrap()
        );
    }

    #[test]
    fn read_pointer_big() {
        let mut buf = Vec::new();
        // insert arbitrary data ('A') so that our test region needs a 2-byte pointer
        let filler_size = 16000;
        buf.resize(filler_size, 65);
        // offset to 'geek' becomes 16000 + 2: ((**254** - 192) << 8) | **130**
        buf.extend_from_slice(&[
            1, 110, // n
            4, 103, 101, 101, 107, // geek
            2, 110, 122, // nz
            0,   // .
            4, 110, 105, 99, 107, 254, 130, // nick->16002
        ]);
        assert_eq!(
            (7, String::from("nick.geek.nz.")),
            read(&buf, filler_size + 11, "").unwrap()
        );
    }

    #[test]
    fn read_pointer_many() {
        // several pointers that just hop incrementally backwards in a chain: www -> n -> geek -> nz -> .
        let buf = vec![
            2, 110, 122, 0, // nz.
            4, 103, 101, 101, 107, 192, 0, // geek->0
            1, 110, 192, 4, // n->4
            3, 119, 119, 119, 192, 11, // www->11
        ];
        // only 'consumes' the 'www' segment, ignores the post-pointer data
        assert_eq!(
            (6, String::from("www.n.geek.nz.")),
            read(&buf, 15, "").unwrap()
        );
    }

    #[test]
    fn read_pointer_tozero() {
        // shouldn't happen since it's less efficient than skipping the pointer, but may as well capture the corner case
        let buf = vec![
            0, // .
            2, 110, 122, 192, 0, // nz->0
        ];
        // only 'consumes' the 'nz' segment, ignores the post-pointer data
        assert_eq!((5, String::from("nz..")), read(&buf, 1, "").unwrap());
    }

    #[test]
    fn read_pointer_inter_loop() {
        // looping A -> B -> C -> A
        let buf = vec![
            3, 119, 119, 119, // www
            1, 110, // n
            4, 103, 101, 101, 107, // geek
            2, 110, 122, 192, 4, // nz->4
        ];
        // any entry segment should result in an error
        read(&buf, 0, "").unwrap_err();
        read(&buf, 4, "").unwrap_err();
        read(&buf, 6, "").unwrap_err();
        read(&buf, 11, "").unwrap_err();
    }

    #[test]
    fn read_pointer_intra_loop() {
        // looping A -> A
        let buf = vec![
            3, 119, 119, 119, // www
            1, 110, // n
            4, 103, 101, 101, 107, 192, 6, // geek->6
        ];
        // any entry segment should result in an error
        read(&buf, 0, "").unwrap_err();
        read(&buf, 4, "").unwrap_err();
        read(&buf, 6, "").unwrap_err();
    }

    #[test]
    fn read_noptr_ptrs() {
        let buf = vec![
            1, 110, // n
            4, 103, 101, 101, 107, // geek
            2, 110, 122, // nz
            0,   // .
            4, 110, 105, 99, 107, 192, 2, // nick->2
        ];
        assert_eq!(
            (11, String::from("n.geek.nz.")),
            read_noptr(&buf, 0, "").unwrap()
        );
        assert_eq!(
            (9, String::from("geek.nz.")),
            read_noptr(&buf, 2, "").unwrap()
        );
        assert_eq!((4, String::from("nz.")), read_noptr(&buf, 7, "").unwrap());
        assert_eq!((1, String::from(".")), read_noptr(&buf, 10, "").unwrap());
        // should return an error about finding an unexpected pointer
        read_noptr(&buf, 11, "").unwrap_err();
    }
}