use iris_abi::{Reader, Tag};
use crate::container::FileHeader;
use crate::digest::Digest;
use crate::error::{Error, Result};
use crate::layout::{
DecoderLocation, FORMAT_MAJOR, HEADER_SIZE, MAGIC, MIN_SIZE, TRAILER_SIZE, tag,
};
use crate::meta::{Dataset, DecoderRef, Schema, Section};
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct Placement {
file_len: u64,
footer_at: u64,
footer_len: u32,
root: Digest,
}
impl Placement {
pub const TRAILER_LEN: usize = TRAILER_SIZE;
pub fn trailer_at(file_len: u64) -> Result<u64> {
if file_len < MIN_SIZE as u64 {
return Err(Error::Truncated {
what: "the container",
needed: MIN_SIZE as u64,
available: file_len,
});
}
Ok(file_len - TRAILER_SIZE as u64)
}
pub fn read(trailer: &[u8], file_len: u64) -> Result<Self> {
let trailer_at = Self::trailer_at(file_len)?;
let t: &[u8; TRAILER_SIZE] = trailer.first_chunk().ok_or(Error::Truncated {
what: "the trailer",
needed: TRAILER_SIZE as u64,
available: trailer.len() as u64,
})?;
if t[48..56] != MAGIC {
return Err(Error::Truncated {
what: "the container, which does not end with the magic",
needed: file_len,
available: file_len,
});
}
let footer_at = u64::from_le_bytes([t[0], t[1], t[2], t[3], t[4], t[5], t[6], t[7]]);
let footer_len = u32::from_le_bytes([t[8], t[9], t[10], t[11]]);
let reserved = u32::from_le_bytes([t[12], t[13], t[14], t[15]]);
if reserved != 0 {
return Err(Error::Reserved { what: "trailer" });
}
let root: [u8; 32] = core::array::from_fn(|i| t[16 + i]);
let footer_end = footer_at
.checked_add(u64::from(footer_len))
.ok_or(Error::Truncated {
what: "the footer",
needed: u64::MAX,
available: file_len,
})?;
if footer_at < HEADER_SIZE as u64 || footer_end > trailer_at {
return Err(Error::Truncated {
what: "the footer",
needed: footer_end,
available: trailer_at,
});
}
Ok(Self {
file_len,
footer_at,
footer_len,
root: Digest(root),
})
}
#[must_use]
pub const fn file_len(&self) -> u64 {
self.file_len
}
#[must_use]
pub const fn footer_at(&self) -> u64 {
self.footer_at
}
#[must_use]
pub const fn footer_len(&self) -> usize {
self.footer_len as usize
}
#[must_use]
pub const fn root_digest(&self) -> Digest {
self.root
}
}
#[derive(Clone, Debug)]
pub struct Directory<'a> {
header: FileHeader,
raw_header: [u8; HEADER_SIZE],
footer: &'a [u8],
placement: Placement,
dataset: Dataset,
schema: Option<Schema<'a>>,
decoder: Option<DecoderRef<'a>>,
sections: Vec<Section>,
}
impl<'a> Directory<'a> {
pub fn parse(header: &[u8], footer: &'a [u8], placement: Placement) -> Result<Self> {
let directory = Self::parse_without_root_digest(header, footer, placement)?;
let actual = directory.compute_root();
if actual != placement.root {
return Err(Error::DigestMismatch {
what: "footer".to_owned(),
expected: placement.root.to_string(),
actual: actual.to_string(),
});
}
Ok(directory)
}
pub fn parse_without_root_digest(
header: &[u8],
footer: &'a [u8],
placement: Placement,
) -> Result<Self> {
let raw_header: [u8; HEADER_SIZE] = *header.first_chunk().ok_or(Error::Truncated {
what: "the header",
needed: HEADER_SIZE as u64,
available: header.len() as u64,
})?;
if raw_header[..MAGIC.len()] != MAGIC {
return Err(Error::NotAContainer {
found: raw_header[..MAGIC.len()].to_vec(),
expected: MAGIC.to_vec(),
});
}
if footer.len() != placement.footer_len() {
return Err(Error::Truncated {
what: "the footer",
needed: placement.footer_len() as u64,
available: footer.len() as u64,
});
}
let mut directory = Self {
header: parse_header(&raw_header)?,
raw_header,
footer,
placement,
dataset: Dataset {
rows: 0,
name: String::new(),
},
schema: None,
decoder: None,
sections: Vec::new(),
};
directory.parse_footer(footer)?;
directory.check_sections()?;
Ok(directory)
}
fn parse_footer(&mut self, footer: &'a [u8]) -> Result<()> {
let mut seen_dataset = false;
let mut p = Reader::new(footer);
while !p.is_empty() {
let (header, mut body) = p.record()?;
match header.tag {
tag::DATASET => {
expect_version(header.tag, header.version, Dataset::VERSION)?;
if seen_dataset {
return Err(Error::RepeatedRecord(tag::DATASET));
}
self.dataset = Dataset::decode(&mut body)?;
seen_dataset = true;
}
tag::SCHEMA => {
expect_version(header.tag, header.version, Schema::VERSION)?;
if self.schema.is_some() {
return Err(Error::RepeatedRecord(tag::SCHEMA));
}
self.schema = Some(Schema::decode(&mut body)?);
}
tag::DECODER => {
expect_version(header.tag, header.version, DecoderRef::VERSION)?;
if self.decoder.is_some() {
return Err(Error::RepeatedRecord(tag::DECODER));
}
self.decoder = Some(DecoderRef::decode(&mut body)?);
}
tag::SECTION => {
expect_version(header.tag, header.version, Section::VERSION)?;
self.sections.push(Section::decode(&mut body)?);
}
_ => {}
}
}
if !seen_dataset {
return Err(Error::MissingRecord(tag::DATASET));
}
Ok(())
}
fn check_sections(&self) -> Result<()> {
let file_len = self.placement.file_len();
for (i, section) in self.sections.iter().enumerate() {
let end = section.end().ok_or(Error::SectionOutOfBounds {
id: section.id,
offset: section.offset,
end: u64::MAX,
file_len,
})?;
if section.offset < HEADER_SIZE as u64 || end > self.placement.footer_at() {
return Err(Error::SectionOutOfBounds {
id: section.id,
offset: section.offset,
end,
file_len,
});
}
if self.sections[..i].iter().any(|s| s.id == section.id) {
return Err(Error::DuplicateSection { id: section.id });
}
}
Ok(())
}
fn compute_root(&self) -> Digest {
let mut hasher = blake3::Hasher::new();
hasher.update(&self.raw_header);
hasher.update(self.footer);
Digest(*hasher.finalize().as_bytes())
}
#[must_use]
pub const fn header(&self) -> FileHeader {
self.header
}
#[must_use]
pub const fn placement(&self) -> Placement {
self.placement
}
#[must_use]
pub const fn root_digest(&self) -> Digest {
self.placement.root_digest()
}
#[must_use]
pub const fn dataset(&self) -> &Dataset {
&self.dataset
}
#[must_use]
pub const fn schema(&self) -> Option<&Schema<'a>> {
self.schema.as_ref()
}
#[must_use]
pub const fn decoder(&self) -> Option<&DecoderRef<'a>> {
self.decoder.as_ref()
}
#[must_use]
pub fn sections(&self) -> &[Section] {
&self.sections
}
#[must_use]
pub fn section(&self, id: u32) -> Option<&Section> {
self.sections.iter().find(|s| s.id == id)
}
#[must_use]
pub fn decoder_section(&self) -> Option<&Section> {
let decoder = self.decoder.as_ref()?;
let DecoderLocation::Embedded { section } = decoder.location else {
return None;
};
self.section(section)
}
}
fn parse_header(header: &[u8; HEADER_SIZE]) -> Result<FileHeader> {
let major = u16::from_le_bytes([header[8], header[9]]);
let minor = u16::from_le_bytes([header[10], header[11]]);
if major != FORMAT_MAJOR {
return Err(Error::UnsupportedFormat {
major,
minor,
supported_major: FORMAT_MAJOR,
});
}
let flags = u32::from_le_bytes([header[12], header[13], header[14], header[15]]);
if flags != 0 {
return Err(Error::Reserved { what: "header" });
}
Ok(FileHeader { major, minor })
}
fn expect_version(tag: Tag, found: u16, supported: u16) -> Result<()> {
if found > supported {
return Err(Error::UnsupportedRecord {
tag,
version: found,
});
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_file_shorter_than_the_smallest_container_has_no_trailer_to_read() {
assert!(matches!(
Placement::trailer_at(MIN_SIZE as u64 - 1),
Err(Error::Truncated { .. })
));
assert_eq!(
Placement::trailer_at(MIN_SIZE as u64).expect("the smallest container has a trailer"),
(MIN_SIZE - TRAILER_SIZE) as u64
);
}
fn trailer(footer_at: u64, footer_len: u32) -> [u8; TRAILER_SIZE] {
let mut t = [0u8; TRAILER_SIZE];
t[0..8].copy_from_slice(&footer_at.to_le_bytes());
t[8..12].copy_from_slice(&footer_len.to_le_bytes());
t[48..56].copy_from_slice(&MAGIC);
t
}
#[test]
fn a_trailer_that_does_not_end_with_the_magic_is_not_a_container() {
let mut t = trailer(HEADER_SIZE as u64, 8);
t[55] = 0;
assert!(matches!(
Placement::read(&t, 1024),
Err(Error::Truncated { .. })
));
}
#[test]
fn a_footer_that_overlaps_the_header_or_the_trailer_is_refused() {
assert!(matches!(
Placement::read(&trailer(HEADER_SIZE as u64 - 1, 8), 1024),
Err(Error::Truncated { .. })
));
assert!(matches!(
Placement::read(&trailer(1024 - TRAILER_SIZE as u64, 8), 1024),
Err(Error::Truncated { .. })
));
assert!(matches!(
Placement::read(&trailer(u64::MAX, 8), 1024),
Err(Error::Truncated { .. })
));
}
#[test]
fn a_reserved_field_that_is_not_zero_is_refused() {
let mut t = trailer(HEADER_SIZE as u64, 8);
t[12] = 1;
assert!(matches!(
Placement::read(&t, 1024),
Err(Error::Reserved { what: "trailer" })
));
}
#[test]
fn a_footer_of_the_wrong_length_is_refused_rather_than_parsed() {
let placement = Placement::read(&trailer(HEADER_SIZE as u64, 8), 1024)
.expect("the placement is well formed");
let mut header = [0u8; HEADER_SIZE];
header[..MAGIC.len()].copy_from_slice(&MAGIC);
header[8..10].copy_from_slice(&FORMAT_MAJOR.to_le_bytes());
assert!(matches!(
Directory::parse_without_root_digest(&header, &[0u8; 4], placement),
Err(Error::Truncated {
what: "the footer",
..
})
));
}
}