1use crate::adc::AdcDataSimple;
46use crate::adcqueue::GraphSampleRate;
47use crate::auth::{
48 AuthCredential, CALIBRATION_ADDRESS, DeviceInfo, HardwareId, PREFERRED_CALIBRATION_ADDRESS,
49 STREAMING_AUTH_CREDENTIAL_SIZE, StreamingAuthResult,
50};
51use crate::error::KMError;
52use crate::message::Packet;
53use crate::offline::{LogMetadata, LogMetadataResponse, OfflineLog};
54use crate::packet::{Attribute, AttributeSet, PacketPeek, PacketType, RawPacket};
55use crate::pd::{PdEventStream, PdStatus};
56use crate::settings::Settings;
57use bytes::Bytes;
58use nusb::Interface;
59use nusb::io::{EndpointRead, EndpointWrite};
60use nusb::transfer::{Bulk, Interrupt};
61use std::collections::VecDeque;
62use std::fmt;
63use std::time::Duration;
64use tokio::io::{AsyncReadExt, AsyncWriteExt};
65use tokio::time::timeout;
66use tracing::{debug, info, trace, warn};
67
68fn calibration_record_is_erased(record: &[u8]) -> bool {
69 record.starts_with(&[0xff; 4])
70}
71
72#[derive(Debug, Clone)]
79pub struct DeviceState {
80 pub info: DeviceInfo,
82 pub hardware_id: HardwareId,
84 pub auth_level: u8,
86 pub adcqueue_enabled: bool,
88}
89
90impl DeviceState {
91 pub fn is_authenticated(&self) -> bool {
93 self.auth_level > 0
94 }
95
96 pub fn model(&self) -> &str {
98 &self.info.model
99 }
100
101 pub fn firmware_version(&self) -> &str {
103 &self.info.fw_version
104 }
105}
106
107impl fmt::Display for DeviceState {
108 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
109 const W: usize = 58; fn line(label: &str, value: &str) -> String {
113 let content = format!(" {}: {}", label, value);
114 format!("│{:<W$}│", content, W = W)
115 }
116
117 let sep = format!("├{:─<W$}┤", "", W = W);
118
119 writeln!(f, "┌{:─<W$}┐", "", W = W)?;
120 writeln!(f, "│{:^W$}│", "POWER-Z KM003C Device Information", W = W)?;
121 writeln!(f, "{}", sep)?;
122 writeln!(f, "{}", line("Model", &self.info.model))?;
123 writeln!(f, "{}", line("Hardware Version", &self.info.hw_version))?;
124 writeln!(f, "{}", line("Mfg Date", &self.info.mfg_date))?;
125 writeln!(f, "{}", sep)?;
126 writeln!(f, "{}", line("Firmware Version", &self.info.fw_version))?;
127 writeln!(f, "{}", line("Firmware Date", &self.info.fw_date))?;
128 writeln!(f, "{}", sep)?;
129 writeln!(f, "{}", line("Serial ID", &self.info.serial_id))?;
130 writeln!(f, "{}", line("UUID", &self.info.uuid))?;
131 writeln!(f, "{}", line("Hardware ID", &format!("{}", self.hardware_id)))?;
132 writeln!(f, "{}", sep)?;
133 writeln!(f, "{}", line("Auth Level", &format!("{}", self.auth_level)))?;
134 writeln!(f, "{}", line("AdcQueue Enabled", &format!("{}", self.adcqueue_enabled)))?;
135 write!(f, "└{:─<W$}┘", "", W = W)
136 }
137}
138
139pub const VID: u16 = 0x5FC9; pub const PID: u16 = 0x0063; pub const INTERFACE_VENDOR: u8 = 0;
145pub const ENDPOINT_OUT_VENDOR: u8 = 0x01;
146pub const ENDPOINT_IN_VENDOR: u8 = 0x81;
147
148pub const INTERFACE_HID: u8 = 3;
150pub const ENDPOINT_OUT_HID: u8 = 0x05;
151pub const ENDPOINT_IN_HID: u8 = 0x85;
152
153const DEFAULT_TIMEOUT: Duration = Duration::from_secs(2);
155const BULK_TRANSFER_SIZE: usize = 2048;
156const INTERRUPT_TRANSFER_SIZE: usize = 64;
157const MAX_PENDING_RESPONSES: usize = 256;
158const AES_BLOCK_SIZE: usize = 16;
159
160async fn find_device(selector: &DeviceSelector) -> Result<nusb::DeviceInfo, KMError> {
162 nusb::list_devices()
163 .await?
164 .find(|device| device.vendor_id() == VID && device.product_id() == PID && selector.matches(device))
165 .ok_or(KMError::DeviceNotFound)
166}
167
168fn response_matches(bytes: &[u8], id: u8, packet_type: PacketType) -> bool {
169 PacketPeek::from_bytes(bytes).is_some_and(|peek| peek.id == id && peek.packet_type == packet_type)
170}
171
172fn response_type_matches(bytes: &[u8], packet_type: PacketType) -> bool {
173 PacketPeek::from_bytes(bytes).is_some_and(|peek| peek.packet_type == packet_type)
174}
175
176fn control_response_matches(bytes: &[u8], id: u8) -> bool {
177 PacketPeek::from_bytes(bytes).is_some_and(|peek| peek.is_ctrl() && peek.id == id)
178}
179
180fn memory_confirmation_matches(bytes: &[u8], id: u8) -> bool {
181 PacketPeek::from_bytes(bytes).is_some_and(|peek| {
182 peek.id == id
183 && matches!(
184 peek.packet_type,
185 PacketType::MemoryRead | PacketType::Rejected | PacketType::NotReadable
186 )
187 })
188}
189
190fn memory_response_size(requested_size: u32) -> usize {
191 (requested_size as usize).div_ceil(AES_BLOCK_SIZE) * AES_BLOCK_SIZE
192}
193
194fn append_memory_chunk(buffer: &mut Vec<u8>, chunk: &[u8], expected_size: usize) -> Result<bool, KMError> {
195 if chunk.is_empty() {
196 return Err(KMError::Protocol(
197 "Received an empty encrypted memory chunk".to_string(),
198 ));
199 }
200
201 let remaining = expected_size.saturating_sub(buffer.len());
202 if chunk.len() > remaining {
203 return Err(KMError::Protocol(format!(
204 "Encrypted memory response exceeded expected size: expected {expected_size} bytes, got at least {}",
205 buffer.len() + chunk.len()
206 )));
207 }
208
209 buffer.extend_from_slice(chunk);
210 Ok(buffer.len() == expected_size)
211}
212
213fn decrypt_memory_response(encrypted: &[u8], requested_size: u32) -> Result<Vec<u8>, KMError> {
214 let expected_size = memory_response_size(requested_size);
215 if encrypted.len() != expected_size {
216 return Err(KMError::Protocol(format!(
217 "Encrypted memory response has the wrong size: expected {expected_size} bytes, got {}",
218 encrypted.len()
219 )));
220 }
221
222 let mut decrypted = crate::auth::aes_ecb_decrypt_blocks(encrypted, crate::auth::MEMORY_READ_KEY)?;
223 decrypted.truncate(requested_size as usize);
224 Ok(decrypted)
225}
226
227fn validate_memory_read_confirmation(
228 packet: &RawPacket,
229 expected_id: u8,
230 expected_address: u32,
231 expected_size: u32,
232) -> Result<(), KMError> {
233 let RawPacket::SimpleData { header, payload } = packet else {
234 return Err(KMError::Protocol(
235 "MemoryRead confirmation has an unexpected packet layout".to_string(),
236 ));
237 };
238
239 if header.id() != expected_id || PacketType::from(header.packet_type()) != PacketType::MemoryRead {
240 return Err(KMError::Protocol(
241 "MemoryRead confirmation does not match the request".to_string(),
242 ));
243 }
244 if payload.len() != crate::auth::MEMORY_READ_CONFIRMATION_SIZE {
245 return Err(KMError::Protocol(format!(
246 "MemoryRead confirmation must contain {} bytes, got {}",
247 crate::auth::MEMORY_READ_CONFIRMATION_SIZE,
248 payload.len()
249 )));
250 }
251
252 let confirmation = crate::auth::MemoryReadConfirmation::from_bytes(payload)
253 .ok_or_else(|| KMError::Protocol("MemoryRead confirmation body is truncated".to_string()))?;
254
255 if confirmation.address != expected_address || confirmation.size != expected_size {
256 return Err(KMError::Protocol(format!(
257 "MemoryRead confirmation echoed address 0x{:08X} and size {}, expected 0x{expected_address:08X} and {expected_size}",
258 confirmation.address, confirmation.size
259 )));
260 }
261 if confirmation.magic != u32::MAX {
262 return Err(KMError::Protocol(format!(
263 "MemoryRead confirmation has invalid magic 0x{:08X}",
264 confirmation.magic
265 )));
266 }
267 if confirmation.checksum != confirmation.expected_checksum {
268 return Err(KMError::Protocol(format!(
269 "MemoryRead confirmation CRC mismatch: expected 0x{:08X}, got 0x{:08X}",
270 confirmation.expected_checksum, confirmation.checksum
271 )));
272 }
273
274 Ok(())
275}
276
277#[derive(Debug, Clone, Copy, PartialEq, Eq)]
279pub enum TransferType {
280 Bulk,
281 Interrupt,
282}
283
284#[derive(Debug, Clone)]
289pub enum ConnectionMode {
290 Basic,
292 Full(DeviceState),
294}
295
296fn ensure_adcqueue_available(mode: &ConnectionMode) -> Result<(), KMError> {
297 match mode {
298 ConnectionMode::Basic => Err(KMError::Protocol(
299 "AdcQueue streaming requires Full mode (vendor interface)".to_string(),
300 )),
301 ConnectionMode::Full(state) if !state.adcqueue_enabled => Err(KMError::Protocol(
302 "StreamingAuth did not enable AdcQueue streaming".to_string(),
303 )),
304 ConnectionMode::Full(_) => Ok(()),
305 }
306}
307
308#[derive(Debug, Clone)]
332pub struct DeviceConfig {
333 interface: u8,
335 endpoint_out: u8,
337 endpoint_in: u8,
339 transfer_type: TransferType,
341 skip_reset: bool,
343 selector: DeviceSelector,
345}
346
347#[derive(Debug, Clone, Default, PartialEq, Eq)]
353pub enum DeviceSelector {
354 #[default]
356 First,
357 SerialNumber(String),
359 BusAddress { bus_id: String, device_address: u8 },
364}
365
366impl DeviceSelector {
367 fn matches(&self, device: &nusb::DeviceInfo) -> bool {
368 match self {
369 Self::First => true,
370 Self::SerialNumber(serial) => device.serial_number() == Some(serial.as_str()),
371 Self::BusAddress { bus_id, device_address } => {
372 device.bus_id() == bus_id && device.device_address() == *device_address
373 }
374 }
375 }
376
377 fn describe(&self) -> String {
378 match self {
379 Self::First => "the first attached KM003C".to_string(),
380 Self::SerialNumber(serial) => format!("KM003C with serial {serial}"),
381 Self::BusAddress { bus_id, device_address } => format!("KM003C at bus {bus_id} address {device_address}"),
382 }
383 }
384}
385
386impl DeviceConfig {
387 pub fn vendor() -> Self {
401 Self {
402 interface: INTERFACE_VENDOR,
403 endpoint_out: ENDPOINT_OUT_VENDOR,
404 endpoint_in: ENDPOINT_IN_VENDOR,
405 transfer_type: TransferType::Bulk,
406 skip_reset: false,
407 selector: DeviceSelector::First,
408 }
409 }
410
411 pub fn hid() -> Self {
427 Self {
428 interface: INTERFACE_HID,
429 endpoint_out: ENDPOINT_OUT_HID,
430 endpoint_in: ENDPOINT_IN_HID,
431 transfer_type: TransferType::Interrupt,
432 skip_reset: false,
433 selector: DeviceSelector::First,
434 }
435 }
436
437 pub fn skip_reset(mut self) -> Self {
442 self.skip_reset = true;
443 self
444 }
445
446 pub fn select(mut self, selector: DeviceSelector) -> Self {
462 self.selector = selector;
463 self
464 }
465
466 pub fn selector(&self) -> &DeviceSelector {
468 &self.selector
469 }
470
471 pub fn is_vendor(&self) -> bool {
473 self.interface == INTERFACE_VENDOR
474 }
475
476 pub fn is_hid(&self) -> bool {
478 self.interface == INTERFACE_HID
479 }
480}
481
482enum EndpointReaderType {
484 Bulk(EndpointRead<Bulk>),
485 Interrupt(EndpointRead<Interrupt>),
486}
487
488enum EndpointWriterType {
490 Bulk(EndpointWrite<Bulk>),
491 Interrupt(EndpointWrite<Interrupt>),
492}
493
494pub struct KM003C {
495 #[allow(dead_code)]
497 interface: Interface,
498 transaction_id: u8,
499 reader: EndpointReaderType,
500 writer: EndpointWriterType,
501 pending_responses: VecDeque<Vec<u8>>,
502 graph_sample_rate: Option<GraphSampleRate>,
504 mode: ConnectionMode,
506}
507
508impl KM003C {
509 pub async fn new(config: DeviceConfig) -> Result<Self, KMError> {
536 let is_vendor = config.is_vendor();
537 let mut device = Self::connect(config).await?;
538
539 if is_vendor {
540 device.run_init().await?;
542 }
543 Ok(device)
546 }
547
548 async fn connect(config: DeviceConfig) -> Result<Self, KMError> {
550 info!("Searching for {}...", config.selector.describe());
551 let device_info = find_device(&config.selector).await?;
552
553 info!(
554 "Found device on bus {} addr {}",
555 device_info.bus_id(),
556 device_info.device_address()
557 );
558
559 let reopen_selector = match &config.selector {
563 DeviceSelector::SerialNumber(_) => config.selector.clone(),
564 _ => device_info
565 .serial_number()
566 .map(|serial| DeviceSelector::SerialNumber(serial.to_string()))
567 .unwrap_or(DeviceSelector::First),
568 };
569
570 let mut device = device_info.open().await?;
571
572 if !config.skip_reset {
574 info!("Resetting device...");
575 device.reset().await?;
576 tokio::time::sleep(Duration::from_millis(1500)).await;
579 device = find_device(&reopen_selector).await?.open().await?;
581 } else {
582 debug!("Skipping USB reset (skip_reset=true)");
583 }
584
585 for interface_num in 0..4 {
592 if let Err(e) = device.detach_kernel_driver(interface_num) {
593 trace!("Could not detach interface {}: {}", interface_num, e);
595 } else {
596 debug!("Detached kernel driver from interface {}", interface_num);
597 }
598 }
599
600 let interface = device.claim_interface(config.interface).await?;
602 info!("Interface {} claimed successfully", config.interface);
603
604 let (reader, writer) = match config.transfer_type {
609 TransferType::Bulk => {
610 let ep_in = interface.endpoint::<Bulk, _>(config.endpoint_in)?;
611 let ep_out = interface.endpoint::<Bulk, _>(config.endpoint_out)?;
612 (
613 EndpointReaderType::Bulk(ep_in.reader(BULK_TRANSFER_SIZE).with_num_transfers(4)),
614 EndpointWriterType::Bulk(ep_out.writer(INTERRUPT_TRANSFER_SIZE).with_num_transfers(4)),
615 )
616 }
617 TransferType::Interrupt => {
618 let ep_in = interface.endpoint::<Interrupt, _>(config.endpoint_in)?;
619 let ep_out = interface.endpoint::<Interrupt, _>(config.endpoint_out)?;
620 (
621 EndpointReaderType::Interrupt(ep_in.reader(INTERRUPT_TRANSFER_SIZE).with_num_transfers(4)),
622 EndpointWriterType::Interrupt(ep_out.writer(INTERRUPT_TRANSFER_SIZE).with_num_transfers(4)),
623 )
624 }
625 };
626
627 let km003c = Self {
628 interface,
629 transaction_id: 0,
630 reader,
631 writer,
632 pending_responses: VecDeque::new(),
633 graph_sample_rate: None,
634 mode: ConnectionMode::Basic,
635 };
636
637 info!("USB connection established");
638 Ok(km003c)
639 }
640
641 fn next_transaction_id(&mut self) -> u8 {
643 let id = self.transaction_id;
644 self.transaction_id = self.transaction_id.wrapping_add(1);
645 self.pending_responses
648 .retain(|response| !PacketPeek::from_bytes(response).is_some_and(|peek| peek.id == id));
649 id
650 }
651
652 pub fn next_tid(&mut self) -> u8 {
654 self.next_transaction_id()
655 }
656
657 pub fn set_transaction_id(&mut self, id: u8) {
659 self.transaction_id = id;
660 }
661
662 pub async fn send(&mut self, packet: Packet) -> Result<(), KMError> {
667 self.send_tracked(packet).await.map(|_| ())
668 }
669
670 async fn send_tracked(&mut self, packet: Packet) -> Result<u8, KMError> {
676 let id = self.next_transaction_id();
677 let raw_packet = packet.to_raw_packet(id)?;
678 self.send_raw_packet(raw_packet).await?;
679 Ok(id)
680 }
681
682 async fn send_raw_packet(&mut self, packet: RawPacket) -> Result<(), KMError> {
684 let (reserved_flag, has_logical_packets) = match &packet {
685 RawPacket::Ctrl { header, .. } => (header.reserved_flag(), false),
686 RawPacket::SimpleData { header, .. } => (header.reserved_flag(), false),
687 RawPacket::Data {
688 header,
689 logical_packets,
690 } => (header.reserved_flag(), !logical_packets.is_empty()),
691 };
692
693 debug!(
694 "Sending packet: type={:?}, reserved_flag={}, has_logical_packets={}, id={}",
695 packet.packet_type(),
696 reserved_flag,
697 has_logical_packets,
698 packet.id(),
699 );
700
701 let message_bytes = Bytes::from(packet);
702 let message = message_bytes.to_vec();
703 trace!("TX [{} bytes]: {:02x?}", message.len(), message);
704
705 match &mut self.writer {
707 EndpointWriterType::Bulk(writer) => {
708 timeout(DEFAULT_TIMEOUT, writer.write_all(&message)).await??;
709 timeout(DEFAULT_TIMEOUT, writer.flush_end_async()).await??;
710 }
711 EndpointWriterType::Interrupt(writer) => {
712 timeout(DEFAULT_TIMEOUT, writer.write_all(&message)).await??;
713 timeout(DEFAULT_TIMEOUT, writer.flush_end_async()).await??;
714 }
715 }
716
717 debug!("Sent successfully");
718 Ok(())
719 }
720
721 pub async fn send_raw(&mut self, data: &[u8]) -> Result<(), KMError> {
723 trace!("TX [{} bytes]: {:02x?}", data.len(), data);
724 match &mut self.writer {
725 EndpointWriterType::Bulk(writer) => {
726 timeout(DEFAULT_TIMEOUT, writer.write_all(data)).await??;
727 timeout(DEFAULT_TIMEOUT, writer.flush_end_async()).await??;
728 }
729 EndpointWriterType::Interrupt(writer) => {
730 timeout(DEFAULT_TIMEOUT, writer.write_all(data)).await??;
731 timeout(DEFAULT_TIMEOUT, writer.flush_end_async()).await??;
732 }
733 }
734 Ok(())
735 }
736
737 async fn read_raw_from_usb(&mut self) -> Result<Vec<u8>, KMError> {
739 let mut buffer = Vec::new();
740 let bytes_read = match &mut self.reader {
741 EndpointReaderType::Bulk(reader) => {
742 let mut message = reader.until_short_packet();
746 let bytes_read = timeout(DEFAULT_TIMEOUT, message.read_to_end(&mut buffer)).await??;
747 message
748 .consume_end()
749 .map_err(|_| KMError::Protocol("Bulk response ended without a short packet".to_string()))?;
750 bytes_read
751 }
752 EndpointReaderType::Interrupt(reader) => {
753 buffer.resize(INTERRUPT_TRANSFER_SIZE, 0);
754 let bytes_read = timeout(DEFAULT_TIMEOUT, reader.read(&mut buffer)).await??;
755 buffer.truncate(bytes_read);
756 bytes_read
757 }
758 };
759 trace!("RX [{} bytes]: {:02x?}", bytes_read, buffer);
760 Ok(buffer)
761 }
762
763 fn queue_pending_response(&mut self, response: Vec<u8>) {
764 if self.pending_responses.len() == MAX_PENDING_RESPONSES {
765 warn!("Dropping oldest unmatched response because the pending queue is full");
766 self.pending_responses.pop_front();
767 }
768 self.pending_responses.push_back(response);
769 }
770
771 async fn receive_matching_raw<F>(&mut self, mut predicate: F) -> Result<Vec<u8>, KMError>
772 where
773 F: FnMut(&[u8]) -> bool,
774 {
775 if let Some(index) = self.pending_responses.iter().position(|response| predicate(response)) {
776 return Ok(self
777 .pending_responses
778 .remove(index)
779 .expect("pending response index is valid"));
780 }
781
782 let deadline = tokio::time::Instant::now() + DEFAULT_TIMEOUT;
783 loop {
784 let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
785 if remaining.is_zero() {
786 return Err(KMError::Protocol(
787 "Timed out waiting for a correlated response".to_string(),
788 ));
789 }
790
791 let response = timeout(remaining, self.read_raw_from_usb()).await??;
792 if predicate(&response) {
793 return Ok(response);
794 }
795
796 if let Some(peek) = PacketPeek::from_bytes(&response) {
797 debug!(
798 "Queueing unmatched response: type={:?}, id={}, len={}",
799 peek.packet_type,
800 peek.id,
801 response.len()
802 );
803 } else {
804 debug!("Queueing unframed response: len={}", response.len());
805 }
806 self.queue_pending_response(response);
807 }
808 }
809
810 pub async fn receive_raw(&mut self) -> Result<Vec<u8>, KMError> {
815 if let Some(response) = self.pending_responses.pop_front() {
816 return Ok(response);
817 }
818 self.read_raw_from_usb().await
819 }
820
821 fn parse_response(raw_bytes: Vec<u8>, graph_rate: Option<GraphSampleRate>) -> Result<Packet, KMError> {
822 if raw_bytes.is_empty() {
823 return Err(KMError::Protocol("Received 0 bytes".to_string()));
824 }
825
826 let raw_packet = RawPacket::try_from(Bytes::from(raw_bytes))?;
827 if let Some(rate) = graph_rate {
828 Packet::from_raw_with_graph_rate(raw_packet, rate)
829 } else {
830 Packet::try_from(raw_packet)
831 }
832 }
833
834 pub async fn receive(&mut self) -> Result<Packet, KMError> {
839 let raw_bytes = self.receive_raw().await?;
840 Self::parse_response(raw_bytes, self.graph_sample_rate)
841 }
842
843 async fn receive_control_response(&mut self, id: u8) -> Result<Packet, KMError> {
844 let raw_bytes = self
845 .receive_matching_raw(|bytes| control_response_matches(bytes, id))
846 .await?;
847 Self::parse_response(raw_bytes, self.graph_sample_rate)
848 }
849
850 async fn expect_accept(&mut self, id: u8, command: &str) -> Result<(), KMError> {
851 match self.receive_control_response(id).await? {
852 Packet::Accept { .. } => Ok(()),
853 Packet::Reject { .. } => Err(KMError::Protocol(format!("{command} was rejected by the device"))),
854 Packet::NotReadable { .. } => Err(KMError::Protocol(format!("{command} is not readable"))),
855 other => Err(KMError::Protocol(format!(
856 "Expected Accept for {command}, got {other:?}"
857 ))),
858 }
859 }
860
861 async fn receive_memory_read_data_exact(&mut self, requested_size: u32) -> Result<Vec<u8>, KMError> {
867 let expected_size = memory_response_size(requested_size);
868 if expected_size == 0 {
869 return Ok(Vec::new());
870 }
871
872 let mut encrypted = Vec::with_capacity(expected_size);
873 while encrypted.len() < expected_size {
874 let chunk = self.read_raw_from_usb().await?;
875 if append_memory_chunk(&mut encrypted, &chunk, expected_size)? {
876 break;
877 }
878 }
879
880 decrypt_memory_response(&encrypted, requested_size)
881 }
882
883 pub async fn request_data(&mut self, mask: AttributeSet) -> Result<Packet, KMError> {
885 let id = self.send_tracked(Packet::GetData { attributes: mask }).await?;
886 let raw_bytes = self
887 .receive_matching_raw(|bytes| response_matches(bytes, id, PacketType::PutData))
888 .await?;
889 let raw_packet = RawPacket::try_from(Bytes::from(raw_bytes))?;
890 raw_packet.validate_correlation(mask)?;
891 if let Some(rate) = self.graph_sample_rate {
892 Packet::from_raw_with_graph_rate(raw_packet, rate)
893 } else {
894 Packet::try_from(raw_packet)
895 }
896 }
897
898 pub async fn request_adc_data(&mut self) -> Result<AdcDataSimple, KMError> {
900 let packet = self.request_data(AttributeSet::single(Attribute::Adc)).await?;
901 packet
902 .get_adc()
903 .cloned()
904 .ok_or_else(|| KMError::Protocol("No ADC data in response".to_string()))
905 }
906
907 pub async fn request_settings(&mut self) -> Result<Settings, KMError> {
909 let packet = self.request_data(AttributeSet::single(Attribute::Settings)).await?;
910 packet
911 .get_settings()
912 .cloned()
913 .ok_or_else(|| KMError::Protocol("No Settings data in response".to_string()))
914 }
915
916 pub async fn request_pd_trace(&mut self) -> Result<crate::pd_trace::PdTrace, KMError> {
918 let packet = self.request_data(AttributeSet::single(Attribute::PdTrace)).await?;
919 packet
920 .get_pd_trace()
921 .cloned()
922 .ok_or_else(|| KMError::Protocol("No PD trace data in response".to_string()))
923 }
924
925 pub async fn request_log_metadata(&mut self) -> Result<Vec<LogMetadata>, KMError> {
929 let packet = self.request_data(AttributeSet::single(Attribute::LogMetadata)).await?;
930 match packet.get_log_metadata() {
931 Some(LogMetadataResponse::Empty) => Ok(Vec::new()),
932 Some(LogMetadataResponse::Available(metadata)) => Ok(metadata.clone()),
933 None => Err(KMError::Protocol("No LogMetadata data in response".to_string())),
934 }
935 }
936
937 pub async fn download_offline_log(&mut self, metadata: LogMetadata) -> Result<OfflineLog, KMError> {
939 let data = self
940 .read_memory_block(metadata.data_address()?, metadata.data_size())
941 .await?;
942 OfflineLog::from_bytes(metadata, &data)
943 }
944
945 pub async fn request_pd_data(&mut self) -> Result<Packet, KMError> {
953 self.request_data(AttributeSet::single(Attribute::PdPacket)).await
954 }
955
956 pub async fn request_adc_with_pd(&mut self) -> Result<Packet, KMError> {
961 let mask = AttributeSet::single(Attribute::Adc).with(Attribute::PdPacket);
962 self.request_data(mask).await
963 }
964
965 pub fn extract_pd_status(packet: &Packet) -> Option<&PdStatus> {
970 packet.get_pd_status()
971 }
972
973 pub fn extract_pd_events(packet: &Packet) -> Option<&PdEventStream> {
978 packet.get_pd_events()
979 }
980
981 pub async fn enable_pd_monitor(&mut self) -> Result<(), KMError> {
989 let id = self.send_tracked(Packet::EnablePdMonitor).await?;
990 self.expect_accept(id, "EnablePdMonitor").await
991 }
992
993 pub async fn disable_pd_monitor(&mut self) -> Result<(), KMError> {
999 let id = self.send_tracked(Packet::DisablePdMonitor).await?;
1000 self.expect_accept(id, "DisablePdMonitor").await
1001 }
1002
1003 async fn run_init(&mut self) -> Result<(), KMError> {
1015 use crate::auth::{HARDWARE_ID_ADDRESS, HARDWARE_ID_SIZE};
1016
1017 const MAX_CONNECT_RETRIES: u8 = 3;
1019 let mut last_error = None;
1020 for attempt in 1..=MAX_CONNECT_RETRIES {
1021 match self.send_tracked(Packet::Connect).await {
1022 Err(err) => last_error = Some(format!("Connect send failed: {err:?}")),
1023 Ok(id) => match self.expect_accept(id, "Connect").await {
1024 Ok(()) => {
1025 last_error = None;
1026 break;
1027 }
1028 Err(err) => last_error = Some(format!("Connect receive failed: {err:?}")),
1029 },
1030 }
1031 if attempt < MAX_CONNECT_RETRIES {
1032 debug!("Connect attempt {} failed, retrying...", attempt);
1033 tokio::time::sleep(Duration::from_millis(100)).await;
1034 }
1035 }
1036 if let Some(err) = last_error {
1037 return Err(KMError::Protocol(err));
1038 }
1039
1040 let info = self.read_device_info_blocks().await;
1041
1042 let data = self
1045 .read_memory_block(HARDWARE_ID_ADDRESS, HARDWARE_ID_SIZE as u32)
1046 .await
1047 .map_err(|error| {
1048 KMError::Protocol(format!(
1049 "Failed to read HardwareID - required for authentication: {error}"
1050 ))
1051 })?;
1052 let hardware_id_bytes: [u8; HARDWARE_ID_SIZE] = data.as_slice().try_into().map_err(|_| {
1053 KMError::Protocol(format!(
1054 "HardwareID read returned {} bytes, expected {HARDWARE_ID_SIZE}",
1055 data.len()
1056 ))
1057 })?;
1058 let hardware_id = HardwareId::from_bytes(hardware_id_bytes);
1059
1060 let auth = self.perform_streaming_auth(AuthCredential::from(&hardware_id)).await?;
1062
1063 let state = DeviceState {
1065 info,
1066 hardware_id,
1067 auth_level: auth.auth_level,
1068 adcqueue_enabled: auth.adcqueue_enabled(),
1069 };
1070
1071 info!("Device initialized: {} (auth_level={})", state.model(), auth.auth_level);
1072
1073 self.mode = ConnectionMode::Full(state);
1074
1075 Ok(())
1076 }
1077
1078 async fn perform_streaming_auth(&mut self, credential: AuthCredential) -> Result<StreamingAuthResult, KMError> {
1079 self.send_tracked(Packet::StreamingAuth { credential }).await?;
1080 let response = self
1081 .receive_matching_raw(|bytes| response_type_matches(bytes, PacketType::StreamingAuth))
1084 .await?;
1085
1086 match Self::parse_response(response, self.graph_sample_rate)? {
1087 Packet::StreamingAuthResponse(result) => Ok(result),
1088 other => Err(KMError::Protocol(format!(
1089 "Expected StreamingAuthResponse, got {other:?}"
1090 ))),
1091 }
1092 }
1093
1094 pub fn mode(&self) -> &ConnectionMode {
1099 &self.mode
1100 }
1101
1102 pub fn is_full_mode(&self) -> bool {
1104 matches!(self.mode, ConnectionMode::Full(_))
1105 }
1106
1107 pub fn is_basic_mode(&self) -> bool {
1109 matches!(self.mode, ConnectionMode::Basic)
1110 }
1111
1112 pub fn state(&self) -> Option<&DeviceState> {
1117 match &self.mode {
1118 ConnectionMode::Full(state) => Some(state),
1119 ConnectionMode::Basic => None,
1120 }
1121 }
1122
1123 pub fn is_authenticated(&self) -> bool {
1127 self.state().map(|s| s.is_authenticated()).unwrap_or(false)
1128 }
1129
1130 pub fn adcqueue_enabled(&self) -> bool {
1134 self.state().map(|s| s.adcqueue_enabled).unwrap_or(false)
1135 }
1136
1137 pub async fn authenticate_calibration(&mut self) -> Result<StreamingAuthResult, KMError> {
1144 if !self.is_full_mode() {
1145 return Err(KMError::Protocol(
1146 "Calibration authentication requires the vendor interface".to_string(),
1147 ));
1148 }
1149
1150 let preferred = self
1151 .read_memory_block(PREFERRED_CALIBRATION_ADDRESS, STREAMING_AUTH_CREDENTIAL_SIZE as u32)
1152 .await?;
1153 let selected = if calibration_record_is_erased(&preferred) {
1154 self.read_memory_block(CALIBRATION_ADDRESS, STREAMING_AUTH_CREDENTIAL_SIZE as u32)
1155 .await?
1156 } else {
1157 preferred
1158 };
1159 let credential = AuthCredential::from_bytes(
1160 selected
1161 .as_slice()
1162 .try_into()
1163 .map_err(|_| KMError::Protocol("Calibration credential has the wrong length".to_string()))?,
1164 );
1165 let result = self.perform_streaming_auth(credential).await?;
1166
1167 if result.auth_level != 2 {
1168 return Err(KMError::Protocol(format!(
1169 "Calibration authentication returned level {} instead of 2",
1170 result.auth_level
1171 )));
1172 }
1173
1174 if let ConnectionMode::Full(state) = &mut self.mode {
1175 state.auth_level = result.auth_level;
1176 state.adcqueue_enabled = result.adcqueue_enabled();
1177 }
1178
1179 Ok(result)
1180 }
1181
1182 pub async fn read_memory_block(&mut self, address: u32, size: u32) -> Result<Vec<u8>, KMError> {
1195 if self.graph_sample_rate.is_some() {
1196 return Err(KMError::Protocol(
1197 "Device memory cannot be read while AdcQueue streaming is active; \
1198 call stop_graph_mode() first"
1199 .to_string(),
1200 ));
1201 }
1202
1203 let id = self.send_tracked(Packet::MemoryRead { address, size }).await?;
1204 let confirmation = self
1205 .receive_matching_raw(|bytes| memory_confirmation_matches(bytes, id))
1206 .await?;
1207
1208 let confirmation = RawPacket::try_from(Bytes::from(confirmation))?;
1209 match confirmation.packet_type() {
1210 PacketType::MemoryRead => {
1211 validate_memory_read_confirmation(&confirmation, id, address, size)?;
1212 self.receive_memory_read_data_exact(size).await
1213 }
1214 PacketType::Rejected => Err(KMError::Protocol(format!("MemoryRead at 0x{address:08X} was rejected"))),
1215 PacketType::NotReadable => Err(KMError::Protocol(format!(
1216 "Memory address 0x{address:08X} is not readable"
1217 ))),
1218 other => Err(KMError::Protocol(format!(
1219 "Expected MemoryRead confirmation, got {other:?}"
1220 ))),
1221 }
1222 }
1223
1224 pub async fn get_device_info(&mut self) -> Result<crate::auth::DeviceInfo, KMError> {
1245 Ok(self.read_device_info_blocks().await)
1246 }
1247
1248 async fn read_device_info_blocks(&mut self) -> crate::auth::DeviceInfo {
1253 use crate::auth::{
1254 CALIBRATION_ADDRESS, DEVICE_INFO_ADDRESS, DeviceInfo, FIRMWARE_INFO_ADDRESS, INFO_BLOCK_SIZE,
1255 };
1256 use tracing::warn;
1257
1258 let mut info = DeviceInfo::default();
1259
1260 match self
1261 .read_memory_block(DEVICE_INFO_ADDRESS, INFO_BLOCK_SIZE as u32)
1262 .await
1263 {
1264 Ok(data) => info.parse_device_info(&data),
1265 Err(e) => warn!("Failed to read DeviceInfo at 0x{:08X}: {}", DEVICE_INFO_ADDRESS, e),
1266 }
1267
1268 match self
1269 .read_memory_block(FIRMWARE_INFO_ADDRESS, INFO_BLOCK_SIZE as u32)
1270 .await
1271 {
1272 Ok(data) => info.parse_firmware_info(&data),
1273 Err(e) => warn!("Failed to read FirmwareInfo at 0x{:08X}: {}", FIRMWARE_INFO_ADDRESS, e),
1274 }
1275
1276 match self
1277 .read_memory_block(CALIBRATION_ADDRESS, INFO_BLOCK_SIZE as u32)
1278 .await
1279 {
1280 Ok(data) => info.parse_calibration(&data),
1281 Err(e) => warn!("Failed to read CalibrationData at 0x{:08X}: {}", CALIBRATION_ADDRESS, e),
1282 }
1283
1284 info
1285 }
1286
1287 pub async fn start_graph_mode(&mut self, rate: GraphSampleRate) -> Result<(), KMError> {
1320 ensure_adcqueue_available(&self.mode)?;
1321
1322 let id = self.send_tracked(Packet::StartGraph { rate }).await?;
1323 self.expect_accept(id, "StartGraph").await?;
1324 self.graph_sample_rate = Some(rate);
1325 Ok(())
1326 }
1327
1328 pub async fn stop_graph_mode(&mut self) -> Result<(), KMError> {
1334 if self.is_basic_mode() {
1335 return Err(KMError::Protocol(
1336 "AdcQueue streaming requires Full mode (vendor interface)".to_string(),
1337 ));
1338 }
1339
1340 let id = self.send_tracked(Packet::StopGraph).await?;
1341 self.expect_accept(id, "StopGraph").await?;
1342 self.graph_sample_rate = None;
1343 Ok(())
1344 }
1345}
1346
1347#[cfg(test)]
1348mod tests {
1349 use super::*;
1350
1351 #[test]
1352 fn calibration_auth_falls_back_only_for_an_erased_preferred_record() {
1353 assert!(calibration_record_is_erased(&[0xff; STREAMING_AUTH_CREDENTIAL_SIZE]));
1354 assert!(!calibration_record_is_erased(b"007965 CDFDD"));
1355 assert!(!calibration_record_is_erased(&[0xff, 0xff, 0xff]));
1356 }
1357
1358 #[test]
1359 fn response_matching_uses_the_typed_header_parser() {
1360 assert!(response_matches(&[0xc1, 7, 0, 0], 7, PacketType::PutData));
1361 }
1362
1363 #[test]
1364 fn graph_mode_requires_streaming_auth_permission() {
1365 assert!(ensure_adcqueue_available(&ConnectionMode::Basic).is_err());
1366
1367 let state = DeviceState {
1368 info: DeviceInfo::default(),
1369 hardware_id: HardwareId::from_bytes([0; 12]),
1370 auth_level: 0,
1371 adcqueue_enabled: false,
1372 };
1373 assert!(ensure_adcqueue_available(&ConnectionMode::Full(state.clone())).is_err());
1374
1375 let enabled = DeviceState {
1376 auth_level: 1,
1377 adcqueue_enabled: true,
1378 ..state
1379 };
1380 assert!(ensure_adcqueue_available(&ConnectionMode::Full(enabled)).is_ok());
1381 }
1382
1383 #[test]
1384 fn response_matching_requires_type_and_transaction_id() {
1385 assert!(!response_matches(&[0xc1, 8, 0, 0], 7, PacketType::PutData));
1386 assert!(!response_matches(&[0xc0, 7, 0, 0], 7, PacketType::PutData));
1387 assert!(!response_matches(&[0xc1], 7, PacketType::PutData));
1388 }
1389
1390 #[test]
1391 fn response_type_matching_supports_fixed_id_protocol_exceptions() {
1392 assert!(response_type_matches(&[0x4c, 0, 3, 2], PacketType::StreamingAuth));
1393 }
1394
1395 #[test]
1396 fn memory_response_size_rounds_to_complete_aes_blocks() {
1397 assert_eq!(memory_response_size(1), 16);
1398 assert_eq!(memory_response_size(12), 16);
1399 assert_eq!(memory_response_size(16), 16);
1400 assert_eq!(memory_response_size(17), 32);
1401 assert_eq!(memory_response_size(64), 64);
1402 }
1403
1404 #[test]
1405 fn decrypted_memory_response_is_trimmed_to_the_requested_size() {
1406 let ciphertext = crate::auth::build_memory_read_payload(0x4001_0450, 12);
1407 let decrypted = decrypt_memory_response(&ciphertext[..16], 12).unwrap();
1408
1409 assert_eq!(decrypted.len(), 12);
1410 assert_eq!(&decrypted[..4], &0x4001_0450_u32.to_le_bytes());
1411 assert_eq!(&decrypted[4..8], &12_u32.to_le_bytes());
1412 assert_eq!(&decrypted[8..12], &u32::MAX.to_le_bytes());
1413 }
1414
1415 #[test]
1416 fn decrypted_memory_response_rejects_an_incorrect_ciphertext_size() {
1417 assert!(decrypt_memory_response(&[0; 32], 12).is_err());
1418 }
1419
1420 #[test]
1421 fn memory_chunks_accumulate_across_usb_transfers() {
1422 let mut response = Vec::new();
1425 assert!(!append_memory_chunk(&mut response, &vec![0x11; 2544], 8336).unwrap());
1426 assert!(!append_memory_chunk(&mut response, &vec![0x22; 2544], 8336).unwrap());
1427 assert!(!append_memory_chunk(&mut response, &vec![0x33; 2544], 8336).unwrap());
1428 assert!(append_memory_chunk(&mut response, &vec![0x44; 704], 8336).unwrap());
1429 assert_eq!(response.len(), 8336);
1430 }
1431
1432 #[test]
1433 fn memory_chunks_reject_empty_and_oversized_transfers() {
1434 let mut response = vec![0; 15];
1435 assert!(append_memory_chunk(&mut response, &[], 16).is_err());
1436 assert!(append_memory_chunk(&mut response, &[1, 2], 16).is_err());
1437 }
1438
1439 #[test]
1440 fn validates_recorded_memory_read_confirmation() {
1441 let bytes = hex::decode("c40201012004000040000000ffffffff1b8c1b24").unwrap();
1443 let packet = RawPacket::try_from(Bytes::from(bytes)).unwrap();
1444
1445 validate_memory_read_confirmation(&packet, 2, 0x420, 64).unwrap();
1446 }
1447
1448 #[test]
1449 fn rejects_corrupted_memory_read_confirmation() {
1450 let mut bytes = hex::decode("c40201012004000040000000ffffffff1b8c1b24").unwrap();
1451 bytes[8] ^= 1;
1452 let packet = RawPacket::try_from(Bytes::from(bytes)).unwrap();
1453
1454 assert!(validate_memory_read_confirmation(&packet, 2, 0x420, 64).is_err());
1455 }
1456}