123 lines
3.0 KiB
Rust
123 lines
3.0 KiB
Rust
// A crummy ad-hoc TLV cause I'm lazy and these are fun
|
|
|
|
#[derive (Debug, PartialEq)]
|
|
pub enum Message {
|
|
IAmOnline {
|
|
peer_id: u32,
|
|
},
|
|
IHaveMessage {
|
|
peer_id: u32,
|
|
msg_index: u32,
|
|
body: Vec <u8>,
|
|
},
|
|
IAcknowledgeYourMessage {
|
|
msg_index: u32,
|
|
},
|
|
}
|
|
|
|
#[derive (Debug, thiserror::Error)]
|
|
pub enum Error {
|
|
#[error ("Bad magic number")]
|
|
BadMagic,
|
|
#[error ("Data is too big to encode / decode")]
|
|
DataTooBig,
|
|
#[error ("Required field missing while decoding")]
|
|
MissingField,
|
|
#[error ("No type in TLV")]
|
|
NoType,
|
|
#[error ("Unrecognized TLV type")]
|
|
UnrecognizedType,
|
|
}
|
|
|
|
const MAGIC: [u8; 4] = [55, 11, 101, 38];
|
|
|
|
pub fn decode (src: &[u8]) -> Result <Message, Error>
|
|
{
|
|
let (magic, src) = take_bytes (src, 4).ok_or (Error::BadMagic)?;
|
|
if magic != MAGIC {
|
|
return Err (Error::BadMagic);
|
|
}
|
|
|
|
let (msg_type, src) = take_bytes(src, 1).ok_or(Error::NoType)?;
|
|
|
|
Ok (match msg_type [0] {
|
|
1 => {
|
|
let (peer_id, src) = take_u32 (src).ok_or (Error::MissingField)?;
|
|
Message::IAmOnline { peer_id, }
|
|
},
|
|
2 => {
|
|
let (peer_id, src) = take_u32 (src).ok_or (Error::MissingField)?;
|
|
let (msg_index, src) = take_u32 (src).ok_or (Error::MissingField)?;
|
|
let (body_len, src) = take_u32 (src).ok_or (Error::MissingField)?;
|
|
let (body_bytes, src) = take_bytes(src, body_len as usize).ok_or(Error::MissingField)?;
|
|
let body = body_bytes.into ();
|
|
Message::IHaveMessage { peer_id, msg_index, body, }
|
|
},
|
|
3 => {
|
|
let (msg_index, src) = take_u32 (src).ok_or (Error::MissingField)?;
|
|
Message::IAcknowledgeYourMessage { msg_index, }
|
|
},
|
|
_ => return Err (Error::UnrecognizedType),
|
|
})
|
|
}
|
|
|
|
fn take_u32 (src: &[u8]) -> Option <(u32, &[u8])>
|
|
{
|
|
let (a, b) = take_bytes (src, 4)?;
|
|
Some ((u32::from_le_bytes([a [0], a [1], a [2], a [3]]), b))
|
|
}
|
|
|
|
fn take_bytes (src: &[u8], n: usize) -> Option <(&[u8], &[u8])>
|
|
{
|
|
if src.len () < n {
|
|
return None;
|
|
}
|
|
|
|
Some (src.split_at (n))
|
|
}
|
|
|
|
pub fn encode (msg: &Message) -> Result <Vec <u8>, Error>
|
|
{
|
|
let mut buf = Vec::from (MAGIC);
|
|
|
|
match msg {
|
|
Message::IAmOnline { peer_id } => {
|
|
buf.push (1);
|
|
buf.extend_from_slice (&peer_id.to_le_bytes());
|
|
},
|
|
Message::IHaveMessage { peer_id, msg_index, body } => {
|
|
buf.push (2);
|
|
buf.extend_from_slice (&peer_id.to_le_bytes());
|
|
buf.extend_from_slice (&msg_index.to_le_bytes());
|
|
|
|
let body_len = u32::try_from (body.len ()).or (Err (Error::DataTooBig))?;
|
|
buf.extend_from_slice (&body_len.to_le_bytes());
|
|
buf.extend_from_slice(body.as_slice());
|
|
},
|
|
Message::IAcknowledgeYourMessage { msg_index } => {
|
|
buf.push (3);
|
|
buf.extend_from_slice (&msg_index.to_le_bytes());
|
|
},
|
|
}
|
|
|
|
Ok (buf)
|
|
}
|
|
|
|
#[cfg (test)]
|
|
mod test {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn roundtrip () {
|
|
for msg in [
|
|
Message::IAmOnline { peer_id: 93 },
|
|
Message::IHaveMessage { peer_id: 93, msg_index: 1000, body: Vec::from (b":V".as_slice()) },
|
|
Message::IAcknowledgeYourMessage { msg_index: 1000 },
|
|
] {
|
|
let encoded = encode (&msg).unwrap ();
|
|
let roundtripped = decode (&encoded).unwrap ();
|
|
assert_eq!(roundtripped, msg);
|
|
}
|
|
}
|
|
}
|