tiny_tree_magic/

lib.rs

use std::collections::{HashMap, VecDeque, hash_map::Entry};
use std::error::Error as StdError;
use std::fmt;
use std::hint::cold_path;
use std::ops::Range;
use std::str::from_utf8;

use memchr::{memchr, memmem::find};

pub struct Database<C> {
    cache: C,
    types: Box<[FrozenType]>,
    matches: FrozenMatches,
    children: FrozenChildren,
    text_idx: u32,
    binary_idx: u32,
}

#[derive(Debug, Default)]
struct Type {
    matches: Vec<Match>,
    children: Vec<u32>,
}

#[derive(Debug)]
struct Match {
    cnt: u32,
    off: u32,
}

#[derive(Debug)]
struct FrozenType {
    mime_type_off: u32,
    matches: Range<u32>,
    children: Range<u32>,
}

type FrozenMatches = Box<[Match]>;

type FrozenChildren = Box<[u32]>;

impl<C> Database<C>
where
    C: AsRef<[u8]>,
{
    /// Open the given MIME-info `cache`
    ///
    /// This can be an in-memory buffer or a memory map of a [mime.cache file](https://specifications.freedesktop.org/shared-mime-info/latest/ar01s02.html#id-1.3.12).
    ///
    /// ```no_run
    /// # use std::error::Error;
    /// use std::fs::File;
    ///
    /// use memmap2::Mmap;
    /// use tiny_tree_magic::Database;
    ///
    /// let cache = unsafe { Mmap::map(&File::open("/usr/share/mime/mime.cache")?)? };
    ///
    /// let database = Database::open(cache)?;
    /// # Ok::<_, Box<dyn Error>>(())
    /// ```
    #[cold]
    pub fn open(cache: C) -> Result<Self, Error> {
        let (parents, types) = parse_offsets(cache.as_ref()).bail(Error::MalformedOffsets)?;

        let mut parents = parse_parents(cache.as_ref(), parents).bail(Error::MalformedParents)?;

        let mut types = parse_types(cache.as_ref(), types).bail(Error::MalformedTypes)?;

        let (text_off, binary_off) = parents_to_children(cache.as_ref(), &mut parents, &mut types)?;

        let types = sort(types, parents)?;

        Ok(Self::freeze(cache, types, text_off, binary_off))
    }

    fn freeze(cache: C, types: Vec<(u32, Type)>, text_off: u32, binary_off: u32) -> Self {
        let mut matches =
            Vec::with_capacity(types.iter().map(|(_, r#type)| r#type.matches.len()).sum());
        let mut children =
            Vec::with_capacity(types.iter().map(|(_, r#type)| r#type.children.len()).sum());

        let types = types
            .into_iter()
            .map(|(mime_type_off, r#type)| {
                let matches_start = u32::try_from(matches.len()).unwrap();
                let matches_len = u32::try_from(r#type.matches.len()).unwrap();
                let matches_end = matches_start.checked_add(matches_len).unwrap();
                matches.extend(r#type.matches);

                let children_start = u32::try_from(children.len()).unwrap();
                let children_len = u32::try_from(r#type.children.len()).unwrap();
                let children_end = children_start.checked_add(children_len).unwrap();
                children.extend(r#type.children);

                FrozenType {
                    mime_type_off,
                    matches: matches_start..matches_end,
                    children: children_start..children_end,
                }
            })
            .collect::<Box<[_]>>();

        let mime_type_off_to_type_idx = types
            .iter()
            .enumerate()
            .map(|(idx, r#type)| (r#type.mime_type_off, u32::try_from(idx).unwrap()))
            .collect::<HashMap<_, _>>();

        for type_idx in &mut children {
            *type_idx = mime_type_off_to_type_idx[type_idx];
        }

        Self {
            cache,
            types,
            matches: matches.into(),
            children: children.into(),
            text_idx: mime_type_off_to_type_idx[&text_off],
            binary_idx: mime_type_off_to_type_idx[&binary_off],
        }
    }

    pub fn r#match<'a>(&'a self, bytes: &[u8]) -> Result<&'a str, Error> {
        let is_text = memchr(0, bytes).is_none();

        let text_type = &self.types[self.text_idx as usize];
        let binary_type = &self.types[self.binary_idx as usize];

        let mime_type_off = if is_text
            && let Some(mime_type_off) = self.match_children(None, &text_type.children, bytes)?
        {
            mime_type_off
        } else if let Some(mime_type_off) =
            self.match_children(None, &binary_type.children, bytes)?
        {
            mime_type_off
        } else if is_text {
            text_type.mime_type_off
        } else {
            binary_type.mime_type_off
        };

        resolve_mime_type(self.cache.as_ref(), mime_type_off).bail(Error::InvalidMimeType)
    }

    fn match_children(
        &self,
        default: Option<u32>,
        children: &Range<u32>,
        bytes: &[u8],
    ) -> Result<Option<u32>, Error> {
        let children = &self.children[children.start as usize..children.end as usize];

        for type_idx in children {
            let r#type = &self.types[*type_idx as usize];

            let matches = &self.matches[r#type.matches.start as usize..r#type.matches.end as usize];

            let mut args = MatchArgs {
                cache: self.cache.as_ref(),
                bytes,
                depth: 0,
            };

            for r#match in matches {
                if args
                    .try_match(r#match.cnt, r#match.off)
                    .bail(Error::MalformedMatch)?
                {
                    return self.match_children(
                        Some(r#type.mime_type_off),
                        &r#type.children,
                        bytes,
                    );
                }
            }
        }

        Ok(default)
    }
}

struct MatchArgs<'a> {
    cache: &'a [u8],
    bytes: &'a [u8],
    depth: u8,
}

impl MatchArgs<'_> {
    fn try_match(&mut self, cnt: u32, off: u32) -> Option<bool> {
        let cnt = cnt as usize;
        let off = off as usize;

        let matchlets = self.cache.get(off..)?.get(..32 * cnt)?;

        for idx in 0..cnt {
            let matchlet = &matchlets[32 * idx..][..32];

            if self.try_matchlet(matchlet)? {
                let cnt = parse_u32(matchlet, 24);
                let off = parse_u32(matchlet, 28);

                if cnt == 0 {
                    return Some(true);
                }

                self.depth += 1;
                if self.depth == 128 {
                    return None;
                }

                if self.try_match(cnt, off)? {
                    return Some(true);
                }

                self.depth -= 1;
            }
        }

        Some(false)
    }

    #[inline]
    fn try_matchlet(&self, matchlet: &[u8]) -> Option<bool> {
        let range_start = parse_u32(matchlet, 0) as usize;
        let range_length = parse_u32(matchlet, 4) as usize;
        let data_length = parse_u32(matchlet, 12) as usize;
        let data_offset = parse_u32(matchlet, 16) as usize;
        let mask_offset = parse_u32(matchlet, 20) as usize;

        let data = self.cache.get(data_offset..)?.get(..data_length)?;

        let mask = if mask_offset != 0 {
            Some(self.cache.get(mask_offset..)?.get(..data_length)?)
        } else {
            None
        };

        for pos in range_start..range_start + range_length {
            if let Some(bytes) = self.bytes.get(pos..pos + data_length) {
                if let Some(mask) = mask {
                    if data
                        .iter()
                        .zip(mask)
                        .zip(bytes)
                        .all(|((data, mask), byte)| data & mask == byte & mask)
                    {
                        return Some(true);
                    }
                } else {
                    if compare(data, bytes) {
                        return Some(true);
                    }
                }
            } else {
                break;
            }
        }

        Some(false)
    }
}

#[inline]
fn compare(lhs: &[u8], rhs: &[u8]) -> bool {
    let len = lhs.len();
    debug_assert_eq!(len, rhs.len());

    if len >= 8 {
        let lhs_lo = u64::from_le_bytes(lhs[..8].try_into().unwrap());
        let rhs_lo = u64::from_le_bytes(rhs[..8].try_into().unwrap());
        if lhs_lo != rhs_lo {
            return false;
        } else if len == 8 {
            return true;
        }

        let lhs_hi = u64::from_le_bytes(lhs[len - 8..].try_into().unwrap());
        let rhs_hi = u64::from_le_bytes(rhs[len - 8..].try_into().unwrap());
        if lhs_hi != rhs_hi {
            return false;
        } else if len <= 16 {
            return true;
        }
    } else if len >= 4 {
        let lhs_lo = u32::from_le_bytes(lhs[..4].try_into().unwrap());
        let rhs_lo = u32::from_le_bytes(rhs[..4].try_into().unwrap());
        if lhs_lo != rhs_lo {
            return false;
        }

        let lhs_hi = u32::from_le_bytes(lhs[len - 4..].try_into().unwrap());
        let rhs_hi = u32::from_le_bytes(rhs[len - 4..].try_into().unwrap());
        return lhs_hi == rhs_hi;
    } else if len > 0 {
        let lhs_lo = lhs[0];
        let rhs_lo = rhs[0];
        if lhs_lo != rhs_lo {
            return false;
        }

        let lhs_mid = lhs[len / 2];
        let rhs_mid = rhs[len / 2];
        if lhs_mid != rhs_mid {
            return false;
        }

        let lhs_hi = lhs[len - 1];
        let rhs_hi = rhs[len - 1];
        return lhs_hi == rhs_hi;
    } else {
        return true;
    }

    lhs == rhs
}

fn sort(
    mut types: HashMap<u32, Type>,
    mut parents: HashMap<u32, Vec<u32>>,
) -> Result<Vec<(u32, Type)>, Error> {
    let mut sorted = Vec::with_capacity(types.len());

    let mut orphans = types
        .keys()
        .copied()
        .filter(|mime_type_off| !parents.contains_key(mime_type_off))
        .collect::<VecDeque<_>>();

    while let Some(orphan_mime_type_off) = orphans.pop_back() {
        let orphan_type = types.remove(&orphan_mime_type_off).unwrap();

        for child_mime_type_off in &orphan_type.children {
            match parents.entry(*child_mime_type_off) {
                Entry::Occupied(mut entry) => {
                    let parents = entry.get_mut();

                    parents.retain(|parent_mime_type_off| {
                        *parent_mime_type_off != orphan_mime_type_off
                    });

                    if parents.is_empty() {
                        entry.remove();

                        orphans.push_front(*child_mime_type_off);
                    }
                }
                Entry::Vacant(_) => unreachable!(),
            }
        }

        sorted.push((orphan_mime_type_off, orphan_type));
    }

    if parents.is_empty() {
        Ok(sorted)
    } else {
        Err(Error::MalformedParents)
    }
}

fn parents_to_children(
    cache: &[u8],
    parents: &mut HashMap<u32, Vec<u32>>,
    types: &mut HashMap<u32, Type>,
) -> Result<(u32, u32), Error> {
    types.retain(|mime_type_off, _type| resolve_mime_type(cache, *mime_type_off).is_some());

    parents.retain(|mime_type_off, parents| {
        if types.contains_key(mime_type_off) {
            parents.retain(|parent_mime_type_off| types.contains_key(parent_mime_type_off));

            !parents.is_empty()
        } else {
            false
        }
    });

    let text_off =
        find_or_insert_mime_type(cache, types, "text/plain\0").bail(Error::MissingTextPlain)?;

    let binary_off = find_or_insert_mime_type(cache, types, "application/octet-stream\0")
        .bail(Error::MissingApplicationOctetStream)?;

    for mime_type_off in types.keys() {
        if *mime_type_off != text_off
            && *mime_type_off != binary_off
            && let Entry::Vacant(entry) = parents.entry(*mime_type_off)
        {
            let mime_type = resolve_mime_type(cache, *mime_type_off).unwrap();

            let parent_mime_type_off = if mime_type.starts_with("text/") {
                text_off
            } else {
                binary_off
            };

            entry.insert(vec![parent_mime_type_off]);
        }
    }

    for (mime_type_off, parents) in parents {
        for parent_mime_type_off in parents {
            types
                .get_mut(parent_mime_type_off)
                .unwrap()
                .children
                .push(*mime_type_off);
        }
    }

    for r#type in types.values_mut() {
        r#type.children.sort_unstable();
        r#type.children.dedup();
    }

    Ok((text_off, binary_off))
}

fn find_or_insert_mime_type(
    cache: &[u8],
    types: &mut HashMap<u32, Type>,
    mime_type0: &str,
) -> Option<u32> {
    let mime_type = &mime_type0[..mime_type0.len() - 1];

    let mime_type_off = types
        .keys()
        .copied()
        .find(|mime_type_off| resolve_mime_type(cache, *mime_type_off) == Some(mime_type));

    mime_type_off.or_else(|| {
        let mime_type_off = find(cache, mime_type0.as_bytes())?.try_into().ok()?;

        types.insert(mime_type_off, Default::default());

        Some(mime_type_off)
    })
}

fn resolve_mime_type(cache: &[u8], mime_type_off: u32) -> Option<&str> {
    let off = mime_type_off as usize;
    let bytes = cache.get(off..)?;

    let pos = memchr(0, bytes)?;
    let mime_type = &bytes[..pos];

    from_utf8(mime_type).ok()
}

fn parse_offsets(cache: &[u8]) -> Option<(&[u8], &[u8])> {
    if cache.len() < 2 * 2 + 9 * 4 {
        return None;
    }

    let parents_off = parse_u32(cache, 2 * 2 + 4) as usize;
    let types_off = parse_u32(cache, 2 * 2 + 5 * 4) as usize;

    let parents = cache.get(parents_off..)?;
    let types = cache.get(types_off..)?;

    Some((parents, types))
}

fn parse_parents<'a>(cache: &'a [u8], bytes: &'a [u8]) -> Option<HashMap<u32, Vec<u32>>> {
    if bytes.len() < 4 {
        return None;
    }

    let cnt = parse_u32(bytes, 0) as usize;

    let bytes = bytes.get(4..)?.get(..8 * cnt)?;

    let mut parents = HashMap::<_, Vec<_>>::with_capacity(cnt);

    for idx in 0..cnt {
        let bytes = &bytes[8 * idx..][..8];

        let mime_type_off = parse_u32(bytes, 0);
        let parents_off = parse_u32(bytes, 4) as usize;

        let bytes = cache.get(parents_off..)?;

        if bytes.len() < 4 {
            return None;
        }

        let cnt = parse_u32(bytes, 0) as usize;

        let bytes = bytes.get(4..)?.get(..4 * cnt)?;

        let parents = parents.entry(mime_type_off).or_default();
        parents.reserve(cnt);

        for idx in 0..cnt {
            let parent_mime_type_off = parse_u32(bytes, 4 * idx);

            parents.push(parent_mime_type_off);
        }
    }

    Some(parents)
}

fn parse_types(cache: &[u8], bytes: &[u8]) -> Option<HashMap<u32, Type>> {
    if bytes.len() < 3 * 4 {
        return None;
    }

    let cnt = parse_u32(bytes, 0) as usize;
    let off = parse_u32(bytes, 8) as usize;

    let bytes = cache.get(off..)?.get(..16 * cnt)?;

    let mut types = HashMap::<_, Type>::with_capacity(cnt);

    for idx in 0..cnt {
        let bytes = &bytes[16 * idx..];

        let mime_type_off = parse_u32(bytes, 4);

        let cnt = parse_u32(bytes, 8);
        let off = parse_u32(bytes, 12);

        types
            .entry(mime_type_off)
            .or_default()
            .matches
            .push(Match { cnt, off });
    }

    Some(types)
}

fn parse_u32(bytes: &[u8], off: usize) -> u32 {
    u32::from_be_bytes(bytes[off..][..4].try_into().unwrap())
}

#[derive(Debug, Clone, Copy)]
pub enum Error {
    MalformedOffsets,
    MalformedParents,
    MalformedTypes,
    MalformedMatch,
    MissingTextPlain,
    MissingApplicationOctetStream,
    InvalidMimeType,
}

trait BailExt<T> {
    fn bail(self, err: Error) -> Result<T, Error>;
}

impl<T> BailExt<T> for Option<T> {
    #[inline(always)]
    fn bail(self, err: Error) -> Result<T, Error> {
        self.ok_or_else(|| {
            cold_path();
            err
        })
    }
}

impl fmt::Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        let msg = match self {
            Self::MalformedOffsets => "Malformed offsets",
            Self::MalformedParents => "Malformed parents",
            Self::MalformedTypes => "Malformed types",
            Self::MalformedMatch => "Malformed match",
            Self::MissingTextPlain => "Missing text/plain MIME type",
            Self::MissingApplicationOctetStream => "Missing application/octet-stream MIME type",
            Self::InvalidMimeType => "Invalid MIME type",
        };

        fmt.write_str(msg)
    }
}

impl StdError for Error {}