Skip to main content

compression/
aa_byte_stream.rs

1use crate::{ffi, util, CompressionError, Result};
2use std::ffi::c_void;
3use std::fs::File;
4use std::ops::{BitAnd, BitAndAssign, BitOr, BitOrAssign};
5use std::os::fd::{IntoRawFd, RawFd};
6use std::ptr::NonNull;
7
8const READ_CHUNK_LEN: usize = 32 * 1024;
9
10pub const OPEN_READ_ONLY: i32 = 0x0000;
11pub const OPEN_WRITE_ONLY: i32 = 0x0001;
12pub const OPEN_READ_WRITE: i32 = 0x0002;
13pub const OPEN_CREATE: i32 = 0x0200;
14pub const OPEN_TRUNCATE: i32 = 0x0400;
15pub const DEFAULT_FILE_MODE: u32 = 0o644;
16
17#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
18pub enum ArchiveCompressionAlgorithm {
19    None,
20    Lz4,
21    Zlib,
22    Lzma,
23    Lzfse,
24    Lzbitmap,
25}
26
27impl ArchiveCompressionAlgorithm {
28    pub const fn from_raw(raw: u32) -> Option<Self> {
29        match raw {
30            0x000 => Some(Self::None),
31            0x100 => Some(Self::Lz4),
32            0x505 => Some(Self::Zlib),
33            0x306 => Some(Self::Lzma),
34            0x801 => Some(Self::Lzfse),
35            0x702 => Some(Self::Lzbitmap),
36            _ => None,
37        }
38    }
39
40    pub const fn as_raw(self) -> u32 {
41        match self {
42            Self::None => 0x000,
43            Self::Lz4 => 0x100,
44            Self::Zlib => 0x505,
45            Self::Lzma => 0x306,
46            Self::Lzfse => 0x801,
47            Self::Lzbitmap => 0x702,
48        }
49    }
50}
51
52#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Hash)]
53pub struct ArchiveFlags(u64);
54
55impl ArchiveFlags {
56    pub const IGNORE_EPERM: Self = Self(1_u64 << 0);
57    pub const ARCHIVE_DEDUPLICATE_DAT: Self = Self(1_u64 << 1);
58    pub const ARCHIVE_NO_RESOLVE_ACL_QUALIFIERS: Self = Self(1_u64 << 2);
59    pub const REPLACE_ATTRIBUTES: Self = Self(1_u64 << 3);
60    pub const EXTRACT_NO_AUTO_DEDUP: Self = Self(1_u64 << 4);
61    pub const EXTRACT_NO_AUTO_SPARSE: Self = Self(1_u64 << 5);
62    pub const CROSS_VOLUME_BOUNDARIES: Self = Self(1_u64 << 6);
63    pub const EXTRACT_AUTO_DEDUP_AS_HARD_LINKS: Self = Self(1_u64 << 7);
64    pub const DECODE_INSERT_IDX: Self = Self(1_u64 << 8);
65    pub const EXCLUDE_METADATA_ENTRIES: Self = Self(1_u64 << 9);
66    pub const PROCESS_RANDOM_ACCESS_OUTPUT: Self = Self(1_u64 << 10);
67    pub const VERBOSITY_0: Self = Self(0_u64 << 62);
68    pub const VERBOSITY_1: Self = Self(1_u64 << 62);
69    pub const VERBOSITY_2: Self = Self(2_u64 << 62);
70    pub const VERBOSITY_3: Self = Self(3_u64 << 62);
71
72    pub const fn empty() -> Self {
73        Self(0)
74    }
75
76    pub const fn bits(self) -> u64 {
77        self.0
78    }
79
80    pub const fn from_bits(bits: u64) -> Self {
81        Self(bits)
82    }
83
84    pub const fn contains(self, other: Self) -> bool {
85        (self.0 & other.0) == other.0
86    }
87
88    pub const fn verbosity(level: u64) -> Self {
89        let clamped = if level > 3 { 3 } else { level };
90        Self(clamped << 62)
91    }
92}
93
94impl BitOr for ArchiveFlags {
95    type Output = Self;
96
97    fn bitor(self, rhs: Self) -> Self::Output {
98        Self(self.0 | rhs.0)
99    }
100}
101
102impl BitOrAssign for ArchiveFlags {
103    fn bitor_assign(&mut self, rhs: Self) {
104        self.0 |= rhs.0;
105    }
106}
107
108impl BitAnd for ArchiveFlags {
109    type Output = Self;
110
111    fn bitand(self, rhs: Self) -> Self::Output {
112        Self(self.0 & rhs.0)
113    }
114}
115
116impl BitAndAssign for ArchiveFlags {
117    fn bitand_assign(&mut self, rhs: Self) {
118        self.0 &= rhs.0;
119    }
120}
121
122#[allow(dead_code)]
123#[derive(Debug)]
124pub enum ByteStreamUpstream {
125    Stream(Box<ByteStream>),
126}
127
128#[derive(Debug)]
129pub struct ByteStream {
130    handle: NonNull<c_void>,
131    _upstream: Option<ByteStreamUpstream>,
132    closed: bool,
133}
134
135impl ByteStream {
136    pub(crate) fn from_handle_with_upstream(
137        handle: *mut c_void,
138        operation: &'static str,
139        upstream: Option<ByteStreamUpstream>,
140    ) -> Result<Self> {
141        Ok(Self {
142            handle: util::nonnull_handle(handle, operation)?,
143            _upstream: upstream,
144            closed: false,
145        })
146    }
147
148    pub fn from_fd(fd: RawFd, automatic_close: bool) -> Result<Self> {
149        let handle = unsafe {
150            ffi::aa_byte_stream::compression_rs_aa_byte_stream_open_with_fd(fd, automatic_close)
151        };
152        Ok(Self {
153            handle: util::nonnull_handle(handle, "AAFileStreamOpenWithFD")?,
154            _upstream: None,
155            closed: false,
156        })
157    }
158
159    pub fn from_file(file: File) -> Result<Self> {
160        Self::from_fd(file.into_raw_fd(), true)
161    }
162
163    pub fn open_with_path(path: &str, open_flags: i32, open_mode: u32) -> Result<Self> {
164        let path = util::cstring("path", path)?;
165        let handle = unsafe {
166            ffi::aa_byte_stream::compression_rs_aa_byte_stream_open_with_path(
167                path.as_ptr(),
168                open_flags,
169                open_mode,
170            )
171        };
172        Ok(Self {
173            handle: util::nonnull_handle(handle, "AAFileStreamOpenWithPath")?,
174            _upstream: None,
175            closed: false,
176        })
177    }
178
179    pub fn temp_file() -> Result<Self> {
180        let handle = unsafe { ffi::aa_byte_stream::compression_rs_aa_temp_file_stream_open() };
181        Ok(Self {
182            handle: util::nonnull_handle(handle, "AATempFileStreamOpen")?,
183            _upstream: None,
184            closed: false,
185        })
186    }
187
188    pub fn shared_buffer_pipe(buffer_capacity: usize) -> Result<(Self, Self)> {
189        let mut ostream = std::ptr::null_mut();
190        let mut istream = std::ptr::null_mut();
191        let status = unsafe {
192            ffi::aa_byte_stream::compression_rs_aa_shared_buffer_pipe_open(
193                &mut ostream,
194                &mut istream,
195                buffer_capacity,
196            )
197        };
198        util::status_result("AASharedBufferPipeOpen", status)?;
199        Ok((
200            Self {
201                handle: util::nonnull_handle(ostream, "AASharedBufferPipeOpen(ostream)")?,
202                _upstream: None,
203                closed: false,
204            },
205            Self {
206                handle: util::nonnull_handle(istream, "AASharedBufferPipeOpen(istream)")?,
207                _upstream: None,
208                closed: false,
209            },
210        ))
211    }
212
213    pub(crate) fn as_ptr(&self) -> *mut c_void {
214        self.handle.as_ptr()
215    }
216
217    pub(crate) fn mark_closed(&mut self) {
218        self.closed = true;
219    }
220
221    fn ensure_open(&self) -> Result<()> {
222        if self.closed {
223            Err(CompressionError::Closed {
224                resource: "byte stream",
225            })
226        } else {
227            Ok(())
228        }
229    }
230
231    pub fn into_compression_output(
232        self,
233        compression_algorithm: ArchiveCompressionAlgorithm,
234        block_size: usize,
235        flags: ArchiveFlags,
236        n_threads: i32,
237    ) -> Result<Self> {
238        let handle = unsafe {
239            ffi::aa_byte_stream::compression_rs_aa_compression_output_stream_open(
240                self.as_ptr(),
241                compression_algorithm.as_raw(),
242                block_size,
243                flags.bits(),
244                n_threads,
245            )
246        };
247        Ok(Self {
248            handle: util::nonnull_handle(handle, "AACompressionOutputStreamOpen")?,
249            _upstream: Some(ByteStreamUpstream::Stream(Box::new(self))),
250            closed: false,
251        })
252    }
253
254    pub fn into_existing_compression_output(
255        self,
256        flags: ArchiveFlags,
257        n_threads: i32,
258    ) -> Result<Self> {
259        let handle = unsafe {
260            ffi::aa_byte_stream::compression_rs_aa_compression_output_stream_open_existing(
261                self.as_ptr(),
262                flags.bits(),
263                n_threads,
264            )
265        };
266        Ok(Self {
267            handle: util::nonnull_handle(handle, "AACompressionOutputStreamOpenExisting")?,
268            _upstream: Some(ByteStreamUpstream::Stream(Box::new(self))),
269            closed: false,
270        })
271    }
272
273    pub fn into_decompression_input(self, flags: ArchiveFlags, n_threads: i32) -> Result<Self> {
274        let handle = unsafe {
275            ffi::aa_byte_stream::compression_rs_aa_decompression_input_stream_open(
276                self.as_ptr(),
277                flags.bits(),
278                n_threads,
279            )
280        };
281        Ok(Self {
282            handle: util::nonnull_handle(handle, "AADecompressionInputStreamOpen")?,
283            _upstream: Some(ByteStreamUpstream::Stream(Box::new(self))),
284            closed: false,
285        })
286    }
287
288    pub fn into_random_access_decompression_input(
289        self,
290        alloc_limit: usize,
291        flags: ArchiveFlags,
292        n_threads: i32,
293    ) -> Result<Self> {
294        let handle = unsafe {
295            ffi::aa_byte_stream::compression_rs_aa_decompression_random_access_input_stream_open(
296                self.as_ptr(),
297                alloc_limit,
298                flags.bits(),
299                n_threads,
300            )
301        };
302        Ok(Self {
303            handle: util::nonnull_handle(handle, "AADecompressionRandomAccessInputStreamOpen")?,
304            _upstream: Some(ByteStreamUpstream::Stream(Box::new(self))),
305            closed: false,
306        })
307    }
308
309    pub fn write(&mut self, buffer: &[u8]) -> Result<usize> {
310        self.ensure_open()?;
311        util::ssize_result("AAByteStreamWrite", unsafe {
312            ffi::aa_byte_stream::compression_rs_aa_byte_stream_write(
313                self.as_ptr(),
314                buffer.as_ptr(),
315                buffer.len(),
316            )
317        })
318    }
319
320    pub fn write_all(&mut self, mut buffer: &[u8]) -> Result<()> {
321        while !buffer.is_empty() {
322            let written = self.write(buffer)?;
323            if written == 0 {
324                return Err(CompressionError::OperationFailed {
325                    operation: "AAByteStreamWrite",
326                    code: -1,
327                });
328            }
329            buffer = &buffer[written..];
330        }
331        Ok(())
332    }
333
334    pub fn pwrite(&mut self, buffer: &[u8], offset: i64) -> Result<usize> {
335        self.ensure_open()?;
336        util::ssize_result("AAByteStreamPWrite", unsafe {
337            ffi::aa_byte_stream::compression_rs_aa_byte_stream_pwrite(
338                self.as_ptr(),
339                buffer.as_ptr(),
340                buffer.len(),
341                offset,
342            )
343        })
344    }
345
346    pub fn read(&mut self, buffer: &mut [u8]) -> Result<usize> {
347        self.ensure_open()?;
348        util::ssize_result("AAByteStreamRead", unsafe {
349            ffi::aa_byte_stream::compression_rs_aa_byte_stream_read(
350                self.as_ptr(),
351                buffer.as_mut_ptr(),
352                buffer.len(),
353            )
354        })
355    }
356
357    pub fn pread(&mut self, buffer: &mut [u8], offset: i64) -> Result<usize> {
358        self.ensure_open()?;
359        util::ssize_result("AAByteStreamPRead", unsafe {
360            ffi::aa_byte_stream::compression_rs_aa_byte_stream_pread(
361                self.as_ptr(),
362                buffer.as_mut_ptr(),
363                buffer.len(),
364                offset,
365            )
366        })
367    }
368
369    pub fn seek(&mut self, offset: i64, whence: i32) -> Result<u64> {
370        self.ensure_open()?;
371        util::off_t_result("AAByteStreamSeek", unsafe {
372            ffi::aa_byte_stream::compression_rs_aa_byte_stream_seek(self.as_ptr(), offset, whence)
373        })
374    }
375
376    pub fn read_to_end(&mut self) -> Result<Vec<u8>> {
377        let mut output = Vec::new();
378        loop {
379            let mut buffer = vec![0_u8; READ_CHUNK_LEN];
380            let read = self.read(&mut buffer)?;
381            if read == 0 {
382                return Ok(output);
383            }
384            output.extend_from_slice(&buffer[..read]);
385        }
386    }
387
388    pub fn cancel(&mut self) -> Result<()> {
389        self.ensure_open()?;
390        unsafe { ffi::aa_byte_stream::compression_rs_aa_byte_stream_cancel(self.as_ptr()) };
391        Ok(())
392    }
393
394    #[deprecated(
395        since = "0.2.2",
396        note = "Use ByteStream::cancel; AAByteStreamAbort is a deprecated AppleArchive compatibility shim."
397    )]
398    pub fn abort(&mut self) -> Result<()> {
399        self.ensure_open()?;
400        unsafe { ffi::aa_byte_stream::compression_rs_aa_byte_stream_abort(self.as_ptr()) };
401        Ok(())
402    }
403
404    pub fn close(&mut self) -> Result<()> {
405        if self.closed {
406            return Ok(());
407        }
408        let status =
409            unsafe { ffi::aa_byte_stream::compression_rs_aa_byte_stream_close(self.as_ptr()) };
410        self.closed = true;
411        util::status_result("AAByteStreamClose", status)
412    }
413
414    pub fn process_into(&mut self, output: &mut Self) -> Result<u64> {
415        self.ensure_open()?;
416        output.ensure_open()?;
417        util::off_t_result("AAByteStreamProcess", unsafe {
418            ffi::aa_byte_stream::compression_rs_aa_byte_stream_process(
419                self.as_ptr(),
420                output.as_ptr(),
421            )
422        })
423    }
424
425    pub fn process_random_access_into(
426        &mut self,
427        output: &mut Self,
428        max_offset: i64,
429        block_size: usize,
430        flags: ArchiveFlags,
431        n_threads: i32,
432    ) -> Result<u64> {
433        self.ensure_open()?;
434        output.ensure_open()?;
435        util::off_t_result("AARandomAccessByteStreamProcess", unsafe {
436            ffi::aa_byte_stream::compression_rs_aa_random_access_byte_stream_process(
437                self.as_ptr(),
438                output.as_ptr(),
439                max_offset,
440                block_size,
441                flags.bits(),
442                n_threads,
443            )
444        })
445    }
446}
447
448impl Drop for ByteStream {
449    fn drop(&mut self) {
450        unsafe { ffi::aa_byte_stream::compression_rs_aa_byte_stream_release(self.as_ptr()) };
451    }
452}
453
454fn custom_byte_stream_error(operation: &'static str) -> CompressionError {
455    CompressionError::OperationFailed {
456        operation,
457        code: -1,
458    }
459}
460
461fn custom_byte_stream_code(error: &CompressionError) -> i32 {
462    match error {
463        CompressionError::OperationFailed { code, .. } if *code < 0 => *code,
464        _ => -1,
465    }
466}
467
468struct CustomByteStreamState {
469    callbacks: Box<dyn CustomByteStreamCallbacks>,
470}
471
472pub trait CustomByteStreamCallbacks {
473    fn write(&mut self, _buffer: &[u8]) -> Result<usize> {
474        Err(custom_byte_stream_error("AAByteStreamWrite"))
475    }
476
477    fn pwrite(&mut self, _buffer: &[u8], _offset: i64) -> Result<usize> {
478        Err(custom_byte_stream_error("AAByteStreamPWrite"))
479    }
480
481    fn read(&mut self, _buffer: &mut [u8]) -> Result<usize> {
482        Err(custom_byte_stream_error("AAByteStreamRead"))
483    }
484
485    fn pread(&mut self, _buffer: &mut [u8], _offset: i64) -> Result<usize> {
486        Err(custom_byte_stream_error("AAByteStreamPRead"))
487    }
488
489    fn seek(&mut self, _offset: i64, _whence: i32) -> Result<i64> {
490        Err(custom_byte_stream_error("AAByteStreamSeek"))
491    }
492
493    fn cancel(&mut self) {}
494
495    fn close(&mut self) -> Result<()> {
496        Ok(())
497    }
498}
499
500unsafe fn custom_byte_stream_slice<'a>(buffer: *const c_void, length: usize) -> Option<&'a [u8]> {
501    if length == 0 {
502        Some(&[])
503    } else if buffer.is_null() {
504        None
505    } else {
506        Some(unsafe { std::slice::from_raw_parts(buffer.cast::<u8>(), length) })
507    }
508}
509
510unsafe fn custom_byte_stream_slice_mut<'a>(
511    buffer: *mut c_void,
512    length: usize,
513) -> Option<&'a mut [u8]> {
514    if length == 0 {
515        Some(&mut [])
516    } else if buffer.is_null() {
517        None
518    } else {
519        Some(unsafe { std::slice::from_raw_parts_mut(buffer.cast::<u8>(), length) })
520    }
521}
522
523unsafe fn custom_byte_stream_state(arg: *mut c_void) -> Option<&'static mut CustomByteStreamState> {
524    if arg.is_null() {
525        None
526    } else {
527        Some(unsafe { &mut *arg.cast::<CustomByteStreamState>() })
528    }
529}
530
531unsafe extern "C" fn custom_byte_stream_write(
532    arg: *mut c_void,
533    buffer: *const c_void,
534    length: usize,
535) -> i64 {
536    let Some(state) = (unsafe { custom_byte_stream_state(arg) }) else {
537        return -1;
538    };
539    let Some(buffer) = (unsafe { custom_byte_stream_slice(buffer, length) }) else {
540        return -1;
541    };
542    match state.callbacks.write(buffer) {
543        Ok(count) => i64::try_from(count).unwrap_or(i64::MAX),
544        Err(error) => i64::from(custom_byte_stream_code(&error)),
545    }
546}
547
548unsafe extern "C" fn custom_byte_stream_pwrite(
549    arg: *mut c_void,
550    buffer: *const c_void,
551    length: usize,
552    offset: i64,
553) -> i64 {
554    let Some(state) = (unsafe { custom_byte_stream_state(arg) }) else {
555        return -1;
556    };
557    let Some(buffer) = (unsafe { custom_byte_stream_slice(buffer, length) }) else {
558        return -1;
559    };
560    match state.callbacks.pwrite(buffer, offset) {
561        Ok(count) => i64::try_from(count).unwrap_or(i64::MAX),
562        Err(error) => i64::from(custom_byte_stream_code(&error)),
563    }
564}
565
566unsafe extern "C" fn custom_byte_stream_read(
567    arg: *mut c_void,
568    buffer: *mut c_void,
569    length: usize,
570) -> i64 {
571    let Some(state) = (unsafe { custom_byte_stream_state(arg) }) else {
572        return -1;
573    };
574    let Some(buffer) = (unsafe { custom_byte_stream_slice_mut(buffer, length) }) else {
575        return -1;
576    };
577    match state.callbacks.read(buffer) {
578        Ok(count) => i64::try_from(count).unwrap_or(i64::MAX),
579        Err(error) => i64::from(custom_byte_stream_code(&error)),
580    }
581}
582
583unsafe extern "C" fn custom_byte_stream_pread(
584    arg: *mut c_void,
585    buffer: *mut c_void,
586    length: usize,
587    offset: i64,
588) -> i64 {
589    let Some(state) = (unsafe { custom_byte_stream_state(arg) }) else {
590        return -1;
591    };
592    let Some(buffer) = (unsafe { custom_byte_stream_slice_mut(buffer, length) }) else {
593        return -1;
594    };
595    match state.callbacks.pread(buffer, offset) {
596        Ok(count) => i64::try_from(count).unwrap_or(i64::MAX),
597        Err(error) => i64::from(custom_byte_stream_code(&error)),
598    }
599}
600
601unsafe extern "C" fn custom_byte_stream_seek(arg: *mut c_void, offset: i64, whence: i32) -> i64 {
602    let Some(state) = (unsafe { custom_byte_stream_state(arg) }) else {
603        return -1;
604    };
605    match state.callbacks.seek(offset, whence) {
606        Ok(position) => position,
607        Err(error) => i64::from(custom_byte_stream_code(&error)),
608    }
609}
610
611unsafe extern "C" fn custom_byte_stream_cancel(arg: *mut c_void) {
612    if let Some(state) = unsafe { custom_byte_stream_state(arg) } {
613        state.callbacks.cancel();
614    }
615}
616
617unsafe extern "C" fn custom_byte_stream_close(arg: *mut c_void) -> i32 {
618    if arg.is_null() {
619        return 0;
620    }
621    let mut state = unsafe { Box::from_raw(arg.cast::<CustomByteStreamState>()) };
622    match state.callbacks.close() {
623        Ok(()) => 0,
624        Err(error) => custom_byte_stream_code(&error),
625    }
626}
627
628impl ByteStream {
629    pub fn custom<T: CustomByteStreamCallbacks + 'static>(callbacks: T) -> Result<Self> {
630        let handle = unsafe { ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_open() };
631        let stream = Self::from_handle_with_upstream(handle, "AACustomByteStreamOpen", None)?;
632        let state = Box::new(CustomByteStreamState {
633            callbacks: Box::new(callbacks),
634        });
635        let data = Box::into_raw(state).cast::<c_void>();
636        unsafe {
637            ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_set_data(
638                stream.as_ptr(),
639                data,
640            );
641            ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_set_write_proc(
642                stream.as_ptr(),
643                Some(custom_byte_stream_write),
644            );
645            ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_set_pwrite_proc(
646                stream.as_ptr(),
647                Some(custom_byte_stream_pwrite),
648            );
649            ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_set_read_proc(
650                stream.as_ptr(),
651                Some(custom_byte_stream_read),
652            );
653            ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_set_pread_proc(
654                stream.as_ptr(),
655                Some(custom_byte_stream_pread),
656            );
657            ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_set_seek_proc(
658                stream.as_ptr(),
659                Some(custom_byte_stream_seek),
660            );
661            ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_set_cancel_proc(
662                stream.as_ptr(),
663                Some(custom_byte_stream_cancel),
664            );
665            ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_set_abort_proc(
666                stream.as_ptr(),
667                Some(custom_byte_stream_cancel),
668            );
669            ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_set_close_proc(
670                stream.as_ptr(),
671                Some(custom_byte_stream_close),
672            );
673        }
674        Ok(stream)
675    }
676}