Skip to main content

smb_server_proto_smb2/
compress.rs

1//! SMB3 compression: the COMPRESSION_TRANSFORM_HEADER ([MS-SMB2] §2.2.42) and
2//! SMB2_COMPRESSION_CAPABILITIES negotiate context ([MS-SMB2] §2.2.3.1.3).
3//!
4//! LZNT1 (compress + decompress) is delegated to the `lznt1` crate; Pattern_V1
5//! (a repeated-byte payload, §2.2.42.1) is expanded inline. Only the unchained
6//! transform form is produced; both forms are accepted on receive.
7
8/// COMPRESSION_TRANSFORM ProtocolId: `0xFC 'S' 'M' 'B'`.
9pub const PROTOCOL_ID: [u8; 4] = [0xFC, b'S', b'M', b'B'];
10
11/// SMB2_COMPRESSION_CAPABILITIES_FLAG_CHAINED ([MS-SMB2] §2.2.3.1.3): advertised
12/// in the negotiate context Flags field when chained compression is supported.
13pub const CAP_FLAG_CHAINED: u32 = 0x0000_0001;
14
15/// SMB2_COMPRESSION_FLAG_CHAINED ([MS-SMB2] §2.2.42.2.1): set in the first
16/// chained payload header. Its position (the Flags field at frame offset 10)
17/// also distinguishes a chained COMPRESSION_TRANSFORM_HEADER from the unchained
18/// form, whose Flags field is zero there ([MS-SMB2] §2.2.42).
19const CHAINED_FLAG: u16 = 0x0001;
20
21/// Minimum single-byte run length worth emitting as a Pattern_V1 payload. A
22/// Pattern_V1 payload costs 16 bytes (8-byte chained header + 8-byte pattern),
23/// so only longer runs shrink the message.
24const CHAINED_PATTERN_MIN: usize = 32;
25
26/// Compression algorithm ids ([MS-SMB2] §2.2.3.1.3).
27pub mod algo {
28    /// No compression.
29    pub const NONE: u16 = 0x0000;
30    /// LZNT1.
31    pub const LZNT1: u16 = 0x0001;
32    /// Plain LZ77 ([MS-XCA]).
33    pub const LZ77: u16 = 0x0002;
34    /// LZ77 + Huffman ([MS-XCA]).
35    pub const LZ77_HUFFMAN: u16 = 0x0003;
36    /// Repeated-byte pattern scrubbing.
37    pub const PATTERN_V1: u16 = 0x0004;
38}
39
40/// Algorithms this server can both send and receive.
41pub const SUPPORTED: &[u16] = &[algo::LZNT1, algo::LZ77, algo::PATTERN_V1];
42
43/// Compress with Plain LZ77 (LZXpress, [MS-XCA]); `None` if the codec errors.
44fn lz77_compress(data: &[u8]) -> Option<Vec<u8>> {
45    xpress_rs::lz77::LZ77Compressor::new(data.to_vec())
46        .compress()
47        .ok()
48}
49
50/// Decompress a Plain LZ77 (LZXpress, [MS-XCA]) stream. Wrapped in `catch_unwind`
51/// so malformed attacker-supplied input that panics the codec is contained as a
52/// decode failure (the dispatcher then disconnects) instead of unwinding the
53/// connection task.
54fn lz77_decompress(data: &[u8]) -> Option<Vec<u8>> {
55    let input = data.to_vec();
56    std::panic::catch_unwind(move || {
57        xpress_rs::lz77::LZ77Decompressor::new(input)
58            .decompress()
59            .ok()
60    })
61    .ok()
62    .flatten()
63}
64
65/// Parse a client `SMB2_COMPRESSION_CAPABILITIES` context (§2.2.3.1.3) into its
66/// advertised algorithm list.
67pub fn parse_compression_caps(data: &[u8]) -> Vec<u16> {
68    if data.len() < 8 {
69        return Vec::new();
70    }
71    let count = u16::from_le_bytes([data[0], data[1]]) as usize;
72    // CompressionAlgorithmCount(2) Padding(2) Flags(4) then the algorithm array.
73    data.get(8..)
74        .map(|rest| {
75            rest.chunks_exact(2)
76                .take(count)
77                .map(|c| u16::from_le_bytes([c[0], c[1]]))
78                .collect()
79        })
80        .unwrap_or_default()
81}
82
83/// Read the Flags field of a client `SMB2_COMPRESSION_CAPABILITIES` context
84/// (§2.2.3.1.3).
85pub fn parse_compression_flags(data: &[u8]) -> u32 {
86    data.get(4..8)
87        .map(|s| u32::from_le_bytes([s[0], s[1], s[2], s[3]]))
88        .unwrap_or(0)
89}
90
91/// Encode the server `SMB2_COMPRESSION_CAPABILITIES` context data advertising
92/// `algos`, setting SMB2_COMPRESSION_CAPABILITIES_FLAG_CHAINED when `chained`.
93pub fn build_compression_caps(algos: &[u16], chained: bool) -> Vec<u8> {
94    let mut d = Vec::with_capacity(8 + algos.len() * 2);
95    d.extend_from_slice(&(algos.len() as u16).to_le_bytes()); // CompressionAlgorithmCount
96    d.extend_from_slice(&0u16.to_le_bytes()); // Padding
97    let flags = if chained { CAP_FLAG_CHAINED } else { 0 };
98    d.extend_from_slice(&flags.to_le_bytes()); // Flags
99    for a in algos {
100        d.extend_from_slice(&a.to_le_bytes());
101    }
102    d
103}
104
105/// Intersect the client's offered algorithms with the server's, preserving the
106/// client's order so the negotiate response echoes the requested CompressionIds
107/// when all are supported ([MS-SMB2] §3.3.5.4).
108pub fn negotiate_algos(client: &[u16]) -> Vec<u16> {
109    client
110        .iter()
111        .copied()
112        .filter(|a| SUPPORTED.contains(a))
113        .collect()
114}
115
116/// Decompress a whole `\xFCSMB` COMPRESSION_TRANSFORM frame into the original
117/// SMB2 message. `None` on malformed input or an unsupported algorithm.
118pub fn decompress_message(frame: &[u8]) -> Option<Vec<u8>> {
119    if frame.len() < 16 || frame[0..4] != PROTOCOL_ID {
120        return None;
121    }
122    // The Flags field at offset 10 distinguishes the chained transform (0x0001)
123    // from the unchained form (0) ([MS-SMB2] §2.2.42).
124    if u16::from_le_bytes([frame[10], frame[11]]) & CHAINED_FLAG != 0 {
125        return decompress_chained(frame);
126    }
127    let orig = u32::from_le_bytes(frame[4..8].try_into().ok()?) as usize;
128    let algorithm = u16::from_le_bytes(frame[8..10].try_into().ok()?);
129    let offset = u32::from_le_bytes(frame[12..16].try_into().ok()?) as usize;
130    let payload = frame.get(16..)?;
131    let prefix = payload.get(..offset)?;
132    let compressed = payload.get(offset..)?;
133    let mut out = Vec::with_capacity(orig.max(prefix.len()));
134    out.extend_from_slice(prefix);
135    match algorithm {
136        algo::LZNT1 => lznt1::decompress(compressed, &mut out).ok()?,
137        algo::LZ77 => out.extend_from_slice(&lz77_decompress(compressed)?),
138        algo::PATTERN_V1 => expand_pattern_v1(compressed, &mut out)?,
139        algo::NONE => out.extend_from_slice(compressed),
140        _ => return None,
141    }
142    Some(out)
143}
144
145/// Decompress a chained COMPRESSION_TRANSFORM_HEADER ([MS-SMB2] \u00a72.2.42.2) by
146/// concatenating each payload's contribution. Per \u00a73.1.5.3 any malformed payload
147/// (an unknown algorithm, a Length reaching past the buffer, a Repetitions or
148/// decompressed size exceeding OriginalCompressedSegmentSize) yields `None`,
149/// which the dispatcher treats as fatal and disconnects the connection.
150fn decompress_chained(frame: &[u8]) -> Option<Vec<u8>> {
151    let orig = u32::from_le_bytes(frame[4..8].try_into().ok()?) as usize;
152    let mut out = Vec::with_capacity(orig);
153    let mut pos = 8;
154    while pos < frame.len() {
155        // SMB2_COMPRESSION_CHAINED_PAYLOAD_HEADER: Algorithm(2) Flags(2) Length(4).
156        let hdr = frame.get(pos..pos + 8)?;
157        let algorithm = u16::from_le_bytes([hdr[0], hdr[1]]);
158        let length = u32::from_le_bytes([hdr[4], hdr[5], hdr[6], hdr[7]]) as usize;
159        pos += 8;
160        let payload = frame.get(pos..pos + length)?;
161        match algorithm {
162            algo::NONE => {
163                if length > orig {
164                    return None;
165                }
166                out.extend_from_slice(payload);
167            }
168            algo::PATTERN_V1 => {
169                // SMB2_COMPRESSION_PATTERN_PAYLOAD_V1: Pattern(1) Reserved1(1)
170                // Reserved2(2) Repetitions(4).
171                let reps = u32::from_le_bytes(payload.get(4..8)?.try_into().ok()?) as usize;
172                if reps > orig {
173                    return None;
174                }
175                out.resize(out.len().checked_add(reps)?, payload[0]);
176            }
177            algo::LZNT1 => {
178                // Length includes the leading 4-byte OriginalPayloadSize field.
179                let orig_payload = u32::from_le_bytes(payload.get(0..4)?.try_into().ok()?) as usize;
180                let before = out.len();
181                lznt1::decompress(&payload[4..], &mut out).ok()?;
182                if out.len() - before != orig_payload {
183                    return None;
184                }
185            }
186            algo::LZ77 => {
187                // Length includes the leading 4-byte OriginalPayloadSize field.
188                let orig_payload = u32::from_le_bytes(payload.get(0..4)?.try_into().ok()?) as usize;
189                let data = lz77_decompress(payload.get(4..)?)?;
190                if data.len() != orig_payload {
191                    return None;
192                }
193                out.extend_from_slice(&data);
194            }
195            _ => return None,
196        }
197        pos += length;
198    }
199    Some(out)
200}
201
202/// Wrap `msg` in an unchained LZNT1 COMPRESSION_TRANSFORM. Returns `None` when
203/// the algorithm is unsupported or compression does not shrink the message (the
204/// caller then sends it uncompressed).
205///
206/// For a READ response only the file data is compressed, leaving the SMB2 header
207/// and READ response structure uncompressed via the transform Offset field. This
208/// mirrors Windows and, per [MS-SMB2] §3.1.4.4 step 5/6, avoids emitting a
209/// transform when compression would not shrink the payload (e.g. a tiny or
210/// incompressible read).
211pub fn compress_message(msg: &[u8], algorithm: u16) -> Option<Vec<u8>> {
212    let offset = read_response_data_offset(msg).unwrap_or(0);
213    compress_message_at(msg, algorithm, offset)
214}
215
216/// Compress the portion of `msg` at or after `offset`, leaving the leading
217/// `offset` bytes uncompressed ([MS-SMB2] §3.1.4.4 step 2).
218fn compress_message_at(msg: &[u8], algorithm: u16, offset: usize) -> Option<Vec<u8>> {
219    let prefix = msg.get(..offset)?;
220    let body = msg.get(offset..)?;
221    let payload = match algorithm {
222        algo::LZNT1 => {
223            let mut p = Vec::new();
224            lznt1::compress(body, &mut p);
225            p
226        }
227        algo::LZ77 => lz77_compress(body)?,
228        _ => return None,
229    };
230    // §3.1.4.4 step 6: send uncompressed when compression does not shrink the
231    // compressed portion.
232    if payload.len() >= body.len() {
233        return None;
234    }
235    let mut f = Vec::with_capacity(16 + prefix.len() + payload.len());
236    f.extend_from_slice(&PROTOCOL_ID);
237    f.extend_from_slice(&(body.len() as u32).to_le_bytes()); // OriginalCompressedSegmentSize
238    f.extend_from_slice(&algorithm.to_le_bytes()); // CompressionAlgorithm
239    f.extend_from_slice(&0u16.to_le_bytes()); // Flags
240    f.extend_from_slice(&(offset as u32).to_le_bytes()); // Offset of the compressed portion
241    f.extend_from_slice(prefix);
242    f.extend_from_slice(&payload);
243    Some(f)
244}
245
246/// Wrap `msg` in a chained COMPRESSION_TRANSFORM_HEADER ([MS-SMB2] \u00a72.2.42.2),
247/// emitting a Pattern_V1 payload for each long single-byte run and an `algorithm`
248/// (or uncompressed NONE) payload for the data between runs ([MS-SMB2] \u00a73.1.4.4).
249/// For a READ response the SMB2 header and READ structure are carried in a
250/// leading NONE payload so only the file data is compressed. Returns `None` when
251/// the result would not shrink the message.
252pub fn compress_message_chained(msg: &[u8], algorithm: u16) -> Option<Vec<u8>> {
253    let offset = read_response_data_offset(msg).unwrap_or(0);
254    let mut payloads = Vec::new();
255    let mut first = true;
256    if offset > 0 {
257        push_chained_payload(&mut payloads, &mut first, algo::NONE, &msg[..offset], None);
258    }
259    let body = msg.get(offset..)?;
260    let mut i = 0;
261    while i < body.len() {
262        let run = run_length(&body[i..]);
263        if run >= CHAINED_PATTERN_MIN {
264            let mut pattern = Vec::with_capacity(8);
265            pattern.push(body[i]); // Pattern
266            pattern.push(0); // Reserved1
267            pattern.extend_from_slice(&0u16.to_le_bytes()); // Reserved2
268            pattern.extend_from_slice(&(run as u32).to_le_bytes()); // Repetitions
269            push_chained_payload(&mut payloads, &mut first, algo::PATTERN_V1, &pattern, None);
270            i += run;
271            continue;
272        }
273        // Accumulate a non-pattern segment up to the next long run.
274        let seg_start = i;
275        while i < body.len() && run_length(&body[i..]) < CHAINED_PATTERN_MIN {
276            i += 1;
277        }
278        let seg = &body[seg_start..i];
279        let comp = match algorithm {
280            algo::LZ77 => lz77_compress(seg),
281            _ => {
282                let mut c = Vec::new();
283                lznt1::compress(seg, &mut c);
284                Some(c)
285            }
286        };
287        match comp {
288            Some(comp) if comp.len() < seg.len() => push_chained_payload(
289                &mut payloads,
290                &mut first,
291                algorithm,
292                &comp,
293                Some(seg.len() as u32),
294            ),
295            _ => push_chained_payload(&mut payloads, &mut first, algo::NONE, seg, None),
296        }
297    }
298    let mut f = Vec::with_capacity(8 + payloads.len());
299    f.extend_from_slice(&PROTOCOL_ID);
300    f.extend_from_slice(&(msg.len() as u32).to_le_bytes()); // OriginalCompressedSegmentSize
301    f.extend_from_slice(&payloads);
302    (f.len() < msg.len()).then_some(f)
303}
304
305/// Append one SMB2_COMPRESSION_CHAINED_PAYLOAD_HEADER plus its payload. The
306/// first payload sets SMB2_COMPRESSION_FLAG_CHAINED; `orig_payload` is the
307/// decompressed size prepended for algorithms that carry it (LZNT1/LZ77/\u2026).
308fn push_chained_payload(
309    out: &mut Vec<u8>,
310    first: &mut bool,
311    algorithm: u16,
312    data: &[u8],
313    orig_payload: Option<u32>,
314) {
315    out.extend_from_slice(&algorithm.to_le_bytes());
316    let flags = if *first { CHAINED_FLAG } else { 0 };
317    *first = false;
318    out.extend_from_slice(&flags.to_le_bytes());
319    let length = data.len() + orig_payload.map_or(0, |_| 4);
320    out.extend_from_slice(&(length as u32).to_le_bytes());
321    if let Some(op) = orig_payload {
322        out.extend_from_slice(&op.to_le_bytes());
323    }
324    out.extend_from_slice(data);
325}
326
327/// Length of the leading run of the first byte in `data`.
328fn run_length(data: &[u8]) -> usize {
329    match data.first() {
330        Some(&b) => data.iter().take_while(|&&x| x == b).count(),
331        None => 0,
332    }
333}
334fn read_response_data_offset(msg: &[u8]) -> Option<usize> {
335    const SMB2_HEADER_LEN: usize = 64;
336    const CMD_READ: u16 = 0x0008;
337    const FLAG_SERVER_TO_REDIR: u32 = 0x0000_0001;
338    if msg.len() < SMB2_HEADER_LEN + 16 || msg[0..4] != crate::SMB2_MAGIC {
339        return None;
340    }
341    let command = u16::from_le_bytes([msg[12], msg[13]]);
342    let flags = u32::from_le_bytes([msg[16], msg[17], msg[18], msg[19]]);
343    let next = u32::from_le_bytes([msg[20], msg[21], msg[22], msg[23]]);
344    if command != CMD_READ || next != 0 || flags & FLAG_SERVER_TO_REDIR == 0 {
345        return None;
346    }
347    // READ Response ([MS-SMB2] §2.2.20): StructureSize(2) then DataOffset(1),
348    // measured from the start of the SMB2 header.
349    let data_offset = msg[SMB2_HEADER_LEN + 2] as usize;
350    (data_offset < msg.len()).then_some(data_offset)
351}
352
353/// Expand an `SMB2_COMPRESSION_PATTERN_PAYLOAD_V1` (§2.2.42.1): a byte repeated
354/// `Repetitions` times.
355fn expand_pattern_v1(data: &[u8], out: &mut Vec<u8>) -> Option<()> {
356    if data.len() < 8 {
357        return None;
358    }
359    let pattern = data[0];
360    let reps = u32::from_le_bytes(data[4..8].try_into().ok()?) as usize;
361    out.resize(out.len() + reps, pattern);
362    Some(())
363}
364
365#[cfg(test)]
366mod tests {
367    use super::*;
368
369    #[test]
370    fn lznt1_transform_round_trips() {
371        let msg = b"SMB2 compression round-trip payload ".repeat(64);
372        let frame = compress_message(&msg, algo::LZNT1).expect("compresses");
373        assert_eq!(&frame[0..4], &PROTOCOL_ID, "transform magic");
374        assert!(frame.len() < msg.len(), "shrunk");
375        let back = decompress_message(&frame).expect("decompresses");
376        assert_eq!(back, msg, "round-trips");
377    }
378
379    #[test]
380    fn lz77_transform_round_trips() {
381        let msg = b"abc".repeat(100);
382        let frame = compress_message(&msg, algo::LZ77).expect("compresses");
383        assert_eq!(&frame[0..4], &PROTOCOL_ID, "transform magic");
384        assert_eq!(
385            u16::from_le_bytes([frame[8], frame[9]]),
386            algo::LZ77,
387            "advertises LZ77"
388        );
389        assert!(frame.len() < msg.len(), "shrunk");
390        let back = decompress_message(&frame).expect("decompresses");
391        assert_eq!(back, msg, "round-trips");
392    }
393
394    #[test]
395    fn incompressible_message_is_not_wrapped() {
396        // Random-ish, tiny input rarely shrinks under LZNT1.
397        let msg: Vec<u8> = (0..32u8).collect();
398        assert!(compress_message(&msg, algo::LZNT1).is_none());
399    }
400
401    #[test]
402    fn pattern_v1_expands() {
403        let mut data = vec![0xAB, 0, 0, 0];
404        data.extend_from_slice(&100u32.to_le_bytes());
405        let mut frame = Vec::new();
406        frame.extend_from_slice(&PROTOCOL_ID);
407        frame.extend_from_slice(&100u32.to_le_bytes());
408        frame.extend_from_slice(&algo::PATTERN_V1.to_le_bytes());
409        frame.extend_from_slice(&0u16.to_le_bytes());
410        frame.extend_from_slice(&0u32.to_le_bytes());
411        frame.extend_from_slice(&data);
412        let out = decompress_message(&frame).expect("pattern");
413        assert_eq!(out, vec![0xABu8; 100]);
414    }
415
416    #[test]
417    fn negotiate_intersects_with_client() {
418        // Order follows the client's request; LZ77+Huffman is not yet supported.
419        let client = [algo::LZ77, algo::LZNT1, algo::LZ77_HUFFMAN];
420        assert_eq!(negotiate_algos(&client), vec![algo::LZ77, algo::LZNT1]);
421        let caps = build_compression_caps(&[algo::LZNT1, algo::PATTERN_V1], true);
422        assert_eq!(
423            parse_compression_caps(&caps),
424            vec![algo::LZNT1, algo::PATTERN_V1]
425        );
426        assert_eq!(parse_compression_flags(&caps), CAP_FLAG_CHAINED);
427    }
428
429    #[test]
430    fn chained_pattern_and_lznt1_round_trip() {
431        // A long single-byte run (Pattern_V1) followed by compressible data.
432        let mut msg = vec![0xAAu8; 256];
433        msg.extend(b"chained compression payload data ".repeat(64));
434        let frame = compress_message_chained(&msg, algo::LZNT1).expect("compresses");
435        assert_eq!(&frame[0..4], &PROTOCOL_ID, "transform magic");
436        assert_eq!(
437            u16::from_le_bytes([frame[10], frame[11]]) & CHAINED_FLAG,
438            CHAINED_FLAG,
439            "first payload carries the chained flag"
440        );
441        assert!(frame.len() < msg.len(), "shrunk");
442        let back = decompress_message(&frame).expect("decompresses");
443        assert_eq!(back, msg, "round-trips");
444    }
445
446    #[test]
447    fn chained_bad_length_is_rejected() {
448        // ProtocolId + OriginalCompressedSegmentSize + one payload header whose
449        // Length runs past the buffer.
450        let mut frame = Vec::new();
451        frame.extend_from_slice(&PROTOCOL_ID);
452        frame.extend_from_slice(&64u32.to_le_bytes());
453        frame.extend_from_slice(&algo::NONE.to_le_bytes());
454        frame.extend_from_slice(&CHAINED_FLAG.to_le_bytes());
455        frame.extend_from_slice(&999u32.to_le_bytes()); // Length past end
456        assert!(decompress_message(&frame).is_none());
457    }
458}