plabble-codec 0.1.0

use crate::{
    abstractions::{
        Serializable,
        SerializationError::{self, InvalidFormat, MissingInfo},
        SerializationInfo, KEY_SIZE, SIGNATURE_SIZE, TYPE_APPEND, TYPE_CONNECT, TYPE_CREATE,
        TYPE_ERROR, TYPE_PUT, TYPE_REQUEST, TYPE_SUBSCRIBE, TYPE_UNSUBSCRIBE, TYPE_WIPE,
    },
    codec::{
        common::{assert_len, SlotBody},
        objects::Certificate,
        ptp_packet::PtpBody,
    },
};

/// The body of a response packet
///
/// # Variants
///
/// * `CONNECT` - the response to a connect request
/// * `CREATE` - the response to a create request
/// * `PUT` - the response to a put request
/// * `APPEND` - the response to an append request
/// * `WIPE` - the response to a wipe request
/// * `REQUEST` - the response to a request request
/// * `SUBSCRIBE` - the response to a subscribe request
/// * `UNSUBSCRIBE` - the response to an unsubscribe request
pub enum ResponseBody {
    CONNECT {
        pub_key: [u8; KEY_SIZE],
        signature: [u8; SIGNATURE_SIZE],
        certificates: Vec<Certificate>,
    },
    CREATE,
    PUT,
    APPEND,
    WIPE,
    REQUEST {
        slots: SlotBody,
    },
    SUBSCRIBE {
        slots: Option<SlotBody>,
    },
    UNSUBSCRIBE,
    ERROR(u8, String),
}

impl PtpBody for ResponseBody {
    fn packet_type(&self) -> u8 {
        match self {
            ResponseBody::CONNECT { .. } => TYPE_CONNECT,
            ResponseBody::CREATE => TYPE_CREATE,
            ResponseBody::PUT => TYPE_PUT,
            ResponseBody::APPEND => TYPE_APPEND,
            ResponseBody::WIPE => TYPE_WIPE,
            ResponseBody::REQUEST { .. } => TYPE_REQUEST,
            ResponseBody::SUBSCRIBE { .. } => TYPE_SUBSCRIBE,
            ResponseBody::UNSUBSCRIBE => TYPE_UNSUBSCRIBE,
            ResponseBody::ERROR(..) => TYPE_ERROR,
        }
    }
}

impl Serializable for ResponseBody {
    fn size(&self) -> usize {
        match self {
            ResponseBody::CONNECT {
                pub_key: _,
                signature: _,
                certificates,
            } => {
                KEY_SIZE
                    + SIGNATURE_SIZE
                    + certificates.iter().map(|c| c.size()).sum::<usize>()
                    + certificates.len()
            }
            ResponseBody::REQUEST { slots } => slots.size(),
            ResponseBody::SUBSCRIBE { slots: Some(slots) } => slots.size(),
            ResponseBody::SUBSCRIBE { slots: None } => 0,
            ResponseBody::ERROR(_, e) => 1 + e.len(),
            _ => 0,
        }
    }

    fn get_bytes(&self) -> Vec<u8> {
        let mut buff = Vec::new();
        match self {
            ResponseBody::CONNECT {
                pub_key,
                signature,
                certificates,
            } => {
                buff.extend_from_slice(pub_key);
                buff.extend_from_slice(signature);
                for c in certificates.iter() {
                    let mut bytes = c.get_bytes();
                    buff.push(bytes.len() as u8);
                    buff.append(&mut bytes);
                }
            }
            ResponseBody::REQUEST { slots } => {
                buff = slots.get_bytes();
            }
            ResponseBody::SUBSCRIBE { slots: Some(slots) } => {
                buff = slots.get_bytes();
            }
            ResponseBody::ERROR(code, err) => {
                buff.push(*code);
                buff.extend_from_slice(err.as_bytes());
            }
            _ => (),
        }

        buff
    }

    fn from_bytes(data: &[u8], info: Option<SerializationInfo>) -> Result<Self, SerializationError>
    where
        Self: Sized,
    {
        if let Some(SerializationInfo::PacketType(packet_type)) = info {
            match packet_type {
                TYPE_CONNECT => {
                    let len = KEY_SIZE + SIGNATURE_SIZE;
                    assert_len(data, len)?;
                    let mut pub_key = [0u8; KEY_SIZE];
                    pub_key.copy_from_slice(&data[..KEY_SIZE]);
                    let mut signature = [0u8; SIGNATURE_SIZE];
                    signature.copy_from_slice(&data[KEY_SIZE..len]);

                    let mut certificates = Vec::new();
                    let mut bytes_read = len;
                    while bytes_read != data.len() {
                        assert_len(data, bytes_read + 1)?;
                        let cert_len = data[bytes_read] as usize;
                        bytes_read += 1;
                        assert_len(data, cert_len)?;
                        certificates.push(Certificate::from_bytes(
                            &data[bytes_read..(bytes_read + cert_len)],
                            None,
                        )?);
                        bytes_read += cert_len;
                    }

                    Ok(ResponseBody::CONNECT {
                        pub_key,
                        signature,
                        certificates,
                    })
                }
                TYPE_CREATE => Ok(ResponseBody::CREATE),
                TYPE_PUT => Ok(ResponseBody::PUT),
                TYPE_APPEND => Ok(ResponseBody::APPEND),
                TYPE_WIPE => Ok(ResponseBody::WIPE),
                TYPE_REQUEST => Ok(ResponseBody::REQUEST {
                    slots: SlotBody::from_bytes(data, None)?,
                }),
                TYPE_SUBSCRIBE => Ok(ResponseBody::SUBSCRIBE {
                    slots: if !data.is_empty() {
                        Some(SlotBody::from_bytes(data, None)?)
                    } else {
                        None
                    },
                }),
                TYPE_UNSUBSCRIBE => Ok(ResponseBody::UNSUBSCRIBE),
                TYPE_ERROR => {
                    assert_len(data, 1)?;
                    let res = String::from_utf8_lossy(&data[1..]).to_string();
                    Ok(ResponseBody::ERROR(data[0], res))
                }
                _ => Err(InvalidFormat(format!(
                    "Packet type {} is not supported",
                    packet_type
                ))),
            }
        } else {
            Err(MissingInfo(String::from("Missing info parameter")))
        }
    }
}

#[cfg(test)]
mod parse_test {
    use time_macros::datetime;

    use crate::abstractions::Serializable;

    use super::*;

    #[test]
    fn can_parse_connect_wo_certs() {
        let connect = &[
            1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
            0, 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0,
        ];
        let connect =
            ResponseBody::from_bytes(connect, Some(SerializationInfo::PacketType(0))).unwrap();

        assert_eq!(96, connect.size());
        assert_eq!(0, connect.packet_type());

        if let ResponseBody::CONNECT {
            pub_key,
            signature,
            certificates,
        } = connect
        {
            assert_eq!(
                &[
                    1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6,
                    7, 8, 9, 0, 1, 2
                ],
                &pub_key
            );
            assert_eq!(&[0u8; 64], &signature);
            assert_eq!(0, certificates.len());
        } else {
            panic!("Not a connect");
        }
    }

    #[test]
    fn can_parse_connect_with_2_certs() {
        let mut connect = Vec::new();
        connect.extend_from_slice(&[1u8; 32]);
        connect.extend_from_slice(&[0u8; 64]);
        let mut cert1 = Certificate {
            public_key: [2u8; 32],
            signature: [3u8; 64],
            valid_from: datetime!(2020-01-01 00:00:00 UTC),
            valid_to: datetime!(2020-01-01 00:00:00 UTC),
            domain_or_ip: String::from("cert.plabble.org"),
        }
        .get_bytes();
        let mut cert2 = Certificate {
            public_key: [3u8; 32],
            signature: [4u8; 64],
            valid_from: datetime!(2020-01-01 00:00:00 UTC),
            valid_to: datetime!(2020-01-01 00:00:00 UTC),
            domain_or_ip: String::from("cert2.plabble.org"),
        }
        .get_bytes();
        connect.push(cert1.len() as u8);
        connect.append(&mut cert1);
        connect.push(cert2.len() as u8);
        connect.append(&mut cert2);

        let conn = ResponseBody::from_bytes(&connect, Some(SerializationInfo::PacketType(0)));
        let conn = conn.unwrap();
        assert_eq!(connect.len(), conn.size());

        if let ResponseBody::CONNECT {
            pub_key,
            signature,
            certificates,
        } = conn
        {
            assert_eq!(&[1u8; 32], &pub_key);
            assert_eq!(&[0u8; 64], &signature);
            assert_eq!(2, certificates.len());
            assert_eq!(&[2u8; 32], &certificates[0].public_key);
            assert_eq!("cert.plabble.org", &certificates[0].domain_or_ip);
            assert_eq!(&[3u8; 32], &certificates[1].public_key);
            assert_eq!("cert2.plabble.org", &certificates[1].domain_or_ip);
        } else {
            panic!("Not a connect");
        }
    }

    #[test]
    fn can_parse_subscribe_without_slots() {
        let bytes = &[];
        let subscribe =
            ResponseBody::from_bytes(bytes, Some(SerializationInfo::PacketType(6))).unwrap();

        assert_eq!(0, subscribe.size());
        if let ResponseBody::SUBSCRIBE { slots } = subscribe {
            assert_eq!(slots, None);
        } else {
            panic!("Not a subscribe");
        }
    }

    #[test]
    fn can_deserialize_error() {
        let bytes = &[7, 97, 110, 32, 101, 114, 114, 111, 114, 33];
        let error =
            ResponseBody::from_bytes(bytes, Some(SerializationInfo::PacketType(15))).unwrap();

        assert_eq!(10, error.size());
        if let ResponseBody::ERROR(code, e) = error {
            assert_eq!("an error!", &e);
            assert_eq!(7, code);
        } else {
            panic!("not an error");
        }
    }
}

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

    #[test]
    fn can_serialize_connect_wo_certificates() {
        let expected = &[
            1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24,
            25, 26, 27, 28, 29, 30, 31, 32, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
            17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38,
            39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60,
            61, 62, 63, 64,
        ];
        let mut pub_key = [0u8; 32];
        pub_key.copy_from_slice((1u8..=32).collect::<Vec<u8>>().as_slice());
        let mut signature = [0u8; 64];
        signature.copy_from_slice((1u8..=64).collect::<Vec<u8>>().as_slice());

        let connect = ResponseBody::CONNECT {
            pub_key,
            signature,
            certificates: Vec::new(),
        };
        assert_eq!(&connect.get_bytes(), expected);
    }

    #[test]
    fn can_serialize_subscribe_without_range() {
        let subscribe = ResponseBody::SUBSCRIBE { slots: None };
        assert_eq!(Vec::<u8>::new(), subscribe.get_bytes());
    }

    #[test]
    fn can_serialize_error() {
        let error = ResponseBody::ERROR(7, String::from("an error!"));
        let bytes = error.get_bytes();
        assert_eq!(vec![7, 97, 110, 32, 101, 114, 114, 111, 114, 33], bytes);
    }
}