use alloc::borrow::Cow;
use std::path::{Path, PathBuf};
use asdf_yaml::{Document, parse_document};
use crate::block::header::CHECKSUM_SIZE;
use crate::compression::Compression;
use crate::error::{Result, err};
use crate::layout::{BlockLocation, Layout, scan};
enum Source {
#[cfg(not(miri))]
Mapped(memmap2::Mmap),
Owned(Vec<u8>),
}
impl core::ops::Deref for Source {
type Target = [u8];
fn deref(&self) -> &[u8] {
match self {
#[cfg(not(miri))]
Source::Mapped(m) => m,
Source::Owned(v) => v,
}
}
}
impl core::fmt::Debug for Source {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
let kind = match self {
#[cfg(not(miri))]
Source::Mapped(_) => "Mapped",
Source::Owned(_) => "Owned",
};
write!(f, "{kind}({} bytes)", self.len())
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum ChecksumStatus {
Absent,
Valid,
Invalid,
}
impl ChecksumStatus {
pub fn is_failure(self) -> bool {
self == ChecksumStatus::Invalid
}
}
fn external_relative_path(uri: &str) -> Result<PathBuf> {
if uri.is_empty() {
return Err(err!(InvalidArgument, "external source is an empty URI"));
}
if uri.contains("://") || uri.starts_with('/') || uri.starts_with('\\') {
return Err(err!(
InvalidArgument,
"external source {uri:?} is not a relative path; only files beside the \
referring one can be resolved"
));
}
let path = Path::new(uri);
for component in path.components() {
use std::path::Component;
match component {
Component::Normal(_) | Component::CurDir => {}
Component::ParentDir => {
return Err(err!(
InvalidArgument,
"external source {uri:?} climbs out of the referring file's directory"
));
}
Component::RootDir | Component::Prefix(_) => {
return Err(err!(InvalidArgument, "external source {uri:?} is not relative"));
}
}
}
Ok(path.to_path_buf())
}
#[derive(Debug)]
pub struct Reader {
source: Source,
layout: Layout,
path: Option<PathBuf>,
}
impl Reader {
pub fn open(path: impl AsRef<Path>) -> Result<Self> {
let path = path.as_ref();
#[cfg(miri)]
{
let bytes = std::fs::read(path)?;
let layout = scan(&bytes)?;
return Ok(Self {
source: Source::Owned(bytes),
layout,
path: Some(path.to_path_buf()),
});
}
#[cfg(not(miri))]
{
let file = std::fs::File::open(path)?;
Self::map(file, path)
}
}
#[cfg(not(miri))]
fn map(file: std::fs::File, path: &Path) -> Result<Self> {
#[allow(unsafe_code)]
let mapped = unsafe { memmap2::Mmap::map(&file) }?;
let layout = scan(&mapped)?;
Ok(Self { source: Source::Mapped(mapped), layout, path: Some(path.to_path_buf()) })
}
pub fn from_bytes(bytes: Vec<u8>) -> Result<Self> {
let layout = scan(&bytes)?;
Ok(Self { source: Source::Owned(bytes), layout, path: None })
}
pub fn path(&self) -> Option<&Path> {
self.path.as_deref()
}
pub fn bytes(&self) -> &[u8] {
&self.source
}
pub fn layout(&self) -> &Layout {
&self.layout
}
pub fn tree_text(&self) -> Option<&str> {
self.layout.tree_str(&self.source)
}
pub fn tree(&self) -> Result<Option<Document>> {
match self.tree_text() {
None => Ok(None),
Some(text) => Ok(Some(parse_document(text)?)),
}
}
pub fn block_count(&self) -> usize {
self.layout.blocks.len()
}
pub fn block(&self, index: usize) -> Result<&BlockLocation> {
self.layout.blocks.get(index).ok_or_else(|| {
err!(
InvalidArgument,
"block index {index} is out of range; the file has {} blocks",
self.layout.blocks.len()
)
})
}
pub fn block_raw(&self, index: usize) -> Result<&[u8]> {
let block = self.block(index)?;
let start = usize::try_from(block.data_pos)
.map_err(|_| err!(UnexpectedEof, "block {index} data offset overflows"))?;
let len = if block.header.is_streamed() {
self.source.len().saturating_sub(start)
} else {
usize::try_from(block.header.used_size)
.map_err(|_| err!(UnexpectedEof, "block {index} used_size overflows"))?
};
let end = start
.checked_add(len)
.ok_or_else(|| err!(UnexpectedEof, "block {index} size overflows"))?;
self.source
.get(start..end)
.ok_or_else(|| err!(UnexpectedEof, "block {index} extends past the end of the file"))
}
pub fn block_compression(&self, index: usize) -> Result<Compression> {
Compression::from_name(self.block(index)?.header.compression_name())
}
pub fn block_data(&self, index: usize) -> Result<Cow<'_, [u8]>> {
let raw = self.block_raw(index)?;
let compression = self.block_compression(index)?;
if compression == Compression::None {
return Ok(Cow::Borrowed(raw));
}
let expected = usize::try_from(self.block(index)?.header.data_size)
.map_err(|_| err!(UnexpectedEof, "block {index} data_size overflows"))?;
Ok(Cow::Owned(compression.decompress(raw, expected)?))
}
pub fn verify_block_checksum(
&self,
index: usize,
) -> Result<(ChecksumStatus, [u8; CHECKSUM_SIZE])> {
let header = &self.block(index)?.header;
if !header.has_checksum() {
return Ok((ChecksumStatus::Absent, [0; CHECKSUM_SIZE]));
}
let expected = header.checksum;
let raw_digest = md5_of(self.block_raw(index)?);
if raw_digest == expected {
return Ok((ChecksumStatus::Valid, raw_digest));
}
if self.block_compression(index)? != Compression::None && self.has_python_checksum_bug() {
let decompressed = md5_of(&self.block_data(index)?);
if decompressed == expected {
return Ok((ChecksumStatus::Valid, decompressed));
}
}
Ok((ChecksumStatus::Invalid, raw_digest))
}
pub fn has_python_checksum_bug(&self) -> bool {
const BUGGY_THROUGH_MAJOR: u32 = 5;
let Ok(Some(doc)) = self.tree() else { return false };
let Some(root) = doc.root() else { return false };
let Some(library) = doc.mapping_get(root, "asdf_library") else {
return false;
};
let name = doc
.mapping_get(library, "name")
.and_then(|id| doc.resolved(id).as_str().map(str::to_string));
if name.as_deref() != Some("asdf") {
return false;
}
doc.mapping_get(library, "version")
.and_then(|id| doc.resolved(id).as_str().map(crate::Version::parse))
.is_some_and(|v| v.major <= BUGGY_THROUGH_MAJOR)
}
}
fn md5_of(data: &[u8]) -> [u8; CHECKSUM_SIZE] {
use md5::{Digest, Md5};
let mut hasher = Md5::new();
hasher.update(data);
hasher.finalize().into()
}
fn find_tagged(doc: &Document, name: &str) -> Vec<asdf_yaml::NodeId> {
use asdf_yaml::NodeData;
let mut out = Vec::new();
let mut seen = std::collections::HashSet::new();
let Some(root) = doc.root() else { return out };
let mut stack = vec![root];
while let Some(id) = stack.pop() {
let resolved = doc.resolve(id);
if !seen.insert(resolved) {
continue;
}
if doc.tag_of(resolved).is_some_and(|t| t.split_version().0 == name) {
out.push(resolved);
}
match &doc.node(resolved).data {
NodeData::Sequence { items, .. } => stack.extend(items.iter().copied()),
NodeData::Mapping { entries, .. } => {
stack.extend(entries.iter().map(|e| e.value));
}
_ => {}
}
}
out.sort();
out
}
impl Reader {
fn block_index_for(&self, source: &crate::core::ndarray::Source) -> Option<usize> {
match source {
crate::core::ndarray::Source::Block(i) => Some(*i),
crate::core::ndarray::Source::LastBlock => self.block_count().checked_sub(1),
_ => None,
}
}
pub fn external_block(&self, uri: &str) -> Result<Vec<u8>> {
let Some(base) = self.path.as_deref().and_then(Path::parent) else {
return Err(err!(
InvalidArgument,
"external source {uri:?} cannot be resolved: this file was not read from disk"
));
};
let relative = external_relative_path(uri)?;
let target = base.join(relative);
let resolved = target.canonicalize().map_err(|e| {
err!(InvalidArgument, "external source {uri:?} ({}): {e}", target.display())
})?;
let root = base.canonicalize().unwrap_or_else(|_| base.to_path_buf());
if !resolved.starts_with(&root) {
return Err(err!(
InvalidArgument,
"external source {uri:?} resolves to {}, outside the referring file's \
directory {}",
resolved.display(),
root.display()
));
}
let referenced = Reader::open(&resolved).map_err(|e| {
err!(InvalidArgument, "external source {uri:?} ({}): {e}", target.display())
})?;
if referenced.block_count() == 0 {
return Err(err!(InvalidArgument, "external source {uri:?} has no blocks"));
}
Ok(referenced.block_data(0)?.into_owned())
}
pub fn tree_inlined(&self) -> Result<Option<(Document, Vec<String>)>> {
use crate::core::elements::{decode_all, inline_ndarray};
use crate::core::ndarray::Ndarray;
let Some(mut doc) = self.tree()? else { return Ok(None) };
let mut skipped = Vec::new();
for id in find_tagged(&doc, "core/ndarray") {
let nd = match Ndarray::parse(&doc, id) {
Ok(nd) => nd,
Err(e) => {
skipped.push(format!("{id:?}: {e}"));
continue;
}
};
if matches!(nd.source, crate::core::ndarray::Source::Inline(_)) {
continue;
}
let data = if let crate::core::ndarray::Source::External(uri) = &nd.source {
match self.external_block(uri) {
Ok(bytes) => Cow::Owned(bytes),
Err(e) => {
skipped.push(format!("{id:?}: {e}"));
continue;
}
}
} else {
let Some(index) = self.block_index_for(&nd.source) else {
skipped.push(format!("{id:?}: data is outside this file ({:?})", nd.source));
continue;
};
match self.block_data(index) {
Ok(d) => d,
Err(e) => {
skipped.push(format!("{id:?}: block {index}: {e}"));
continue;
}
}
};
let shape = match nd.resolved_shape(Some(data.len() as u64)) {
Ok(s) => s,
Err(e) => {
skipped.push(format!("{id:?}: {e}"));
continue;
}
};
match decode_all(&nd, &shape, &data) {
Ok(elements) => inline_ndarray(&mut doc, id, &elements, &shape)?,
Err(e) => skipped.push(format!("{id:?}: {e}")),
}
}
Ok(Some((doc, skipped)))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::block::header::BlockHeader;
use crate::layout::write_block_index;
fn build(payload: &[u8], compression: Compression, checksum_over: Option<&[u8]>) -> Vec<u8> {
let stored = compression.compress(payload).unwrap();
let mut buf = Vec::new();
buf.extend_from_slice(b"#ASDF 1.0.0\n#ASDF_STANDARD 1.6.0\n");
buf.extend_from_slice(
b"%YAML 1.1\n%TAG ! tag:stsci.edu:asdf/\n--- !core/asdf-1.1.0\nx: 1\n...\n",
);
let mut header = BlockHeader {
allocated_size: stored.len() as u64,
used_size: stored.len() as u64,
data_size: payload.len() as u64,
..Default::default()
};
header.set_compression(compression.name()).unwrap();
if let Some(over) = checksum_over {
header.checksum = md5_of(over);
}
let offset = buf.len() as u64;
header.write(&mut buf);
buf.extend_from_slice(&stored);
buf.extend_from_slice(&write_block_index(&[offset]));
buf
}
#[test]
fn external_source_uris_may_not_escape_the_directory() {
assert!(external_relative_path("exploded0000.asdf").is_ok());
assert!(external_relative_path("data/block0.asdf").is_ok());
assert!(external_relative_path("./here.asdf").is_ok());
for bad in [
"",
"/etc/passwd",
"../secrets.asdf",
"data/../../secrets.asdf",
"file:///etc/passwd",
"https://example.invalid/x.asdf",
] {
assert!(
external_relative_path(bad).is_err(),
"{bad:?} should be rejected as an external source"
);
}
}
#[test]
fn a_memory_backed_file_resolves_no_external_sources() {
let file = build(b"whatever", Compression::None, None);
let r = Reader::from_bytes(file);
let r = r.unwrap();
assert!(r.path().is_none());
assert!(r.external_block("other.asdf").is_err());
}
#[test]
fn an_external_source_is_read_from_the_neighbouring_file() {
let dir = std::env::temp_dir().join(format!("asdf-exploded-{}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
let payload: Vec<u8> = [1i32, 2, 3, 4].iter().flat_map(|v| v.to_le_bytes()).collect();
let data_file = dir.join("holder0000.asdf");
std::fs::write(&data_file, build(&payload, Compression::None, None)).unwrap();
let mut buf = Vec::new();
buf.extend_from_slice(b"#ASDF 1.0.0\n#ASDF_STANDARD 1.6.0\n");
buf.extend_from_slice(b"%YAML 1.1\n%TAG ! tag:stsci.edu:asdf/\n--- !core/asdf-1.1.0\n");
buf.extend_from_slice(
b"data: !core/ndarray-1.1.0\n source: holder0000.asdf\n \
datatype: int32\n byteorder: little\n shape: [4]\n",
);
buf.extend_from_slice(b"...\n");
let referring = dir.join("holder.asdf");
std::fs::write(&referring, buf).unwrap();
let r = Reader::open(&referring).unwrap();
assert_eq!(r.block_count(), 0, "the referring file holds no blocks itself");
assert_eq!(r.external_block("holder0000.asdf").unwrap(), payload);
let (doc, skipped) = r.tree_inlined().unwrap().unwrap();
assert!(skipped.is_empty(), "nothing should be left un-inlined: {skipped:?}");
let root = doc.root().unwrap();
let array = doc.mapping_get(root, "data").unwrap();
let values = doc.mapping_get(array, "data").unwrap();
let items = doc.sequence_items(values).unwrap();
let read: Vec<&str> = items.iter().map(|i| doc.resolved(*i).as_str().unwrap()).collect();
assert_eq!(read, ["1", "2", "3", "4"]);
assert!(doc.mapping_get(array, "source").is_none());
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn a_missing_external_file_is_reported_not_silently_skipped() {
let dir =
std::env::temp_dir().join(format!("asdf-exploded-missing-{}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
let referring = dir.join("dangling.asdf");
let mut buf = Vec::new();
buf.extend_from_slice(b"#ASDF 1.0.0\n#ASDF_STANDARD 1.6.0\n");
buf.extend_from_slice(b"%YAML 1.1\n%TAG ! tag:stsci.edu:asdf/\n--- !core/asdf-1.1.0\n");
buf.extend_from_slice(
b"data: !core/ndarray-1.1.0\n source: nowhere0000.asdf\n \
datatype: int32\n byteorder: little\n shape: [4]\n",
);
buf.extend_from_slice(b"...\n");
std::fs::write(&referring, buf).unwrap();
let r = Reader::open(&referring).unwrap();
let (_, skipped) = r.tree_inlined().unwrap().unwrap();
assert_eq!(skipped.len(), 1);
assert!(skipped[0].contains("nowhere0000.asdf"), "{skipped:?}");
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn reads_tree_and_block_data() {
let payload = b"hello block data".to_vec();
let file = build(&payload, Compression::None, None);
let r = Reader::from_bytes(file).unwrap();
assert_eq!(r.block_count(), 1);
assert_eq!(&*r.block_data(0).unwrap(), &payload[..]);
let doc = r.tree().unwrap().unwrap();
let root = doc.root().unwrap();
assert!(doc.mapping_get(root, "x").is_some());
}
#[test]
fn uncompressed_data_is_borrowed_not_copied() {
let file = build(b"borrow me", Compression::None, None);
let r = Reader::from_bytes(file).unwrap();
assert!(matches!(r.block_data(0).unwrap(), Cow::Borrowed(_)));
}
#[test]
fn compressed_data_round_trips() {
let payload = vec![7u8; 4096];
for c in crate::compression::available() {
let file = build(&payload, c, None);
let r = Reader::from_bytes(file).unwrap();
assert_eq!(r.block_compression(0).unwrap(), c);
assert_eq!(&*r.block_data(0).unwrap(), &payload[..], "{c:?}");
assert!(r.block_raw(0).unwrap().len() < payload.len(), "{c:?}");
}
}
#[test]
fn valid_checksums_verify() {
let payload = b"checksum me".to_vec();
let file = build(&payload, Compression::None, Some(&payload));
let r = Reader::from_bytes(file).unwrap();
let (status, _) = r.verify_block_checksum(0).unwrap();
assert_eq!(status, ChecksumStatus::Valid);
}
#[test]
fn invalid_checksums_are_reported() {
let payload = b"checksum me".to_vec();
let file = build(&payload, Compression::None, Some(b"something else"));
let r = Reader::from_bytes(file).unwrap();
let (status, computed) = r.verify_block_checksum(0).unwrap();
assert_eq!(status, ChecksumStatus::Invalid);
assert_eq!(computed, md5_of(&payload), "the digest of the real data is reported");
}
#[test]
fn an_absent_checksum_is_not_a_failure() {
let file = build(b"no checksum", Compression::None, None);
let r = Reader::from_bytes(file).unwrap();
let (status, _) = r.verify_block_checksum(0).unwrap();
assert_eq!(status, ChecksumStatus::Absent);
assert!(!status.is_failure());
}
#[cfg(feature = "zlib")]
#[test]
fn compressed_checksums_cover_the_stored_bytes() {
let payload = vec![3u8; 2048];
let stored = Compression::Zlib.compress(&payload).unwrap();
let file = build(&payload, Compression::Zlib, Some(&stored));
let r = Reader::from_bytes(file).unwrap();
assert_eq!(r.verify_block_checksum(0).unwrap().0, ChecksumStatus::Valid);
}
#[cfg(feature = "zlib")]
#[test]
fn the_python_checksum_bug_is_worked_around() {
let payload = vec![9u8; 2048];
let stored = Compression::Zlib.compress(&payload).unwrap();
let make = |library_version: &str| {
let mut buf = Vec::new();
buf.extend_from_slice(b"#ASDF 1.0.0\n#ASDF_STANDARD 1.6.0\n");
buf.extend_from_slice(b"%YAML 1.1\n%TAG ! tag:stsci.edu:asdf/\n--- !core/asdf-1.1.0\n");
buf.extend_from_slice(
format!(
"asdf_library: !core/software-1.0.0 {{name: asdf, version: {library_version}}}\n"
)
.as_bytes(),
);
buf.extend_from_slice(b"...\n");
let mut header = BlockHeader {
allocated_size: stored.len() as u64,
used_size: stored.len() as u64,
data_size: payload.len() as u64,
checksum: md5_of(&payload),
..Default::default()
};
header.set_compression("zlib").unwrap();
header.write(&mut buf);
buf.extend_from_slice(&stored);
buf
};
let r = Reader::from_bytes(make("4.1.0")).unwrap();
assert!(r.has_python_checksum_bug());
assert_eq!(
r.verify_block_checksum(0).unwrap().0,
ChecksumStatus::Valid,
"an affected writer's checksum should verify against the uncompressed data"
);
let r = Reader::from_bytes(make("6.0.0")).unwrap();
assert!(!r.has_python_checksum_bug());
assert_eq!(
r.verify_block_checksum(0).unwrap().0,
ChecksumStatus::Invalid,
"the workaround must not apply to writers that are not affected"
);
}
#[test]
fn out_of_range_block_indices_error() {
let file = build(b"one block", Compression::None, None);
let r = Reader::from_bytes(file).unwrap();
assert!(r.block(1).is_err());
assert!(r.block_data(99).is_err());
}
#[test]
fn a_file_without_a_tree_reads_cleanly() {
let mut buf = Vec::new();
buf.extend_from_slice(b"#ASDF 1.0.0\n#ASDF_STANDARD 1.6.0\n");
let header =
BlockHeader { allocated_size: 4, used_size: 4, data_size: 4, ..Default::default() };
header.write(&mut buf);
buf.extend_from_slice(b"data");
let r = Reader::from_bytes(buf).unwrap();
assert!(r.tree().unwrap().is_none());
assert_eq!(&*r.block_data(0).unwrap(), b"data");
}
}