use crate::adc::AdcDataSimple;
use crate::adcqueue::GraphSampleRate;
use crate::auth::{
AuthCredential, CALIBRATION_ADDRESS, DeviceInfo, HardwareId, PREFERRED_CALIBRATION_ADDRESS,
STREAMING_AUTH_CREDENTIAL_SIZE, StreamingAuthResult,
};
use crate::error::KMError;
use crate::message::Packet;
use crate::offline::{LogMetadata, LogMetadataResponse, OfflineLog};
use crate::packet::{Attribute, AttributeSet, PacketPeek, PacketType, RawPacket};
use crate::pd::{PdEventStream, PdStatus};
use crate::settings::Settings;
use bytes::Bytes;
use nusb::Interface;
use nusb::io::{EndpointRead, EndpointWrite};
use nusb::transfer::{Bulk, Interrupt};
use std::collections::VecDeque;
use std::fmt;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::time::timeout;
use tracing::{debug, info, trace, warn};
fn calibration_record_is_erased(record: &[u8]) -> bool {
record.starts_with(&[0xff; 4])
}
#[derive(Debug, Clone)]
pub struct DeviceState {
pub info: DeviceInfo,
pub hardware_id: HardwareId,
pub auth_level: u8,
pub adcqueue_enabled: bool,
}
impl DeviceState {
pub fn is_authenticated(&self) -> bool {
self.auth_level > 0
}
pub fn model(&self) -> &str {
&self.info.model
}
pub fn firmware_version(&self) -> &str {
&self.info.fw_version
}
}
impl fmt::Display for DeviceState {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
const W: usize = 58;
fn line(label: &str, value: &str) -> String {
let content = format!(" {}: {}", label, value);
format!("│{:<W$}│", content, W = W)
}
let sep = format!("├{:─<W$}┤", "", W = W);
writeln!(f, "┌{:─<W$}┐", "", W = W)?;
writeln!(f, "│{:^W$}│", "POWER-Z KM003C Device Information", W = W)?;
writeln!(f, "{}", sep)?;
writeln!(f, "{}", line("Model", &self.info.model))?;
writeln!(f, "{}", line("Hardware Version", &self.info.hw_version))?;
writeln!(f, "{}", line("Mfg Date", &self.info.mfg_date))?;
writeln!(f, "{}", sep)?;
writeln!(f, "{}", line("Firmware Version", &self.info.fw_version))?;
writeln!(f, "{}", line("Firmware Date", &self.info.fw_date))?;
writeln!(f, "{}", sep)?;
writeln!(f, "{}", line("Serial ID", &self.info.serial_id))?;
writeln!(f, "{}", line("UUID", &self.info.uuid))?;
writeln!(f, "{}", line("Hardware ID", &format!("{}", self.hardware_id)))?;
writeln!(f, "{}", sep)?;
writeln!(f, "{}", line("Auth Level", &format!("{}", self.auth_level)))?;
writeln!(f, "{}", line("AdcQueue Enabled", &format!("{}", self.adcqueue_enabled)))?;
write!(f, "└{:─<W$}┘", "", W = W)
}
}
pub const VID: u16 = 0x5FC9; pub const PID: u16 = 0x0063;
pub const INTERFACE_VENDOR: u8 = 0;
pub const ENDPOINT_OUT_VENDOR: u8 = 0x01;
pub const ENDPOINT_IN_VENDOR: u8 = 0x81;
pub const INTERFACE_HID: u8 = 3;
pub const ENDPOINT_OUT_HID: u8 = 0x05;
pub const ENDPOINT_IN_HID: u8 = 0x85;
const DEFAULT_TIMEOUT: Duration = Duration::from_secs(2);
const BULK_TRANSFER_SIZE: usize = 2048;
const INTERRUPT_TRANSFER_SIZE: usize = 64;
const MAX_PENDING_RESPONSES: usize = 256;
const AES_BLOCK_SIZE: usize = 16;
async fn find_device(selector: &DeviceSelector) -> Result<nusb::DeviceInfo, KMError> {
nusb::list_devices()
.await?
.find(|device| device.vendor_id() == VID && device.product_id() == PID && selector.matches(device))
.ok_or(KMError::DeviceNotFound)
}
fn response_matches(bytes: &[u8], id: u8, packet_type: PacketType) -> bool {
PacketPeek::from_bytes(bytes).is_some_and(|peek| peek.id == id && peek.packet_type == packet_type)
}
fn response_type_matches(bytes: &[u8], packet_type: PacketType) -> bool {
PacketPeek::from_bytes(bytes).is_some_and(|peek| peek.packet_type == packet_type)
}
fn control_response_matches(bytes: &[u8], id: u8) -> bool {
PacketPeek::from_bytes(bytes).is_some_and(|peek| peek.is_ctrl() && peek.id == id)
}
fn memory_confirmation_matches(bytes: &[u8], id: u8) -> bool {
PacketPeek::from_bytes(bytes).is_some_and(|peek| {
peek.id == id
&& matches!(
peek.packet_type,
PacketType::MemoryRead | PacketType::Rejected | PacketType::NotReadable
)
})
}
fn memory_response_size(requested_size: u32) -> usize {
(requested_size as usize).div_ceil(AES_BLOCK_SIZE) * AES_BLOCK_SIZE
}
fn append_memory_chunk(buffer: &mut Vec<u8>, chunk: &[u8], expected_size: usize) -> Result<bool, KMError> {
if chunk.is_empty() {
return Err(KMError::Protocol(
"Received an empty encrypted memory chunk".to_string(),
));
}
let remaining = expected_size.saturating_sub(buffer.len());
if chunk.len() > remaining {
return Err(KMError::Protocol(format!(
"Encrypted memory response exceeded expected size: expected {expected_size} bytes, got at least {}",
buffer.len() + chunk.len()
)));
}
buffer.extend_from_slice(chunk);
Ok(buffer.len() == expected_size)
}
fn decrypt_memory_response(encrypted: &[u8], requested_size: u32) -> Result<Vec<u8>, KMError> {
let expected_size = memory_response_size(requested_size);
if encrypted.len() != expected_size {
return Err(KMError::Protocol(format!(
"Encrypted memory response has the wrong size: expected {expected_size} bytes, got {}",
encrypted.len()
)));
}
let mut decrypted = crate::auth::aes_ecb_decrypt_blocks(encrypted, crate::auth::MEMORY_READ_KEY)?;
decrypted.truncate(requested_size as usize);
Ok(decrypted)
}
fn validate_memory_read_confirmation(
packet: &RawPacket,
expected_id: u8,
expected_address: u32,
expected_size: u32,
) -> Result<(), KMError> {
let RawPacket::SimpleData { header, payload } = packet else {
return Err(KMError::Protocol(
"MemoryRead confirmation has an unexpected packet layout".to_string(),
));
};
if header.id() != expected_id || PacketType::from(header.packet_type()) != PacketType::MemoryRead {
return Err(KMError::Protocol(
"MemoryRead confirmation does not match the request".to_string(),
));
}
if payload.len() != crate::auth::MEMORY_READ_CONFIRMATION_SIZE {
return Err(KMError::Protocol(format!(
"MemoryRead confirmation must contain {} bytes, got {}",
crate::auth::MEMORY_READ_CONFIRMATION_SIZE,
payload.len()
)));
}
let confirmation = crate::auth::MemoryReadConfirmation::from_bytes(payload)
.ok_or_else(|| KMError::Protocol("MemoryRead confirmation body is truncated".to_string()))?;
if confirmation.address != expected_address || confirmation.size != expected_size {
return Err(KMError::Protocol(format!(
"MemoryRead confirmation echoed address 0x{:08X} and size {}, expected 0x{expected_address:08X} and {expected_size}",
confirmation.address, confirmation.size
)));
}
if confirmation.magic != u32::MAX {
return Err(KMError::Protocol(format!(
"MemoryRead confirmation has invalid magic 0x{:08X}",
confirmation.magic
)));
}
if confirmation.checksum != confirmation.expected_checksum {
return Err(KMError::Protocol(format!(
"MemoryRead confirmation CRC mismatch: expected 0x{:08X}, got 0x{:08X}",
confirmation.expected_checksum, confirmation.checksum
)));
}
Ok(())
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TransferType {
Bulk,
Interrupt,
}
#[derive(Debug, Clone)]
pub enum ConnectionMode {
Basic,
Full(DeviceState),
}
fn ensure_adcqueue_available(mode: &ConnectionMode) -> Result<(), KMError> {
match mode {
ConnectionMode::Basic => Err(KMError::Protocol(
"AdcQueue streaming requires Full mode (vendor interface)".to_string(),
)),
ConnectionMode::Full(state) if !state.adcqueue_enabled => Err(KMError::Protocol(
"StreamingAuth did not enable AdcQueue streaming".to_string(),
)),
ConnectionMode::Full(_) => Ok(()),
}
}
#[derive(Debug, Clone)]
pub struct DeviceConfig {
interface: u8,
endpoint_out: u8,
endpoint_in: u8,
transfer_type: TransferType,
skip_reset: bool,
selector: DeviceSelector,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub enum DeviceSelector {
#[default]
First,
SerialNumber(String),
BusAddress { bus_id: String, device_address: u8 },
}
impl DeviceSelector {
fn matches(&self, device: &nusb::DeviceInfo) -> bool {
match self {
Self::First => true,
Self::SerialNumber(serial) => device.serial_number() == Some(serial.as_str()),
Self::BusAddress { bus_id, device_address } => {
device.bus_id() == bus_id && device.device_address() == *device_address
}
}
}
fn describe(&self) -> String {
match self {
Self::First => "the first attached KM003C".to_string(),
Self::SerialNumber(serial) => format!("KM003C with serial {serial}"),
Self::BusAddress { bus_id, device_address } => format!("KM003C at bus {bus_id} address {device_address}"),
}
}
}
impl DeviceConfig {
pub fn vendor() -> Self {
Self {
interface: INTERFACE_VENDOR,
endpoint_out: ENDPOINT_OUT_VENDOR,
endpoint_in: ENDPOINT_IN_VENDOR,
transfer_type: TransferType::Bulk,
skip_reset: false,
selector: DeviceSelector::First,
}
}
pub fn hid() -> Self {
Self {
interface: INTERFACE_HID,
endpoint_out: ENDPOINT_OUT_HID,
endpoint_in: ENDPOINT_IN_HID,
transfer_type: TransferType::Interrupt,
skip_reset: false,
selector: DeviceSelector::First,
}
}
pub fn skip_reset(mut self) -> Self {
self.skip_reset = true;
self
}
pub fn select(mut self, selector: DeviceSelector) -> Self {
self.selector = selector;
self
}
pub fn selector(&self) -> &DeviceSelector {
&self.selector
}
pub fn is_vendor(&self) -> bool {
self.interface == INTERFACE_VENDOR
}
pub fn is_hid(&self) -> bool {
self.interface == INTERFACE_HID
}
}
enum EndpointReaderType {
Bulk(EndpointRead<Bulk>),
Interrupt(EndpointRead<Interrupt>),
}
enum EndpointWriterType {
Bulk(EndpointWrite<Bulk>),
Interrupt(EndpointWrite<Interrupt>),
}
pub struct KM003C {
#[allow(dead_code)]
interface: Interface,
transaction_id: u8,
reader: EndpointReaderType,
writer: EndpointWriterType,
pending_responses: VecDeque<Vec<u8>>,
graph_sample_rate: Option<GraphSampleRate>,
mode: ConnectionMode,
}
impl KM003C {
pub async fn new(config: DeviceConfig) -> Result<Self, KMError> {
let is_vendor = config.is_vendor();
let mut device = Self::connect(config).await?;
if is_vendor {
device.run_init().await?;
}
Ok(device)
}
async fn connect(config: DeviceConfig) -> Result<Self, KMError> {
info!("Searching for {}...", config.selector.describe());
let device_info = find_device(&config.selector).await?;
info!(
"Found device on bus {} addr {}",
device_info.bus_id(),
device_info.device_address()
);
let reopen_selector = match &config.selector {
DeviceSelector::SerialNumber(_) => config.selector.clone(),
_ => device_info
.serial_number()
.map(|serial| DeviceSelector::SerialNumber(serial.to_string()))
.unwrap_or(DeviceSelector::First),
};
let mut device = device_info.open().await?;
if !config.skip_reset {
info!("Resetting device...");
device.reset().await?;
tokio::time::sleep(Duration::from_millis(1500)).await;
device = find_device(&reopen_selector).await?.open().await?;
} else {
debug!("Skipping USB reset (skip_reset=true)");
}
for interface_num in 0..4 {
if let Err(e) = device.detach_kernel_driver(interface_num) {
trace!("Could not detach interface {}: {}", interface_num, e);
} else {
debug!("Detached kernel driver from interface {}", interface_num);
}
}
let interface = device.claim_interface(config.interface).await?;
info!("Interface {} claimed successfully", config.interface);
let (reader, writer) = match config.transfer_type {
TransferType::Bulk => {
let ep_in = interface.endpoint::<Bulk, _>(config.endpoint_in)?;
let ep_out = interface.endpoint::<Bulk, _>(config.endpoint_out)?;
(
EndpointReaderType::Bulk(ep_in.reader(BULK_TRANSFER_SIZE).with_num_transfers(4)),
EndpointWriterType::Bulk(ep_out.writer(INTERRUPT_TRANSFER_SIZE).with_num_transfers(4)),
)
}
TransferType::Interrupt => {
let ep_in = interface.endpoint::<Interrupt, _>(config.endpoint_in)?;
let ep_out = interface.endpoint::<Interrupt, _>(config.endpoint_out)?;
(
EndpointReaderType::Interrupt(ep_in.reader(INTERRUPT_TRANSFER_SIZE).with_num_transfers(4)),
EndpointWriterType::Interrupt(ep_out.writer(INTERRUPT_TRANSFER_SIZE).with_num_transfers(4)),
)
}
};
let km003c = Self {
interface,
transaction_id: 0,
reader,
writer,
pending_responses: VecDeque::new(),
graph_sample_rate: None,
mode: ConnectionMode::Basic,
};
info!("USB connection established");
Ok(km003c)
}
fn next_transaction_id(&mut self) -> u8 {
let id = self.transaction_id;
self.transaction_id = self.transaction_id.wrapping_add(1);
self.pending_responses
.retain(|response| !PacketPeek::from_bytes(response).is_some_and(|peek| peek.id == id));
id
}
pub fn next_tid(&mut self) -> u8 {
self.next_transaction_id()
}
pub fn set_transaction_id(&mut self, id: u8) {
self.transaction_id = id;
}
pub async fn send(&mut self, packet: Packet) -> Result<(), KMError> {
self.send_tracked(packet).await.map(|_| ())
}
async fn send_tracked(&mut self, packet: Packet) -> Result<u8, KMError> {
let id = self.next_transaction_id();
let raw_packet = packet.to_raw_packet(id)?;
self.send_raw_packet(raw_packet).await?;
Ok(id)
}
async fn send_raw_packet(&mut self, packet: RawPacket) -> Result<(), KMError> {
let (reserved_flag, has_logical_packets) = match &packet {
RawPacket::Ctrl { header, .. } => (header.reserved_flag(), false),
RawPacket::SimpleData { header, .. } => (header.reserved_flag(), false),
RawPacket::Data {
header,
logical_packets,
} => (header.reserved_flag(), !logical_packets.is_empty()),
};
debug!(
"Sending packet: type={:?}, reserved_flag={}, has_logical_packets={}, id={}",
packet.packet_type(),
reserved_flag,
has_logical_packets,
packet.id(),
);
let message_bytes = Bytes::from(packet);
let message = message_bytes.to_vec();
trace!("TX [{} bytes]: {:02x?}", message.len(), message);
match &mut self.writer {
EndpointWriterType::Bulk(writer) => {
timeout(DEFAULT_TIMEOUT, writer.write_all(&message)).await??;
timeout(DEFAULT_TIMEOUT, writer.flush_end_async()).await??;
}
EndpointWriterType::Interrupt(writer) => {
timeout(DEFAULT_TIMEOUT, writer.write_all(&message)).await??;
timeout(DEFAULT_TIMEOUT, writer.flush_end_async()).await??;
}
}
debug!("Sent successfully");
Ok(())
}
pub async fn send_raw(&mut self, data: &[u8]) -> Result<(), KMError> {
trace!("TX [{} bytes]: {:02x?}", data.len(), data);
match &mut self.writer {
EndpointWriterType::Bulk(writer) => {
timeout(DEFAULT_TIMEOUT, writer.write_all(data)).await??;
timeout(DEFAULT_TIMEOUT, writer.flush_end_async()).await??;
}
EndpointWriterType::Interrupt(writer) => {
timeout(DEFAULT_TIMEOUT, writer.write_all(data)).await??;
timeout(DEFAULT_TIMEOUT, writer.flush_end_async()).await??;
}
}
Ok(())
}
async fn read_raw_from_usb(&mut self) -> Result<Vec<u8>, KMError> {
let mut buffer = Vec::new();
let bytes_read = match &mut self.reader {
EndpointReaderType::Bulk(reader) => {
let mut message = reader.until_short_packet();
let bytes_read = timeout(DEFAULT_TIMEOUT, message.read_to_end(&mut buffer)).await??;
message
.consume_end()
.map_err(|_| KMError::Protocol("Bulk response ended without a short packet".to_string()))?;
bytes_read
}
EndpointReaderType::Interrupt(reader) => {
buffer.resize(INTERRUPT_TRANSFER_SIZE, 0);
let bytes_read = timeout(DEFAULT_TIMEOUT, reader.read(&mut buffer)).await??;
buffer.truncate(bytes_read);
bytes_read
}
};
trace!("RX [{} bytes]: {:02x?}", bytes_read, buffer);
Ok(buffer)
}
fn queue_pending_response(&mut self, response: Vec<u8>) {
if self.pending_responses.len() == MAX_PENDING_RESPONSES {
warn!("Dropping oldest unmatched response because the pending queue is full");
self.pending_responses.pop_front();
}
self.pending_responses.push_back(response);
}
async fn receive_matching_raw<F>(&mut self, mut predicate: F) -> Result<Vec<u8>, KMError>
where
F: FnMut(&[u8]) -> bool,
{
if let Some(index) = self.pending_responses.iter().position(|response| predicate(response)) {
return Ok(self
.pending_responses
.remove(index)
.expect("pending response index is valid"));
}
let deadline = tokio::time::Instant::now() + DEFAULT_TIMEOUT;
loop {
let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
if remaining.is_zero() {
return Err(KMError::Protocol(
"Timed out waiting for a correlated response".to_string(),
));
}
let response = timeout(remaining, self.read_raw_from_usb()).await??;
if predicate(&response) {
return Ok(response);
}
if let Some(peek) = PacketPeek::from_bytes(&response) {
debug!(
"Queueing unmatched response: type={:?}, id={}, len={}",
peek.packet_type,
peek.id,
response.len()
);
} else {
debug!("Queueing unframed response: len={}", response.len());
}
self.queue_pending_response(response);
}
}
pub async fn receive_raw(&mut self) -> Result<Vec<u8>, KMError> {
if let Some(response) = self.pending_responses.pop_front() {
return Ok(response);
}
self.read_raw_from_usb().await
}
fn parse_response(raw_bytes: Vec<u8>, graph_rate: Option<GraphSampleRate>) -> Result<Packet, KMError> {
if raw_bytes.is_empty() {
return Err(KMError::Protocol("Received 0 bytes".to_string()));
}
let raw_packet = RawPacket::try_from(Bytes::from(raw_bytes))?;
if let Some(rate) = graph_rate {
Packet::from_raw_with_graph_rate(raw_packet, rate)
} else {
Packet::try_from(raw_packet)
}
}
pub async fn receive(&mut self) -> Result<Packet, KMError> {
let raw_bytes = self.receive_raw().await?;
Self::parse_response(raw_bytes, self.graph_sample_rate)
}
async fn receive_control_response(&mut self, id: u8) -> Result<Packet, KMError> {
let raw_bytes = self
.receive_matching_raw(|bytes| control_response_matches(bytes, id))
.await?;
Self::parse_response(raw_bytes, self.graph_sample_rate)
}
async fn expect_accept(&mut self, id: u8, command: &str) -> Result<(), KMError> {
match self.receive_control_response(id).await? {
Packet::Accept { .. } => Ok(()),
Packet::Reject { .. } => Err(KMError::Protocol(format!("{command} was rejected by the device"))),
Packet::NotReadable { .. } => Err(KMError::Protocol(format!("{command} is not readable"))),
other => Err(KMError::Protocol(format!(
"Expected Accept for {command}, got {other:?}"
))),
}
}
async fn receive_memory_read_data_exact(&mut self, requested_size: u32) -> Result<Vec<u8>, KMError> {
let expected_size = memory_response_size(requested_size);
if expected_size == 0 {
return Ok(Vec::new());
}
let mut encrypted = Vec::with_capacity(expected_size);
while encrypted.len() < expected_size {
let chunk = self.read_raw_from_usb().await?;
if append_memory_chunk(&mut encrypted, &chunk, expected_size)? {
break;
}
}
decrypt_memory_response(&encrypted, requested_size)
}
pub async fn request_data(&mut self, mask: AttributeSet) -> Result<Packet, KMError> {
let id = self.send_tracked(Packet::GetData { attributes: mask }).await?;
let raw_bytes = self
.receive_matching_raw(|bytes| response_matches(bytes, id, PacketType::PutData))
.await?;
let raw_packet = RawPacket::try_from(Bytes::from(raw_bytes))?;
raw_packet.validate_correlation(mask)?;
if let Some(rate) = self.graph_sample_rate {
Packet::from_raw_with_graph_rate(raw_packet, rate)
} else {
Packet::try_from(raw_packet)
}
}
pub async fn request_adc_data(&mut self) -> Result<AdcDataSimple, KMError> {
let packet = self.request_data(AttributeSet::single(Attribute::Adc)).await?;
packet
.get_adc()
.cloned()
.ok_or_else(|| KMError::Protocol("No ADC data in response".to_string()))
}
pub async fn request_settings(&mut self) -> Result<Settings, KMError> {
let packet = self.request_data(AttributeSet::single(Attribute::Settings)).await?;
packet
.get_settings()
.cloned()
.ok_or_else(|| KMError::Protocol("No Settings data in response".to_string()))
}
pub async fn request_pd_trace(&mut self) -> Result<crate::pd_trace::PdTrace, KMError> {
let packet = self.request_data(AttributeSet::single(Attribute::PdTrace)).await?;
packet
.get_pd_trace()
.cloned()
.ok_or_else(|| KMError::Protocol("No PD trace data in response".to_string()))
}
pub async fn request_log_metadata(&mut self) -> Result<Vec<LogMetadata>, KMError> {
let packet = self.request_data(AttributeSet::single(Attribute::LogMetadata)).await?;
match packet.get_log_metadata() {
Some(LogMetadataResponse::Empty) => Ok(Vec::new()),
Some(LogMetadataResponse::Available(metadata)) => Ok(metadata.clone()),
None => Err(KMError::Protocol("No LogMetadata data in response".to_string())),
}
}
pub async fn download_offline_log(&mut self, metadata: LogMetadata) -> Result<OfflineLog, KMError> {
let data = self
.read_memory_block(metadata.data_address()?, metadata.data_size())
.await?;
OfflineLog::from_bytes(metadata, &data)
}
pub async fn request_pd_data(&mut self) -> Result<Packet, KMError> {
self.request_data(AttributeSet::single(Attribute::PdPacket)).await
}
pub async fn request_adc_with_pd(&mut self) -> Result<Packet, KMError> {
let mask = AttributeSet::single(Attribute::Adc).with(Attribute::PdPacket);
self.request_data(mask).await
}
pub fn extract_pd_status(packet: &Packet) -> Option<&PdStatus> {
packet.get_pd_status()
}
pub fn extract_pd_events(packet: &Packet) -> Option<&PdEventStream> {
packet.get_pd_events()
}
pub async fn enable_pd_monitor(&mut self) -> Result<(), KMError> {
let id = self.send_tracked(Packet::EnablePdMonitor).await?;
self.expect_accept(id, "EnablePdMonitor").await
}
pub async fn disable_pd_monitor(&mut self) -> Result<(), KMError> {
let id = self.send_tracked(Packet::DisablePdMonitor).await?;
self.expect_accept(id, "DisablePdMonitor").await
}
async fn run_init(&mut self) -> Result<(), KMError> {
use crate::auth::{HARDWARE_ID_ADDRESS, HARDWARE_ID_SIZE};
const MAX_CONNECT_RETRIES: u8 = 3;
let mut last_error = None;
for attempt in 1..=MAX_CONNECT_RETRIES {
match self.send_tracked(Packet::Connect).await {
Err(err) => last_error = Some(format!("Connect send failed: {err:?}")),
Ok(id) => match self.expect_accept(id, "Connect").await {
Ok(()) => {
last_error = None;
break;
}
Err(err) => last_error = Some(format!("Connect receive failed: {err:?}")),
},
}
if attempt < MAX_CONNECT_RETRIES {
debug!("Connect attempt {} failed, retrying...", attempt);
tokio::time::sleep(Duration::from_millis(100)).await;
}
}
if let Some(err) = last_error {
return Err(KMError::Protocol(err));
}
let info = self.read_device_info_blocks().await;
let data = self
.read_memory_block(HARDWARE_ID_ADDRESS, HARDWARE_ID_SIZE as u32)
.await
.map_err(|error| {
KMError::Protocol(format!(
"Failed to read HardwareID - required for authentication: {error}"
))
})?;
let hardware_id_bytes: [u8; HARDWARE_ID_SIZE] = data.as_slice().try_into().map_err(|_| {
KMError::Protocol(format!(
"HardwareID read returned {} bytes, expected {HARDWARE_ID_SIZE}",
data.len()
))
})?;
let hardware_id = HardwareId::from_bytes(hardware_id_bytes);
let auth = self.perform_streaming_auth(AuthCredential::from(&hardware_id)).await?;
let state = DeviceState {
info,
hardware_id,
auth_level: auth.auth_level,
adcqueue_enabled: auth.adcqueue_enabled(),
};
info!("Device initialized: {} (auth_level={})", state.model(), auth.auth_level);
self.mode = ConnectionMode::Full(state);
Ok(())
}
async fn perform_streaming_auth(&mut self, credential: AuthCredential) -> Result<StreamingAuthResult, KMError> {
self.send_tracked(Packet::StreamingAuth { credential }).await?;
let response = self
.receive_matching_raw(|bytes| response_type_matches(bytes, PacketType::StreamingAuth))
.await?;
match Self::parse_response(response, self.graph_sample_rate)? {
Packet::StreamingAuthResponse(result) => Ok(result),
other => Err(KMError::Protocol(format!(
"Expected StreamingAuthResponse, got {other:?}"
))),
}
}
pub fn mode(&self) -> &ConnectionMode {
&self.mode
}
pub fn is_full_mode(&self) -> bool {
matches!(self.mode, ConnectionMode::Full(_))
}
pub fn is_basic_mode(&self) -> bool {
matches!(self.mode, ConnectionMode::Basic)
}
pub fn state(&self) -> Option<&DeviceState> {
match &self.mode {
ConnectionMode::Full(state) => Some(state),
ConnectionMode::Basic => None,
}
}
pub fn is_authenticated(&self) -> bool {
self.state().map(|s| s.is_authenticated()).unwrap_or(false)
}
pub fn adcqueue_enabled(&self) -> bool {
self.state().map(|s| s.adcqueue_enabled).unwrap_or(false)
}
pub async fn authenticate_calibration(&mut self) -> Result<StreamingAuthResult, KMError> {
if !self.is_full_mode() {
return Err(KMError::Protocol(
"Calibration authentication requires the vendor interface".to_string(),
));
}
let preferred = self
.read_memory_block(PREFERRED_CALIBRATION_ADDRESS, STREAMING_AUTH_CREDENTIAL_SIZE as u32)
.await?;
let selected = if calibration_record_is_erased(&preferred) {
self.read_memory_block(CALIBRATION_ADDRESS, STREAMING_AUTH_CREDENTIAL_SIZE as u32)
.await?
} else {
preferred
};
let credential = AuthCredential::from_bytes(
selected
.as_slice()
.try_into()
.map_err(|_| KMError::Protocol("Calibration credential has the wrong length".to_string()))?,
);
let result = self.perform_streaming_auth(credential).await?;
if result.auth_level != 2 {
return Err(KMError::Protocol(format!(
"Calibration authentication returned level {} instead of 2",
result.auth_level
)));
}
if let ConnectionMode::Full(state) = &mut self.mode {
state.auth_level = result.auth_level;
state.adcqueue_enabled = result.adcqueue_enabled();
}
Ok(result)
}
pub async fn read_memory_block(&mut self, address: u32, size: u32) -> Result<Vec<u8>, KMError> {
if self.graph_sample_rate.is_some() {
return Err(KMError::Protocol(
"Device memory cannot be read while AdcQueue streaming is active; \
call stop_graph_mode() first"
.to_string(),
));
}
let id = self.send_tracked(Packet::MemoryRead { address, size }).await?;
let confirmation = self
.receive_matching_raw(|bytes| memory_confirmation_matches(bytes, id))
.await?;
let confirmation = RawPacket::try_from(Bytes::from(confirmation))?;
match confirmation.packet_type() {
PacketType::MemoryRead => {
validate_memory_read_confirmation(&confirmation, id, address, size)?;
self.receive_memory_read_data_exact(size).await
}
PacketType::Rejected => Err(KMError::Protocol(format!("MemoryRead at 0x{address:08X} was rejected"))),
PacketType::NotReadable => Err(KMError::Protocol(format!(
"Memory address 0x{address:08X} is not readable"
))),
other => Err(KMError::Protocol(format!(
"Expected MemoryRead confirmation, got {other:?}"
))),
}
}
pub async fn get_device_info(&mut self) -> Result<crate::auth::DeviceInfo, KMError> {
Ok(self.read_device_info_blocks().await)
}
async fn read_device_info_blocks(&mut self) -> crate::auth::DeviceInfo {
use crate::auth::{
CALIBRATION_ADDRESS, DEVICE_INFO_ADDRESS, DeviceInfo, FIRMWARE_INFO_ADDRESS, INFO_BLOCK_SIZE,
};
use tracing::warn;
let mut info = DeviceInfo::default();
match self
.read_memory_block(DEVICE_INFO_ADDRESS, INFO_BLOCK_SIZE as u32)
.await
{
Ok(data) => info.parse_device_info(&data),
Err(e) => warn!("Failed to read DeviceInfo at 0x{:08X}: {}", DEVICE_INFO_ADDRESS, e),
}
match self
.read_memory_block(FIRMWARE_INFO_ADDRESS, INFO_BLOCK_SIZE as u32)
.await
{
Ok(data) => info.parse_firmware_info(&data),
Err(e) => warn!("Failed to read FirmwareInfo at 0x{:08X}: {}", FIRMWARE_INFO_ADDRESS, e),
}
match self
.read_memory_block(CALIBRATION_ADDRESS, INFO_BLOCK_SIZE as u32)
.await
{
Ok(data) => info.parse_calibration(&data),
Err(e) => warn!("Failed to read CalibrationData at 0x{:08X}: {}", CALIBRATION_ADDRESS, e),
}
info
}
pub async fn start_graph_mode(&mut self, rate: GraphSampleRate) -> Result<(), KMError> {
ensure_adcqueue_available(&self.mode)?;
let id = self.send_tracked(Packet::StartGraph { rate }).await?;
self.expect_accept(id, "StartGraph").await?;
self.graph_sample_rate = Some(rate);
Ok(())
}
pub async fn stop_graph_mode(&mut self) -> Result<(), KMError> {
if self.is_basic_mode() {
return Err(KMError::Protocol(
"AdcQueue streaming requires Full mode (vendor interface)".to_string(),
));
}
let id = self.send_tracked(Packet::StopGraph).await?;
self.expect_accept(id, "StopGraph").await?;
self.graph_sample_rate = None;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn calibration_auth_falls_back_only_for_an_erased_preferred_record() {
assert!(calibration_record_is_erased(&[0xff; STREAMING_AUTH_CREDENTIAL_SIZE]));
assert!(!calibration_record_is_erased(b"007965 CDFDD"));
assert!(!calibration_record_is_erased(&[0xff, 0xff, 0xff]));
}
#[test]
fn response_matching_uses_the_typed_header_parser() {
assert!(response_matches(&[0xc1, 7, 0, 0], 7, PacketType::PutData));
}
#[test]
fn graph_mode_requires_streaming_auth_permission() {
assert!(ensure_adcqueue_available(&ConnectionMode::Basic).is_err());
let state = DeviceState {
info: DeviceInfo::default(),
hardware_id: HardwareId::from_bytes([0; 12]),
auth_level: 0,
adcqueue_enabled: false,
};
assert!(ensure_adcqueue_available(&ConnectionMode::Full(state.clone())).is_err());
let enabled = DeviceState {
auth_level: 1,
adcqueue_enabled: true,
..state
};
assert!(ensure_adcqueue_available(&ConnectionMode::Full(enabled)).is_ok());
}
#[test]
fn response_matching_requires_type_and_transaction_id() {
assert!(!response_matches(&[0xc1, 8, 0, 0], 7, PacketType::PutData));
assert!(!response_matches(&[0xc0, 7, 0, 0], 7, PacketType::PutData));
assert!(!response_matches(&[0xc1], 7, PacketType::PutData));
}
#[test]
fn response_type_matching_supports_fixed_id_protocol_exceptions() {
assert!(response_type_matches(&[0x4c, 0, 3, 2], PacketType::StreamingAuth));
}
#[test]
fn memory_response_size_rounds_to_complete_aes_blocks() {
assert_eq!(memory_response_size(1), 16);
assert_eq!(memory_response_size(12), 16);
assert_eq!(memory_response_size(16), 16);
assert_eq!(memory_response_size(17), 32);
assert_eq!(memory_response_size(64), 64);
}
#[test]
fn decrypted_memory_response_is_trimmed_to_the_requested_size() {
let ciphertext = crate::auth::build_memory_read_payload(0x4001_0450, 12);
let decrypted = decrypt_memory_response(&ciphertext[..16], 12).unwrap();
assert_eq!(decrypted.len(), 12);
assert_eq!(&decrypted[..4], &0x4001_0450_u32.to_le_bytes());
assert_eq!(&decrypted[4..8], &12_u32.to_le_bytes());
assert_eq!(&decrypted[8..12], &u32::MAX.to_le_bytes());
}
#[test]
fn decrypted_memory_response_rejects_an_incorrect_ciphertext_size() {
assert!(decrypt_memory_response(&[0; 32], 12).is_err());
}
#[test]
fn memory_chunks_accumulate_across_usb_transfers() {
let mut response = Vec::new();
assert!(!append_memory_chunk(&mut response, &vec![0x11; 2544], 8336).unwrap());
assert!(!append_memory_chunk(&mut response, &vec![0x22; 2544], 8336).unwrap());
assert!(!append_memory_chunk(&mut response, &vec![0x33; 2544], 8336).unwrap());
assert!(append_memory_chunk(&mut response, &vec![0x44; 704], 8336).unwrap());
assert_eq!(response.len(), 8336);
}
#[test]
fn memory_chunks_reject_empty_and_oversized_transfers() {
let mut response = vec![0; 15];
assert!(append_memory_chunk(&mut response, &[], 16).is_err());
assert!(append_memory_chunk(&mut response, &[1, 2], 16).is_err());
}
#[test]
fn validates_recorded_memory_read_confirmation() {
let bytes = hex::decode("c40201012004000040000000ffffffff1b8c1b24").unwrap();
let packet = RawPacket::try_from(Bytes::from(bytes)).unwrap();
validate_memory_read_confirmation(&packet, 2, 0x420, 64).unwrap();
}
#[test]
fn rejects_corrupted_memory_read_confirmation() {
let mut bytes = hex::decode("c40201012004000040000000ffffffff1b8c1b24").unwrap();
bytes[8] ^= 1;
let packet = RawPacket::try_from(Bytes::from(bytes)).unwrap();
assert!(validate_memory_read_confirmation(&packet, 2, 0x420, 64).is_err());
}
}