lpc55prog/src/packet.rs

use nom::bytes::complete::take;
use nom::error::{Error, ErrorKind};
use nom::multi::many0;
use nom::number::complete;
use nom::IResult;
use nom::{Err, InputTake};

use nom::sequence::Tuple;
use num_enum::{FromPrimitive, IntoPrimitive, TryFromPrimitive};
use std::convert::TryFrom;
use std::fmt::Display;

use bitflags::bitflags;
use byteorder::{LittleEndian, WriteBytesExt};

use strum_macros::{EnumIter, IntoStaticStr};

#[repr(u8)]
#[derive(Copy, Clone, TryFromPrimitive, IntoPrimitive, Debug)]
pub enum ReportId {
    CommandOut = 1,
    DataOut = 2,
    CommandIn = 3,
    DataIn = 4,
}

#[repr(u8)]
#[derive(Copy, Clone, TryFromPrimitive, IntoPrimitive, Debug, EnumIter)]
pub enum CommandTag {
    FlashEraseAll = 0x01,
    FlashEraseRegion = 0x02,
    ReadMemory = 0x03,
    WriteMemory = 0x04,
    FillMemory = 0x05,
    GetProperty = 0x07,
    ReceiveSbFile = 0x08,
    Execute = 0x09,
    Call = 0x0a,
    Reset = 0x0b,
    SetProperty = 0x0c,
    ConfigureMemory = 0x11,
    KeyProvision = 0x15,
}

#[repr(u8)]
#[derive(Copy, Clone, TryFromPrimitive, IntoPrimitive, Debug)]
pub enum ResponseTag {
    GenericResponse = 0xa0,
    ReadMemoryResponse = 0xa3,
    GetPropertyResponse = 0xa7,
    FlashReadOnceResponse = 0xaf,
    KeyProvisionResponse = 0xb5,
}
#[repr(u32)]
#[derive(Copy, Clone, PartialEq, IntoStaticStr, FromPrimitive, IntoPrimitive, Debug)]
pub enum StatusCode {
    Success = 0,
    StatusFlashAlignmentError = 101,
    StatusFlashEccError = 116,
    StatusMemoryRangeInvalid = 10200,
    StatusMemoryBlankPageReadDisallowed = 10211,
    #[num_enum(catch_all)]
    Unknown(u32),
}

impl Display for StatusCode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            StatusCode::Success => f.write_str("Success"),
            StatusCode::Unknown(code) => f.write_fmt(format_args!("Error({})", code)),
            code => {
                let name: &'static str = code.into();
                f.write_fmt(format_args!("{}({})", name, u32::from(code.to_owned())))
            }
        }
    }
}

#[derive(Debug)]
pub enum ResponseParameters {
    GenericResponse(GenericResponseParams),
    ReadMemoryResponse(ReadMemoryResponseParams),
    GetPropertyResponse(GetPropertyResponseParams),
    FlashReadOnceResponse(FlashReadOnceResponseParams),
    KeyProvisionResponse(KeyProvisionResponseParams),
}
#[derive(Debug)]
pub struct GenericResponseParams {
    pub status: StatusCode,
    pub command: CommandTag,
}
#[derive(Debug)]
pub struct ReadMemoryResponseParams {
    pub status: StatusCode,
    pub data_bytes: u32,
}
#[derive(Debug)]
pub struct GetPropertyResponseParams {
    pub status: StatusCode,
    pub properties: Vec<u32>,
}
#[derive(Debug)]
pub struct FlashReadOnceResponseParams {
    status: StatusCode,
    byte_count: u32,
    read_data: [u8; 20],
}
#[derive(Debug)]
pub struct KeyProvisionResponseParams {
    status: StatusCode,
    data_bytes: u32,
}

bitflags! {
    pub struct CommandFlags: u8 {
        const DATA_FOLLOWS = 0b00000001;
    }
}

impl Default for CommandFlags {
    fn default() -> Self {
        CommandFlags { bits: 0 }
    }
}

#[derive(Debug, Clone)]
pub struct CommandPacket {
    pub tag: CommandTag,
    pub flags: CommandFlags,
    pub reserved: u8,
    pub param_count: u8,
    pub params: Vec<u32>,
}

#[derive(Debug)]
pub struct ResponsePacket {
    pub tag: ResponseTag,
    pub flags: CommandFlags,
    pub reserved: u8,
    pub param_count: u8,
    pub params: ResponseParameters,
}

#[derive(Debug)]
pub struct DataPacket {
    pub bytes: Vec<u8>,
}

pub struct UsbPacket {
    report_id: ReportId,
    _padding: u8,
    packet_length: u16,
    packet: Packet,
}

#[derive(Debug)]
//TODO: Refactor this to use traits and borrows
pub enum Packet {
    CommandPacket(CommandPacket),
    ResponsePacket(ResponsePacket),
    DataPacket(DataPacket),
}

fn status_code(input: &[u8]) -> IResult<&[u8], StatusCode> {
    let (input, code) = complete::le_u32(input)?;
    Ok((input, StatusCode::from(code)))
}

fn response_tag(input: &[u8]) -> IResult<&[u8], ResponseTag> {
    let res = take(1u8)(input);
    match res {
        Ok(ir) => match ResponseTag::try_from(ir.1[0]) {
            Ok(val) => Ok((ir.0, val)),
            Err(e) => {
                println!("{:?}", e);
                unimplemented!()
            }
        },
        Err(e) => Err(e),
    }
}

fn command_tag(input: &[u8]) -> IResult<&[u8], CommandTag> {
    let res = take(1u8)(input);
    match res {
        Ok(ir) => match CommandTag::try_from(ir.1[0]) {
            Ok(val) => Ok((ir.0, val)),
            Err(e) => {
                println!("{:?}", e);
                unimplemented!()
            }
        },
        Err(e) => Err(e),
    }
}

fn generic_response_params(input: &[u8]) -> IResult<&[u8], ResponseParameters> {
    let (input, (status, command)) = (status_code, complete::u8).parse(input)?;
    Ok((
        input,
        ResponseParameters::GenericResponse(GenericResponseParams {
            status,
            command: CommandTag::try_from(command).unwrap(),
        }),
    ))
}

fn get_property_response_params(input: &[u8]) -> IResult<&[u8], ResponseParameters> {
    let (input, (status, properties)) = (status_code, many0(complete::le_u32)).parse(input)?;
    Ok((
        input,
        ResponseParameters::GetPropertyResponse(GetPropertyResponseParams { status, properties }),
    ))
}

fn get_read_memory_response_params(input: &[u8]) -> IResult<&[u8], ResponseParameters> {
    let (input, (status, properties)) = (status_code, many0(complete::le_u32)).parse(input)?;
    assert!(properties.len() == 1);
    Ok((
        input,
        ResponseParameters::ReadMemoryResponse(ReadMemoryResponseParams {
            status,
            data_bytes: properties[0],
        }),
    ))
}

fn report_id(input: &[u8]) -> IResult<&[u8], ReportId> {
    let (input, id) = complete::u8(input)?;
    match ReportId::try_from(id) {
        Ok(val) => Ok((input, val)),
        Err(_) => Err(Err::Error(Error::new(input, ErrorKind::Fail))),
    }
}

pub fn usb_packet(input: &[u8]) -> IResult<&[u8], Packet> {
    let (input, (reportid, _, packet_length)) =
        (report_id, complete::u8, complete::le_u16).parse(input)?;
    let (input, buf) = take(packet_length)(input)?;

    let pack = match reportid {
        ReportId::CommandOut => command_packet(buf)?,
        ReportId::DataOut => unimplemented!(),
        ReportId::CommandIn => response_packet(buf)?,
        ReportId::DataIn => data_packet(buf, packet_length)?,
    };

    Ok((input, pack.1))
}

pub fn response_packet(input: &[u8]) -> IResult<&[u8], Packet> {
    let (input, (tag, flags, reserved, param_count)) =
        (response_tag, complete::u8, complete::u8, complete::u8).parse(input)?;
    let (input, params) = match tag {
        ResponseTag::GenericResponse => generic_response_params(input)?,
        ResponseTag::GetPropertyResponse => get_property_response_params(input)?,
        ResponseTag::ReadMemoryResponse => get_read_memory_response_params(input)?,
        ResponseTag::FlashReadOnceResponse => unimplemented!(),
        ResponseTag::KeyProvisionResponse => unimplemented!(),
    };
    Ok((
        input,
        Packet::ResponsePacket(ResponsePacket {
            tag,
            flags: CommandFlags::from_bits_truncate(flags),
            reserved,
            param_count,
            params,
        }),
    ))
}

pub fn data_packet(input: &[u8], packet_length: u16) -> IResult<&[u8], Packet> {
    let data = input.take(packet_length as usize);
    Ok((
        input,
        Packet::DataPacket(DataPacket {
            bytes: data.to_vec(),
        }),
    ))
}

pub fn command_packet(input: &[u8]) -> IResult<&[u8], Packet> {
    let (input, (tag, flags, reserved, param_count)) =
        (command_tag, complete::u8, complete::u8, complete::u8).parse(input)?;
    Ok((
        input,
        Packet::CommandPacket(CommandPacket {
            tag,
            flags: CommandFlags::from_bits_truncate(flags),
            reserved,
            param_count,
            params: Vec::new(),
        }),
    ))
}

impl Packet {
    pub fn write<W: std::io::Write>(&self, w: &mut W) -> Result<(), std::io::Error> {
        //TODO: Refactor to a trait
        match self {
            Packet::CommandPacket(p) => {
                w.write(&[
                    p.tag as u8,
                    p.flags.bits() as u8,
                    p.reserved as u8,
                    p.param_count as u8,
                ])?;
                for param in &p.params {
                    w.write_u32::<LittleEndian>(*param)?;
                }
            }
            Packet::DataPacket(p) => {
                w.write(p.bytes.as_slice())?;
            }
            Packet::ResponsePacket(p) => {
                unimplemented!("Serializing response packets is not implemented");
            }
        }
        Ok(())
    }
    pub fn length(&self) -> u16 {
        let header_len = 4u16;
        match self {
            Packet::CommandPacket(p) => header_len + (p.param_count as u16) * 4,
            Packet::ResponsePacket(p) => header_len + (p.param_count as u16) * 4,
            Packet::DataPacket(p) => p.bytes.len() as u16,
        }
    }
}

impl UsbPacket {
    pub fn new(report_id: ReportId, packet: Packet) -> UsbPacket {
        UsbPacket {
            report_id,
            _padding: 0,
            packet_length: packet.length(),
            packet,
        }
    }

    pub fn length(&self) -> usize {
        self.packet_length as usize + 4
    }

    pub fn write<W: std::io::Write>(&self, w: &mut W) -> Result<(), std::io::Error> {
        w.write(&[self.report_id as u8, 0])?;
        w.write_u16::<LittleEndian>(self.packet_length)?;
        self.packet.write(w)?;
        Ok(())
    }
}