rust-url/src/lib.rs

#[macro_use] extern crate matches;
pub extern crate mime;

pub enum DataUrlError {
    NotADataUrl,
    NoComma,
}

pub struct DataUrl<'a> {
    mime_type: mime::Mime,
    base64: bool,
    encoded_body_plus_fragment: &'a str,
}

pub enum DecodeError<E> {
    InvalidBase64(InvalidBase64),
    WriteError(E),
}

pub struct InvalidBase64(());

impl<E> From<InvalidBase64> for DecodeError<E> {
    fn from(e: InvalidBase64) -> Self { DecodeError::InvalidBase64(e) }
}

/// The URL’s fragment identifier (after `#`) encoded as in the original input.
///
/// It needs to be either percent-encoded to obtain the same string as in a parsed URL,
/// or percent-decoded to interpret it as text.
pub struct UrlFragmentIdentifier<'a>(pub &'a str);

impl<'a> DataUrl<'a> {
    /// <https://fetch.spec.whatwg.org/#data-url-processor>
    /// but starting from a string rather than a Url, to avoid extra string copies.
    pub fn process(input: &'a str) -> Result<Self, DataUrlError> {
        use DataUrlError::*;

        let after_colon = pretend_parse_data_url(input).ok_or(NotADataUrl)?;

        let (from_colon_to_comma, encoded_body_plus_fragment) =
            find_comma_before_fragment(after_colon).ok_or(NoComma)?;

        let (mime_type, base64) = parse_header(from_colon_to_comma);

        Ok(DataUrl { mime_type, base64, encoded_body_plus_fragment })
    }

    pub fn mime_type(&self) -> &mime::Mime {
        &self.mime_type
    }

    /// Streaming-decode the data URL’s body to `write_body_bytes`,
    /// and return the URL’s fragment identifier is returned if it has one.
    pub fn decode<F, E>(&self, write_body_bytes: F)
                        -> Result<Option<UrlFragmentIdentifier<'a>>, DecodeError<E>>
        where F: FnMut(&[u8]) -> Result<(), E>
    {
        if self.base64 {
            decode_with_base64(self.encoded_body_plus_fragment, write_body_bytes)
        } else {
            decode_without_base64(self.encoded_body_plus_fragment, write_body_bytes)
                .map_err(DecodeError::WriteError)
        }
    }

    /// Return the decoded body and the URL’s fragment identifier
    pub fn decode_to_vec(&self)
        -> Result<(Vec<u8>, Option<UrlFragmentIdentifier<'a>>), InvalidBase64>
    {
        enum Impossible {}
        let mut body = Vec::new();
        let result = self.decode::<_, Impossible>(|bytes| Ok(body.extend_from_slice(bytes)));
        match result {
            Ok(url_fragment) => Ok((body, url_fragment)),
            Err(DecodeError::InvalidBase64(e)) => Err(e),
            Err(DecodeError::WriteError(e)) => match e {}
        }
    }
}

macro_rules! require {
    ($condition: expr) => {
        if !$condition {
            return None
        }
    }
}

/// Similar to <https://url.spec.whatwg.org/#concept-basic-url-parser>
/// followed by <https://url.spec.whatwg.org/#concept-url-serializer>
///
/// * `None`: not a data URL.
///
/// * `Some(s)`: sort of the result of serialization, except:
///
///   - `data:` prefix removed
///   - The fragment is included
///   - Other components are **not** UTF-8 percent-encoded
///   - ASCII tabs and newlines in the middle are **not** removed
fn pretend_parse_data_url(input: &str) -> Option<&str> {
    // Trim C0 control or space
    let left_trimmed = input.trim_left_matches(|ch| ch <= ' ');

    let mut bytes = left_trimmed.bytes();
    {
        // Ignore ASCII tabs or newlines
        let mut iter = bytes.by_ref().filter(|&byte| !matches!(byte, b'\t' | b'\n' | b'\r'));
        require!(iter.next()?.to_ascii_lowercase() == b'd');
        require!(iter.next()?.to_ascii_lowercase() == b'a');
        require!(iter.next()?.to_ascii_lowercase() == b't');
        require!(iter.next()?.to_ascii_lowercase() == b'a');
        require!(iter.next()? == b':');
    }
    let bytes_consumed = left_trimmed.len() - bytes.len();
    let after_colon = &left_trimmed[bytes_consumed..];

    // Trim C0 control or space
    Some(after_colon.trim_right_matches(|ch| ch <= ' '))
}

fn find_comma_before_fragment(after_colon: &str) -> Option<(&str, &str)> {
    for (i, byte) in after_colon.bytes().enumerate() {
        if byte == b',' {
            return Some((&after_colon[..i], &after_colon[i + 1..]))
        }
        if byte == b'#' {
            break
        }
    }
    None
}

fn parse_header(from_colon_to_comma: &str) -> (mime::Mime, bool) {
    let input = from_colon_to_comma.chars()
        .filter(|&c| !matches!(c, '\t' | '\n' | '\r'))  // Removed by the URL parser
        .collect::<String>();
    let mut string;

    let input = input.trim_matches(' ');

    let (mut input, base64) = match without_base64_suffix(input) {
        Some(s) => (s, true),
        None => (input, false),
    };

    // FIXME: percent-encode

    if input.starts_with(';') {
        string = String::from("text/plain");
        string.push_str(input);
        input = &*string;
    }

    // FIXME: does Mime::from_str match the MIME Sniffing Standard’s parsing algorithm?
    // <https://mimesniff.spec.whatwg.org/#parse-a-mime-type>
    let mime_type = input.parse()
        .unwrap_or_else(|_| "text/plain;charset=US-ASCII".parse().unwrap());

    (mime_type, base64)
}

/// None: no base64 suffix
fn without_base64_suffix(s: &str) -> Option<&str> {
    remove_suffix(
        remove_suffix(s, "base64", str::eq_ignore_ascii_case)?
            .trim_right_matches(' '),
        ";", str::eq
    )
}

fn remove_suffix<'a, Eq>(haystack: &'a str, needle: &str, eq: Eq) -> Option<&'a str>
    where Eq: Fn(&str, &str) -> bool
{
    let start_index = haystack.len().checked_sub(needle.len())?;
    let (before, after) = haystack.split_at(start_index);
    if eq(after, needle) {
        Some(before)
    } else {
        None
    }
}

/// This is <https://url.spec.whatwg.org/#string-percent-decode> while also:
///
/// * Ignoring ASCII tab or newlines
/// * Stopping at the first '#' (which indicates the start of the fragment)
///
/// Anything that would have been UTF-8 percent-encoded by the URL parser
/// would be percent-decoded here.
/// We skip that round-trip and pass it through unchanged.
fn decode_without_base64<F, E>(encoded_body_plus_fragment: &str, mut write_bytes: F)
                               -> Result<Option<UrlFragmentIdentifier>, E>
    where F: FnMut(&[u8]) -> Result<(), E>
{
    let bytes = encoded_body_plus_fragment.as_bytes();
    let mut slice_start = 0;
    for (i, &byte) in bytes.iter().enumerate() {
        // We only need to look for 5 different "special" byte values.
        // For everything else we make slices as large as possible, borrowing the input,
        // in order to make fewer write_all() calls.
        if matches!(byte, b'%' | b'#' | b'\t' | b'\n' | b'\r') {
            // Write everything (if anything) "non-special" we’ve accumulated
            // before this special byte
            if i > slice_start {
                write_bytes(&bytes[slice_start..i])?;
            }
            // Then deal with the special byte.
            match byte {
                b'%' => {
                    let l = bytes.get(i + 2).and_then(|&b| (b as char).to_digit(16));
                    let h = bytes.get(i + 1).and_then(|&b| (b as char).to_digit(16));
                    if let (Some(h), Some(l)) = (h, l) {
                        // '%' followed by two ASCII hex digits
                        let one_byte = h as u8 * 0x10 + l as u8;
                        write_bytes(&[one_byte])?;
                        slice_start = i + 3;
                    } else {
                        // Do nothing. Leave slice_start unchanged.
                        // The % sign will be part of the next slice.
                    }
                }

                b'#' => {
                    let fragment_start = i + 1;
                    let fragment = &encoded_body_plus_fragment[fragment_start..];
                    return Ok(Some(UrlFragmentIdentifier(fragment)))
                }

                // Ignore over '\t' | '\n' | '\r'
                _ => slice_start = i + 1
            }
        }
    }
    write_bytes(&bytes[slice_start..])?;
    Ok(None)
}

/// `decode_without_base64()` composed with
/// <https://infra.spec.whatwg.org/#isomorphic-decode> composed with
/// <https://infra.spec.whatwg.org/#forgiving-base64-decode>.
fn decode_with_base64<F, E>(encoded_body_plus_fragment: &str, mut write_bytes: F)
                            -> Result<Option<UrlFragmentIdentifier>, DecodeError<E>>
    where F: FnMut(&[u8]) -> Result<(), E>
{
    let mut bit_buffer: u32 = 0;
    let mut buffer_bit_length: u8 = 0;
    let mut padding_symbols: u8 = 0;

    let fragment = decode_without_base64::<_, DecodeError<E>>(encoded_body_plus_fragment, |bytes| {
        for &byte in bytes.iter() {
            let value = BASE64_DECODE_TABLE[byte as usize];
            if value < 0 {
                // A character that’s not part of the alphabet

                // Remove ASCII whitespace
                // '\t' | '\n' | '\r' was already filtered by decode_without_base64()
                if byte == b' ' || byte == b'\x0C' {
                    continue
                }

                if byte == b'=' {
                    padding_symbols = padding_symbols.saturating_add(8);
                    continue
                }

                Err(InvalidBase64(()))?
            }
            if padding_symbols > 0 {
                // Alphabet symbols after padding
                Err(InvalidBase64(()))?
            }
            bit_buffer <<= 6;
            bit_buffer |= value as u32;
            if buffer_bit_length < 24 {
                buffer_bit_length += 6;
            } else {
                // We’ve accumulated four times 6 bits, which equals three times 8 bits.
                let byte_buffer = [
                    (bit_buffer >> 16) as u8,
                    (bit_buffer >> 8) as u8,
                    bit_buffer as u8,
                ];
                write_bytes(&byte_buffer).map_err(DecodeError::WriteError)?;
                buffer_bit_length = 0;
                // No need to reset bit_buffer,
                // since next time we’re only gonna read relevant bits.
            }
        }
        Ok(())
    })?;
    match (buffer_bit_length, padding_symbols) {
        (0, 0) => {
            // A multiple of four of alphabet symbols, and nothing else.
        }
        (12, 2) | (12, 0) => {
            // A multiple of four of alphabet symbols, followed by two more symbols,
            // optionally followed by two padding characters (which make a total multiple of four).
            let byte_buffer = [
                (bit_buffer >> 4) as u8,
            ];
            write_bytes(&byte_buffer).map_err(DecodeError::WriteError)?;
        }
        (18, 1) | (18, 0) => {
            // A multiple of four of alphabet symbols, followed by three more symbols,
            // optionally followed by one padding character (which make a total multiple of four).
            let byte_buffer = [
                (bit_buffer >> 10) as u8,
                (bit_buffer >> 2) as u8,
            ];
            write_bytes(&byte_buffer).map_err(DecodeError::WriteError)?;
        }
        _ => {
            // No other combination is acceptable
            Err(InvalidBase64(()))?
        }
    }
    Ok(fragment)
}

/// Generated by `make_base64_decode_table.py` based on "Table 1: The Base 64 Alphabet"
/// at <https://tools.ietf.org/html/rfc4648#section-4>
///
/// Array indices are the byte value of symbols.
/// Array values are their positions in the base64 alphabet,
/// or -1 for symbols not in the alphabet.
/// The position contributes 6 bits to the decoded bytes.
const BASE64_DECODE_TABLE: [i8; 256] = [
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 62, -1, -1, -1, 63,
    52, 53, 54, 55, 56, 57, 58, 59, 60, 61, -1, -1, -1, -1, -1, -1,
    -1,  0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14,
    15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, -1, -1, -1, -1, -1,
    -1, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40,
    41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
];