km003c_lib/
packet.rs

1#![allow(unused_parens)]
2
3use crate::constants::*;
4use crate::error::KMError;
5use bytes::Bytes;
6use modular_bitfield::prelude::*;
7use num_enum::{FromPrimitive, IntoPrimitive};
8
9#[bitfield(bytes = 4)]
10#[derive(Debug, Clone, Copy, PartialEq)]
11pub struct CtrlHeader {
12    pub packet_type: B7,
13    /// Reserved flag bit in the first byte, purpose unknown.
14    /// IMPORTANT: This is NOT an indicator that an "extended header" is present.
15    pub reserved_flag: bool,
16    pub id: u8,
17    #[skip]
18    unused: bool,
19    pub attribute: B15,
20}
21
22#[bitfield(bytes = 4)]
23#[derive(Debug, Clone, Copy, PartialEq)]
24pub struct DataHeader {
25    pub packet_type: B7,
26    /// Reserved flag bit in the first byte, purpose unknown.
27    /// IMPORTANT: This is NOT an indicator that an "extended header" is present.
28    /// As per protocol research, PutData (0x41) packets always include a 4-byte
29    /// extended header regardless of this bit.
30    pub reserved_flag: bool,
31    pub id: u8,
32    #[skip]
33    unused: B6,
34    pub obj_count_words: B10,
35}
36
37/// Header used by StreamingAuth, whose final two bytes are a raw result value.
38#[bitfield(bytes = 4)]
39#[derive(Debug, Clone, Copy, PartialEq)]
40pub(crate) struct StreamingAuthHeader {
41    pub packet_type: B7,
42    pub reserved_flag: bool,
43    pub id: u8,
44    pub attribute: u16,
45}
46
47#[bitfield(bytes = 4)]
48#[derive(Debug, Clone, Copy, PartialEq)]
49pub struct ExtendedHeader {
50    pub attribute: B15,
51    pub next: bool,
52    pub chunk: B6,
53    pub size: B10,
54}
55
56/// KM003C protocol packet types.
57///
58/// Values < 0x40 are control packet types, >= 0x40 are data packet types.
59#[derive(Debug, Clone, Copy, PartialEq, Eq, IntoPrimitive, FromPrimitive)]
60#[repr(u8)]
61pub enum PacketType {
62    // 0 is reserved
63    // less than 0x40 is ctrl type
64    Sync = 0x01,
65    Connect = 0x02,
66    Disconnect = 0x03,
67    Reset = 0x04,
68    Accept = 0x05,
69    Rejected = 0x06,
70    Finished = 0x07,
71    JumpAprom = 0x08,
72    JumpDfu = 0x09,
73    GetStatus = 0x0A,
74    Error = 0x0B,
75    GetData = 0x0C,
76    GetFile = 0x0D,
77    StartGraph = 0x0E,       // Start AdcQueue streaming with rate selector
78    StopGraph = 0x0F,        // Stop AdcQueue streaming
79    EnablePdMonitor = 0x10,  // Enable PD sniffer/capture
80    DisablePdMonitor = 0x11, // Disable PD sniffer/capture
81
82    // Authentication commands
83    MemoryRead = 0x44,    // 68 - Read device memory (encrypted)
84    StreamingAuth = 0x4C, // 76 - Authenticate for AdcQueue streaming
85
86    // Error responses
87    NotReadable = 0x27, // Memory address not accessible
88
89    // >= 0x40 is data type
90    Head = 64,
91    PutData = 65,
92    #[num_enum(catch_all)]
93    Unknown(u8),
94}
95
96impl PacketType {
97    pub fn is_ctrl_type(&self) -> bool {
98        let value: u8 = (*self).into();
99        value < 0x40
100    }
101}
102
103// Python support for PacketType
104#[cfg(feature = "python")]
105impl<'py> pyo3::IntoPyObject<'py> for PacketType {
106    type Target = pyo3::PyAny;
107    type Output = pyo3::Bound<'py, Self::Target>;
108    type Error = pyo3::PyErr;
109
110    fn into_pyobject(self, py: pyo3::Python<'py>) -> Result<Self::Output, <Self as pyo3::IntoPyObject<'py>>::Error> {
111        // Convert to u8 value for Python
112        let value: u8 = self.into();
113        Ok(value.into_pyobject(py).unwrap().into_any())
114    }
115}
116
117/// Attribute values used in command headers and extended headers.
118///
119/// These values specify the type of data or command being sent.
120#[derive(Debug, Clone, Copy, PartialEq, Eq, IntoPrimitive, FromPrimitive)]
121#[repr(u16)]
122pub enum Attribute {
123    None = 0,
124    Adc = 0x1,
125    AdcQueue = 0x2,
126    AdcQueue10k = 0x4,
127    Settings = 0x8,
128    /// PD protocol data. Response is either 12-byte PdStatus (measurements only)
129    /// or longer PdEventStream (12-byte preamble + PD wire events).
130    PdPacket = 0x10,
131    /// Internal USB PD state-machine trace queues.
132    PdTrace = 0x20,
133
134    /// Offline recording metadata.
135    LogMetadata = 0x200,
136
137    #[num_enum(catch_all)]
138    Unknown(u16),
139}
140
141// Python support for Attribute
142#[cfg(feature = "python")]
143impl<'py> pyo3::IntoPyObject<'py> for Attribute {
144    type Target = pyo3::PyAny;
145    type Output = pyo3::Bound<'py, Self::Target>;
146    type Error = pyo3::PyErr;
147
148    fn into_pyobject(self, py: pyo3::Python<'py>) -> Result<Self::Output, Self::Error> {
149        // Convert to u16 value for Python
150        let value: u16 = self.into();
151        Ok(value.into_pyobject(py).unwrap().into_any())
152    }
153}
154
155/// Set of attributes for use in request masks.
156/// Can represent single or multiple attributes combined with bitwise OR.
157#[derive(Debug, Clone, Copy, PartialEq, Eq)]
158#[cfg_attr(feature = "python", derive(pyo3::IntoPyObject))]
159pub struct AttributeSet {
160    mask: u16,
161}
162
163impl AttributeSet {
164    pub const fn empty() -> Self {
165        Self { mask: 0 }
166    }
167
168    pub const fn from_raw(mask: u16) -> Self {
169        Self { mask }
170    }
171
172    pub fn single(attr: Attribute) -> Self {
173        Self { mask: attr.into() }
174    }
175
176    pub fn from_attributes<I>(attrs: I) -> Self
177    where
178        I: IntoIterator<Item = Attribute>,
179    {
180        Self {
181            mask: attrs
182                .into_iter()
183                .map(|a| {
184                    let val: u16 = a.into();
185                    val
186                })
187                .fold(0, |acc, x| acc | x),
188        }
189    }
190
191    pub fn with(mut self, attr: Attribute) -> Self {
192        let val: u16 = attr.into();
193        self.mask |= val;
194        self
195    }
196
197    pub fn contains(&self, attr: Attribute) -> bool {
198        let val: u16 = attr.into();
199        self.mask & val != 0
200    }
201
202    /// Check if any of the given attributes are present
203    pub fn contains_any<I>(&self, attrs: I) -> bool
204    where
205        I: IntoIterator<Item = Attribute>,
206    {
207        attrs.into_iter().any(|a| self.contains(a))
208    }
209
210    /// Check if all of the given attributes are present
211    pub fn contains_all<I>(&self, attrs: I) -> bool
212    where
213        I: IntoIterator<Item = Attribute>,
214    {
215        attrs.into_iter().all(|a| self.contains(a))
216    }
217
218    /// Remove an attribute from the set
219    pub fn without(mut self, attr: Attribute) -> Self {
220        let val: u16 = attr.into();
221        self.mask &= !val;
222        self
223    }
224
225    /// Iterate over all attributes in the set
226    pub fn iter(&self) -> impl Iterator<Item = Attribute> + '_ {
227        (0..16).filter_map(move |bit| {
228            let value = 1u16 << bit;
229            if self.mask & value != 0 {
230                Some(Attribute::from_primitive(value))
231            } else {
232                None
233            }
234        })
235    }
236
237    pub fn to_vec(&self) -> Vec<Attribute> {
238        self.iter().collect()
239    }
240
241    pub const fn raw(&self) -> u16 {
242        self.mask
243    }
244
245    pub fn is_empty(&self) -> bool {
246        self.mask == 0
247    }
248
249    pub fn len(&self) -> usize {
250        self.mask.count_ones() as usize
251    }
252}
253
254impl From<Attribute> for AttributeSet {
255    fn from(attr: Attribute) -> Self {
256        Self::single(attr)
257    }
258}
259
260impl FromIterator<Attribute> for AttributeSet {
261    fn from_iter<I: IntoIterator<Item = Attribute>>(iter: I) -> Self {
262        Self::from_attributes(iter)
263    }
264}
265
266/// Represents a single logical packet within a PutData response.
267/// PutData packets can contain multiple chained logical packets,
268/// each with its own extended header and payload.
269#[derive(Debug, Clone, PartialEq)]
270#[cfg_attr(
271    feature = "python",
272    pyo3::pyclass(get_all, skip_from_py_object, name = "LogicalPacket")
273)]
274pub struct LogicalPacket {
275    pub attribute: Attribute,
276    pub next: bool,
277    pub chunk: u8,
278    pub size: u16,
279    pub payload: Vec<u8>,
280}
281
282#[derive(Debug, Clone, PartialEq)]
283pub enum RawPacket {
284    Ctrl {
285        header: CtrlHeader,
286        payload: Vec<u8>,
287    },
288    SimpleData {
289        header: DataHeader,
290        payload: Vec<u8>,
291    },
292    Data {
293        header: DataHeader,
294        logical_packets: Vec<LogicalPacket>,
295    },
296}
297
298impl RawPacket {
299    pub fn id(&self) -> u8 {
300        match self {
301            RawPacket::Ctrl { header, .. } => header.id(),
302            RawPacket::SimpleData { header, .. } => header.id(),
303            RawPacket::Data { header, .. } => header.id(),
304        }
305    }
306
307    pub fn packet_type(&self) -> PacketType {
308        match self {
309            RawPacket::Ctrl { header, .. } => PacketType::from_primitive(header.packet_type()),
310            RawPacket::SimpleData { header, .. } => PacketType::from_primitive(header.packet_type()),
311            RawPacket::Data { header, .. } => PacketType::from_primitive(header.packet_type()),
312        }
313    }
314
315    /// Get the attribute for control packets
316    pub fn get_attribute(&self) -> Option<Attribute> {
317        match self {
318            RawPacket::Ctrl { header, .. } => Some(Attribute::from_primitive(header.attribute())),
319            _ => None,
320        }
321    }
322
323    /// Get the attribute set for control packets (for GetData requests)
324    pub fn get_attribute_set(&self) -> Option<AttributeSet> {
325        match self {
326            RawPacket::Ctrl { header, .. } => Some(AttributeSet::from_raw(header.attribute())),
327            _ => None,
328        }
329    }
330
331    /// Get logical packets for Data variant
332    pub fn logical_packets(&self) -> Option<&[LogicalPacket]> {
333        match self {
334            RawPacket::Data { logical_packets, .. } => Some(logical_packets),
335            _ => None,
336        }
337    }
338
339    /// Validate that response attributes match the request mask
340    ///
341    /// Returns Ok(()) if all response attributes were requested in the mask,
342    /// or Err if there's a mismatch.
343    pub fn validate_correlation(&self, request_mask: u16) -> Result<(), KMError> {
344        match self {
345            RawPacket::Data { logical_packets, .. } => {
346                let request_set = AttributeSet::from_raw(request_mask);
347
348                // Check each logical packet's attribute
349                for lp in logical_packets {
350                    if !request_set.contains(lp.attribute) {
351                        let expected: Vec<u16> = request_set.iter().map(|a| a.into()).collect();
352                        let actual: Vec<u16> = logical_packets.iter().map(|lp| lp.attribute.into()).collect();
353
354                        return Err(KMError::AttributeMismatch { expected, actual });
355                    }
356                }
357
358                Ok(())
359            }
360            _ => Ok(()), // Non-data packets don't need validation
361        }
362    }
363
364    /// Check if this is an empty PutData response (no logical packets)
365    pub fn is_empty_response(&self) -> bool {
366        matches!(self, RawPacket::Data { logical_packets, .. } if logical_packets.is_empty())
367    }
368}
369
370impl TryFrom<Bytes> for RawPacket {
371    type Error = KMError;
372
373    fn try_from(mut bytes: Bytes) -> Result<Self, Self::Error> {
374        // Check minimum length first to prevent panic in split_to
375        if bytes.len() < MAIN_HEADER_SIZE {
376            return Err(KMError::InvalidPacket(format!(
377                "Packet too short for header: expected {}, got {}",
378                MAIN_HEADER_SIZE,
379                bytes.len()
380            )));
381        }
382
383        // the first byte contains packet type (7 bits) + header flag bit
384        let first_byte = bytes[0]; // Safe now that we know len >= 4
385        // Extract only the packet type (lower 7 bits), ignoring the header flag bit
386        let package_type_byte = first_byte & 0x7F;
387        let is_ctrl_packet = PacketType::from_primitive(package_type_byte).is_ctrl_type();
388
389        let header_bytes: [u8; 4] = bytes
390            .split_to(4) // Safe now - we know there are at least 4 bytes
391            .as_ref()
392            .try_into()
393            .unwrap(); // Safe to unwrap since we know the slice is exactly 4 bytes
394        let mut payload = bytes;
395
396        if is_ctrl_packet {
397            let header = CtrlHeader::from_bytes(header_bytes);
398            Ok(RawPacket::Ctrl {
399                header,
400                payload: payload.to_vec(),
401            })
402        } else {
403            let header = DataHeader::from_bytes(header_bytes);
404            let packet_type = PacketType::from_primitive(header.packet_type());
405
406            // Only PutData packets have chained logical packets with extended headers
407            if packet_type == PacketType::PutData {
408                // The firmware's top-level count encoding can be zero even
409                // when a valid logical packet follows (for example, a
410                // 15-byte PdTrace response). Only the actual absence of bytes
411                // means that PutData has no logical packets.
412                if payload.is_empty() {
413                    return Ok(RawPacket::Data {
414                        header,
415                        logical_packets: vec![],
416                    });
417                }
418
419                if payload.len() < EXTENDED_HEADER_SIZE {
420                    // Spec indicates all PutData (0x41) packets  must carry a 4-byte extended header.
421                    // We currently fall back to SimpleData when payload < 4 for
422                    // robustness against malformed frames.
423                    return Ok(RawPacket::SimpleData {
424                        header,
425                        payload: payload.to_vec(),
426                    });
427                }
428
429                // Parse chained logical packets
430                let mut logical_packets = Vec::new();
431
432                loop {
433                    if payload.len() < EXTENDED_HEADER_SIZE {
434                        return Err(KMError::InvalidPacket(format!(
435                            "Insufficient bytes for extended header: need {}, got {}",
436                            EXTENDED_HEADER_SIZE,
437                            payload.len()
438                        )));
439                    }
440
441                    // Parse extended header
442                    let ext_header_bytes = payload.split_to(4);
443                    let ext_header_array: [u8; 4] = ext_header_bytes
444                        .as_ref()
445                        .try_into()
446                        .map_err(|_| KMError::InvalidPacket("Failed to extract extended header bytes".to_string()))?;
447                    let ext = ExtendedHeader::from_bytes(ext_header_array);
448
449                    let has_next = ext.next();
450                    let attribute = Attribute::from_primitive(ext.attribute());
451                    // AdcQueue uses chunk as its sample count and size as bytes per sample.
452                    let payload_size = if attribute == Attribute::PdTrace {
453                        crate::pd_trace::payload_size(&payload)?
454                    } else {
455                        ext.size() as usize
456                            * if attribute == Attribute::AdcQueue {
457                                ext.chunk() as usize
458                            } else {
459                                1
460                            }
461                    };
462
463                    // For AdcQueue, the size field indicates sample size (20 bytes),
464                    // but the actual payload contains multiple samples.
465                    // Take all remaining payload if this is the last logical packet.
466                    let logical_payload = if !has_next && attribute == Attribute::AdcQueue {
467                        // Last packet and AdcQueue: take all remaining bytes
468                        let all = payload.clone();
469                        payload = Bytes::new();
470                        all
471                    } else {
472                        // Normal case: take exactly size bytes
473                        if payload.len() < payload_size {
474                            return Err(KMError::InvalidPacket(format!(
475                                "Insufficient payload bytes: expected {}, got {}",
476                                payload_size,
477                                payload.len()
478                            )));
479                        }
480                        let chunk = payload.slice(..payload_size);
481                        payload = payload.slice(payload_size..);
482                        chunk
483                    };
484
485                    logical_packets.push(LogicalPacket {
486                        attribute,
487                        next: has_next,
488                        chunk: ext.chunk(),
489                        size: ext.size(),
490                        payload: logical_payload.to_vec(),
491                    });
492
493                    // Check if there are more logical packets
494                    if !has_next {
495                        break;
496                    }
497                }
498
499                if logical_packets.is_empty() {
500                    return Err(KMError::InvalidPacket(
501                        "PutData packet must have at least one logical packet".to_string(),
502                    ));
503                }
504
505                Ok(RawPacket::Data {
506                    header,
507                    logical_packets,
508                })
509            } else {
510                Ok(RawPacket::SimpleData {
511                    header,
512                    payload: payload.to_vec(),
513                })
514            }
515        }
516    }
517}
518
519// Python support
520#[cfg(feature = "python")]
521impl<'py> pyo3::IntoPyObject<'py> for RawPacket {
522    type Target = pyo3::PyAny;
523    type Output = pyo3::Bound<'py, Self::Target>;
524    type Error = pyo3::PyErr;
525
526    fn into_pyobject(self, py: pyo3::Python<'py>) -> Result<Self::Output, Self::Error> {
527        use pyo3::types::{PyDict, PyDictMethods, PyListMethods};
528
529        let dict = PyDict::new(py);
530        match self {
531            RawPacket::Ctrl { header, payload } => {
532                let inner = PyDict::new(py);
533                let header_dict = PyDict::new(py);
534                header_dict.set_item("packet_type", header.packet_type())?;
535                header_dict.set_item("reserved_flag", header.reserved_flag())?;
536                header_dict.set_item("id", header.id())?;
537                header_dict.set_item("attribute", header.attribute())?;
538                inner.set_item("header", header_dict)?;
539                inner.set_item("payload", payload)?;
540                dict.set_item("Ctrl", inner)?;
541            }
542            RawPacket::SimpleData { header, payload } => {
543                let inner = PyDict::new(py);
544                let header_dict = PyDict::new(py);
545                header_dict.set_item("packet_type", header.packet_type())?;
546                header_dict.set_item("reserved_flag", header.reserved_flag())?;
547                header_dict.set_item("id", header.id())?;
548                header_dict.set_item("obj_count_words", header.obj_count_words())?;
549                inner.set_item("header", header_dict)?;
550                inner.set_item("payload", payload)?;
551                dict.set_item("SimpleData", inner)?;
552            }
553            RawPacket::Data {
554                header,
555                logical_packets,
556            } => {
557                let inner = PyDict::new(py);
558                let header_dict = PyDict::new(py);
559                header_dict.set_item("packet_type", header.packet_type())?;
560                header_dict.set_item("reserved_flag", header.reserved_flag())?;
561                header_dict.set_item("id", header.id())?;
562                header_dict.set_item("obj_count_words", header.obj_count_words())?;
563                inner.set_item("header", header_dict)?;
564                // Convert LogicalPacket vec -> list of dicts for Python
565                let lp_list = pyo3::types::PyList::empty(py);
566                for lp in logical_packets {
567                    let lp_dict = PyDict::new(py);
568                    lp_dict.set_item("attribute", u16::from(lp.attribute))?;
569                    lp_dict.set_item("next", lp.next)?;
570                    lp_dict.set_item("chunk", lp.chunk)?;
571                    lp_dict.set_item("size", lp.size)?;
572                    lp_dict.set_item("payload", lp.payload)?;
573                    lp_list.append(lp_dict)?;
574                }
575                inner.set_item("logical_packets", lp_list)?;
576                dict.set_item("Data", inner)?;
577            }
578        }
579        Ok(dict.into_any())
580    }
581}
582
583impl From<RawPacket> for Bytes {
584    fn from(packet: RawPacket) -> Self {
585        let (header_bytes, payload) = match packet {
586            RawPacket::Ctrl { header, payload } => (header.into_bytes(), payload),
587            RawPacket::SimpleData { header, payload } => (header.into_bytes(), payload),
588            RawPacket::Data {
589                header,
590                logical_packets,
591            } => {
592                // Reconstruct chained logical packets
593                let mut full_payload = Vec::new();
594
595                for logical_packet in logical_packets {
596                    // Build extended header
597                    let ext = ExtendedHeader::new()
598                        .with_attribute(logical_packet.attribute.into())
599                        .with_next(logical_packet.next)
600                        .with_chunk(logical_packet.chunk)
601                        .with_size(logical_packet.size);
602
603                    full_payload.extend_from_slice(&ext.into_bytes());
604                    full_payload.extend_from_slice(&logical_packet.payload);
605                }
606
607                (header.into_bytes(), full_payload)
608            }
609        };
610
611        // Create the full message by combining header and payload
612        let mut message = Vec::with_capacity(4 + payload.len());
613        message.extend_from_slice(&header_bytes);
614        message.extend_from_slice(payload.as_ref());
615
616        Bytes::from(message)
617    }
618}
km003c_lib/packet.rs

km003c_lib/
packet.rs
