use std::io::{self, Read, Seek, SeekFrom};
use crate::tree::TreeNode;
const KOLY_MAGIC: &[u8; 4] = b"koly";
const KOLY_SIZE: usize = 512;
const KOLY_VERSION: u32 = 4;
#[derive(Debug)]
pub enum Error {
TooShort,
BadMagic,
BadVersion(u32),
Io(io::Error),
}
impl std::fmt::Display for Error {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Error::TooShort => write!(f, "DMG file is shorter than the 512-byte koly trailer"),
Error::BadMagic => write!(f, "DMG koly magic b\"koly\" not found at file end − 512"),
Error::BadVersion(v) => write!(f, "DMG koly version {v} is not supported (expected 4)"),
Error::Io(e) => write!(f, "DMG I/O error: {e}"),
}
}
}
impl std::error::Error for Error {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Error::Io(e) => Some(e),
_ => None,
}
}
}
impl From<io::Error> for Error {
fn from(e: io::Error) -> Self {
Error::Io(e)
}
}
struct Koly {
xml_offset: u64,
xml_length: u64,
sector_count: u64,
}
fn read_koly<R: Read + Seek>(r: &mut R) -> Result<Koly, Error> {
let file_len = r.seek(SeekFrom::End(0))?;
if file_len < KOLY_SIZE as u64 {
return Err(Error::TooShort);
}
r.seek(SeekFrom::Start(file_len - KOLY_SIZE as u64))?;
let mut buf = [0u8; KOLY_SIZE];
match r.read_exact(&mut buf) {
Ok(()) => {}
Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Err(Error::TooShort),
Err(e) => return Err(Error::Io(e)),
}
if &buf[0..4] != KOLY_MAGIC {
return Err(Error::BadMagic);
}
let version = u32::from_be_bytes([buf[4], buf[5], buf[6], buf[7]]);
if version != KOLY_VERSION {
return Err(Error::BadVersion(version));
}
let xml_offset = u64::from_be_bytes([
buf[216], buf[217], buf[218], buf[219], buf[220], buf[221], buf[222], buf[223],
]);
let xml_length = u64::from_be_bytes([
buf[224], buf[225], buf[226], buf[227], buf[228], buf[229], buf[230], buf[231],
]);
let sector_count = u64::from_be_bytes([
buf[492], buf[493], buf[494], buf[495], buf[496], buf[497], buf[498], buf[499],
]);
Ok(Koly {
xml_offset,
xml_length,
sector_count,
})
}
pub fn detect<R: Read + Seek>(r: &mut R) -> Result<(), Error> {
let pos = r.stream_position()?;
let result = detect_inner(r);
let _ = r.seek(SeekFrom::Start(pos));
result
}
fn detect_inner<R: Read + Seek>(r: &mut R) -> Result<(), Error> {
let file_len = r.seek(SeekFrom::End(0))?;
if file_len < KOLY_SIZE as u64 {
return Err(Error::TooShort);
}
r.seek(SeekFrom::Start(file_len - KOLY_SIZE as u64))?;
let mut magic = [0u8; 4];
match r.read_exact(&mut magic) {
Ok(()) => {}
Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Err(Error::TooShort),
Err(e) => return Err(Error::Io(e)),
}
if &magic != KOLY_MAGIC {
return Err(Error::BadMagic);
}
Ok(())
}
#[derive(Debug)]
struct BlkxEntry {
name: String,
}
fn parse_plist_xml(xml: &str) -> Vec<BlkxEntry> {
let mut entries = Vec::new();
let blkx_key = "<key>blkx</key>";
let blkx_pos = match xml.find(blkx_key) {
Some(p) => p + blkx_key.len(),
None => return entries,
};
let array_open = "<array>";
let array_close = "</array>";
let array_start = match xml[blkx_pos..].find(array_open) {
Some(p) => blkx_pos + p + array_open.len(),
None => return entries,
};
let array_end = match xml[array_start..].find(array_close) {
Some(p) => array_start + p,
None => return entries,
};
let array_body = &xml[array_start..array_end];
let mut pos = 0;
while let Some(dict_rel) = array_body[pos..].find("<dict>") {
let dict_start = pos + dict_rel + "<dict>".len();
let dict_end = match array_body[dict_start..].find("</dict>") {
Some(e) => dict_start + e,
None => break,
};
let dict_body = &array_body[dict_start..dict_end];
let cf_name = extract_keyed_string(dict_body, "CFName");
let plain_name = extract_keyed_string(dict_body, "Name");
let name = cf_name.or(plain_name).unwrap_or_default();
if !name.is_empty() {
entries.push(BlkxEntry { name });
}
pos = dict_end + "</dict>".len();
}
entries
}
fn extract_keyed_string(text: &str, key: &str) -> Option<String> {
let key_tag = format!("<key>{key}</key>");
let key_pos = text.find(&key_tag)? + key_tag.len();
let rest = &text[key_pos..];
let string_open = "<string>";
let string_close = "</string>";
let val_start = rest.find(string_open)? + string_open.len();
let val_end = rest[val_start..].find(string_close)? + val_start;
Some(rest[val_start..val_end].trim().to_string())
}
pub fn detect_and_parse<R: Read + Seek>(r: &mut R) -> Result<TreeNode, Error> {
let koly = read_koly(r)?;
let file_len = r.seek(SeekFrom::End(0))?;
let xml_sane = koly.xml_length > 0
&& koly.xml_length <= 32 * 1024 * 1024
&& koly.xml_offset < file_len
&& koly.xml_offset + koly.xml_length <= file_len;
let xml_text = if xml_sane {
r.seek(SeekFrom::Start(koly.xml_offset))?;
let read_len = koly.xml_length as usize;
let mut raw = vec![0u8; read_len];
match r.read_exact(&mut raw) {
Ok(()) => String::from_utf8_lossy(&raw).into_owned(),
Err(_) => String::new(),
}
} else {
String::new()
};
let entries = parse_plist_xml(&xml_text);
let mut root = TreeNode::new_directory("/".to_string());
if entries.is_empty() {
let disk_size = koly.sector_count.saturating_mul(512);
let mut disk_node = TreeNode::new_file("disk.dmg".to_string(), disk_size);
disk_node.file_length = Some(disk_size);
root.add_child(disk_node);
} else {
for entry in &entries {
let node = TreeNode::new_directory(entry.name.clone());
root.add_child(node);
}
}
root.calculate_directory_size();
Ok(root)
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
fn build_dmg(prefix_len: usize, xml: &str, sector_count: u64) -> Vec<u8> {
let xml_bytes = xml.as_bytes();
let xml_len = xml_bytes.len();
let total = prefix_len + xml_len + KOLY_SIZE;
let mut buf = vec![0u8; total];
let xml_start = prefix_len;
buf[xml_start..xml_start + xml_len].copy_from_slice(xml_bytes);
let koly_start = total - KOLY_SIZE;
buf[koly_start..koly_start + 4].copy_from_slice(KOLY_MAGIC);
buf[koly_start + 4..koly_start + 8].copy_from_slice(&4u32.to_be_bytes());
buf[koly_start + 8..koly_start + 12].copy_from_slice(&512u32.to_be_bytes());
if xml_len > 0 {
let xml_offset = xml_start as u64;
buf[koly_start + 216..koly_start + 224].copy_from_slice(&xml_offset.to_be_bytes());
buf[koly_start + 224..koly_start + 232]
.copy_from_slice(&(xml_len as u64).to_be_bytes());
}
buf[koly_start + 492..koly_start + 500].copy_from_slice(§or_count.to_be_bytes());
buf
}
#[test]
fn detect_valid_dmg_ok() {
let dmg = build_dmg(0, "", 2048);
let mut c = Cursor::new(&dmg);
assert!(
detect(&mut c).is_ok(),
"detect() should succeed on a valid koly trailer"
);
}
#[test]
fn detect_restores_position() {
let dmg = build_dmg(0, "", 2048);
let mut c = Cursor::new(&dmg);
c.seek(SeekFrom::Start(42)).unwrap();
detect(&mut c).unwrap();
assert_eq!(
c.stream_position().unwrap(),
42,
"detect() must restore the stream position"
);
}
#[test]
fn detect_rejects_bad_magic() {
let data = vec![0u8; 1024];
let mut c = Cursor::new(&data);
assert!(
matches!(detect(&mut c), Err(Error::BadMagic)),
"all-zeros buffer should fail with BadMagic"
);
}
#[test]
fn detect_rejects_too_short() {
let data = vec![0u8; 64];
let mut c = Cursor::new(&data);
assert!(
matches!(detect(&mut c), Err(Error::TooShort)),
"64-byte buffer should fail with TooShort"
);
}
#[test]
fn detect_rejects_bad_version() {
let mut buf = vec![0u8; KOLY_SIZE];
buf[0..4].copy_from_slice(KOLY_MAGIC);
buf[4..8].copy_from_slice(&99u32.to_be_bytes());
let mut c = Cursor::new(&buf);
assert!(
detect(&mut c).is_ok(),
"detect() checks only the magic, not the version"
);
let mut c2 = Cursor::new(&buf);
assert!(
matches!(detect_and_parse(&mut c2), Err(Error::BadVersion(99))),
"detect_and_parse() should return BadVersion(99) for version 99"
);
}
#[test]
fn parse_plist_no_blkx() {
let xml = r#"<?xml version="1.0"?><plist version="1.0"><dict></dict></plist>"#;
let entries = parse_plist_xml(xml);
assert!(entries.is_empty(), "no blkx key → empty entries");
}
#[test]
fn parse_plist_single_partition_cfname() {
let xml = r#"<plist><dict>
<key>resource-fork</key>
<dict>
<key>blkx</key>
<array>
<dict>
<key>CFName</key><string>Apple_HFS : macOS</string>
<key>Name</key><string>Driver Descriptor Map</string>
<key>Data</key><data>AAAA</data>
</dict>
</array>
</dict>
</dict></plist>"#;
let entries = parse_plist_xml(xml);
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].name, "Apple_HFS : macOS");
}
#[test]
fn parse_plist_multiple_partitions() {
let xml = r#"<plist><dict>
<key>resource-fork</key>
<dict>
<key>blkx</key>
<array>
<dict><key>CFName</key><string>Driver Descriptor Map</string></dict>
<dict><key>CFName</key><string>Apple_partition_map</string></dict>
<dict><key>Name</key><string>Apple_HFS : macOS</string></dict>
</array>
</dict>
</dict></plist>"#;
let entries = parse_plist_xml(xml);
assert_eq!(entries.len(), 3);
assert_eq!(entries[0].name, "Driver Descriptor Map");
assert_eq!(entries[1].name, "Apple_partition_map");
assert_eq!(entries[2].name, "Apple_HFS : macOS");
}
#[test]
fn parse_plist_falls_back_to_name_when_no_cfname() {
let xml = r#"<plist><dict>
<key>resource-fork</key><dict>
<key>blkx</key><array>
<dict><key>Name</key><string>free</string></dict>
</array></dict></dict></plist>"#;
let entries = parse_plist_xml(xml);
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].name, "free");
}
#[test]
fn parse_no_xml_returns_synthetic_disk_node() {
let dmg = build_dmg(0, "", 4096);
let mut c = Cursor::new(&dmg);
let root = detect_and_parse(&mut c).expect("detect_and_parse should succeed");
assert_eq!(root.name, "/");
assert!(root.is_directory);
assert_eq!(root.children.len(), 1);
assert_eq!(root.children[0].name, "disk.dmg");
assert!(!root.children[0].is_directory);
assert_eq!(root.children[0].file_length, Some(4096 * 512));
}
#[test]
fn parse_xml_with_partitions_returns_directory_children() {
let xml = r#"<plist><dict>
<key>resource-fork</key><dict>
<key>blkx</key><array>
<dict><key>CFName</key><string>Apple_HFS : macOS</string></dict>
<dict><key>CFName</key><string>free</string></dict>
</array></dict></dict></plist>"#;
let dmg = build_dmg(512, xml, 8192);
let mut c = Cursor::new(&dmg);
let root = detect_and_parse(&mut c).expect("detect_and_parse should succeed");
assert_eq!(root.children.len(), 2);
assert_eq!(root.children[0].name, "Apple_HFS : macOS");
assert!(root.children[0].is_directory);
assert_eq!(root.children[1].name, "free");
}
#[test]
fn parse_rejects_bad_magic() {
let mut buf = vec![0u8; KOLY_SIZE];
buf[0..4].copy_from_slice(b"XXXX");
let mut c = Cursor::new(&buf);
assert!(
matches!(detect_and_parse(&mut c), Err(Error::BadMagic)),
"bad magic should return Error::BadMagic"
);
}
#[test]
fn koly_sector_count_read_correctly() {
let sector_count: u64 = 123_456;
let dmg = build_dmg(0, "", sector_count);
let mut c = Cursor::new(&dmg);
let root = detect_and_parse(&mut c).expect("detect_and_parse should succeed");
assert_eq!(root.children[0].file_length, Some(sector_count * 512));
}
#[test]
fn error_display_too_short() {
let msg = format!("{}", Error::TooShort);
assert!(msg.contains("512") || msg.contains("short"), "got: {msg}");
}
#[test]
fn error_display_bad_magic() {
let msg = format!("{}", Error::BadMagic);
assert!(msg.contains("koly") || msg.contains("magic"), "got: {msg}");
}
#[test]
fn error_display_bad_version() {
let msg = format!("{}", Error::BadVersion(99));
assert!(msg.contains("99"), "got: {msg}");
}
#[test]
fn error_display_io() {
let io = io::Error::other("disk");
let msg = format!("{}", Error::Io(io));
assert!(msg.contains("disk"), "got: {msg}");
}
#[test]
fn error_source_io() {
use std::error::Error as StdError;
assert!(Error::Io(io::Error::other("s")).source().is_some());
}
#[test]
fn error_source_non_io() {
use std::error::Error as StdError;
assert!(Error::TooShort.source().is_none());
assert!(Error::BadMagic.source().is_none());
assert!(Error::BadVersion(4).source().is_none());
}
#[test]
fn error_from_io_error() {
let e = Error::from(io::Error::other("dmg test"));
assert!(matches!(e, Error::Io(_)));
}
#[test]
fn read_koly_too_short_returns_error() {
let data = vec![0u8; 100];
let mut c = Cursor::new(data);
assert!(matches!(read_koly(&mut c), Err(Error::TooShort)));
}
#[test]
fn parse_plist_xml_no_array_after_blkx_key() {
let xml = "<key>blkx</key><key>something_else</key>";
let entries = parse_plist_xml(xml);
assert!(entries.is_empty(), "no <array> → should return empty");
}
#[test]
fn parse_plist_xml_no_closing_array_tag() {
let xml = "<key>blkx</key><array><dict><key>CFName</key><string>EFI</string></dict>";
let entries = parse_plist_xml(xml);
assert!(entries.is_empty(), "no </array> → should return empty");
}
#[test]
fn parse_plist_xml_no_closing_dict_tag() {
let xml = "<key>blkx</key><array><dict><key>CFName</key><string>EFI</string></array>";
let entries = parse_plist_xml(xml);
assert!(entries.is_empty(), "unclosed <dict> → should return empty");
}
}