donglora/src/usb/task.rs

//! USB CDC-ACM task: COBS-framed fixed-size LE command/response protocol.
//!
//! Reads COBS frames from USB, decodes them, parses commands, routes them
//! to the radio or display task, and sends COBS-framed responses back.

use defmt::warn;
use embassy_executor::task;
use embassy_futures::join::join;
use embassy_usb::class::cdc_acm::{CdcAcmClass, Receiver, Sender, State};
use embassy_usb::Builder;
use static_cell::StaticCell;

use crate::board::UsbParts;
use crate::channel::{
    CommandChannel, DisplayCommand, DisplayCommandChannel, ResponseChannel,
};
use crate::protocol::{Command, ErrorCode, Response};

const MAX_FRAME: usize = 512;
const MAX_PACKET: usize = 64;

// Worst case: tag(1) + rssi(2) + snr(2) + len(2) + payload(256) = 263 bytes.
// COBS adds ceil(263/254)+1 = 3 bytes. 266 < 512. Plenty of room.
const _: () = assert!(
    MAX_FRAME >= crate::protocol::MAX_PAYLOAD + 64,
    "MAX_FRAME too small for max payload + COBS overhead"
);

#[task]
pub async fn usb_task(
    parts: UsbParts,
    commands: &'static CommandChannel,
    responses: &'static ResponseChannel,
    display_commands: &'static DisplayCommandChannel,
    has_display: bool,
) {
    // ── Build embassy-usb device with CDC-ACM class ────────────────
    let mut config = embassy_usb::Config::new(0x1209, 0x5741);
    config.manufacturer = Some("DongLoRa");
    config.product = Some("DongLoRa LoRa Radio");
    config.serial_number = Some("001");
    config.device_class = 0xEF;
    config.device_sub_class = 0x02;
    config.device_protocol = 0x01;
    config.composite_with_iads = true;

    static DESC_BUF: StaticCell<[u8; 256]> = StaticCell::new();
    static CONF_BUF: StaticCell<[u8; 256]> = StaticCell::new();
    static BOS_BUF: StaticCell<[u8; 256]> = StaticCell::new();
    static CTRL_BUF: StaticCell<[u8; 64]> = StaticCell::new();
    static CDC_STATE: StaticCell<State> = StaticCell::new();

    let desc_buf = DESC_BUF.init([0; 256]);
    let conf_buf = CONF_BUF.init([0; 256]);
    let bos_buf = BOS_BUF.init([0; 256]);
    let ctrl_buf = CTRL_BUF.init([0; 64]);
    let cdc_state = CDC_STATE.init(State::new());

    let mut builder = Builder::new(
        parts.driver,
        config,
        desc_buf,
        conf_buf,
        bos_buf,
        ctrl_buf,
    );

    let class = CdcAcmClass::new(&mut builder, cdc_state, MAX_PACKET as u16);
    let mut usb_dev = builder.build();
    let (sender, receiver) = class.split();

    // ── Run USB device + protocol loop concurrently ────────────────
    join(
        usb_dev.run(),
        protocol_loop(sender, receiver, commands, responses, display_commands, has_display),
    )
    .await;
}

async fn protocol_loop<'d, D: embassy_usb_driver::Driver<'d>>(
    mut sender: Sender<'d, D>,
    mut receiver: Receiver<'d, D>,
    commands: &'static CommandChannel,
    responses: &'static ResponseChannel,
    display_commands: &'static DisplayCommandChannel,
    has_display: bool,
) {
    use embassy_futures::select::select3;
    use embassy_futures::select::Either3;

    let mut read_buf = [0u8; MAX_PACKET];
    let mut write_buf = [0u8; MAX_FRAME];
    let mut cobs_encode_buf = [0u8; MAX_FRAME];

    // COBS accumulator: collect bytes until 0x00 sentinel
    let mut frame_buf = [0u8; MAX_FRAME];
    let mut frame_len: usize = 0;

    let mut was_connected = false;

    loop {
        match select3(
            receiver.read_packet(&mut read_buf),
            responses.receive(),
            embassy_time::Timer::after_millis(250),
        )
        .await
        {
            Either3::First(result) => {
                let n = match result {
                    Ok(0) => continue,
                    Ok(n) => n,
                    Err(_) => {
                        frame_len = 0; // reset accumulator
                        embassy_time::Timer::after_millis(100).await;
                        continue;
                    }
                };

                // Feed bytes into the COBS accumulator
                for &byte in &read_buf[..n] {
                    if byte == 0x00 {
                        // End of COBS frame — decode and parse
                        if frame_len > 0 {
                            let mut decode_buf = [0u8; MAX_FRAME];
                            if let Some(decoded_len) =
                                cobs_decode(&frame_buf[..frame_len], &mut decode_buf)
                            {
                                if let Some(cmd) = Command::from_bytes(&decode_buf[..decoded_len]) {
                                    route_command(
                                        cmd, commands, responses, display_commands, has_display,
                                    )
                                    .await;
                                }
                            }
                        }
                        frame_len = 0;
                    } else if frame_len < MAX_FRAME {
                        frame_buf[frame_len] = byte;
                        frame_len += 1;
                    } else {
                        // Frame too large — discard
                        frame_len = 0;
                    }
                }
            }
            Either3::Second(response) => {
                // Serialize response to fixed-size LE bytes, then COBS encode
                let raw_len = response.to_bytes(&mut write_buf);
                let encoded_len = cobs::encode(&write_buf[..raw_len], &mut cobs_encode_buf);
                // Append 0x00 sentinel
                if encoded_len < cobs_encode_buf.len() {
                    cobs_encode_buf[encoded_len] = 0x00;
                    let frame = &cobs_encode_buf[..encoded_len + 1];
                    for chunk in frame.chunks(MAX_PACKET) {
                        if sender.write_packet(chunk).await.is_err() {
                            warn!("USB write failed, response dropped");
                            break;
                        }
                    }
                } else {
                    warn!("COBS encode buffer overflow");
                }
            }
            Either3::Third(()) => {
                // Poll DTR every 250ms to detect host disconnect
            }
        }

        // Check DTR after every select wake
        let connected = receiver.dtr();
        if was_connected && !connected && has_display {
            display_commands.send(DisplayCommand::Reset).await;
        }
        if !was_connected && connected && has_display {
            display_commands.send(DisplayCommand::On).await;
        }
        was_connected = connected;
    }
}

async fn route_command(
    cmd: Command,
    commands: &CommandChannel,
    responses: &ResponseChannel,
    display_commands: &DisplayCommandChannel,
    has_display: bool,
) {
    match cmd {
        Command::DisplayOn => {
            if has_display {
                display_commands.send(DisplayCommand::On).await;
                responses.send(Response::Ok).await;
            } else {
                responses.send(Response::Error(ErrorCode::NoDisplay)).await;
            }
        }
        Command::DisplayOff => {
            if has_display {
                display_commands.send(DisplayCommand::Off).await;
                responses.send(Response::Ok).await;
            } else {
                responses.send(Response::Error(ErrorCode::NoDisplay)).await;
            }
        }
        other => {
            commands.send(other).await;
        }
    }
}

/// Decode a COBS-encoded buffer (without the trailing 0x00 sentinel).
/// Returns the number of decoded bytes, or None on error.
fn cobs_decode(src: &[u8], dest: &mut [u8]) -> Option<usize> {
    let mut si = 0;
    let mut di = 0;
    while si < src.len() {
        let code = src[si] as usize;
        si += 1;
        if code == 0 {
            return None; // unexpected zero in COBS data
        }
        for _ in 1..code {
            if si >= src.len() || di >= dest.len() {
                return None;
            }
            dest[di] = src[si];
            di += 1;
            si += 1;
        }
        if code < 0xFF && si < src.len() {
            if di >= dest.len() {
                return None;
            }
            dest[di] = 0;
            di += 1;
        }
    }
    // Strip trailing zero if present (COBS artifact)
    if di > 0 && dest[di - 1] == 0 {
        di -= 1;
    }
    Some(di)
}