1pub const PROTOCOL_ID: [u8; 4] = [0xFC, b'S', b'M', b'B'];
10
11pub const CAP_FLAG_CHAINED: u32 = 0x0000_0001;
14
15const CHAINED_FLAG: u16 = 0x0001;
20
21const CHAINED_PATTERN_MIN: usize = 32;
25
26pub mod algo {
28 pub const NONE: u16 = 0x0000;
30 pub const LZNT1: u16 = 0x0001;
32 pub const LZ77: u16 = 0x0002;
34 pub const LZ77_HUFFMAN: u16 = 0x0003;
36 pub const PATTERN_V1: u16 = 0x0004;
38}
39
40pub const SUPPORTED: &[u16] = &[algo::LZNT1, algo::LZ77, algo::PATTERN_V1];
42
43fn lz77_compress(data: &[u8]) -> Option<Vec<u8>> {
45 xpress_rs::lz77::LZ77Compressor::new(data.to_vec())
46 .compress()
47 .ok()
48}
49
50fn 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
65pub 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 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
83pub 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
91pub 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()); d.extend_from_slice(&0u16.to_le_bytes()); let flags = if chained { CAP_FLAG_CHAINED } else { 0 };
98 d.extend_from_slice(&flags.to_le_bytes()); for a in algos {
100 d.extend_from_slice(&a.to_le_bytes());
101 }
102 d
103}
104
105pub fn negotiate_algos(client: &[u16]) -> Vec<u16> {
109 client
110 .iter()
111 .copied()
112 .filter(|a| SUPPORTED.contains(a))
113 .collect()
114}
115
116pub fn decompress_message(frame: &[u8]) -> Option<Vec<u8>> {
119 if frame.len() < 16 || frame[0..4] != PROTOCOL_ID {
120 return None;
121 }
122 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
145fn 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 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 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 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 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
202pub 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
216fn 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 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()); f.extend_from_slice(&algorithm.to_le_bytes()); f.extend_from_slice(&0u16.to_le_bytes()); f.extend_from_slice(&(offset as u32).to_le_bytes()); f.extend_from_slice(prefix);
242 f.extend_from_slice(&payload);
243 Some(f)
244}
245
246pub 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.push(0); pattern.extend_from_slice(&0u16.to_le_bytes()); pattern.extend_from_slice(&(run as u32).to_le_bytes()); push_chained_payload(&mut payloads, &mut first, algo::PATTERN_V1, &pattern, None);
270 i += run;
271 continue;
272 }
273 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()); f.extend_from_slice(&payloads);
302 (f.len() < msg.len()).then_some(f)
303}
304
305fn 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
327fn 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 let data_offset = msg[SMB2_HEADER_LEN + 2] as usize;
350 (data_offset < msg.len()).then_some(data_offset)
351}
352
353fn 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 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 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 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 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()); assert!(decompress_message(&frame).is_none());
457 }
458}