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            loop {
55                let (dcm_count, dcm_data) = decompressor.update(&self.decrypt_buffer)?;
56                self.result_buffer.extend(dcm_data);
57                if dcm_count == 0 {
58                    break;
59                }
60                self.decrypt_buffer.drain(0..dcm_count);
61            }
62
63            Ok((dcr_count, &self.result_buffer))
64        } else if let Some(decryptor) = self.decryptor.as_deref_mut() {
65            Ok(decryptor.update(data)?)
66        } else if let Some(decompressor) = self.decompressor.as_deref_mut() {
67            Ok(decompressor.update(data)?)
68        } else {
69            // FIXME maybe it is worth to avoid this copy...
70            self.result_buffer = Vec::from(data);
71            Ok((data.len(), &self.result_buffer))
72        }
73    }
74}