1use anyhow::Result;
2use crc32fast::Hasher;
3use fatfs::{Dir, FileSystem, FormatVolumeOptions, FsOptions, ReadWriteSeek};
4use flate2::bufread::ZlibEncoder;
5use flate2::read::ZlibDecoder;
6use flate2::Compression;
7use fscommon::BufStream;
8use gpt::mbr::{PartRecord, ProtectiveMBR};
9use std::fs::File;
10use std::io::{BufReader, BufWriter, Cursor, Read, Seek, SeekFrom, Write};
11use std::path::Path;
12
13mod blkx;
14mod koly;
15mod xml;
16
17pub use crate::blkx::*;
18pub use crate::koly::*;
19pub use crate::xml::*;
20
21pub struct DmgReader<R: Read + Seek> {
22 koly: KolyTrailer,
23 xml: Plist,
24 r: R,
25}
26
27impl DmgReader<BufReader<File>> {
28 pub fn open(path: &Path) -> Result<Self> {
29 let r = BufReader::new(File::open(path)?);
30 Self::new(r)
31 }
32}
33
34impl<R: Read + Seek> DmgReader<R> {
35 pub fn new(mut r: R) -> Result<Self> {
36 let koly = KolyTrailer::read_from(&mut r)?;
37 r.seek(SeekFrom::Start(koly.plist_offset))?;
38 let mut xml = Vec::with_capacity(koly.plist_length as usize);
39 (&mut r).take(koly.plist_length).read_to_end(&mut xml)?;
40 let xml: Plist = plist::from_reader_xml(&xml[..])?;
41 Ok(Self { koly, xml, r })
42 }
43
44 pub fn koly(&self) -> &KolyTrailer {
45 &self.koly
46 }
47
48 pub fn plist(&self) -> &Plist {
49 &self.xml
50 }
51
52 pub fn sector(&mut self, chunk: &BlkxChunk) -> Result<impl Read + '_> {
53 self.r.seek(SeekFrom::Start(chunk.compressed_offset))?;
54 let compressed_chunk = (&mut self.r).take(chunk.compressed_length);
55 match chunk.ty().expect("unknown chunk type") {
56 ChunkType::Ignore | ChunkType::Zero | ChunkType::Comment => {
57 Ok(Box::new(std::io::repeat(0).take(chunk.compressed_length)) as Box<dyn Read>)
58 }
59 ChunkType::Raw => Ok(Box::new(compressed_chunk)),
60 ChunkType::Zlib => Ok(Box::new(ZlibDecoder::new(compressed_chunk))),
61 ChunkType::Adc | ChunkType::Bzlib | ChunkType::Lzfse => unimplemented!(),
62 ChunkType::Term => Ok(Box::new(std::io::empty())),
63 }
64 }
65
66 pub fn data_checksum(&mut self) -> Result<u32> {
67 self.r.seek(SeekFrom::Start(self.koly.data_fork_offset))?;
68 let mut data_fork = Vec::with_capacity(self.koly.data_fork_length as usize);
69 (&mut self.r)
70 .take(self.koly.data_fork_length)
71 .read_to_end(&mut data_fork)?;
72 Ok(crc32fast::hash(&data_fork))
73 }
74
75 pub fn partition_table(&self, i: usize) -> Result<BlkxTable> {
76 self.plist().partitions()[i].table()
77 }
78
79 pub fn partition_name(&self, i: usize) -> &str {
80 &self.plist().partitions()[i].name
81 }
82
83 pub fn partition_data(&mut self, i: usize) -> Result<Vec<u8>> {
84 let table = self.plist().partitions()[i].table()?;
85 let mut partition = vec![];
86 for chunk in &table.chunks {
87 std::io::copy(&mut self.sector(chunk)?, &mut partition)?;
88 }
89 Ok(partition)
90 }
91}
92
93pub struct DmgWriter<W: Write + Seek> {
94 xml: Plist,
95 w: W,
96 data_hasher: Hasher,
97 main_hasher: Hasher,
98 sector_number: u64,
99 compressed_offset: u64,
100}
101
102impl DmgWriter<BufWriter<File>> {
103 pub fn create(path: &Path) -> Result<Self> {
104 let w = BufWriter::new(File::create(path)?);
105 Ok(Self::new(w))
106 }
107}
108
109impl<W: Write + Seek> DmgWriter<W> {
110 pub fn new(w: W) -> Self {
111 Self {
112 xml: Default::default(),
113 w,
114 data_hasher: Hasher::new(),
115 main_hasher: Hasher::new(),
116 sector_number: 0,
117 compressed_offset: 0,
118 }
119 }
120
121 pub fn create_fat32(mut self, fat32: &[u8]) -> Result<()> {
122 anyhow::ensure!(fat32.len() % 512 == 0);
123 let sector_count = fat32.len() as u64 / 512;
124 let mut mbr = ProtectiveMBR::new();
125 let mut partition = PartRecord::new_protective(Some(sector_count.try_into()?));
126 partition.os_type = 11;
127 mbr.set_partition(0, partition);
128 let mbr = mbr.as_bytes()?;
129 self.add_partition("Master Boot Record (MBR : 0)", &mbr)?;
130 self.add_partition("FAT32 (FAT32 : 1)", fat32)?;
131 self.finish()?;
132 Ok(())
133 }
134
135 pub fn add_partition(&mut self, name: &str, bytes: &[u8]) -> Result<()> {
136 anyhow::ensure!(bytes.len() % 512 == 0);
137 let id = self.xml.partitions().len() as u32;
138 let name = name.to_string();
139 let mut table = BlkxTable::new(id, self.sector_number, crc32fast::hash(bytes));
140 for chunk in bytes.chunks(2048 * 512) {
141 let mut encoder = ZlibEncoder::new(chunk, Compression::best());
142 let mut compressed = vec![];
143 encoder.read_to_end(&mut compressed)?;
144 let compressed_length = compressed.len() as u64;
145 let sector_count = chunk.len() as u64 / 512;
146 self.w.write_all(&compressed)?;
147 self.data_hasher.update(&compressed);
148 table.add_chunk(BlkxChunk::new(
149 ChunkType::Zlib,
150 self.sector_number,
151 sector_count,
152 self.compressed_offset,
153 compressed_length,
154 ));
155 self.sector_number += sector_count;
156 self.compressed_offset += compressed_length;
157 }
158 table.add_chunk(BlkxChunk::term(self.sector_number, self.compressed_offset));
159 self.main_hasher.update(&table.checksum.data[..4]);
160 self.xml
161 .add_partition(Partition::new(id as i32 - 1, name, table));
162 Ok(())
163 }
164
165 pub fn finish(mut self) -> Result<()> {
166 let mut xml = vec![];
167 plist::to_writer_xml(&mut xml, &self.xml)?;
168 let pos = self.w.seek(SeekFrom::Current(0))?;
169 let data_digest = self.data_hasher.finalize();
170 let main_digest = self.main_hasher.finalize();
171 let koly = KolyTrailer::new(
172 pos,
173 self.sector_number,
174 pos,
175 xml.len() as _,
176 data_digest,
177 main_digest,
178 );
179 self.w.write_all(&xml)?;
180 koly.write_to(&mut self.w)?;
181 Ok(())
182 }
183}
184
185fn symlink(target: &str) -> Vec<u8> {
186 let xsym = format!(
187 "XSym\n{:04}\n{:x}\n{}\n",
188 target.as_bytes().len(),
189 md5::compute(target.as_bytes()),
190 target,
191 );
192 let mut xsym = xsym.into_bytes();
193 xsym.resize(1067, b' ');
194 xsym
195}
196
197fn add_dir<T: ReadWriteSeek>(src: &Path, dest: &Dir<'_, T>) -> Result<()> {
198 for entry in std::fs::read_dir(src)? {
199 let entry = entry?;
200 let file_name = entry.file_name();
201 let file_name = file_name.to_str().unwrap();
202 let source = src.join(&file_name);
203 let file_type = entry.file_type()?;
204 if file_type.is_dir() {
205 let d = dest.create_dir(file_name)?;
206 add_dir(&source, &d)?;
207 } else if file_type.is_file() {
208 let mut f = dest.create_file(file_name)?;
209 std::io::copy(&mut File::open(source)?, &mut f)?;
210 } else if file_type.is_symlink() {
211 let target = std::fs::read_link(&source)?;
212 let xsym = symlink(target.to_str().unwrap());
213 let mut f = dest.create_file(file_name)?;
214 std::io::copy(&mut &xsym[..], &mut f)?;
215 }
216 }
217 Ok(())
218}
219
220pub fn create_dmg(dir: &Path, dmg: &Path, volume_label: &str, total_sectors: u32) -> Result<()> {
221 let mut fat32 = vec![0; total_sectors as usize * 512];
222 {
223 let mut volume_label_bytes = [0; 11];
224 let end = std::cmp::min(volume_label_bytes.len(), volume_label.len());
225 volume_label_bytes[..end].copy_from_slice(&volume_label.as_bytes()[..end]);
226 let volume_options = FormatVolumeOptions::new()
227 .volume_label(volume_label_bytes)
228 .bytes_per_sector(512)
229 .total_sectors(total_sectors);
230 let mut disk = BufStream::new(Cursor::new(&mut fat32));
231 fatfs::format_volume(&mut disk, volume_options)?;
232 let fs = FileSystem::new(disk, FsOptions::new())?;
233 let file_name = dir.file_name().unwrap().to_str().unwrap();
234 let dest = fs.root_dir().create_dir(file_name)?;
235 add_dir(dir, &dest)?;
236 }
237 DmgWriter::create(dmg)?.create_fat32(&fat32)
238}
239
240#[cfg(test)]
241mod tests {
242 use super::*;
243 use gpt::disk::LogicalBlockSize;
244
245 static DMG: &[u8] = include_bytes!("../assets/raqote-winit.dmg");
246
247 fn print_dmg<R: Read + Seek>(dmg: &DmgReader<R>) -> Result<()> {
248 println!("{:?}", dmg.koly());
249 println!("{:?}", dmg.plist());
250 for partition in dmg.plist().partitions() {
251 let table = partition.table()?;
252 println!("{:?}", table);
253 println!("table checksum 0x{:x}", u32::from(table.checksum));
254 for (i, chunk) in table.chunks.iter().enumerate() {
255 println!("{} {:?}", i, chunk);
256 }
257 }
258 Ok(())
259 }
260
261 #[test]
262 fn read_koly_trailer() -> Result<()> {
263 let koly = KolyTrailer::read_from(&mut Cursor::new(DMG))?;
264 let mut bytes = [0; 512];
266 koly.write_to(&mut &mut bytes[..])?;
267 let koly2 = KolyTrailer::read_from(&mut Cursor::new(&bytes))?;
268 assert_eq!(koly, koly2);
269 Ok(())
270 }
271
272 #[test]
273 fn only_read_dmg() -> Result<()> {
274 let mut dmg = DmgReader::new(Cursor::new(DMG))?;
275 print_dmg(&dmg)?;
276 assert_eq!(
277 UdifChecksum::new(dmg.data_checksum()?),
278 dmg.koly().data_fork_digest
279 );
280 let mut buffer = vec![];
281 let mut dmg2 = DmgWriter::new(Cursor::new(&mut buffer));
282 for i in 0..dmg.plist().partitions().len() {
283 let data = dmg.partition_data(i)?;
284 let name = dmg.partition_name(i);
285 dmg2.add_partition(name, &data)?;
286 }
287 dmg2.finish()?;
288 let mut dmg2 = DmgReader::new(Cursor::new(buffer))?;
289 print_dmg(&dmg2)?;
290 assert_eq!(
291 UdifChecksum::new(dmg2.data_checksum()?),
292 dmg2.koly().data_fork_digest
293 );
294 for i in 0..dmg.plist().partitions().len() {
295 let table = dmg.partition_table(i)?;
296 let data = dmg.partition_data(i)?;
297 let expected = u32::from(table.checksum);
298 let calculated = crc32fast::hash(&data);
299 assert_eq!(expected, calculated);
300 }
301 assert_eq!(dmg.koly().main_digest, dmg2.koly().main_digest);
302 println!("data crc32 0x{:x}", u32::from(dmg.koly().data_fork_digest));
303 println!("main crc32 0x{:x}", u32::from(dmg.koly().main_digest));
304 Ok(())
305 }
306
307 #[test]
308 fn read_dmg_partition_mbr() -> Result<()> {
309 let mut dmg = DmgReader::new(Cursor::new(DMG))?;
310 let mbr = dmg.partition_data(0)?;
311 println!("{:?}", mbr);
312 let mbr = ProtectiveMBR::from_bytes(&mbr, LogicalBlockSize::Lb512)?;
313 println!("{:?}", mbr);
314 Ok(())
315 }
316
317 #[test]
318 fn read_dmg_partition_fat32() -> Result<()> {
319 let mut dmg = DmgReader::new(Cursor::new(DMG))?;
320 let fat32 = dmg.partition_data(1)?;
321 let fs = FileSystem::new(Cursor::new(fat32), FsOptions::new())?;
322 println!("volume: {}", fs.volume_label());
323 for entry in fs.root_dir().iter() {
324 let entry = entry?;
325 println!("{}", entry.file_name());
326 }
327 Ok(())
328 }
329
330 #[test]
331 fn checksum() -> Result<()> {
332 let mut dmg = DmgReader::new(Cursor::new(DMG))?;
333 assert_eq!(
334 UdifChecksum::new(dmg.data_checksum()?),
335 dmg.koly().data_fork_digest
336 );
337 for i in 0..dmg.plist().partitions().len() {
338 let table = dmg.partition_table(i)?;
339 let data = dmg.partition_data(i)?;
340 let expected = u32::from(table.checksum);
341 let calculated = crc32fast::hash(&data);
342 assert_eq!(expected, calculated);
343 }
344 Ok(())
345 }
346}