CS3700-project3/hptp/src/msg.rs

use super::seg::{SegData, SegIdx};
pub use super::seg::{DOWN_HEADER_SIZE, MAX_TOTAL_PACKET_SIZE, UP_HEADER_SIZE, MAX_SEG_SIZE};
use byteorder::ByteOrder;

#[derive(Clone)]
pub enum UpMsg {
    Data { payload: SegData, seg_idx: SegIdx },
}

#[derive(Clone)]
pub enum DownMsg {
    /// `idxs` must be distinct and in increasing order.
    Ack { idxs: Vec<SegIdx> },
}

#[derive(Error, Debug)]
#[error("deserialization failed; malformed packet")]
pub struct DesError;

pub trait SerDes: Sized {
    // `buf.len()` must be <= `MAX_TOTAL_PACKET_SIZE`
    fn des(buf: &[u8]) -> Result<Self, DesError>;

    // `buf.len()` must be >= `MAX_TOTAL_PACKET_SIZE`
    fn ser_to(&self, buf: &mut [u8]) -> usize;
}

type BO = byteorder::LE;

const LAST_SEG_MASK: u32 = 1 << 31;
const MEOW_CODED_MASK: u32 = 1 << 30;
const CUT_MASK: u32 = 1 << 29;

impl SerDes for UpMsg {
    fn des(buf: &[u8]) -> Result<Self, DesError> {
        if buf.len() < UP_HEADER_SIZE {
            Err(DesError)
        } else {
            let hdr = BO::read_u32(&buf[0..4]);
            Ok(UpMsg::Data {
                seg_idx: hdr & !LAST_SEG_MASK & !MEOW_CODED_MASK & !CUT_MASK,
                payload: SegData {
                    bytes: buf[4..].into(),
                    is_last_segment: (hdr & LAST_SEG_MASK) != 0,
                    is_meow_encoded: (hdr & MEOW_CODED_MASK) != 0,
                    is_cut: (hdr & CUT_MASK) != 0,
                },
            })
        }
    }

    fn ser_to(&self, buf: &mut [u8]) -> usize {
        match self {
            UpMsg::Data {
                payload:
                    SegData {
                        bytes,
                        is_last_segment,
                        is_meow_encoded,
                        is_cut,
                    },
                seg_idx,
            } => {
                let hdr = *seg_idx
                    | if *is_last_segment { LAST_SEG_MASK } else { 0 }
                    | if *is_meow_encoded { MEOW_CODED_MASK } else { 0 }
                    | if *is_cut { CUT_MASK } else { 0 };
                BO::write_u32(&mut buf[0..4], hdr);
                let len = bytes.len();
                buf[4..4 + len].copy_from_slice(&bytes[..]);
                4 + len
            }
        }
    }
}

impl SerDes for DownMsg {
    fn des(buf: &[u8]) -> Result<Self, DesError> {
        if buf.len() < DOWN_HEADER_SIZE {
            Err(DesError)
        } else {
            let first_unacked = BO::read_u32(&buf[0..4]);
            let mut idxs = vec![];
            for i in 0..first_unacked {
                idxs.push(i);
            }
            for (i, b) in buf.iter().skip(4).cloned().enumerate() {
                for j in 0..8 {
                    if b & (1 << j) != 0 {
                        idxs.push((i * 8 + j + (first_unacked as usize) + 1) as SegIdx);
                    }
                }
            }
            Ok(DownMsg::Ack { idxs })
        }
    }

    fn ser_to(&self, buf: &mut [u8]) -> usize {
        match self {
            DownMsg::Ack { idxs } => {
                let mut first_unacked: u32 = 0;
                for idx in idxs.iter().cloned() {
                    if idx == first_unacked {
                        first_unacked += 1;
                    } else {
                        break;
                    }
                }
                BO::write_u32(&mut buf[0..4], first_unacked);
                let mut max_i = 0;
                for seg_idx in idxs.iter().cloned() {
                    if seg_idx < first_unacked {
                        continue;
                    }
                    let seg_idx_delta: usize = (seg_idx - first_unacked - 1) as usize;
                    let i = seg_idx_delta / 8;
                    let j = (seg_idx_delta % 8) as u8;
                    buf[i + 4] |= 1 << j;
                    max_i = std::cmp::max(i, max_i);
                }
                max_i + 4 + 1
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn serdes_downmsg(init_idxs: &Vec<SegIdx>, expected_hdr: u32) {
        let mut buf = [0u8; MAX_SEG_SIZE];
        DownMsg::Ack{idxs: init_idxs.clone()}.ser_to(&mut buf);
        assert_eq!(BO::read_u32(&buf[0..4]), expected_hdr);
        let des = DownMsg::des(&buf).unwrap();
        match des {
            DownMsg::Ack { idxs } => {
                assert_eq!(*init_idxs, idxs);
            }
        }
    }

    #[test]
    fn test_serdes_downmsg() {
        let init_idxs: Vec<SegIdx> = vec![0, 1, 2, 3];
        serdes_downmsg(&init_idxs, 4);
    }

    #[test]
    fn test_serdes_downmsg_empty() {
        let init_idxs: Vec<SegIdx> = vec![];
        serdes_downmsg(&init_idxs, 0);
    }

    #[test]
    fn test_serdes_downmsg_missing0() {
        let init_idxs: Vec<SegIdx> = vec![1, 2, 4];
        serdes_downmsg(&init_idxs, 0);
    }

    #[test]
    fn test_serdes_downmsg_missing1() {
        let init_idxs: Vec<SegIdx> = vec![0, 2, 7];
        serdes_downmsg(&init_idxs, 1);
    }

    #[test]
    fn test_serdes_downmsg_missing2() {
        let init_idxs: Vec<SegIdx> = vec![0, 1, 2, 3, 4, 7, 8, 10];
        serdes_downmsg(&init_idxs, 5);
    }
}