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    pub fn close(&mut self) -> Result<()> {
395        if self.closed {
396            return Ok(());
397        }
398        let status =
399            unsafe { ffi::aa_byte_stream::compression_rs_aa_byte_stream_close(self.as_ptr()) };
400        self.closed = true;
401        util::status_result("AAByteStreamClose", status)
402    }
403
404    pub fn process_into(&mut self, output: &mut Self) -> Result<u64> {
405        self.ensure_open()?;
406        output.ensure_open()?;
407        util::off_t_result("AAByteStreamProcess", unsafe {
408            ffi::aa_byte_stream::compression_rs_aa_byte_stream_process(
409                self.as_ptr(),
410                output.as_ptr(),
411            )
412        })
413    }
414
415    pub fn process_random_access_into(
416        &mut self,
417        output: &mut Self,
418        max_offset: i64,
419        block_size: usize,
420        flags: ArchiveFlags,
421        n_threads: i32,
422    ) -> Result<u64> {
423        self.ensure_open()?;
424        output.ensure_open()?;
425        util::off_t_result("AARandomAccessByteStreamProcess", unsafe {
426            ffi::aa_byte_stream::compression_rs_aa_random_access_byte_stream_process(
427                self.as_ptr(),
428                output.as_ptr(),
429                max_offset,
430                block_size,
431                flags.bits(),
432                n_threads,
433            )
434        })
435    }
436}
437
438impl Drop for ByteStream {
439    fn drop(&mut self) {
440        unsafe { ffi::aa_byte_stream::compression_rs_aa_byte_stream_release(self.as_ptr()) };
441    }
442}
443
444fn custom_byte_stream_error(operation: &'static str) -> CompressionError {
445    CompressionError::OperationFailed {
446        operation,
447        code: -1,
448    }
449}
450
451fn custom_byte_stream_code(error: &CompressionError) -> i32 {
452    match error {
453        CompressionError::OperationFailed { code, .. } if *code < 0 => *code,
454        _ => -1,
455    }
456}
457
458struct CustomByteStreamState {
459    callbacks: Box<dyn CustomByteStreamCallbacks>,
460}
461
462pub trait CustomByteStreamCallbacks {
463    fn write(&mut self, _buffer: &[u8]) -> Result<usize> {
464        Err(custom_byte_stream_error("AAByteStreamWrite"))
465    }
466
467    fn pwrite(&mut self, _buffer: &[u8], _offset: i64) -> Result<usize> {
468        Err(custom_byte_stream_error("AAByteStreamPWrite"))
469    }
470
471    fn read(&mut self, _buffer: &mut [u8]) -> Result<usize> {
472        Err(custom_byte_stream_error("AAByteStreamRead"))
473    }
474
475    fn pread(&mut self, _buffer: &mut [u8], _offset: i64) -> Result<usize> {
476        Err(custom_byte_stream_error("AAByteStreamPRead"))
477    }
478
479    fn seek(&mut self, _offset: i64, _whence: i32) -> Result<i64> {
480        Err(custom_byte_stream_error("AAByteStreamSeek"))
481    }
482
483    fn cancel(&mut self) {}
484
485    fn close(&mut self) -> Result<()> {
486        Ok(())
487    }
488}
489
490unsafe fn custom_byte_stream_slice<'a>(buffer: *const c_void, length: usize) -> Option<&'a [u8]> {
491    if length == 0 {
492        Some(&[])
493    } else if buffer.is_null() {
494        None
495    } else {
496        Some(unsafe { std::slice::from_raw_parts(buffer.cast::<u8>(), length) })
497    }
498}
499
500unsafe fn custom_byte_stream_slice_mut<'a>(
501    buffer: *mut c_void,
502    length: usize,
503) -> Option<&'a mut [u8]> {
504    if length == 0 {
505        Some(&mut [])
506    } else if buffer.is_null() {
507        None
508    } else {
509        Some(unsafe { std::slice::from_raw_parts_mut(buffer.cast::<u8>(), length) })
510    }
511}
512
513unsafe fn custom_byte_stream_state(arg: *mut c_void) -> Option<&'static mut CustomByteStreamState> {
514    if arg.is_null() {
515        None
516    } else {
517        Some(unsafe { &mut *arg.cast::<CustomByteStreamState>() })
518    }
519}
520
521unsafe extern "C" fn custom_byte_stream_write(
522    arg: *mut c_void,
523    buffer: *const c_void,
524    length: usize,
525) -> i64 {
526    let Some(state) = (unsafe { custom_byte_stream_state(arg) }) else {
527        return -1;
528    };
529    let Some(buffer) = (unsafe { custom_byte_stream_slice(buffer, length) }) else {
530        return -1;
531    };
532    match state.callbacks.write(buffer) {
533        Ok(count) => i64::try_from(count).unwrap_or(i64::MAX),
534        Err(error) => i64::from(custom_byte_stream_code(&error)),
535    }
536}
537
538unsafe extern "C" fn custom_byte_stream_pwrite(
539    arg: *mut c_void,
540    buffer: *const c_void,
541    length: usize,
542    offset: i64,
543) -> i64 {
544    let Some(state) = (unsafe { custom_byte_stream_state(arg) }) else {
545        return -1;
546    };
547    let Some(buffer) = (unsafe { custom_byte_stream_slice(buffer, length) }) else {
548        return -1;
549    };
550    match state.callbacks.pwrite(buffer, offset) {
551        Ok(count) => i64::try_from(count).unwrap_or(i64::MAX),
552        Err(error) => i64::from(custom_byte_stream_code(&error)),
553    }
554}
555
556unsafe extern "C" fn custom_byte_stream_read(
557    arg: *mut c_void,
558    buffer: *mut c_void,
559    length: usize,
560) -> i64 {
561    let Some(state) = (unsafe { custom_byte_stream_state(arg) }) else {
562        return -1;
563    };
564    let Some(buffer) = (unsafe { custom_byte_stream_slice_mut(buffer, length) }) else {
565        return -1;
566    };
567    match state.callbacks.read(buffer) {
568        Ok(count) => i64::try_from(count).unwrap_or(i64::MAX),
569        Err(error) => i64::from(custom_byte_stream_code(&error)),
570    }
571}
572
573unsafe extern "C" fn custom_byte_stream_pread(
574    arg: *mut c_void,
575    buffer: *mut c_void,
576    length: usize,
577    offset: i64,
578) -> i64 {
579    let Some(state) = (unsafe { custom_byte_stream_state(arg) }) else {
580        return -1;
581    };
582    let Some(buffer) = (unsafe { custom_byte_stream_slice_mut(buffer, length) }) else {
583        return -1;
584    };
585    match state.callbacks.pread(buffer, offset) {
586        Ok(count) => i64::try_from(count).unwrap_or(i64::MAX),
587        Err(error) => i64::from(custom_byte_stream_code(&error)),
588    }
589}
590
591unsafe extern "C" fn custom_byte_stream_seek(arg: *mut c_void, offset: i64, whence: i32) -> i64 {
592    let Some(state) = (unsafe { custom_byte_stream_state(arg) }) else {
593        return -1;
594    };
595    match state.callbacks.seek(offset, whence) {
596        Ok(position) => position,
597        Err(error) => i64::from(custom_byte_stream_code(&error)),
598    }
599}
600
601unsafe extern "C" fn custom_byte_stream_cancel(arg: *mut c_void) {
602    if let Some(state) = unsafe { custom_byte_stream_state(arg) } {
603        state.callbacks.cancel();
604    }
605}
606
607unsafe extern "C" fn custom_byte_stream_close(arg: *mut c_void) -> i32 {
608    if arg.is_null() {
609        return 0;
610    }
611    let mut state = unsafe { Box::from_raw(arg.cast::<CustomByteStreamState>()) };
612    match state.callbacks.close() {
613        Ok(()) => 0,
614        Err(error) => custom_byte_stream_code(&error),
615    }
616}
617
618impl ByteStream {
619    pub fn custom<T: CustomByteStreamCallbacks + 'static>(callbacks: T) -> Result<Self> {
620        let handle =
621            unsafe { ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_open() };
622        let stream = Self::from_handle_with_upstream(handle, "AACustomByteStreamOpen", None)?;
623        let state = Box::new(CustomByteStreamState {
624            callbacks: Box::new(callbacks),
625        });
626        let data = Box::into_raw(state).cast::<c_void>();
627        unsafe {
628            ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_set_data(
629                stream.as_ptr(),
630                data,
631            );
632            ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_set_write_proc(
633                stream.as_ptr(),
634                Some(custom_byte_stream_write),
635            );
636            ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_set_pwrite_proc(
637                stream.as_ptr(),
638                Some(custom_byte_stream_pwrite),
639            );
640            ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_set_read_proc(
641                stream.as_ptr(),
642                Some(custom_byte_stream_read),
643            );
644            ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_set_pread_proc(
645                stream.as_ptr(),
646                Some(custom_byte_stream_pread),
647            );
648            ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_set_seek_proc(
649                stream.as_ptr(),
650                Some(custom_byte_stream_seek),
651            );
652            ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_set_cancel_proc(
653                stream.as_ptr(),
654                Some(custom_byte_stream_cancel),
655            );
656            ffi::aa_byte_stream::compression_rs_aa_custom_byte_stream_set_close_proc(
657                stream.as_ptr(),
658                Some(custom_byte_stream_close),
659            );
660        }
661        Ok(stream)
662    }
663}
compression/aa_byte_stream.rs

compression/
aa_byte_stream.rs
