ssh/algorithm/compression/
zlib.rs

1use flate2;
2
3use crate::SshError;
4
5use super::Compression;
6
7/// The "zlib" compression is described in [RFC1950] and in [RFC1951].
8/// The compression context is initialized after each key exchange, and
9/// is passed from one packet to the next, with only a partial flush
10/// being performed at the end of each packet.  A partial flush means
11/// that the current compressed block is ended and all data will be
12/// output.  If the current block is not a stored block, one or more
13/// empty blocks are added after the current block to ensure that there
14/// are at least 8 bits, counting from the start of the end-of-block code
15/// of the current block to the end of the packet payload.
16///
17/// <https://www.openssh.com/txt/draft-miller-secsh-compression-delayed-00.txt>
18/// The "zlib@openssh.com" method operates identically to the "zlib"
19/// method described in [RFC4252] except that packet compression does not
20/// start until the server sends a SSH_MSG_USERAUTH_SUCCESS packet,
21/// replacing the "zlib" method's start of compression when the server
22/// sends SSH_MSG_NEWKEYS.
23pub(super) struct CompressZlib {
24    decompressor: flate2::Decompress,
25    compressor: flate2::Compress,
26}
27
28impl Compression for CompressZlib {
29    fn new() -> Self
30    where
31        Self: Sized,
32    {
33        Self {
34            decompressor: flate2::Decompress::new(true),
35            compressor: flate2::Compress::new(flate2::Compression::fast(), true),
36        }
37    }
38
39    fn decompress(&mut self, buf: &[u8]) -> crate::SshResult<Vec<u8>> {
40        let mut buf_in = buf;
41        let mut buf_once = [0; 4096];
42        let mut buf_out = vec![];
43        loop {
44            let in_before = self.decompressor.total_in();
45            let out_before = self.decompressor.total_out();
46
47            let result =
48                self.decompressor
49                    .decompress(buf_in, &mut buf_once, flate2::FlushDecompress::Sync);
50
51            let consumed = (self.decompressor.total_in() - in_before) as usize;
52            let produced = (self.decompressor.total_out() - out_before) as usize;
53
54            match result {
55                Ok(flate2::Status::Ok) => {
56                    buf_in = &buf_in[consumed..];
57                    buf_out.extend(&buf_once[..produced]);
58                }
59                Ok(flate2::Status::StreamEnd) => {
60                    return Err(SshError::CompressionError(
61                        "Stream ends during the decompress".to_owned(),
62                    ));
63                }
64                Ok(flate2::Status::BufError) => {
65                    break;
66                }
67                Err(e) => return Err(SshError::CompressionError(e.to_string())),
68            }
69        }
70
71        Ok(buf_out)
72    }
73
74    fn compress(&mut self, buf: &[u8]) -> crate::SshResult<Vec<u8>> {
75        let mut buf_in = buf;
76        let mut buf_once = [0; 4096];
77        let mut buf_out = vec![];
78        loop {
79            let in_before = self.compressor.total_in();
80            let out_before = self.compressor.total_out();
81
82            let result =
83                self.compressor
84                    .compress(buf_in, &mut buf_once, flate2::FlushCompress::Partial);
85
86            let consumed = (self.compressor.total_in() - in_before) as usize;
87            let produced = (self.compressor.total_out() - out_before) as usize;
88
89            // tracing::info!(consumed);
90            // tracing::info!(produced);
91
92            // means an empty compress
93            // 2 bytes ZLIB header at the start of the stream
94            // 4 bytes CRC checksum at the end of the stream
95            if produced == 6 {
96                break;
97            }
98
99            match result {
100                Ok(flate2::Status::Ok) => {
101                    buf_in = &buf_in[consumed..];
102                    buf_out.extend(&buf_once[..produced]);
103                }
104                Ok(flate2::Status::StreamEnd) => {
105                    return Err(SshError::CompressionError(
106                        "Stream ends during the compress".to_owned(),
107                    ));
108                }
109                Ok(flate2::Status::BufError) => {
110                    break;
111                }
112                Err(e) => return Err(SshError::CompressionError(e.to_string())),
113            }
114        }
115
116        Ok(buf_out)
117    }
118}
ssh/algorithm/compression/zlib.rs

ssh/algorithm/compression/
zlib.rs
