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/// Identity of a framed packet, read without parsing its body.
104///
105/// Correlating a response to a request only needs the first two header bytes.
106/// Going through [`RawPacket::try_from`] for that would copy the whole transfer
107/// and allocate a `Vec` per logical packet, which is wasteful in a polling
108/// loop that runs hundreds of times per second.
109#[derive(Debug, Clone, Copy, PartialEq, Eq)]
110pub struct PacketPeek {
111    pub packet_type: PacketType,
112    pub id: u8,
113}
114
115impl PacketPeek {
116    /// Read the packet type and transaction ID from a framed transfer.
117    ///
118    /// Returns `None` when the transfer is too short to contain a header.
119    pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
120        let header: &[u8; MAIN_HEADER_SIZE] = bytes.first_chunk()?;
121        Some(Self {
122            // The top bit of byte 0 is the reserved flag, not part of the type.
123            packet_type: PacketType::from_primitive(header[0] & 0x7F),
124            id: header[1],
125        })
126    }
127
128    /// Whether this transfer uses the control header layout.
129    pub fn is_ctrl(&self) -> bool {
130        self.packet_type.is_ctrl_type()
131    }
132}
133
134// Python support for PacketType
135#[cfg(feature = "python")]
136impl<'py> pyo3::IntoPyObject<'py> for PacketType {
137    type Target = pyo3::PyAny;
138    type Output = pyo3::Bound<'py, Self::Target>;
139    type Error = pyo3::PyErr;
140
141    fn into_pyobject(self, py: pyo3::Python<'py>) -> Result<Self::Output, <Self as pyo3::IntoPyObject<'py>>::Error> {
142        // Convert to u8 value for Python
143        let value: u8 = self.into();
144        Ok(value.into_pyobject(py).unwrap().into_any())
145    }
146}
147
148/// Attribute values used in command headers and extended headers.
149///
150/// These values specify the type of data or command being sent.
151#[derive(Debug, Clone, Copy, PartialEq, Eq, IntoPrimitive, FromPrimitive)]
152#[repr(u16)]
153pub enum Attribute {
154    None = 0,
155    Adc = 0x1,
156    AdcQueue = 0x2,
157    AdcQueue10k = 0x4,
158    Settings = 0x8,
159    /// PD protocol data. Response is either 12-byte PdStatus (measurements only)
160    /// or longer PdEventStream (12-byte preamble + PD wire events).
161    PdPacket = 0x10,
162    /// Internal USB PD state-machine trace queues.
163    PdTrace = 0x20,
164
165    /// Offline recording metadata.
166    LogMetadata = 0x200,
167
168    #[num_enum(catch_all)]
169    Unknown(u16),
170}
171
172// Python support for Attribute
173#[cfg(feature = "python")]
174impl<'py> pyo3::IntoPyObject<'py> for Attribute {
175    type Target = pyo3::PyAny;
176    type Output = pyo3::Bound<'py, Self::Target>;
177    type Error = pyo3::PyErr;
178
179    fn into_pyobject(self, py: pyo3::Python<'py>) -> Result<Self::Output, Self::Error> {
180        // Convert to u16 value for Python
181        let value: u16 = self.into();
182        Ok(value.into_pyobject(py).unwrap().into_any())
183    }
184}
185
186/// Set of attributes for use in request masks.
187/// Can represent single or multiple attributes combined with bitwise OR.
188#[derive(Debug, Clone, Copy, PartialEq, Eq)]
189#[cfg_attr(feature = "python", derive(pyo3::IntoPyObject))]
190pub struct AttributeSet {
191    mask: u16,
192}
193
194impl AttributeSet {
195    /// Bits actually carried by the 15-bit wire attribute field.
196    const WIRE_MASK: u16 = 0x7FFF;
197
198    pub const fn empty() -> Self {
199        Self { mask: 0 }
200    }
201
202    /// Build a set from a raw wire mask.
203    ///
204    /// The header field is 15 bits wide, so bit 15 is dropped rather than
205    /// silently lost later during serialization.
206    pub const fn from_raw(mask: u16) -> Self {
207        Self {
208            mask: mask & Self::WIRE_MASK,
209        }
210    }
211
212    pub fn single(attr: Attribute) -> Self {
213        Self { mask: attr.into() }
214    }
215
216    pub fn from_attributes<I>(attrs: I) -> Self
217    where
218        I: IntoIterator<Item = Attribute>,
219    {
220        Self {
221            mask: attrs
222                .into_iter()
223                .map(|a| {
224                    let val: u16 = a.into();
225                    val
226                })
227                .fold(0, |acc, x| acc | x),
228        }
229    }
230
231    pub fn with(mut self, attr: Attribute) -> Self {
232        let val: u16 = attr.into();
233        self.mask |= val;
234        self
235    }
236
237    pub fn contains(&self, attr: Attribute) -> bool {
238        let val: u16 = attr.into();
239        self.mask & val != 0
240    }
241
242    /// Check if any of the given attributes are present
243    pub fn contains_any<I>(&self, attrs: I) -> bool
244    where
245        I: IntoIterator<Item = Attribute>,
246    {
247        attrs.into_iter().any(|a| self.contains(a))
248    }
249
250    /// Check if all of the given attributes are present
251    pub fn contains_all<I>(&self, attrs: I) -> bool
252    where
253        I: IntoIterator<Item = Attribute>,
254    {
255        attrs.into_iter().all(|a| self.contains(a))
256    }
257
258    /// Remove an attribute from the set
259    pub fn without(mut self, attr: Attribute) -> Self {
260        let val: u16 = attr.into();
261        self.mask &= !val;
262        self
263    }
264
265    /// Iterate over all attributes in the set
266    ///
267    /// The wire field is 15 bits wide, so bit 15 is never part of a set.
268    pub fn iter(&self) -> impl Iterator<Item = Attribute> + '_ {
269        (0..15).filter_map(move |bit| {
270            let value = 1u16 << bit;
271            if self.mask & value != 0 {
272                Some(Attribute::from_primitive(value))
273            } else {
274                None
275            }
276        })
277    }
278
279    pub fn to_vec(&self) -> Vec<Attribute> {
280        self.iter().collect()
281    }
282
283    pub const fn raw(&self) -> u16 {
284        self.mask
285    }
286
287    pub fn is_empty(&self) -> bool {
288        self.mask == 0
289    }
290
291    pub fn len(&self) -> usize {
292        self.mask.count_ones() as usize
293    }
294}
295
296impl From<Attribute> for AttributeSet {
297    fn from(attr: Attribute) -> Self {
298        Self::single(attr)
299    }
300}
301
302impl FromIterator<Attribute> for AttributeSet {
303    fn from_iter<I: IntoIterator<Item = Attribute>>(iter: I) -> Self {
304        Self::from_attributes(iter)
305    }
306}
307
308/// Represents a single logical packet within a PutData response.
309/// PutData packets can contain multiple chained logical packets,
310/// each with its own extended header and payload.
311#[derive(Debug, Clone, PartialEq)]
312#[cfg_attr(
313    feature = "python",
314    pyo3::pyclass(get_all, skip_from_py_object, name = "LogicalPacket")
315)]
316pub struct LogicalPacket {
317    pub attribute: Attribute,
318    pub next: bool,
319    pub chunk: u8,
320    pub size: u16,
321    pub payload: Vec<u8>,
322}
323
324#[derive(Debug, Clone, PartialEq)]
325pub enum RawPacket {
326    Ctrl {
327        header: CtrlHeader,
328        payload: Vec<u8>,
329    },
330    SimpleData {
331        header: DataHeader,
332        payload: Vec<u8>,
333    },
334    Data {
335        header: DataHeader,
336        logical_packets: Vec<LogicalPacket>,
337    },
338}
339
340impl RawPacket {
341    pub fn id(&self) -> u8 {
342        match self {
343            RawPacket::Ctrl { header, .. } => header.id(),
344            RawPacket::SimpleData { header, .. } => header.id(),
345            RawPacket::Data { header, .. } => header.id(),
346        }
347    }
348
349    pub fn packet_type(&self) -> PacketType {
350        match self {
351            RawPacket::Ctrl { header, .. } => PacketType::from_primitive(header.packet_type()),
352            RawPacket::SimpleData { header, .. } => PacketType::from_primitive(header.packet_type()),
353            RawPacket::Data { header, .. } => PacketType::from_primitive(header.packet_type()),
354        }
355    }
356
357    /// Get the attribute for control packets
358    pub fn get_attribute(&self) -> Option<Attribute> {
359        match self {
360            RawPacket::Ctrl { header, .. } => Some(Attribute::from_primitive(header.attribute())),
361            _ => None,
362        }
363    }
364
365    /// Get the attribute set for control packets (for GetData requests)
366    pub fn get_attribute_set(&self) -> Option<AttributeSet> {
367        match self {
368            RawPacket::Ctrl { header, .. } => Some(AttributeSet::from_raw(header.attribute())),
369            _ => None,
370        }
371    }
372
373    /// Get logical packets for Data variant
374    pub fn logical_packets(&self) -> Option<&[LogicalPacket]> {
375        match self {
376            RawPacket::Data { logical_packets, .. } => Some(logical_packets),
377            _ => None,
378        }
379    }
380
381    /// Validate that response attributes match the request mask
382    ///
383    /// Returns Ok(()) if all response attributes were requested in the mask,
384    /// or Err if there's a mismatch.
385    pub fn validate_correlation(&self, request_set: AttributeSet) -> Result<(), KMError> {
386        match self {
387            RawPacket::Data { logical_packets, .. } => {
388                // Check each logical packet's attribute
389                for lp in logical_packets {
390                    if !request_set.contains(lp.attribute) {
391                        let expected: Vec<u16> = request_set.iter().map(|a| a.into()).collect();
392                        let actual: Vec<u16> = logical_packets.iter().map(|lp| lp.attribute.into()).collect();
393
394                        return Err(KMError::AttributeMismatch { expected, actual });
395                    }
396                }
397
398                Ok(())
399            }
400            _ => Ok(()), // Non-data packets don't need validation
401        }
402    }
403
404    /// Check if this is an empty PutData response (no logical packets)
405    pub fn is_empty_response(&self) -> bool {
406        matches!(self, RawPacket::Data { logical_packets, .. } if logical_packets.is_empty())
407    }
408}
409
410impl TryFrom<Bytes> for RawPacket {
411    type Error = KMError;
412
413    fn try_from(mut bytes: Bytes) -> Result<Self, Self::Error> {
414        // Check minimum length first to prevent panic in split_to
415        if bytes.len() < MAIN_HEADER_SIZE {
416            return Err(KMError::InvalidPacket(format!(
417                "Packet too short for header: expected {}, got {}",
418                MAIN_HEADER_SIZE,
419                bytes.len()
420            )));
421        }
422
423        // the first byte contains packet type (7 bits) + header flag bit
424        let first_byte = bytes[0]; // Safe now that we know len >= 4
425        // Extract only the packet type (lower 7 bits), ignoring the header flag bit
426        let package_type_byte = first_byte & 0x7F;
427        let is_ctrl_packet = PacketType::from_primitive(package_type_byte).is_ctrl_type();
428
429        let header_bytes: [u8; 4] = bytes
430            .split_to(4) // Safe now - we know there are at least 4 bytes
431            .as_ref()
432            .try_into()
433            .unwrap(); // Safe to unwrap since we know the slice is exactly 4 bytes
434        let mut payload = bytes;
435
436        if is_ctrl_packet {
437            let header = CtrlHeader::from_bytes(header_bytes);
438            Ok(RawPacket::Ctrl {
439                header,
440                payload: payload.to_vec(),
441            })
442        } else {
443            let header = DataHeader::from_bytes(header_bytes);
444            let packet_type = PacketType::from_primitive(header.packet_type());
445
446            // Only PutData packets have chained logical packets with extended headers
447            if packet_type == PacketType::PutData {
448                // The firmware's top-level count encoding can be zero even
449                // when a valid logical packet follows (for example, a
450                // 15-byte PdTrace response). Only the actual absence of bytes
451                // means that PutData has no logical packets.
452                if payload.is_empty() {
453                    return Ok(RawPacket::Data {
454                        header,
455                        logical_packets: vec![],
456                    });
457                }
458
459                if payload.len() < EXTENDED_HEADER_SIZE {
460                    // Spec indicates all PutData (0x41) packets  must carry a 4-byte extended header.
461                    // We currently fall back to SimpleData when payload < 4 for
462                    // robustness against malformed frames.
463                    return Ok(RawPacket::SimpleData {
464                        header,
465                        payload: payload.to_vec(),
466                    });
467                }
468
469                // Parse chained logical packets
470                let mut logical_packets = Vec::new();
471
472                loop {
473                    if payload.len() < EXTENDED_HEADER_SIZE {
474                        return Err(KMError::InvalidPacket(format!(
475                            "Insufficient bytes for extended header: need {}, got {}",
476                            EXTENDED_HEADER_SIZE,
477                            payload.len()
478                        )));
479                    }
480
481                    // Parse extended header
482                    let ext_header_bytes = payload.split_to(4);
483                    let ext_header_array: [u8; 4] = ext_header_bytes
484                        .as_ref()
485                        .try_into()
486                        .map_err(|_| KMError::InvalidPacket("Failed to extract extended header bytes".to_string()))?;
487                    let ext = ExtendedHeader::from_bytes(ext_header_array);
488
489                    let has_next = ext.next();
490                    let attribute = Attribute::from_primitive(ext.attribute());
491                    // AdcQueue uses chunk as its sample count and size as bytes per sample.
492                    let payload_size = if attribute == Attribute::PdTrace {
493                        crate::pd_trace::payload_size(&payload)?
494                    } else {
495                        ext.size() as usize
496                            * if attribute == Attribute::AdcQueue {
497                                ext.chunk() as usize
498                            } else {
499                                1
500                            }
501                    };
502
503                    // For AdcQueue, the size field indicates sample size (20 bytes),
504                    // but the actual payload contains multiple samples.
505                    // Take all remaining payload if this is the last logical packet.
506                    let logical_payload = if !has_next && attribute == Attribute::AdcQueue {
507                        // Last packet and AdcQueue: take all remaining bytes
508                        let all = payload.clone();
509                        payload = Bytes::new();
510                        all
511                    } else {
512                        // Normal case: take exactly size bytes
513                        if payload.len() < payload_size {
514                            return Err(KMError::InvalidPacket(format!(
515                                "Insufficient payload bytes: expected {}, got {}",
516                                payload_size,
517                                payload.len()
518                            )));
519                        }
520                        let chunk = payload.slice(..payload_size);
521                        payload = payload.slice(payload_size..);
522                        chunk
523                    };
524
525                    logical_packets.push(LogicalPacket {
526                        attribute,
527                        next: has_next,
528                        chunk: ext.chunk(),
529                        size: ext.size(),
530                        payload: logical_payload.to_vec(),
531                    });
532
533                    // Check if there are more logical packets
534                    if !has_next {
535                        break;
536                    }
537                }
538
539                if logical_packets.is_empty() {
540                    return Err(KMError::InvalidPacket(
541                        "PutData packet must have at least one logical packet".to_string(),
542                    ));
543                }
544
545                Ok(RawPacket::Data {
546                    header,
547                    logical_packets,
548                })
549            } else {
550                Ok(RawPacket::SimpleData {
551                    header,
552                    payload: payload.to_vec(),
553                })
554            }
555        }
556    }
557}
558
559// Python support
560#[cfg(feature = "python")]
561impl<'py> pyo3::IntoPyObject<'py> for RawPacket {
562    type Target = pyo3::PyAny;
563    type Output = pyo3::Bound<'py, Self::Target>;
564    type Error = pyo3::PyErr;
565
566    fn into_pyobject(self, py: pyo3::Python<'py>) -> Result<Self::Output, Self::Error> {
567        use pyo3::types::{PyDict, PyDictMethods, PyListMethods};
568
569        let dict = PyDict::new(py);
570        match self {
571            RawPacket::Ctrl { header, payload } => {
572                let inner = PyDict::new(py);
573                let header_dict = PyDict::new(py);
574                header_dict.set_item("packet_type", header.packet_type())?;
575                header_dict.set_item("reserved_flag", header.reserved_flag())?;
576                header_dict.set_item("id", header.id())?;
577                header_dict.set_item("attribute", header.attribute())?;
578                inner.set_item("header", header_dict)?;
579                inner.set_item("payload", payload)?;
580                dict.set_item("Ctrl", inner)?;
581            }
582            RawPacket::SimpleData { header, payload } => {
583                let inner = PyDict::new(py);
584                let header_dict = PyDict::new(py);
585                header_dict.set_item("packet_type", header.packet_type())?;
586                header_dict.set_item("reserved_flag", header.reserved_flag())?;
587                header_dict.set_item("id", header.id())?;
588                header_dict.set_item("obj_count_words", header.obj_count_words())?;
589                inner.set_item("header", header_dict)?;
590                inner.set_item("payload", payload)?;
591                dict.set_item("SimpleData", inner)?;
592            }
593            RawPacket::Data {
594                header,
595                logical_packets,
596            } => {
597                let inner = PyDict::new(py);
598                let header_dict = PyDict::new(py);
599                header_dict.set_item("packet_type", header.packet_type())?;
600                header_dict.set_item("reserved_flag", header.reserved_flag())?;
601                header_dict.set_item("id", header.id())?;
602                header_dict.set_item("obj_count_words", header.obj_count_words())?;
603                inner.set_item("header", header_dict)?;
604                // Convert LogicalPacket vec -> list of dicts for Python
605                let lp_list = pyo3::types::PyList::empty(py);
606                for lp in logical_packets {
607                    let lp_dict = PyDict::new(py);
608                    lp_dict.set_item("attribute", u16::from(lp.attribute))?;
609                    lp_dict.set_item("next", lp.next)?;
610                    lp_dict.set_item("chunk", lp.chunk)?;
611                    lp_dict.set_item("size", lp.size)?;
612                    lp_dict.set_item("payload", lp.payload)?;
613                    lp_list.append(lp_dict)?;
614                }
615                inner.set_item("logical_packets", lp_list)?;
616                dict.set_item("Data", inner)?;
617            }
618        }
619        Ok(dict.into_any())
620    }
621}
622
623impl From<RawPacket> for Bytes {
624    fn from(packet: RawPacket) -> Self {
625        let (header_bytes, payload) = match packet {
626            RawPacket::Ctrl { header, payload } => (header.into_bytes(), payload),
627            RawPacket::SimpleData { header, payload } => (header.into_bytes(), payload),
628            RawPacket::Data {
629                header,
630                logical_packets,
631            } => {
632                // Reconstruct chained logical packets
633                let mut full_payload = Vec::new();
634
635                for logical_packet in logical_packets {
636                    // Build extended header
637                    let ext = ExtendedHeader::new()
638                        .with_attribute(logical_packet.attribute.into())
639                        .with_next(logical_packet.next)
640                        .with_chunk(logical_packet.chunk)
641                        .with_size(logical_packet.size);
642
643                    full_payload.extend_from_slice(&ext.into_bytes());
644                    full_payload.extend_from_slice(&logical_packet.payload);
645                }
646
647                (header.into_bytes(), full_payload)
648            }
649        };
650
651        // Create the full message by combining header and payload
652        let mut message = Vec::with_capacity(4 + payload.len());
653        message.extend_from_slice(&header_bytes);
654        message.extend_from_slice(payload.as_ref());
655
656        Bytes::from(message)
657    }
658}
km003c_lib/packet.rs

km003c_lib/
packet.rs
