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, 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
160fn parse_framed_response(bytes: &[u8]) -> Option<RawPacket> {
161 RawPacket::try_from(Bytes::copy_from_slice(bytes)).ok()
162}
163
164fn response_matches(bytes: &[u8], id: u8, packet_type: PacketType) -> bool {
165 parse_framed_response(bytes).is_some_and(|packet| packet.id() == id && packet.packet_type() == packet_type)
166}
167
168fn response_type_matches(bytes: &[u8], packet_type: PacketType) -> bool {
169 parse_framed_response(bytes).is_some_and(|packet| packet.packet_type() == packet_type)
170}
171
172fn control_response_matches(bytes: &[u8], id: u8) -> bool {
173 matches!(
174 parse_framed_response(bytes),
175 Some(RawPacket::Ctrl { header, .. }) if header.id() == id
176 )
177}
178
179fn memory_confirmation_matches(bytes: &[u8], id: u8) -> bool {
180 parse_framed_response(bytes).is_some_and(|packet| {
181 packet.id() == id
182 && matches!(
183 packet.packet_type(),
184 PacketType::MemoryRead | PacketType::Rejected | PacketType::NotReadable
185 )
186 })
187}
188
189fn memory_response_size(requested_size: u32) -> usize {
190 (requested_size as usize).div_ceil(AES_BLOCK_SIZE) * AES_BLOCK_SIZE
191}
192
193fn append_memory_chunk(buffer: &mut Vec<u8>, chunk: &[u8], expected_size: usize) -> Result<bool, KMError> {
194 if chunk.is_empty() {
195 return Err(KMError::Protocol(
196 "Received an empty encrypted memory chunk".to_string(),
197 ));
198 }
199
200 let remaining = expected_size.saturating_sub(buffer.len());
201 if chunk.len() > remaining {
202 return Err(KMError::Protocol(format!(
203 "Encrypted memory response exceeded expected size: expected {expected_size} bytes, got at least {}",
204 buffer.len() + chunk.len()
205 )));
206 }
207
208 buffer.extend_from_slice(chunk);
209 Ok(buffer.len() == expected_size)
210}
211
212fn decrypt_memory_response(encrypted: &[u8], requested_size: u32) -> Result<Vec<u8>, KMError> {
213 let expected_size = memory_response_size(requested_size);
214 if encrypted.len() != expected_size {
215 return Err(KMError::Protocol(format!(
216 "Encrypted memory response has the wrong size: expected {expected_size} bytes, got {}",
217 encrypted.len()
218 )));
219 }
220
221 let mut decrypted = crate::auth::aes_ecb_decrypt_blocks(encrypted, crate::auth::MEMORY_READ_KEY)?;
222 decrypted.truncate(requested_size as usize);
223 Ok(decrypted)
224}
225
226fn validate_memory_read_confirmation(
227 packet: &RawPacket,
228 expected_id: u8,
229 expected_address: u32,
230 expected_size: u32,
231) -> Result<(), KMError> {
232 let RawPacket::SimpleData { header, payload } = packet else {
233 return Err(KMError::Protocol(
234 "MemoryRead confirmation has an unexpected packet layout".to_string(),
235 ));
236 };
237
238 if header.id() != expected_id || PacketType::from(header.packet_type()) != PacketType::MemoryRead {
239 return Err(KMError::Protocol(
240 "MemoryRead confirmation does not match the request".to_string(),
241 ));
242 }
243 if payload.len() != 16 {
244 return Err(KMError::Protocol(format!(
245 "MemoryRead confirmation must contain 16 bytes, got {}",
246 payload.len()
247 )));
248 }
249
250 let address = u32::from_le_bytes(payload[0..4].try_into()?);
251 let size = u32::from_le_bytes(payload[4..8].try_into()?);
252 let magic = u32::from_le_bytes(payload[8..12].try_into()?);
253 let crc = u32::from_le_bytes(payload[12..16].try_into()?);
254 let expected_crc = crc32fast::hash(&payload[..12]);
255
256 if address != expected_address || size != expected_size {
257 return Err(KMError::Protocol(format!(
258 "MemoryRead confirmation echoed address 0x{address:08X} and size {size}, expected 0x{expected_address:08X} and {expected_size}"
259 )));
260 }
261 if magic != u32::MAX {
262 return Err(KMError::Protocol(format!(
263 "MemoryRead confirmation has invalid magic 0x{magic:08X}"
264 )));
265 }
266 if crc != expected_crc {
267 return Err(KMError::Protocol(format!(
268 "MemoryRead confirmation CRC mismatch: expected 0x{expected_crc:08X}, got 0x{crc:08X}"
269 )));
270 }
271
272 Ok(())
273}
274
275#[derive(Debug, Clone, Copy, PartialEq, Eq)]
277pub enum TransferType {
278 Bulk,
279 Interrupt,
280}
281
282#[derive(Debug, Clone)]
287pub enum ConnectionMode {
288 Basic,
290 Full(DeviceState),
292}
293
294fn ensure_adcqueue_available(mode: &ConnectionMode) -> Result<(), KMError> {
295 match mode {
296 ConnectionMode::Basic => Err(KMError::Protocol(
297 "AdcQueue streaming requires Full mode (vendor interface)".to_string(),
298 )),
299 ConnectionMode::Full(state) if !state.adcqueue_enabled => Err(KMError::Protocol(
300 "StreamingAuth did not enable AdcQueue streaming".to_string(),
301 )),
302 ConnectionMode::Full(_) => Ok(()),
303 }
304}
305
306#[derive(Debug, Clone, Copy)]
330pub struct DeviceConfig {
331 interface: u8,
333 endpoint_out: u8,
335 endpoint_in: u8,
337 transfer_type: TransferType,
339 skip_reset: bool,
341}
342
343impl DeviceConfig {
344 pub fn vendor() -> Self {
358 Self {
359 interface: INTERFACE_VENDOR,
360 endpoint_out: ENDPOINT_OUT_VENDOR,
361 endpoint_in: ENDPOINT_IN_VENDOR,
362 transfer_type: TransferType::Bulk,
363 skip_reset: false,
364 }
365 }
366
367 pub fn hid() -> Self {
383 Self {
384 interface: INTERFACE_HID,
385 endpoint_out: ENDPOINT_OUT_HID,
386 endpoint_in: ENDPOINT_IN_HID,
387 transfer_type: TransferType::Interrupt,
388 skip_reset: false,
389 }
390 }
391
392 pub fn skip_reset(mut self) -> Self {
397 self.skip_reset = true;
398 self
399 }
400
401 pub fn is_vendor(&self) -> bool {
403 self.interface == INTERFACE_VENDOR
404 }
405
406 pub fn is_hid(&self) -> bool {
408 self.interface == INTERFACE_HID
409 }
410}
411
412enum EndpointReaderType {
414 Bulk(EndpointRead<Bulk>),
415 Interrupt(EndpointRead<Interrupt>),
416}
417
418enum EndpointWriterType {
420 Bulk(EndpointWrite<Bulk>),
421 Interrupt(EndpointWrite<Interrupt>),
422}
423
424pub struct KM003C {
425 #[allow(dead_code)]
427 interface: Interface,
428 transaction_id: u8,
429 reader: EndpointReaderType,
430 writer: EndpointWriterType,
431 pending_responses: VecDeque<Vec<u8>>,
432 graph_sample_rate: Option<GraphSampleRate>,
434 mode: ConnectionMode,
436}
437
438impl KM003C {
439 pub async fn new(config: DeviceConfig) -> Result<Self, KMError> {
466 let mut device = Self::connect(config).await?;
467
468 if config.is_vendor() {
469 device.run_init().await?;
471 }
472 Ok(device)
475 }
476
477 async fn connect(config: DeviceConfig) -> Result<Self, KMError> {
479 info!("Searching for POWER-Z KM003C...");
480 let device_info = nusb::list_devices()
481 .await?
482 .find(|d| d.vendor_id() == VID && d.product_id() == PID)
483 .ok_or(KMError::DeviceNotFound)?;
484
485 info!(
486 "Found device on bus {} addr {}",
487 device_info.bus_id(),
488 device_info.device_address()
489 );
490
491 let mut device = device_info.open().await?;
492
493 if !config.skip_reset {
495 info!("Resetting device...");
496 device.reset().await?;
497 tokio::time::sleep(Duration::from_millis(1500)).await;
500 let device_info = nusb::list_devices()
502 .await?
503 .find(|d| d.vendor_id() == VID && d.product_id() == PID)
504 .ok_or(KMError::DeviceNotFound)?;
505 device = device_info.open().await?;
506 } else {
507 debug!("Skipping USB reset (skip_reset=true)");
508 }
509
510 for interface_num in 0..4 {
517 if let Err(e) = device.detach_kernel_driver(interface_num) {
518 trace!("Could not detach interface {}: {}", interface_num, e);
520 } else {
521 debug!("Detached kernel driver from interface {}", interface_num);
522 }
523 }
524
525 let interface = device.claim_interface(config.interface).await?;
527 info!("Interface {} claimed successfully", config.interface);
528
529 let (reader, writer) = match config.transfer_type {
534 TransferType::Bulk => {
535 let ep_in = interface.endpoint::<Bulk, _>(config.endpoint_in)?;
536 let ep_out = interface.endpoint::<Bulk, _>(config.endpoint_out)?;
537 (
538 EndpointReaderType::Bulk(ep_in.reader(BULK_TRANSFER_SIZE).with_num_transfers(4)),
539 EndpointWriterType::Bulk(ep_out.writer(INTERRUPT_TRANSFER_SIZE).with_num_transfers(4)),
540 )
541 }
542 TransferType::Interrupt => {
543 let ep_in = interface.endpoint::<Interrupt, _>(config.endpoint_in)?;
544 let ep_out = interface.endpoint::<Interrupt, _>(config.endpoint_out)?;
545 (
546 EndpointReaderType::Interrupt(ep_in.reader(INTERRUPT_TRANSFER_SIZE).with_num_transfers(4)),
547 EndpointWriterType::Interrupt(ep_out.writer(INTERRUPT_TRANSFER_SIZE).with_num_transfers(4)),
548 )
549 }
550 };
551
552 let km003c = Self {
553 interface,
554 transaction_id: 0,
555 reader,
556 writer,
557 pending_responses: VecDeque::new(),
558 graph_sample_rate: None,
559 mode: ConnectionMode::Basic,
560 };
561
562 info!("USB connection established");
563 Ok(km003c)
564 }
565
566 fn next_transaction_id(&mut self) -> u8 {
568 let id = self.transaction_id;
569 self.transaction_id = self.transaction_id.wrapping_add(1);
570 id
571 }
572
573 pub fn next_tid(&mut self) -> u8 {
575 self.next_transaction_id()
576 }
577
578 pub fn set_transaction_id(&mut self, id: u8) {
580 self.transaction_id = id;
581 }
582
583 pub async fn send(&mut self, packet: Packet) -> Result<(), KMError> {
588 self.send_tracked(packet).await.map(|_| ())
589 }
590
591 async fn send_tracked(&mut self, packet: Packet) -> Result<u8, KMError> {
593 use crate::auth;
594
595 let id = self.next_transaction_id();
596
597 match &packet {
600 Packet::MemoryRead { address, size } => {
601 let raw = auth::build_memory_read_packet(*address, *size, id);
602 self.send_raw(&raw).await?;
603 return Ok(id);
604 }
605 Packet::StreamingAuth { credential } => {
606 let raw = auth::build_streaming_auth_packet(credential, id);
607 self.send_raw(&raw).await?;
608 return Ok(id);
609 }
610 _ => {}
611 }
612
613 let raw_packet = packet.to_raw_packet(id)?;
614 self.send_raw_packet(raw_packet).await?;
615 Ok(id)
616 }
617
618 async fn send_raw_packet(&mut self, packet: RawPacket) -> Result<(), KMError> {
620 let (reserved_flag, has_logical_packets) = match &packet {
621 RawPacket::Ctrl { header, .. } => (header.reserved_flag(), false),
622 RawPacket::SimpleData { header, .. } => (header.reserved_flag(), false),
623 RawPacket::Data {
624 header,
625 logical_packets,
626 } => (header.reserved_flag(), !logical_packets.is_empty()),
627 };
628
629 debug!(
630 "Sending packet: type={:?}, reserved_flag={}, has_logical_packets={}, id={}",
631 packet.packet_type(),
632 reserved_flag,
633 has_logical_packets,
634 packet.id(),
635 );
636
637 let message_bytes = Bytes::from(packet);
638 let message = message_bytes.to_vec();
639 trace!("TX [{} bytes]: {:02x?}", message.len(), message);
640
641 match &mut self.writer {
643 EndpointWriterType::Bulk(writer) => {
644 timeout(DEFAULT_TIMEOUT, writer.write_all(&message)).await??;
645 timeout(DEFAULT_TIMEOUT, writer.flush_end_async()).await??;
646 }
647 EndpointWriterType::Interrupt(writer) => {
648 timeout(DEFAULT_TIMEOUT, writer.write_all(&message)).await??;
649 timeout(DEFAULT_TIMEOUT, writer.flush_end_async()).await??;
650 }
651 }
652
653 debug!("Sent successfully");
654 Ok(())
655 }
656
657 pub async fn send_raw(&mut self, data: &[u8]) -> Result<(), KMError> {
659 trace!("TX [{} bytes]: {:02x?}", data.len(), data);
660 match &mut self.writer {
661 EndpointWriterType::Bulk(writer) => {
662 timeout(DEFAULT_TIMEOUT, writer.write_all(data)).await??;
663 timeout(DEFAULT_TIMEOUT, writer.flush_end_async()).await??;
664 }
665 EndpointWriterType::Interrupt(writer) => {
666 timeout(DEFAULT_TIMEOUT, writer.write_all(data)).await??;
667 timeout(DEFAULT_TIMEOUT, writer.flush_end_async()).await??;
668 }
669 }
670 Ok(())
671 }
672
673 async fn read_raw_from_usb(&mut self) -> Result<Vec<u8>, KMError> {
675 let mut buffer = Vec::new();
676 let bytes_read = match &mut self.reader {
677 EndpointReaderType::Bulk(reader) => {
678 let mut message = reader.until_short_packet();
682 let bytes_read = timeout(DEFAULT_TIMEOUT, message.read_to_end(&mut buffer)).await??;
683 message
684 .consume_end()
685 .map_err(|_| KMError::Protocol("Bulk response ended without a short packet".to_string()))?;
686 bytes_read
687 }
688 EndpointReaderType::Interrupt(reader) => {
689 buffer.resize(INTERRUPT_TRANSFER_SIZE, 0);
690 let bytes_read = timeout(DEFAULT_TIMEOUT, reader.read(&mut buffer)).await??;
691 buffer.truncate(bytes_read);
692 bytes_read
693 }
694 };
695 trace!("RX [{} bytes]: {:02x?}", bytes_read, buffer);
696 Ok(buffer)
697 }
698
699 fn queue_pending_response(&mut self, response: Vec<u8>) {
700 if self.pending_responses.len() == MAX_PENDING_RESPONSES {
701 warn!("Dropping oldest unmatched response because the pending queue is full");
702 self.pending_responses.pop_front();
703 }
704 self.pending_responses.push_back(response);
705 }
706
707 async fn receive_matching_raw<F>(&mut self, mut predicate: F) -> Result<Vec<u8>, KMError>
708 where
709 F: FnMut(&[u8]) -> bool,
710 {
711 if let Some(index) = self.pending_responses.iter().position(|response| predicate(response)) {
712 return Ok(self
713 .pending_responses
714 .remove(index)
715 .expect("pending response index is valid"));
716 }
717
718 let deadline = tokio::time::Instant::now() + DEFAULT_TIMEOUT;
719 loop {
720 let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
721 if remaining.is_zero() {
722 return Err(KMError::Protocol(
723 "Timed out waiting for a correlated response".to_string(),
724 ));
725 }
726
727 let response = timeout(remaining, self.read_raw_from_usb()).await??;
728 if predicate(&response) {
729 return Ok(response);
730 }
731
732 if let Some(packet) = parse_framed_response(&response) {
733 debug!(
734 "Queueing unmatched response: type={:?}, id={}, len={}",
735 packet.packet_type(),
736 packet.id(),
737 response.len()
738 );
739 } else {
740 debug!("Queueing unframed response: len={}", response.len());
741 }
742 self.queue_pending_response(response);
743 }
744 }
745
746 pub async fn receive_raw(&mut self) -> Result<Vec<u8>, KMError> {
751 if let Some(response) = self.pending_responses.pop_front() {
752 return Ok(response);
753 }
754 self.read_raw_from_usb().await
755 }
756
757 fn parse_response(raw_bytes: Vec<u8>, graph_rate: Option<GraphSampleRate>) -> Result<Packet, KMError> {
758 if raw_bytes.is_empty() {
759 return Err(KMError::Protocol("Received 0 bytes".to_string()));
760 }
761
762 let raw_packet = RawPacket::try_from(Bytes::from(raw_bytes))?;
763 if let Some(rate) = graph_rate {
764 Packet::from_raw_with_graph_rate(raw_packet, rate)
765 } else {
766 Packet::try_from(raw_packet)
767 }
768 }
769
770 pub async fn receive(&mut self) -> Result<Packet, KMError> {
775 let raw_bytes = self.receive_raw().await?;
776 Self::parse_response(raw_bytes, self.graph_sample_rate)
777 }
778
779 async fn receive_control_response(&mut self, id: u8) -> Result<Packet, KMError> {
780 let raw_bytes = self
781 .receive_matching_raw(|bytes| control_response_matches(bytes, id))
782 .await?;
783 Self::parse_response(raw_bytes, self.graph_sample_rate)
784 }
785
786 async fn expect_accept(&mut self, id: u8, command: &str) -> Result<(), KMError> {
787 match self.receive_control_response(id).await? {
788 Packet::Accept { .. } => Ok(()),
789 Packet::Reject { .. } => Err(KMError::Protocol(format!("{command} was rejected by the device"))),
790 Packet::NotReadable { .. } => Err(KMError::Protocol(format!("{command} is not readable"))),
791 other => Err(KMError::Protocol(format!(
792 "Expected Accept for {command}, got {other:?}"
793 ))),
794 }
795 }
796
797 async fn receive_memory_read_data_exact(&mut self, requested_size: u32) -> Result<Vec<u8>, KMError> {
798 let expected_size = memory_response_size(requested_size);
799 if expected_size == 0 {
800 return Ok(Vec::new());
801 }
802
803 let mut encrypted = Vec::with_capacity(expected_size);
804 while encrypted.len() < expected_size {
805 let chunk = self.read_raw_from_usb().await?;
806 if append_memory_chunk(&mut encrypted, &chunk, expected_size)? {
807 break;
808 }
809 }
810
811 decrypt_memory_response(&encrypted, requested_size)
812 }
813
814 pub async fn request_data(&mut self, mask: AttributeSet) -> Result<Packet, KMError> {
816 let id = self
817 .send_tracked(Packet::GetData {
818 attribute_mask: mask.raw(),
819 })
820 .await?;
821 let raw_bytes = self
822 .receive_matching_raw(|bytes| response_matches(bytes, id, PacketType::PutData))
823 .await?;
824 let raw_packet = RawPacket::try_from(Bytes::from(raw_bytes))?;
825 raw_packet.validate_correlation(mask.raw())?;
826 if let Some(rate) = self.graph_sample_rate {
827 Packet::from_raw_with_graph_rate(raw_packet, rate)
828 } else {
829 Packet::try_from(raw_packet)
830 }
831 }
832
833 pub async fn request_adc_data(&mut self) -> Result<AdcDataSimple, KMError> {
835 let packet = self.request_data(AttributeSet::single(Attribute::Adc)).await?;
836 packet
837 .get_adc()
838 .cloned()
839 .ok_or_else(|| KMError::Protocol("No ADC data in response".to_string()))
840 }
841
842 pub async fn request_settings(&mut self) -> Result<Settings, KMError> {
844 let packet = self.request_data(AttributeSet::single(Attribute::Settings)).await?;
845 packet
846 .get_settings()
847 .cloned()
848 .ok_or_else(|| KMError::Protocol("No Settings data in response".to_string()))
849 }
850
851 pub async fn request_pd_trace(&mut self) -> Result<crate::pd_trace::PdTrace, KMError> {
853 let packet = self.request_data(AttributeSet::single(Attribute::PdTrace)).await?;
854 packet
855 .get_pd_trace()
856 .cloned()
857 .ok_or_else(|| KMError::Protocol("No PD trace data in response".to_string()))
858 }
859
860 pub async fn request_log_metadata(&mut self) -> Result<Vec<LogMetadata>, KMError> {
864 let packet = self.request_data(AttributeSet::single(Attribute::LogMetadata)).await?;
865 match packet.get_log_metadata() {
866 Some(LogMetadataResponse::Empty) => Ok(Vec::new()),
867 Some(LogMetadataResponse::Available(metadata)) => Ok(metadata.clone()),
868 None => Err(KMError::Protocol("No LogMetadata data in response".to_string())),
869 }
870 }
871
872 pub async fn download_offline_log(&mut self, metadata: LogMetadata) -> Result<OfflineLog, KMError> {
874 let data = self
875 .read_memory_block(metadata.data_address()?, metadata.data_size())
876 .await?;
877 OfflineLog::from_bytes(metadata, &data)
878 }
879
880 pub async fn request_pd_data(&mut self) -> Result<Packet, KMError> {
888 self.request_data(AttributeSet::single(Attribute::PdPacket)).await
889 }
890
891 pub async fn request_adc_with_pd(&mut self) -> Result<Packet, KMError> {
896 let mask = AttributeSet::single(Attribute::Adc).with(Attribute::PdPacket);
897 self.request_data(mask).await
898 }
899
900 pub fn extract_pd_status(packet: &Packet) -> Option<&PdStatus> {
905 packet.get_pd_status()
906 }
907
908 pub fn extract_pd_events(packet: &Packet) -> Option<&PdEventStream> {
913 packet.get_pd_events()
914 }
915
916 pub async fn enable_pd_monitor(&mut self) -> Result<(), KMError> {
924 let id = self.send_tracked(Packet::EnablePdMonitor).await?;
925 self.expect_accept(id, "EnablePdMonitor").await
926 }
927
928 pub async fn disable_pd_monitor(&mut self) -> Result<(), KMError> {
934 let id = self.send_tracked(Packet::DisablePdMonitor).await?;
935 self.expect_accept(id, "DisablePdMonitor").await
936 }
937
938 async fn run_init(&mut self) -> Result<(), KMError> {
950 use crate::auth::{HARDWARE_ID_ADDRESS, HARDWARE_ID_SIZE};
951
952 const MAX_CONNECT_RETRIES: u8 = 3;
954 let mut last_error = None;
955 for attempt in 1..=MAX_CONNECT_RETRIES {
956 match self.send_tracked(Packet::Connect).await {
957 Err(err) => last_error = Some(format!("Connect send failed: {err:?}")),
958 Ok(id) => match self.expect_accept(id, "Connect").await {
959 Ok(()) => {
960 last_error = None;
961 break;
962 }
963 Err(err) => last_error = Some(format!("Connect receive failed: {err:?}")),
964 },
965 }
966 if attempt < MAX_CONNECT_RETRIES {
967 debug!("Connect attempt {} failed, retrying...", attempt);
968 tokio::time::sleep(Duration::from_millis(100)).await;
969 }
970 }
971 if let Some(err) = last_error {
972 return Err(KMError::Protocol(err));
973 }
974
975 let info = self.read_device_info_blocks().await;
976 let mut hardware_id_bytes = [0u8; HARDWARE_ID_SIZE];
977
978 let hardware_id = if let Ok(data) = self
980 .read_memory_block(HARDWARE_ID_ADDRESS, HARDWARE_ID_SIZE as u32)
981 .await
982 {
983 if data.len() >= HARDWARE_ID_SIZE {
984 hardware_id_bytes.copy_from_slice(&data[..HARDWARE_ID_SIZE]);
985 }
986 HardwareId::from_bytes(hardware_id_bytes)
987 } else {
988 return Err(KMError::Protocol(
989 "Failed to read HardwareID - required for authentication".to_string(),
990 ));
991 };
992
993 let auth = self.perform_streaming_auth(AuthCredential::from(&hardware_id)).await?;
995
996 let state = DeviceState {
998 info,
999 hardware_id,
1000 auth_level: auth.auth_level,
1001 adcqueue_enabled: auth.adcqueue_enabled(),
1002 };
1003
1004 info!("Device initialized: {} (auth_level={})", state.model(), auth.auth_level);
1005
1006 self.mode = ConnectionMode::Full(state);
1007
1008 Ok(())
1009 }
1010
1011 async fn perform_streaming_auth(&mut self, credential: AuthCredential) -> Result<StreamingAuthResult, KMError> {
1012 self.send_tracked(Packet::StreamingAuth { credential }).await?;
1013 let response = self
1014 .receive_matching_raw(|bytes| response_type_matches(bytes, PacketType::StreamingAuth))
1017 .await?;
1018
1019 match Self::parse_response(response, self.graph_sample_rate)? {
1020 Packet::StreamingAuthResponse(result) => Ok(result),
1021 other => Err(KMError::Protocol(format!(
1022 "Expected StreamingAuthResponse, got {other:?}"
1023 ))),
1024 }
1025 }
1026
1027 pub fn mode(&self) -> &ConnectionMode {
1032 &self.mode
1033 }
1034
1035 pub fn is_full_mode(&self) -> bool {
1037 matches!(self.mode, ConnectionMode::Full(_))
1038 }
1039
1040 pub fn is_basic_mode(&self) -> bool {
1042 matches!(self.mode, ConnectionMode::Basic)
1043 }
1044
1045 pub fn state(&self) -> Option<&DeviceState> {
1050 match &self.mode {
1051 ConnectionMode::Full(state) => Some(state),
1052 ConnectionMode::Basic => None,
1053 }
1054 }
1055
1056 pub fn is_authenticated(&self) -> bool {
1060 self.state().map(|s| s.is_authenticated()).unwrap_or(false)
1061 }
1062
1063 pub fn adcqueue_enabled(&self) -> bool {
1067 self.state().map(|s| s.adcqueue_enabled).unwrap_or(false)
1068 }
1069
1070 pub async fn authenticate_calibration(&mut self) -> Result<StreamingAuthResult, KMError> {
1077 if !self.is_full_mode() {
1078 return Err(KMError::Protocol(
1079 "Calibration authentication requires the vendor interface".to_string(),
1080 ));
1081 }
1082
1083 let preferred = self
1084 .read_memory_block(PREFERRED_CALIBRATION_ADDRESS, STREAMING_AUTH_CREDENTIAL_SIZE as u32)
1085 .await?;
1086 let selected = if calibration_record_is_erased(&preferred) {
1087 self.read_memory_block(CALIBRATION_ADDRESS, STREAMING_AUTH_CREDENTIAL_SIZE as u32)
1088 .await?
1089 } else {
1090 preferred
1091 };
1092 let credential = AuthCredential::from_bytes(
1093 selected
1094 .as_slice()
1095 .try_into()
1096 .map_err(|_| KMError::Protocol("Calibration credential has the wrong length".to_string()))?,
1097 );
1098 let result = self.perform_streaming_auth(credential).await?;
1099
1100 if result.auth_level != 2 {
1101 return Err(KMError::Protocol(format!(
1102 "Calibration authentication returned level {} instead of 2",
1103 result.auth_level
1104 )));
1105 }
1106
1107 if let ConnectionMode::Full(state) = &mut self.mode {
1108 state.auth_level = result.auth_level;
1109 state.adcqueue_enabled = result.adcqueue_enabled();
1110 }
1111
1112 Ok(result)
1113 }
1114
1115 pub async fn read_memory_block(&mut self, address: u32, size: u32) -> Result<Vec<u8>, KMError> {
1121 let id = self.send_tracked(Packet::MemoryRead { address, size }).await?;
1122 let confirmation = self
1123 .receive_matching_raw(|bytes| memory_confirmation_matches(bytes, id))
1124 .await?;
1125
1126 let confirmation = RawPacket::try_from(Bytes::from(confirmation))?;
1127 match confirmation.packet_type() {
1128 PacketType::MemoryRead => {
1129 validate_memory_read_confirmation(&confirmation, id, address, size)?;
1130 self.receive_memory_read_data_exact(size).await
1131 }
1132 PacketType::Rejected => Err(KMError::Protocol(format!("MemoryRead at 0x{address:08X} was rejected"))),
1133 PacketType::NotReadable => Err(KMError::Protocol(format!(
1134 "Memory address 0x{address:08X} is not readable"
1135 ))),
1136 other => Err(KMError::Protocol(format!(
1137 "Expected MemoryRead confirmation, got {other:?}"
1138 ))),
1139 }
1140 }
1141
1142 pub async fn get_device_info(&mut self) -> Result<crate::auth::DeviceInfo, KMError> {
1163 Ok(self.read_device_info_blocks().await)
1164 }
1165
1166 async fn read_device_info_blocks(&mut self) -> crate::auth::DeviceInfo {
1171 use crate::auth::{
1172 CALIBRATION_ADDRESS, DEVICE_INFO_ADDRESS, DeviceInfo, FIRMWARE_INFO_ADDRESS, INFO_BLOCK_SIZE,
1173 };
1174 use tracing::warn;
1175
1176 let mut info = DeviceInfo::default();
1177
1178 match self
1179 .read_memory_block(DEVICE_INFO_ADDRESS, INFO_BLOCK_SIZE as u32)
1180 .await
1181 {
1182 Ok(data) => info.parse_device_info(&data),
1183 Err(e) => warn!("Failed to read DeviceInfo at 0x{:08X}: {}", DEVICE_INFO_ADDRESS, e),
1184 }
1185
1186 match self
1187 .read_memory_block(FIRMWARE_INFO_ADDRESS, INFO_BLOCK_SIZE as u32)
1188 .await
1189 {
1190 Ok(data) => info.parse_firmware_info(&data),
1191 Err(e) => warn!("Failed to read FirmwareInfo at 0x{:08X}: {}", FIRMWARE_INFO_ADDRESS, e),
1192 }
1193
1194 match self
1195 .read_memory_block(CALIBRATION_ADDRESS, INFO_BLOCK_SIZE as u32)
1196 .await
1197 {
1198 Ok(data) => info.parse_calibration(&data),
1199 Err(e) => warn!("Failed to read CalibrationData at 0x{:08X}: {}", CALIBRATION_ADDRESS, e),
1200 }
1201
1202 info
1203 }
1204
1205 pub async fn start_graph_mode(&mut self, rate: GraphSampleRate) -> Result<(), KMError> {
1238 ensure_adcqueue_available(&self.mode)?;
1239
1240 let id = self
1242 .send_tracked(Packet::StartGraph {
1243 rate_index: rate as u16,
1244 })
1245 .await?;
1246 self.expect_accept(id, "StartGraph").await?;
1247 self.graph_sample_rate = Some(rate);
1248 Ok(())
1249 }
1250
1251 pub async fn stop_graph_mode(&mut self) -> Result<(), KMError> {
1257 if self.is_basic_mode() {
1258 return Err(KMError::Protocol(
1259 "AdcQueue streaming requires Full mode (vendor interface)".to_string(),
1260 ));
1261 }
1262
1263 let id = self.send_tracked(Packet::StopGraph).await?;
1264 self.expect_accept(id, "StopGraph").await?;
1265 self.graph_sample_rate = None;
1266 Ok(())
1267 }
1268}
1269
1270#[cfg(test)]
1271mod tests {
1272 use super::*;
1273
1274 #[test]
1275 fn calibration_auth_falls_back_only_for_an_erased_preferred_record() {
1276 assert!(calibration_record_is_erased(&[0xff; STREAMING_AUTH_CREDENTIAL_SIZE]));
1277 assert!(!calibration_record_is_erased(b"007965 CDFDD"));
1278 assert!(!calibration_record_is_erased(&[0xff, 0xff, 0xff]));
1279 }
1280
1281 #[test]
1282 fn response_matching_uses_the_typed_header_parser() {
1283 assert!(response_matches(&[0xc1, 7, 0, 0], 7, PacketType::PutData));
1284 }
1285
1286 #[test]
1287 fn graph_mode_requires_streaming_auth_permission() {
1288 assert!(ensure_adcqueue_available(&ConnectionMode::Basic).is_err());
1289
1290 let state = DeviceState {
1291 info: DeviceInfo::default(),
1292 hardware_id: HardwareId::from_bytes([0; 12]),
1293 auth_level: 0,
1294 adcqueue_enabled: false,
1295 };
1296 assert!(ensure_adcqueue_available(&ConnectionMode::Full(state.clone())).is_err());
1297
1298 let enabled = DeviceState {
1299 auth_level: 1,
1300 adcqueue_enabled: true,
1301 ..state
1302 };
1303 assert!(ensure_adcqueue_available(&ConnectionMode::Full(enabled)).is_ok());
1304 }
1305
1306 #[test]
1307 fn response_matching_requires_type_and_transaction_id() {
1308 assert!(!response_matches(&[0xc1, 8, 0, 0], 7, PacketType::PutData));
1309 assert!(!response_matches(&[0xc0, 7, 0, 0], 7, PacketType::PutData));
1310 assert!(!response_matches(&[0xc1], 7, PacketType::PutData));
1311 }
1312
1313 #[test]
1314 fn response_type_matching_supports_fixed_id_protocol_exceptions() {
1315 assert!(response_type_matches(&[0x4c, 0, 3, 2], PacketType::StreamingAuth));
1316 }
1317
1318 #[test]
1319 fn memory_response_size_rounds_to_complete_aes_blocks() {
1320 assert_eq!(memory_response_size(1), 16);
1321 assert_eq!(memory_response_size(12), 16);
1322 assert_eq!(memory_response_size(16), 16);
1323 assert_eq!(memory_response_size(17), 32);
1324 assert_eq!(memory_response_size(64), 64);
1325 }
1326
1327 #[test]
1328 fn decrypted_memory_response_is_trimmed_to_the_requested_size() {
1329 let ciphertext = crate::auth::build_memory_read_payload(0x4001_0450, 12);
1330 let decrypted = decrypt_memory_response(&ciphertext[..16], 12).unwrap();
1331
1332 assert_eq!(decrypted.len(), 12);
1333 assert_eq!(&decrypted[..4], &0x4001_0450_u32.to_le_bytes());
1334 assert_eq!(&decrypted[4..8], &12_u32.to_le_bytes());
1335 assert_eq!(&decrypted[8..12], &u32::MAX.to_le_bytes());
1336 }
1337
1338 #[test]
1339 fn decrypted_memory_response_rejects_an_incorrect_ciphertext_size() {
1340 assert!(decrypt_memory_response(&[0; 32], 12).is_err());
1341 }
1342
1343 #[test]
1344 fn memory_chunks_accumulate_across_usb_transfers() {
1345 let mut response = Vec::new();
1348 assert!(!append_memory_chunk(&mut response, &vec![0x11; 2544], 8336).unwrap());
1349 assert!(!append_memory_chunk(&mut response, &vec![0x22; 2544], 8336).unwrap());
1350 assert!(!append_memory_chunk(&mut response, &vec![0x33; 2544], 8336).unwrap());
1351 assert!(append_memory_chunk(&mut response, &vec![0x44; 704], 8336).unwrap());
1352 assert_eq!(response.len(), 8336);
1353 }
1354
1355 #[test]
1356 fn memory_chunks_reject_empty_and_oversized_transfers() {
1357 let mut response = vec![0; 15];
1358 assert!(append_memory_chunk(&mut response, &[], 16).is_err());
1359 assert!(append_memory_chunk(&mut response, &[1, 2], 16).is_err());
1360 }
1361
1362 #[test]
1363 fn validates_recorded_memory_read_confirmation() {
1364 let bytes = hex::decode("c40201012004000040000000ffffffff1b8c1b24").unwrap();
1366 let packet = RawPacket::try_from(Bytes::from(bytes)).unwrap();
1367
1368 validate_memory_read_confirmation(&packet, 2, 0x420, 64).unwrap();
1369 }
1370
1371 #[test]
1372 fn rejects_corrupted_memory_read_confirmation() {
1373 let mut bytes = hex::decode("c40201012004000040000000ffffffff1b8c1b24").unwrap();
1374 bytes[8] ^= 1;
1375 let packet = RawPacket::try_from(Bytes::from(bytes)).unwrap();
1376
1377 assert!(validate_memory_read_confirmation(&packet, 2, 0x420, 64).is_err());
1378 }
1379}