vendor/measureme/src/file_header.rs

//! All binary files generated by measureme have a simple file header that
//! consists of a 4 byte file magic string and a 4 byte little-endian version
//! number.
use std::convert::TryInto;
use std::error::Error;
use std::path::Path;

pub const CURRENT_FILE_FORMAT_VERSION: u32 = 9;

pub const FILE_MAGIC_TOP_LEVEL: &[u8; 4] = b"MMPD";
pub const FILE_MAGIC_EVENT_STREAM: &[u8; 4] = b"MMES";
pub const FILE_MAGIC_STRINGTABLE_DATA: &[u8; 4] = b"MMSD";
pub const FILE_MAGIC_STRINGTABLE_INDEX: &[u8; 4] = b"MMSI";

pub const FILE_EXTENSION: &str = "mm_profdata";

/// The size of the file header in bytes. Note that functions in this module
/// rely on this size to be `8`.
pub const FILE_HEADER_SIZE: usize = 8;

pub fn write_file_header(
    s: &mut dyn std::io::Write,
    file_magic: &[u8; 4],
) -> Result<(), Box<dyn Error + Send + Sync>> {
    // The implementation here relies on FILE_HEADER_SIZE to have the value 8.
    // Let's make sure this assumption cannot be violated without being noticed.
    assert_eq!(FILE_HEADER_SIZE, 8);

    s.write_all(file_magic).map_err(Box::new)?;
    s.write_all(&CURRENT_FILE_FORMAT_VERSION.to_le_bytes())
        .map_err(Box::new)?;

    Ok(())
}

#[must_use]
pub fn verify_file_header(
    bytes: &[u8],
    expected_magic: &[u8; 4],
    diagnostic_file_path: Option<&Path>,
    stream_tag: &str,
) -> Result<(), Box<dyn Error + Send + Sync>> {
    // The implementation here relies on FILE_HEADER_SIZE to have the value 8.
    // Let's make sure this assumption cannot be violated without being noticed.
    assert_eq!(FILE_HEADER_SIZE, 8);

    let diagnostic_file_path = diagnostic_file_path.unwrap_or(Path::new("<in-memory>"));

    if bytes.len() < FILE_HEADER_SIZE {
        let msg = format!(
            "Error reading {} stream in file `{}`: Expected file to contain at least `{:?}` bytes but found `{:?}` bytes",
            stream_tag,
            diagnostic_file_path.display(),
            FILE_HEADER_SIZE,
            bytes.len()
        );

        return Err(From::from(msg));
    }

    let actual_magic = &bytes[0..4];

    if actual_magic != expected_magic {
        let msg = format!(
            "Error reading {} stream in file `{}`: Expected file magic `{:?}` but found `{:?}`",
            stream_tag,
            diagnostic_file_path.display(),
            expected_magic,
            actual_magic
        );

        return Err(From::from(msg));
    }

    let file_format_version = u32::from_le_bytes(bytes[4..8].try_into().unwrap());

    if file_format_version != CURRENT_FILE_FORMAT_VERSION {
        let msg = format!(
            "Error reading {} stream in file `{}`: Expected file format version {} but found `{}`",
            stream_tag,
            diagnostic_file_path.display(),
            CURRENT_FILE_FORMAT_VERSION,
            file_format_version
        );

        return Err(From::from(msg));
    }

    Ok(())
}

pub fn strip_file_header(data: &[u8]) -> &[u8] {
    &data[FILE_HEADER_SIZE..]
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{PageTag, SerializationSinkBuilder};

    #[test]
    fn roundtrip() {
        let data_sink = SerializationSinkBuilder::new_in_memory().new_sink(PageTag::Events);

        write_file_header(&mut data_sink.as_std_write(), FILE_MAGIC_EVENT_STREAM).unwrap();

        let data = data_sink.into_bytes();

        verify_file_header(&data, FILE_MAGIC_EVENT_STREAM, None, "test").unwrap();
    }

    #[test]
    fn invalid_magic() {
        let data_sink = SerializationSinkBuilder::new_in_memory().new_sink(PageTag::Events);
        write_file_header(&mut data_sink.as_std_write(), FILE_MAGIC_STRINGTABLE_DATA).unwrap();
        let mut data = data_sink.into_bytes();

        // Invalidate the filemagic
        data[2] = 0;
        assert!(verify_file_header(&data, FILE_MAGIC_STRINGTABLE_DATA, None, "test").is_err());
    }

    #[test]
    fn other_version() {
        let data_sink = SerializationSinkBuilder::new_in_memory().new_sink(PageTag::Events);

        write_file_header(&mut data_sink.as_std_write(), FILE_MAGIC_STRINGTABLE_INDEX).unwrap();

        let mut data = data_sink.into_bytes();

        // Change version
        data[4] = 0xFF;
        data[5] = 0xFF;
        data[6] = 0xFF;
        data[7] = 0xFF;
        assert!(verify_file_header(&data, FILE_MAGIC_STRINGTABLE_INDEX, None, "test").is_err());
    }

    #[test]
    fn empty_file() {
        let data: [u8; 0] = [];

        assert!(verify_file_header(&data, FILE_MAGIC_STRINGTABLE_DATA, None, "test").is_err());
    }
}