Skip to main content

stream_unpack/
pipeline.rs

1use thiserror::Error;
2
3use crate::{decompress::{DecompressionError, Decompressor}, decrypt::{DecryptionError, Decryptor}};
4
5#[derive(Debug, Error)]
6pub enum PipelineError {
7    #[error("decryption error: {0}")]
8    Decryption(#[from] DecryptionError),
9
10    #[error("decompression error: {0}")]
11    Decompression(#[from] DecompressionError)
12}
13
14/// Represents a generalized decompression pipeline for a file (decrypt then decompress)
15/// as a single entity.
16///
17/// Either of the steps can be skipped.
18#[derive(Debug)]
19pub struct Pipeline {
20    decryptor: Option<Box<dyn Decryptor>>,
21    decompressor: Option<Box<dyn Decompressor>>,
22
23    decrypt_buffer: Vec<u8>,
24    result_buffer: Vec<u8>
25}
26
27impl Pipeline {
28    /// Creates a new Pipeilne. Providing None for either step skips it. None can be provided for both steps.
29    pub fn new(decryptor: Option<Box<dyn Decryptor>>, decompressor: Option<Box<dyn Decompressor>>) -> Self {
30        Self {
31            decryptor,
32            decompressor,
33
34            decrypt_buffer: Vec::new(),
35            result_buffer: Vec::new()
36        }
37    }
38
39    /// Updates the current pipeline with new data. Returns the amount of data actually consumed, and the
40    /// pipeline result. The input data should be advanced exactly the amount that this method returned.
41    ///
42    /// The returned result data may be empty, which should not be of concern to the caller.
43    /// The returned amount can be zero, which means that there is not enough data in the input to do
44    /// anything meaningful, and the caller should retry with a larger slice when it becomes available.
45    pub fn update(&mut self, data: &[u8]) -> Result<(usize, &[u8]), PipelineError> {
46        if self.decryptor.is_some() && self.decompressor.is_some() {
47            let decryptor = self.decryptor.as_deref_mut().unwrap();
48            let decompressor = self.decompressor.as_deref_mut().unwrap();
49
50            let (dcr_count, dcr_data) = decryptor.update(data)?;
51            self.decrypt_buffer.extend(dcr_data);
52
53            self.result_buffer.clear();
54            let decrypt_slice = self.decrypt_buffer.as_slice();
55            let mut offset = 0;
56            loop {
57                let (dcm_count, dcm_data) = decompressor.update(&decrypt_slice[offset..])?;
58                self.result_buffer.extend(dcm_data);
59                if dcm_count == 0 {
60                    break;
61                }
62                offset += dcm_count;
63            }
64            self.decrypt_buffer.drain(0..offset);
65
66            Ok((dcr_count, &self.result_buffer))
67        } else if let Some(decryptor) = self.decryptor.as_deref_mut() {
68            Ok(decryptor.update(data)?)
69        } else if let Some(decompressor) = self.decompressor.as_deref_mut() {
70            Ok(decompressor.update(data)?)
71        } else {
72            // FIXME maybe it is worth to avoid this copy...
73            self.result_buffer = Vec::from(data);
74            Ok((data.len(), &self.result_buffer))
75        }
76    }
77}