maarcadetweet/crates/at-blob/src/mime.rs

//! MIME type detection for uploaded blobs.
//!
//! The PDS serves blobs through `com.atproto.sync.getBlob` and
//! `com.atproto.uploadBlob`. The wire protocol for `uploadBlob` carries
//! the MIME type as a request header (`Content-Type`), so the happy
//! path doesn't need any sniffing at write time — the client tells us
//! what they uploaded.
//!
//! On the read side we may not always have the header preserved (e.g.
//! blobs uploaded by older clients, or blobs referenced from a record
//! without their original MIME type available). [`detect_mime`] sniffs
//! the magic bytes of the payload to recover a sensible
//! `Content-Type` for the response.
//!
//! We use the [`infer`] crate for the common image / media formats
//! (PNG, JPEG, GIF, WebP, …) and a tiny inline ASCII heuristic for
//! plain text. Anything unknown returns `None` so the caller can fall
//! back to `application/octet-stream`.

/// The set of MIME types we can detect from content sniffing.
///
/// Kept as an enum (not a `&'static str` alias) so callers can exhaust
/// over the supported set when they want to — e.g. the
/// `mime_type_str` mapping below is the single source of truth for
/// the wire-level string form.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum MimeType {
    Png,
    Jpeg,
    Gif,
    Webp,
    Text,
}

impl MimeType {
    /// Wire-level MIME type string (e.g. `image/png`). Always ASCII
    /// and safe to use as an HTTP `Content-Type` value.
    pub fn as_str(&self) -> &'static str {
        match self {
            MimeType::Png => "image/png",
            MimeType::Jpeg => "image/jpeg",
            MimeType::Gif => "image/gif",
            MimeType::Webp => "image/webp",
            MimeType::Text => "text/plain; charset=utf-8",
        }
    }
}

impl std::fmt::Display for MimeType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.as_str())
    }
}

/// Sniff the magic bytes of `data` to recover a MIME type. Returns
/// `None` when the bytes don't match any known signature — the caller
/// is expected to fall back to a generic `application/octet-stream`.
///
/// Detection is deliberately conservative: we'd rather return `None`
/// than guess wrong. The matching logic, in order:
/// 1. PNG signature (`89 50 4E 47 0D 0A 1A 0A`)
/// 2. JPEG signature (`FF D8 FF`)
/// 3. GIF signature (`47 49 46 38 …` — `GIF8` prefix)
/// 4. WebP signature (`RIFF…WEBP`)
///
/// The `infer` crate is used for step 1–4 because its matchers are
/// well-maintained and we get proper `image/png`, `image/jpeg`, etc.
/// strings for free. The plain-text check at the end is inline because
/// `infer` doesn't classify text and the heuristic is a one-liner:
/// every byte must be printable ASCII, or a common whitespace / line
/// ending.
pub fn detect_mime(data: &[u8]) -> Option<MimeType> {
    if data.is_empty() {
        return None;
    }
    if let Some(kind) = infer::get(data) {
        return match kind.mime_type() {
            "image/png" => Some(MimeType::Png),
            "image/jpeg" => Some(MimeType::Jpeg),
            "image/gif" => Some(MimeType::Gif),
            "image/webp" => Some(MimeType::Webp),
            _ => None,
        };
    }
    if looks_like_text(data) {
        return Some(MimeType::Text);
    }
    None
}

/// True if `data` is non-empty printable ASCII (allowing tab and the
/// usual line endings). We use this as a last-ditch sniff for blobs
/// that aren't tagged as anything by `infer` but that *look* like
/// text — useful when an older client uploaded, say, a JSON string
/// blob with `Content-Type: text/plain` but we don't have the header
/// any more.
///
/// We deliberately don't accept UTF-8 multi-byte sequences here —
/// keeping it ASCII means we won't false-positive on, e.g., a tiny
/// PNG-prefixed binary blob. Real text blobs that need a UTF-8
/// charset should be uploaded with the explicit `Content-Type`
/// header.
fn looks_like_text(data: &[u8]) -> bool {
    if data.is_empty() {
        return false;
    }
    data.iter().all(|&b| {
        b == b'\n'
            || b == b'\r'
            || b == b'\t'
            || (0x20..=0x7e).contains(&b)
    })
}

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

    /// A minimal but valid PNG signature followed by enough bytes
    /// that `infer::get` accepts it (the full file would have more
    /// chunks, but the magic is in the first 8 bytes).
    fn png_signature() -> Vec<u8> {
        vec![0x89, b'P', b'N', b'G', 0x0d, 0x0a, 0x1a, 0x0a, 0, 0, 0, 0]
    }

    /// JPEG starts with `FF D8 FF`. We add an arbitrary fourth byte
    /// (`E0` = JFIF marker) to make it more realistic — `infer`
    /// matches on the first three.
    fn jpeg_signature() -> Vec<u8> {
        vec![0xff, 0xd8, 0xff, 0xe0, 0, 0]
    }

    /// GIF87a prefix is `47 49 46 38 37 61`; GIF89a is `47 49 46 38 39 61`.
    /// Either is enough for `infer`.
    fn gif_signature() -> Vec<u8> {
        vec![b'G', b'I', b'F', b'8', b'9', b'a', 0, 0]
    }

    /// WebP is `RIFF…WEBP`. The size field (4 bytes LE) between
    /// `RIFF` and `WEBP` must be present but its value doesn't matter
    /// for the signature check.
    fn webp_signature() -> Vec<u8> {
        let mut v = vec![b'R', b'I', b'F', b'F'];
        v.extend_from_slice(&[0, 0, 0, 0]);
        v.extend_from_slice(b"WEBP");
        v.extend_from_slice(&[0; 8]);
        v
    }

    #[test]
    fn detects_png() {
        assert_eq!(detect_mime(&png_signature()), Some(MimeType::Png));
        assert_eq!(MimeType::Png.as_str(), "image/png");
    }

    #[test]
    fn detects_jpeg() {
        assert_eq!(detect_mime(&jpeg_signature()), Some(MimeType::Jpeg));
        assert_eq!(MimeType::Jpeg.as_str(), "image/jpeg");
    }

    #[test]
    fn detects_gif() {
        assert_eq!(detect_mime(&gif_signature()), Some(MimeType::Gif));
        assert_eq!(MimeType::Gif.as_str(), "image/gif");
    }

    #[test]
    fn detects_webp() {
        assert_eq!(detect_mime(&webp_signature()), Some(MimeType::Webp));
        assert_eq!(MimeType::Webp.as_str(), "image/webp");
    }

    #[test]
    fn detects_plain_text() {
        let txt = b"hello world\nthis is plain text, with punctuation: !@#$%^&*()\n";
        assert_eq!(detect_mime(txt), Some(MimeType::Text));
        assert_eq!(
            MimeType::Text.as_str(),
            "text/plain; charset=utf-8"
        );
    }

    #[test]
    fn text_allows_tabs_and_crlf() {
        let txt = b"line1\r\nline2\tindented\n";
        assert_eq!(detect_mime(txt), Some(MimeType::Text));
    }

    #[test]
    fn unknown_binary_returns_none() {
        // Random bytes that don't match any known signature.
        let blob = vec![0x00, 0x01, 0x02, 0x03, 0xff, 0xfe, 0xfd];
        assert_eq!(detect_mime(blob.as_slice()), None);
    }

    #[test]
    fn empty_input_returns_none() {
        assert_eq!(detect_mime(b""), None);
    }

    #[test]
    fn non_ascii_bytes_not_classified_as_text() {
        // High-bit bytes aren't ASCII; the text heuristic must skip
        // them. This is intentional: it prevents false-positives on
        // tiny binary blobs (e.g. a 4-byte integer that happens to
        // spell "ABCD").
        let blob = vec![b'A', b'B', 0x80, 0x81];
        assert_eq!(detect_mime(blob.as_slice()), None);
    }

    #[test]
    fn mime_type_display_matches_as_str() {
        for m in [
            MimeType::Png,
            MimeType::Jpeg,
            MimeType::Gif,
            MimeType::Webp,
            MimeType::Text,
        ] {
            assert_eq!(format!("{m}"), m.as_str());
        }
    }
}