parent
7b4eeea12c
commit
5843432bfe
@ -0,0 +1,122 @@
|
||||
// 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);
|
||||
}
|
||||
}
|
||||
}
|
Loading…
Reference in new issue