flatline/cipher/
compress.rs

1use super::*;
2use crate::error::{builder, Result};
3use flate2::{Compress, Compression, Decompress, Status};
4use indexmap::IndexMap;
5use std::mem;
6
7algo_list!(
8    encode_all,
9    new_encode_all,
10    new_encode_by_name,
11    dyn Encode + Send,
12    "zlib" => ZEncoder::new(true),
13    "zlib@openssh.com" => ZEncoder::new(false),
14    "none" => Never::default(),
15);
16
17algo_list!(
18    decode_all,
19    new_decode_all,
20    new_decode_by_name,
21    dyn Decode + Send,
22    "zlib" => ZDecoder::new(true),
23    "zlib@openssh.com" => ZDecoder::new(false),
24    "none" => Never::default(),
25);
26
27pub fn none_encode() -> Boxtory<dyn Encode + Send> {
28    create_boxtory!(Never::default())
29}
30
31pub fn none_decode() -> Boxtory<dyn Decode + Send> {
32    create_boxtory!(Never::default())
33}
34
35pub trait Encode {
36    fn compress_in_auth(&self) -> bool;
37    fn update(&mut self, data: &[u8]) -> Result<()>;
38    fn finalize(&mut self) -> Result<Vec<u8>>;
39}
40
41pub trait Decode {
42    fn compress_in_auth(&self) -> bool;
43    fn update(&mut self, data: &[u8]) -> Result<()>;
44    fn finalize(&mut self) -> Result<Vec<u8>>;
45}
46
47impl Encode for Never {
48    fn compress_in_auth(&self) -> bool {
49        false
50    }
51
52    fn update(&mut self, data: &[u8]) -> Result<()> {
53        self.buf.extend(data);
54        Ok(())
55    }
56
57    fn finalize(&mut self) -> Result<Vec<u8>> {
58        Ok(mem::take(&mut self.buf))
59    }
60}
61
62impl Decode for Never {
63    fn compress_in_auth(&self) -> bool {
64        false
65    }
66
67    fn update(&mut self, data: &[u8]) -> Result<()> {
68        self.buf.extend(data);
69        Ok(())
70    }
71
72    fn finalize(&mut self) -> Result<Vec<u8>> {
73        Ok(mem::take(&mut self.buf))
74    }
75}
76
77#[derive(Default)]
78struct Never {
79    buf: Vec<u8>,
80}
81
82struct ZEncoder {
83    compress_in_auth: bool,
84    encoder: Compress,
85    buf: Vec<u8>,
86}
87
88impl ZEncoder {
89    fn new(compress_in_auth: bool) -> Self {
90        Self {
91            compress_in_auth,
92            encoder: Compress::new(Compression::default(), true),
93            buf: vec![],
94        }
95    }
96}
97
98impl Encode for ZEncoder {
99    fn compress_in_auth(&self) -> bool {
100        self.compress_in_auth
101    }
102
103    fn update(&mut self, data: &[u8]) -> Result<()> {
104        let before = self.encoder.total_in() as usize;
105        let data_len = data.len();
106        let mut pos = 0;
107        loop {
108            let cap = ((data_len - pos) / 1024 + 1) * 1024;
109            let mut tmp = Vec::with_capacity(cap);
110            let status =
111                self.encoder
112                    .compress_vec(&data[pos..], &mut tmp, flate2::FlushCompress::Partial);
113            return match status {
114                Ok(Status::Ok) => {
115                    self.buf.extend(&tmp);
116                    let after = self.encoder.total_in() as usize;
117                    pos = after - before;
118                    if pos < data_len || tmp.len() == cap {
119                        continue;
120                    }
121                    Ok(())
122                }
123                _ => builder::CompressFailed.fail(), //Err(Error::CompressFailed),
124            };
125        }
126    }
127
128    fn finalize(&mut self) -> Result<Vec<u8>> {
129        Ok(mem::take(&mut self.buf))
130    }
131}
132
133struct ZDecoder {
134    compress_in_auth: bool,
135    decoder: Decompress,
136    buf: Vec<u8>,
137}
138
139impl ZDecoder {
140    fn new(compress_in_auth: bool) -> Self {
141        Self {
142            compress_in_auth,
143            decoder: Decompress::new(true),
144            buf: vec![],
145        }
146    }
147}
148
149impl Decode for ZDecoder {
150    fn compress_in_auth(&self) -> bool {
151        self.compress_in_auth
152    }
153
154    fn update(&mut self, data: &[u8]) -> Result<()> {
155        let before = self.decoder.total_in();
156        let mut pos = 0;
157        loop {
158            let cap = 1024 * 4;
159            let mut tmp = Vec::with_capacity(cap);
160            let status =
161                self.decoder
162                    .decompress_vec(&data[pos..], &mut tmp, flate2::FlushDecompress::Sync);
163            pos = (self.decoder.total_in() - before) as usize;
164            return match status {
165                Ok(Status::Ok) => {
166                    self.buf.extend(tmp);
167                    continue;
168                }
169                Ok(_) => {
170                    self.buf.extend(tmp);
171                    if pos < data.len() {
172                        continue;
173                    }
174                    Ok(())
175                }
176                _ => builder::CompressFailed.fail(), //Err(Error::CompressFailed),
177            };
178        }
179
180        // let mut len = (data.len() / 1024 + 1) * 1024;
181        // let out_before = self.decoder.total_out() as usize;
182        // let in_before = self.decoder.total_in() as usize;
183        // let mut tmp = vec![0; len];
184        // loop {
185        //     // let status = self
186        //     //     .decoder
187        //     //     .decompress_vec(data, &mut tmp, flate2::FlushDecompress::None);
188
189        //     let input_index = self.decoder.total_in() as usize - in_before;
190        //     let output_index = self.decoder.total_out() as usize - out_before;
191        //     let status = self.decoder.decompress(
192        //         &data[input_index..],
193        //         &mut tmp[output_index..],
194        //         flate2::FlushDecompress::Sync,
195        //     );
196
197        //     return match status {
198        //         Ok(Status::Ok) => {
199        //             len *= 2;
200        //             tmp.resize(len, 0);
201        //             continue;
202        //         }
203        //         Ok(_) => {
204        //             self.buf.extend(tmp);
205        //             Ok(())
206        //         }
207        //         _ => Err(Error::CompressFailed),
208        //     };
209        // }
210    }
211
212    fn finalize(&mut self) -> Result<Vec<u8>> {
213        Ok(mem::take(&mut self.buf))
214    }
215}