iptr_decoder/raw_packet_handler/
level1.rs

1use core::{hint::unreachable_unchecked, num::NonZero};
2
3use derive_more::Display;
4
5use crate::{
6    DecoderContext, HandlePacket, TraceeMode,
7    error::{DecoderError, DecoderResult},
8    raw_packet_handler::{RawPacketHandler, RawPacketHandlers},
9};
10
11impl<H: HandlePacket> RawPacketHandlers<H> {
12    const LEVEL1_HANDLERS: [RawPacketHandler<H>; 256] = const {
13        let mut handlers: [RawPacketHandler<H>; 256] = [handle_wrong_packet::<H>; 256];
14
15        let mut index = 0;
16
17        loop {
18            if index >= 256 {
19                break;
20            }
21            let cur_index = index;
22            index += 1;
23
24            let handler = if cur_index == 0b0000_0000 {
25                // 00000000
26                handle_pad_packet::<H>
27            } else if cur_index & 0b0001_1111 == 0b0000_0001 {
28                // xxx00001
29                handle_tip_pgd_packet::<H>
30            } else if cur_index == 0b0000_0010 {
31                // 00000010
32                handle_level2_packet::<H>
33            } else if cur_index & 0b0000_0011 == 0b0000_0011 {
34                // xxxxxx11
35                handle_cyc_packet::<H>
36            } else if cur_index & 0b0000_0001 == 0b0000_0000 {
37                // xxxxxxx0 but not 00000000 and 00000010
38                handle_short_tnt_packet::<H>
39            } else if cur_index & 0b0001_1111 == 0b0000_1101 {
40                // xxx01101
41                handle_tip_packet::<H>
42            } else if cur_index & 0b0001_1111 == 0b0001_0001 {
43                // xxx10001
44                handle_tip_pge_packet::<H>
45            } else if cur_index == 0b0001_1001 {
46                // 00011001
47                handle_tsc_packet::<H>
48            } else if cur_index & 0b0001_1111 == 0b0001_1101 {
49                // xxx11101
50                handle_fup_packet::<H>
51            } else if cur_index == 0b0101_1001 {
52                // 01011001
53                handle_mtc_packet::<H>
54            } else if cur_index == 0b1001_1001 {
55                // 10011001
56                handle_mode_packet::<H>
57            } else {
58                // Anything else
59                handle_wrong_packet::<H>
60            };
61
62            handlers[cur_index] = handler;
63        }
64
65        handlers
66    };
67}
68
69#[inline]
70fn handle_pad_packet<H: HandlePacket>(
71    buf: &[u8],
72    _byte: u8,
73    context: &mut DecoderContext,
74    packet_handler: &mut H,
75) -> DecoderResult<(), H> {
76    let packet_length = 1;
77
78    loop {
79        packet_handler
80            .on_pad_packet(context)
81            .map_err(DecoderError::PacketHandler)?;
82
83        context.pos += packet_length;
84        let Some(byte) = buf.get(context.pos) else {
85            break;
86        };
87        if *byte != 0b0000_0000 {
88            break;
89        }
90        // Fast path for continuous PAD packet
91    }
92
93    Ok(())
94}
95
96#[inline]
97fn handle_short_tnt_packet<H: HandlePacket>(
98    buf: &[u8],
99    byte: u8,
100    context: &mut DecoderContext,
101    packet_handler: &mut H,
102) -> DecoderResult<(), H> {
103    if let Some(packet_block) = context.packet_block {
104        // Special handling if we are between BBP and BEP
105
106        if (byte & 0b0000_0111) != 0b0000_0100 {
107            // BIP's first byte is end with 100
108            return Err(DecoderError::InvalidPacket);
109        }
110
111        let packet_length = packet_block.size.size() + 1;
112        let id = byte >> 3;
113        let Some(bytes) = buf
114            .get((context.pos + 1)..)
115            .and_then(|buf| buf.chunks_exact(packet_block.size.size()).next())
116        else {
117            return Err(DecoderError::UnexpectedEOF);
118        };
119        packet_handler
120            .on_bip_packet(context, id, bytes, packet_block.r#type)
121            .map_err(DecoderError::PacketHandler)?;
122
123        context.pos += packet_length;
124
125        return Ok(());
126    }
127
128    // SAFETY: byte will never be zero
129    debug_assert_ne!(byte, 0, "0b0000_0000 should be PAD packet!");
130    let byte = unsafe { NonZero::new_unchecked(byte) };
131    // The short TNT packets always ends with 0, so leading zeros will never be 7;
132    // The 0b00000000 is PAD packet, so leading zeros will never be 8, so no need
133    // to check the trailing 1
134    debug_assert!(byte.leading_zeros() <= 6, "Unexpected short TNT packet!");
135
136    let packet_length = 1;
137
138    let highest_bit = 6 - byte.leading_zeros();
139    packet_handler
140        .on_short_tnt_packet(context, byte, highest_bit)
141        .map_err(DecoderError::PacketHandler)?;
142
143    context.pos += packet_length;
144
145    Ok(())
146}
147
148#[inline]
149fn handle_tip_packet<H: HandlePacket>(
150    buf: &[u8],
151    byte: u8,
152    context: &mut DecoderContext,
153    packet_handler: &mut H,
154) -> DecoderResult<(), H> {
155    context.pos += 1; // Header
156
157    let ip_bytes = byte >> 5;
158    // SAFETY: ip_bytes is not greater than 0b111
159    let ip_reconstruction_pattern = unsafe { ip_reconstruction(buf, ip_bytes, context)? };
160
161    packet_handler
162        .on_tip_packet(context, ip_reconstruction_pattern)
163        .map_err(DecoderError::PacketHandler)?;
164
165    Ok(())
166}
167
168#[inline]
169fn handle_tip_pgd_packet<H: HandlePacket>(
170    buf: &[u8],
171    byte: u8,
172    context: &mut DecoderContext,
173    packet_handler: &mut H,
174) -> DecoderResult<(), H> {
175    context.pos += 1; // Header
176
177    let ip_bytes = byte >> 5;
178    // SAFETY: ip_bytes is not greater than 0b111
179    let ip_reconstruction_pattern = unsafe { ip_reconstruction(buf, ip_bytes, context)? };
180
181    packet_handler
182        .on_tip_pgd_packet(context, ip_reconstruction_pattern)
183        .map_err(DecoderError::PacketHandler)?;
184
185    Ok(())
186}
187
188#[inline]
189fn handle_tip_pge_packet<H: HandlePacket>(
190    buf: &[u8],
191    byte: u8,
192    context: &mut DecoderContext,
193    packet_handler: &mut H,
194) -> DecoderResult<(), H> {
195    context.pos += 1; // Header
196
197    let ip_bytes = byte >> 5;
198    // SAFETY: ip_bytes is not greater than 0b111
199    let ip_reconstruction_pattern = unsafe { ip_reconstruction(buf, ip_bytes, context)? };
200
201    packet_handler
202        .on_tip_pge_packet(context, ip_reconstruction_pattern)
203        .map_err(DecoderError::PacketHandler)?;
204
205    Ok(())
206}
207
208#[inline]
209fn handle_fup_packet<H: HandlePacket>(
210    buf: &[u8],
211    byte: u8,
212    context: &mut DecoderContext,
213    packet_handler: &mut H,
214) -> DecoderResult<(), H> {
215    context.pos += 1; // Header
216
217    let ip_bytes = byte >> 5;
218    // SAFETY: ip_bytes is not greater than 0b111
219    let ip_reconstruction_pattern = unsafe { ip_reconstruction(buf, ip_bytes, context)? };
220
221    packet_handler
222        .on_fup_packet(context, ip_reconstruction_pattern)
223        .map_err(DecoderError::PacketHandler)?;
224
225    Ok(())
226}
227
228/// Pattern for IP reconstruction
229///
230/// You can use utility function [`reconstruct_ip_and_update_last`][crate::utils::reconstruct_ip_and_update_last]
231/// to use this enumerate.
232#[derive(Debug, Display, Clone, Copy)]
233pub enum IpReconstructionPattern {
234    /// None, IP is out of context
235    OutOfContext,
236    /// `IP Payload[15:0]`
237    #[display("TwoBytesWithLastIp({_0:#x})")]
238    TwoBytesWithLastIp(u16),
239    /// `IP Payload[31:0]`
240    #[display("FourBytesWithLastIp({_0:#x})")]
241    FourBytesWithLastIp(u32),
242    /// `IP Payload[47:0]`, the upper 2 bytes are guaranteed to be cleared
243    #[display("SixBytesExtended({_0:#x})")]
244    SixBytesExtended(u64),
245    /// `IP Payload[47:0]`, the upper 2 bytes are guaranteed to be cleared
246    #[display("SixBytesWithLastIp({_0:#x})")]
247    SixBytesWithLastIp(u64),
248    /// `IP Payload[63:0]`
249    #[display("EightBytes({_0:#x})")]
250    EightBytes(u64),
251}
252
253/// pos should be updated by 1 (header) before calling the function
254///
255/// # SAFETY
256///
257/// `ip_bytes` should be no greater than 0b111
258unsafe fn ip_reconstruction<H: HandlePacket>(
259    buf: &[u8],
260    ip_bytes: u8,
261    context: &mut DecoderContext,
262) -> DecoderResult<IpReconstructionPattern, H> {
263    debug_assert!(ip_bytes <= 0b111, "Unexpected ip bytes.");
264    let pattern = match ip_bytes {
265        // Header only, no IP payload
266        0b000 => IpReconstructionPattern::OutOfContext,
267        0b001 => {
268            let Some(bytes) = buf
269                .get(context.pos..)
270                .and_then(|buf| buf.first_chunk::<2>())
271            else {
272                return Err(DecoderError::UnexpectedEOF);
273            };
274            let ip_payload = u16::from_le_bytes(*bytes);
275
276            context.pos += 2;
277
278            IpReconstructionPattern::TwoBytesWithLastIp(ip_payload)
279        }
280        0b010 => {
281            let Some(bytes) = buf
282                .get(context.pos..)
283                .and_then(|buf| buf.first_chunk::<4>())
284            else {
285                return Err(DecoderError::UnexpectedEOF);
286            };
287            let ip_payload = u32::from_le_bytes(*bytes);
288
289            context.pos += 4;
290
291            IpReconstructionPattern::FourBytesWithLastIp(ip_payload)
292        }
293        0b011 => {
294            let Some([byte1, byte2, byte3, byte4, byte5, byte6]) = buf
295                .get(context.pos..)
296                .and_then(|buf| buf.first_chunk::<6>())
297            else {
298                return Err(DecoderError::UnexpectedEOF);
299            };
300            let ip_payload =
301                u64::from_le_bytes([*byte1, *byte2, *byte3, *byte4, *byte5, *byte6, 0, 0]);
302
303            context.pos += 6;
304
305            IpReconstructionPattern::SixBytesExtended(ip_payload)
306        }
307        0b100 => {
308            let Some([byte1, byte2, byte3, byte4, byte5, byte6]) = buf
309                .get(context.pos..)
310                .and_then(|buf| buf.first_chunk::<6>())
311            else {
312                return Err(DecoderError::UnexpectedEOF);
313            };
314            let ip_payload =
315                u64::from_le_bytes([*byte1, *byte2, *byte3, *byte4, *byte5, *byte6, 0, 0]);
316
317            context.pos += 6;
318
319            IpReconstructionPattern::SixBytesWithLastIp(ip_payload)
320        }
321        0b110 => {
322            let Some(bytes) = buf
323                .get(context.pos..)
324                .and_then(|buf| buf.first_chunk::<8>())
325            else {
326                return Err(DecoderError::UnexpectedEOF);
327            };
328            let ip_payload = u64::from_le_bytes(*bytes);
329
330            context.pos += 8;
331
332            IpReconstructionPattern::EightBytes(ip_payload)
333        }
334        0b101 | 0b111 => {
335            return Err(DecoderError::InvalidPacket);
336        }
337        _ => {
338            // SAFETY: ip_bytes should be no greater than than 0b111
339            unsafe {
340                unreachable_unchecked();
341            }
342        }
343    };
344
345    Ok(pattern)
346}
347
348#[inline]
349fn handle_cyc_packet<H: HandlePacket>(
350    buf: &[u8],
351    byte: u8,
352    context: &mut DecoderContext,
353    packet_handler: &mut H,
354) -> DecoderResult<(), H> {
355    let mut exp = (byte & 0b0000_0100) != 0;
356    let mut end_pos = context.pos + 1;
357
358    loop {
359        if !exp {
360            break;
361        }
362        let Some(byte) = buf.get(end_pos) else {
363            return Err(DecoderError::UnexpectedEOF);
364        };
365        exp = byte % 2 != 0;
366        end_pos += 1;
367    }
368
369    // SAFETY: All bytes are accessed before.
370    debug_assert!(buf.len() > end_pos, "Unexpected");
371    packet_handler
372        .on_cyc_packet(context, unsafe { buf.get_unchecked(context.pos..end_pos) })
373        .map_err(DecoderError::PacketHandler)?;
374
375    context.pos = end_pos;
376
377    Ok(())
378}
379
380#[inline]
381fn handle_tsc_packet<H: HandlePacket>(
382    buf: &[u8],
383    _byte: u8,
384    context: &mut DecoderContext,
385    packet_handler: &mut H,
386) -> DecoderResult<(), H> {
387    let packet_length = 8;
388
389    let Some([byte1, byte2, byte3, byte4, byte5, byte6, byte7]) = buf
390        .get((context.pos + 1)..)
391        .and_then(|buf| buf.first_chunk::<7>())
392    else {
393        return Err(DecoderError::UnexpectedEOF);
394    };
395    let tsc_bytes = [*byte1, *byte2, *byte3, *byte4, *byte5, *byte6, *byte7, 0];
396    let tsc_value = u64::from_le_bytes(tsc_bytes);
397
398    packet_handler
399        .on_tsc_packet(context, tsc_value)
400        .map_err(DecoderError::PacketHandler)?;
401
402    context.pos += packet_length;
403
404    Ok(())
405}
406
407#[inline]
408fn handle_mtc_packet<H: HandlePacket>(
409    buf: &[u8],
410    _byte: u8,
411    context: &mut DecoderContext,
412    packet_handler: &mut H,
413) -> DecoderResult<(), H> {
414    let packet_length = 2;
415
416    let Some(byte) = buf.get(context.pos + 1) else {
417        return Err(DecoderError::UnexpectedEOF);
418    };
419    let ctc_payload = *byte;
420
421    packet_handler
422        .on_mtc_packet(context, ctc_payload)
423        .map_err(DecoderError::PacketHandler)?;
424
425    context.pos += packet_length;
426
427    Ok(())
428}
429
430#[inline]
431fn handle_mode_packet<H: HandlePacket>(
432    buf: &[u8],
433    _byte: u8,
434    context: &mut DecoderContext,
435    packet_handler: &mut H,
436) -> DecoderResult<(), H> {
437    let packet_length = 2;
438
439    let Some(byte) = buf.get(context.pos + 1) else {
440        return Err(DecoderError::UnexpectedEOF);
441    };
442    let byte = *byte;
443    let leaf_id = (byte & 0b1110_0000) >> 5;
444    let mode = byte & 0b0001_1111;
445
446    if leaf_id == 0b000 {
447        // MODE.exec packet
448        match mode & 0b0000_0011 {
449            0b00 => context.tracee_mode = TraceeMode::Mode16,
450            0b01 => context.tracee_mode = TraceeMode::Mode64,
451            0b10 => context.tracee_mode = TraceeMode::Mode32,
452            _ => {}
453        }
454    }
455
456    packet_handler
457        .on_mode_packet(context, leaf_id, mode)
458        .map_err(DecoderError::PacketHandler)?;
459
460    context.pos += packet_length;
461
462    Ok(())
463}
464
465#[inline]
466fn handle_wrong_packet<H: HandlePacket>(
467    _buf: &[u8],
468    _byte: u8,
469    _context: &mut DecoderContext,
470    _packet_handler: &mut H,
471) -> DecoderResult<(), H> {
472    Err(DecoderError::InvalidPacket)
473}
474
475#[inline]
476fn handle_level2_packet<H: HandlePacket>(
477    buf: &[u8],
478    _byte: u8,
479    context: &mut DecoderContext,
480    packet_handler: &mut H,
481) -> DecoderResult<(), H> {
482    // All pos should be updated by level2's decode
483    super::level2::decode(buf, context, packet_handler)?;
484
485    Ok(())
486}
487
488macro_rules! h {
489    ($byte: ident, $buf: ident, $context: ident, $packet_handler: ident : $($val:literal),*) => {
490        match $byte {
491            $(
492                $val => RawPacketHandlers::<H>::LEVEL1_HANDLERS[$val]($buf, $byte, $context, $packet_handler),
493            )*
494        }
495    };
496}
497
498pub fn decode<H: HandlePacket>(
499    buf: &[u8],
500    context: &mut DecoderContext,
501    packet_handler: &mut H,
502) -> DecoderResult<(), H> {
503    loop {
504        let Some(byte) = buf.get(context.pos) else {
505            break;
506        };
507        let byte = *byte;
508        // Note that context.pos has not been updated before calling dispatch functions
509        h!(byte, buf, context, packet_handler: 0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112,113,114,115,116,117,118,119,120,121,122,123,124,125,126,127,128,129,130,131,132,133,134,135,136,137,138,139,140,141,142,143,144,145,146,147,148,149,150,151,152,153,154,155,156,157,158,159,160,161,162,163,164,165,166,167,168,169,170,171,172,173,174,175,176,177,178,179,180,181,182,183,184,185,186,187,188,189,190,191,192,193,194,195,196,197,198,199,200,201,202,203,204,205,206,207,208,209,210,211,212,213,214,215,216,217,218,219,220,221,222,223,224,225,226,227,228,229,230,231,232,233,234,235,236,237,238,239,240,241,242,243,244,245,246,247,248,249,250,251,252,253,254,255)?;
510    }
511
512    Ok(())
513}