russh/
compression.rs

1use std::convert::TryFrom;
2
3use delegate::delegate;
4use ssh_encoding::Encode;
5
6#[cfg(feature = "flate2")]
7use crate::cipher::MAXIMUM_DECOMPRESSED_PACKET_LEN;
8
9#[derive(Debug, Clone)]
10pub enum Compression {
11    None,
12    #[cfg(feature = "flate2")]
13    Zlib,
14    #[cfg(feature = "flate2")]
15    ZlibOpenSSH,
16}
17
18#[derive(Debug)]
19pub enum Compress {
20    None,
21    #[cfg(feature = "flate2")]
22    Zlib(flate2::Compress),
23}
24
25#[derive(Debug)]
26pub enum Decompress {
27    None,
28    #[cfg(feature = "flate2")]
29    Zlib(flate2::Decompress),
30}
31
32#[derive(Debug, PartialEq, Eq, Copy, Clone, Hash)]
33pub struct Name(&'static str);
34impl AsRef<str> for Name {
35    fn as_ref(&self) -> &str {
36        self.0
37    }
38}
39
40impl Encode for Name {
41    delegate! { to self.as_ref() {
42        fn encoded_len(&self) -> Result<usize, ssh_encoding::Error>;
43        fn encode(&self, writer: &mut impl ssh_encoding::Writer) -> Result<(), ssh_encoding::Error>;
44    }}
45}
46
47impl TryFrom<&str> for Name {
48    type Error = ();
49    fn try_from(s: &str) -> Result<Name, ()> {
50        ALL_COMPRESSION_ALGORITHMS
51            .iter()
52            .find(|x| x.0 == s)
53            .map(|x| **x)
54            .ok_or(())
55    }
56}
57
58pub const NONE: Name = Name("none");
59#[cfg(feature = "flate2")]
60pub const ZLIB: Name = Name("zlib");
61#[cfg(feature = "flate2")]
62pub const ZLIB_LEGACY: Name = Name("zlib@openssh.com");
63
64pub const ALL_COMPRESSION_ALGORITHMS: &[&Name] = &[
65    &NONE,
66    #[cfg(feature = "flate2")]
67    &ZLIB,
68    #[cfg(feature = "flate2")]
69    &ZLIB_LEGACY,
70];
71
72#[cfg(feature = "flate2")]
73impl Compression {
74    pub fn new(name: &Name) -> Self {
75        if name == &ZLIB {
76            Compression::Zlib
77        } else if name == &ZLIB_LEGACY {
78            Compression::ZlibOpenSSH
79        } else {
80            Compression::None
81        }
82    }
83
84    pub fn init_compress(&self, comp: &mut Compress) {
85        match *self {
86            Compression::Zlib | Compression::ZlibOpenSSH => {
87                if let Compress::Zlib(ref mut c) = *comp {
88                    c.reset()
89                } else {
90                    *comp =
91                        Compress::Zlib(flate2::Compress::new(flate2::Compression::fast(), true))
92                }
93            }
94            Compression::None => {
95                *comp = Compress::None;
96            }
97        }
98    }
99
100    pub fn init_decompress(&self, comp: &mut Decompress) {
101        match *self {
102            Compression::Zlib | Compression::ZlibOpenSSH => {
103                if let Decompress::Zlib(ref mut c) = *comp {
104                    c.reset(true)
105                } else {
106                    *comp = Decompress::Zlib(flate2::Decompress::new(true))
107                }
108            }
109            Compression::None => {
110                *comp = Decompress::None;
111            }
112        }
113    }
114}
115
116impl Compression {
117    /// Returns true if compression should be deferred until after authentication.
118    /// "zlib@openssh.com" defers; RFC 4253 "zlib" does not.
119    pub fn is_deferred(&self) -> bool {
120        match self {
121            #[cfg(feature = "flate2")]
122            Compression::ZlibOpenSSH => true,
123            _ => false,
124        }
125    }
126}
127
128#[cfg(not(feature = "flate2"))]
129impl Compression {
130    pub fn new(_name: &Name) -> Self {
131        Compression::None
132    }
133
134    pub fn init_compress(&self, _: &mut Compress) {}
135
136    pub fn init_decompress(&self, _: &mut Decompress) {}
137}
138
139#[cfg(not(feature = "flate2"))]
140impl Compress {
141    pub fn compress<'a>(
142        &mut self,
143        input: &'a [u8],
144        _: &'a mut Vec<u8>,
145    ) -> Result<&'a [u8], crate::Error> {
146        Ok(input)
147    }
148
149    pub fn compress_into(
150        &mut self,
151        input: &[u8],
152        output: &mut Vec<u8>,
153        start_len: usize,
154    ) -> Result<usize, crate::Error> {
155        output.truncate(start_len);
156        output.extend_from_slice(input);
157        Ok(input.len())
158    }
159}
160
161#[cfg(not(feature = "flate2"))]
162impl Decompress {
163    pub fn decompress<'a>(
164        &mut self,
165        input: &'a [u8],
166        _: &'a mut Vec<u8>,
167    ) -> Result<&'a [u8], crate::Error> {
168        Ok(input)
169    }
170}
171
172#[cfg(all(test, feature = "flate2"))]
173mod tests {
174    use std::io::Write;
175
176    use flate2::write::ZlibEncoder;
177
178    use super::*;
179
180    #[test]
181    fn decompressed_packet_at_limit_is_accepted() {
182        let payload = vec![b'A'; MAXIMUM_DECOMPRESSED_PACKET_LEN];
183        let mut encoder = ZlibEncoder::new(Vec::new(), flate2::Compression::best());
184        encoder.write_all(&payload).unwrap();
185        let compressed = encoder.finish().unwrap();
186
187        let mut decompressor = Decompress::Zlib(flate2::Decompress::new(true));
188        let mut output = Vec::new();
189
190        let decompressed = decompressor.decompress(&compressed, &mut output).unwrap();
191        assert_eq!(decompressed.len(), MAXIMUM_DECOMPRESSED_PACKET_LEN);
192    }
193
194    #[test]
195    fn oversized_decompressed_packet_is_rejected() {
196        let payload = vec![b'A'; MAXIMUM_DECOMPRESSED_PACKET_LEN + 1024];
197        let mut encoder = ZlibEncoder::new(Vec::new(), flate2::Compression::best());
198        encoder.write_all(&payload).unwrap();
199        let compressed = encoder.finish().unwrap();
200
201        let mut decompressor = Decompress::Zlib(flate2::Decompress::new(true));
202        let mut output = Vec::new();
203
204        let err = decompressor.decompress(&compressed, &mut output).unwrap_err();
205        assert!(
206            matches!(err, crate::Error::PacketSize(len) if len > MAXIMUM_DECOMPRESSED_PACKET_LEN)
207        );
208    }
209
210    #[test]
211    fn empty_compressed_packet_does_not_spin() {
212        let compressed = Vec::new();
213        let mut decompressor = Decompress::Zlib(flate2::Decompress::new(true));
214        let mut output = Vec::new();
215
216        let decompressed = decompressor.decompress(&compressed, &mut output).unwrap();
217        assert!(decompressed.is_empty());
218    }
219}
220
221#[cfg(feature = "flate2")]
222impl Compress {
223    fn zlib_output_reserve_bound(input_len: usize) -> usize {
224        input_len.saturating_add(10)
225    }
226
227    pub fn compress<'a>(
228        &mut self,
229        input: &'a [u8],
230        output: &'a mut Vec<u8>,
231    ) -> Result<&'a [u8], crate::Error> {
232        match *self {
233            Compress::None => Ok(input),
234            Compress::Zlib(ref mut z) => {
235                output.clear();
236                let n_in = z.total_in() as usize;
237                let n_out = z.total_out() as usize;
238                output.resize(input.len() + 10, 0);
239                let flush = flate2::FlushCompress::Partial;
240                loop {
241                    let n_in_ = z.total_in() as usize - n_in;
242                    let n_out_ = z.total_out() as usize - n_out;
243                    #[allow(clippy::indexing_slicing)] // length checked
244                    let c = z.compress(&input[n_in_..], &mut output[n_out_..], flush)?;
245                    match c {
246                        flate2::Status::BufError => {
247                            output.resize(output.len() * 2, 0);
248                        }
249                        _ => break,
250                    }
251                }
252                let n_out_ = z.total_out() as usize - n_out;
253                #[allow(clippy::indexing_slicing)] // length checked
254                Ok(&output[..n_out_])
255            }
256        }
257    }
258
259    pub fn compress_into(
260        &mut self,
261        input: &[u8],
262        output: &mut Vec<u8>,
263        start_len: usize,
264    ) -> Result<usize, crate::Error> {
265        match *self {
266            Compress::None => {
267                output.truncate(start_len);
268                output.extend_from_slice(input);
269                Ok(input.len())
270            }
271            Compress::Zlib(ref mut z) => {
272                output.truncate(start_len);
273                let n_in = z.total_in() as usize;
274                let n_out = z.total_out() as usize;
275                let reserve = Self::zlib_output_reserve_bound(input.len());
276                output.resize(start_len + reserve, 0);
277                let flush = flate2::FlushCompress::Partial;
278                loop {
279                    let n_in_ = z.total_in() as usize - n_in;
280                    let n_out_ = z.total_out() as usize - n_out;
281                    #[allow(clippy::indexing_slicing)] // length checked
282                    let c = z.compress(&input[n_in_..], &mut output[start_len + n_out_..], flush)?;
283                    match c {
284                        flate2::Status::BufError => {
285                            let growth = output.len().saturating_sub(start_len).max(1);
286                            output.resize(output.len() + growth, 0);
287                        }
288                        _ => break,
289                    }
290                }
291                let n_out_ = z.total_out() as usize - n_out;
292                output.truncate(start_len + n_out_);
293                Ok(n_out_)
294            }
295        }
296    }
297}
298
299#[cfg(feature = "flate2")]
300impl Decompress {
301    pub fn decompress<'a>(
302        &mut self,
303        input: &'a [u8],
304        output: &'a mut Vec<u8>,
305    ) -> Result<&'a [u8], crate::Error> {
306        match *self {
307            Decompress::None => Ok(input),
308            Decompress::Zlib(ref mut z) => {
309                output.clear();
310                let n_in = z.total_in() as usize;
311                let n_out = z.total_out() as usize;
312                let max_output_len = MAXIMUM_DECOMPRESSED_PACKET_LEN
313                    .checked_add(1)
314                    .ok_or(crate::Error::PacketSize(usize::MAX))?;
315                output.resize(input.len().clamp(1, max_output_len), 0);
316                let flush = flate2::FlushDecompress::None;
317                loop {
318                    let n_in_ = z.total_in() as usize - n_in;
319                    let n_out_ = z.total_out() as usize - n_out;
320                    #[allow(clippy::indexing_slicing)] // length checked
321                    let d = z.decompress(&input[n_in_..], &mut output[n_out_..], flush);
322                    match d? {
323                        flate2::Status::Ok | flate2::Status::BufError => {
324                            let consumed_all_input = n_in_ == input.len();
325                            let output_full = n_out_ == output.len();
326
327                            if !output_full && consumed_all_input {
328                                break;
329                            }
330
331                            if output_full {
332                                if output.len() == max_output_len {
333                                    break;
334                                }
335                                let next_len = output
336                                    .len()
337                                    .checked_mul(2)
338                                    .map(|len| len.min(max_output_len))
339                                    .ok_or(crate::Error::PacketSize(usize::MAX))?;
340                                output.resize(next_len, 0);
341                            }
342                        }
343                        _ => break,
344                    }
345                }
346                let n_out_ = z.total_out() as usize - n_out;
347                if n_out_ > MAXIMUM_DECOMPRESSED_PACKET_LEN {
348                    return Err(crate::Error::PacketSize(n_out_));
349                }
350                #[allow(clippy::indexing_slicing)] // length checked
351                Ok(&output[..n_out_])
352            }
353        }
354    }
355}
russh/compression.rs

russh/
compression.rs
