1 files changed, 145 insertions, 0 deletions
diff --git a/vendor/measureme/src/file_header.rs b/vendor/measureme/src/file_header.rs
new file mode 100644
index 000000000..8ad192895
--- /dev/null
+++ b/vendor/measureme/src/file_header.rs
@@ -0,0 +1,145 @@
+//! 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 = 8;
+
+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());
+    }
+}