rmodem/src/cli.rs

use std::io::{self, Read, Seek, Write};
use std::time;

use clap::Parser;
#[cfg(feature = "crc")]
use rmodem::XmodemCrcPacket;
use rmodem::{Control, Error, Result, Sequence, XmodemData, XmodemPacket};
#[cfg(feature = "1k")]
use rmodem::{OneKData, XmodemOneKPacket};

const TIMEOUT_MS: u64 = 10_0000;
const RETRIES: usize = 16;

#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
    /// Path to the serial device for XMODEM communication.
    #[arg(short, long)]
    device: String,
    /// Baud rate for serial communication.
    #[arg(short, long, default_value_t = 115200)]
    baud: u32,
    /// Path to the file to send using XMODEM.
    #[arg(short, long)]
    file: String,
    /// Request that files be sent using this XMODEM protocol variant. Valid options: ["xmodem", "crc", "1k"].
    ///
    /// If the receiver does not support the requested protocol, fall-back will try to use other
    /// implemented protocols.
    #[arg(short, long, default_value_t = String::from("1k"))]
    protocol: String,
}

fn send_xmodem<R: Read>(
    port: &mut serialport::TTYPort,
    reader: &mut io::BufReader<R>,
) -> Result<()> {
    let mut block = [rmodem::CPMEOF; XmodemData::LEN];
    let mut sequence = Sequence::new();
    let mut packet = 1usize;
    let mut is_eof = false;

    while !is_eof {
        match reader.read(&mut block) {
            Ok(n) if n > 0 => {
                log::debug!("Sending XMODEM, packet: {packet}, sequence: {sequence}");

                retry_packet(
                    port,
                    XmodemPacket::new()
                        .with_sequence(sequence)
                        .with_data(XmodemData::from_block(&block[..n]))
                        .into_bytes()
                        .as_ref(),
                    RETRIES,
                )?;

                log::debug!("Received ACK for packet: {packet}, sequence: {sequence}");

                block.iter_mut().take(n).for_each(|b| *b = rmodem::CPMEOF);
                sequence = sequence.next();
                packet = packet.saturating_add(1);
            }
            Ok(_n) => {
                log::debug!("Reached EOF");
                is_eof = true;
            }
            Err(err) => {
                log::error!("Error reading block {sequence} from file: {err}");
                return Err(Error::Cli);
            }
        }
    }

    Ok(())
}

#[cfg(feature = "crc")]
fn send_xmodem_crc<R: Read + Seek>(
    port: &mut serialport::TTYPort,
    reader: &mut io::BufReader<R>,
) -> Result<()> {
    let mut block = [rmodem::CPMEOF; XmodemData::LEN];
    let n = reader.read(&mut block).map_err(|err| {
        log::error!("Error reading next block from the file: {err}");
        Error::Cli
    })?;

    log::debug!("Sending XMODEM-CRC, packet: 1, sequence: 1");

    match send_packet(
        port,
        XmodemCrcPacket::new()
            .with_data(XmodemData::from_block(&block[..n]))
            .into_bytes()
            .as_ref(),
    )? {
        Control::Nak => {
            // fall-back to XMODEM, rewind read bytes from the input file
            reader.seek(io::SeekFrom::Start(0)).map_err(|err| {
                log::error!("Error seeking input file: {err}");
                Error::Cli
            })?;
            send_xmodem(port, reader)
        }
        Control::Ack => {
            log::debug!("Received ACK for packet: 1, sequence: 1");
            send_xmodem_crc_inner(port, reader)
        }
        control => {
            log::error!("Unexpected control byte: {control:?}");
            Err(Error::Cli)
        }
    }
}

#[cfg(not(feature = "crc"))]
#[allow(dead_code)]
fn send_xmodem_crc<R: Read + Seek>(
    port: &mut serialport::TTYPort,
    reader: &mut io::BufReader<R>,
) -> Result<()> {
    send_xmodem(port, reader)
}

#[cfg(feature = "crc")]
fn send_xmodem_crc_inner<R: Read + Seek>(
    port: &mut serialport::TTYPort,
    reader: &mut io::BufReader<R>,
) -> Result<()> {
    let mut block = [rmodem::CPMEOF; XmodemData::LEN];
    let mut sequence = Sequence::from(2u8);
    let mut packet = 1usize;
    let mut is_eof = false;

    while !is_eof {
        match reader.read(&mut block) {
            Ok(n) if n > 0 => {
                log::debug!("Sending XMODEM-CRC, packet: {packet}, sequence: {sequence}");

                retry_packet(
                    port,
                    XmodemCrcPacket::new()
                        .with_sequence(sequence)
                        .with_data(XmodemData::from_block(&block[..n]))
                        .into_bytes()
                        .as_ref(),
                    RETRIES,
                )?;

                log::debug!("Received ACK for packet: {packet}, sequence: {sequence}");

                block.iter_mut().take(n).for_each(|b| *b = rmodem::CPMEOF);
                sequence = sequence.next();
                packet = packet.saturating_add(1);
            }
            Ok(_n) => {
                log::debug!("Reached EOF");
                is_eof = true;
            }
            Err(err) => {
                log::error!("Error reading block {sequence} from file: {err}");
                return Err(Error::Cli);
            }
        }
    }

    Ok(())
}

#[cfg(feature = "1k")]
fn send_xmodem_1k<R: Read + Seek>(
    port: &mut serialport::TTYPort,
    reader: &mut io::BufReader<R>,
) -> Result<()> {
    let mut block = [rmodem::CPMEOF; OneKData::LEN];
    let n = reader.read(&mut block).map_err(|err| {
        log::error!("Error reading next block from the file: {err}");
        Error::Cli
    })?;

    log::debug!("Sending XMODEM-1K, packet: 1, sequence: 1");

    match send_packet(
        port,
        XmodemOneKPacket::new()
            .with_data(OneKData::from_block(&block[..n]))
            .into_bytes()
            .as_ref(),
    )? {
        Control::Nak => {
            // fall-back to XMODEM/CRC, rewind read bytes from the input file
            reader.seek(io::SeekFrom::Start(0)).map_err(|err| {
                log::error!("Error seeking to start of input file: {err}");
                Error::Cli
            })?;
            send_xmodem_crc(port, reader)
        }
        Control::Ack => {
            log::debug!("Received ACK for packet: 1, sequence: 1");
            send_xmodem_1k_inner(port, reader)
        }
        control => {
            log::error!("Unexpected control byte: {control:?}");
            Err(Error::Cli)
        }
    }
}

#[cfg(all(not(feature = "1k"), feature = "crc"))]
fn send_xmodem_1k<R: Read + Seek>(
    port: &mut serialport::TTYPort,
    reader: &mut io::BufReader<R>,
) -> Result<()> {
    send_xmodem_crc(port, reader)
}

#[cfg(all(not(feature = "1k"), not(feature = "crc")))]
fn send_xmodem_1k<R: Read + Seek>(
    port: &mut serialport::TTYPort,
    reader: &mut io::BufReader<R>,
) -> Result<()> {
    send_xmodem(port, reader)
}

#[cfg(feature = "1k")]
fn send_xmodem_1k_inner<R: Read + Seek>(
    port: &mut serialport::TTYPort,
    reader: &mut io::BufReader<R>,
) -> Result<()> {
    let mut block = [rmodem::CPMEOF; OneKData::LEN];
    let mut sequence = Sequence::from(2u8);
    let mut packet = 1usize;
    let mut is_eof = false;

    while !is_eof {
        match reader.read(&mut block) {
            Ok(n) if n > 0 => {
                log::debug!("Sending XMODEM-1K, packet: {packet}, sequence: {sequence}");

                retry_packet(
                    port,
                    XmodemOneKPacket::new()
                        .with_sequence(sequence)
                        .with_data(OneKData::from_block(&block[..n]))
                        .into_bytes()
                        .as_ref(),
                    RETRIES,
                )?;

                log::debug!("Received ACK for packet: {packet}, sequence: {sequence}");

                block.iter_mut().take(n).for_each(|b| *b = rmodem::CPMEOF);
                sequence = sequence.next();
                packet = packet.saturating_add(1);
            }
            Ok(_n) => {
                log::debug!("Reached EOF");
                is_eof = true;
            }
            Err(err) => {
                log::error!("Error reading block {sequence} from file: {err}");
                return Err(Error::Cli);
            }
        }
    }

    Ok(())
}

fn retry_packet(port: &mut serialport::TTYPort, packet: &[u8], retries: usize) -> Result<()> {
    for retry in 1..=retries {
        log::debug!("Sending packet to device, attempt: {retry}");

        match send_packet(port, packet)? {
            Control::Ack => return Ok(()),
            Control::Nak => (),
            _ => return Err(Error::Cli),
        }
    }

    log::error!("Resend failed after {retries} retries");
    Err(Error::Cli)
}

fn send_packet(port: &mut serialport::TTYPort, packet: &[u8]) -> Result<Control> {
    port.write_all(packet).map_err(|err| {
        log::error!("Error writing next packet: {err}");
        Error::Cli
    })?;

    let mut res = [0u8; 1];
    port.read_exact(&mut res).map_err(|err| {
        log::error!("Error reading next response: {err}");
        Error::Cli
    })?;

    match Control::try_from(res[0])? {
        ctrl @ Control::Nak | ctrl @ Control::Ack => Ok(ctrl),
        ctrl => {
            log::error!("Unexpected control byte: {ctrl:?}");
            Err(Error::Cli)
        }
    }
}

pub fn cli_main() -> Result<()> {
    let mut builder = env_logger::Builder::new();
    builder.filter_level(log::LevelFilter::Info);
    builder.parse_default_env();
    builder.init();

    let args = Args::parse();

    let device = args.device;
    let baud = args.baud;
    let file_path = args.file;
    let protocol = rmodem::Protocol::try_from(args.protocol.as_str())?;

    log::debug!("Reading file from path: {file_path}");

    let file = std::fs::File::open(file_path.as_str()).map_err(|err| {
        log::error!("Error opening {file_path}: {err}");
        Error::Cli
    })?;

    let mut reader = io::BufReader::new(file);

    log::debug!("Opening serial port, path: {device}, baud: {baud}");

    let mut port = serialport::new(device.as_str(), baud)
        .data_bits(serialport::DataBits::Eight)
        .parity(serialport::Parity::None)
        .stop_bits(serialport::StopBits::One)
        .timeout(time::Duration::from_millis(TIMEOUT_MS))
        .open_native()
        .map_err(|err| {
            log::error!("Error opening serial device: {err}");
            Error::Cli
        })?;

    let mut idle = [0u8; 1];

    log::debug!("Reading XMODEM idle byte...");
    port.read_exact(&mut idle).map_err(|err| {
        log::error!("Error reading IDLE byte from the receiver: {err}");
        Error::Cli
    })?;

    log::debug!("Read XMODEM idle byte: {:#x}", idle[0]);

    match Control::try_from(idle[0])? {
        Control::Nak => send_xmodem(&mut port, &mut reader),
        Control::Idle => match protocol {
            rmodem::Protocol::Xmodem => send_xmodem(&mut port, &mut reader),
            rmodem::Protocol::Crc => send_xmodem_crc(&mut port, &mut reader),
            rmodem::Protocol::OneK => send_xmodem_1k(&mut port, &mut reader),
        },
        control => {
            log::error!("Unexpected initialization control byte: {control:?}");
            Err(Error::Cli)
        }
    }?;

    log::debug!("Sending EOT");
    retry_packet(&mut port, &[rmodem::EOT], RETRIES)?;

    log::info!("XMODEM transmission successful");

    Ok(())
}