use bencode::bencode_serialize_to_writer; use bencode::BencodeDeserializer; use buffers::ByteBufT; use bytes::Bytes; use clone_to_owned::CloneToOwned; use serde::Deserialize; use serde::Serialize; use std::io::Write; use crate::MessageDeserializeError; #[derive(Debug)] pub enum UtMetadata { Request(u32), Data { piece: u32, total_size: u32, data: ByteBuf, }, Reject(u32), } impl CloneToOwned for UtMetadata { type Target = UtMetadata<::Target>; fn clone_to_owned(&self, within_buffer: Option<&Bytes>) -> Self::Target { match self { UtMetadata::Request(req) => UtMetadata::Request(*req), UtMetadata::Data { piece, total_size, data, } => UtMetadata::Data { piece: *piece, total_size: *total_size, data: data.clone_to_owned(within_buffer), }, UtMetadata::Reject(piece) => UtMetadata::Reject(*piece), } } } impl UtMetadata { pub fn serialize(&self, buf: &mut Vec) where ByteBuf: AsRef<[u8]>, { #[derive(Serialize)] struct Message { msg_type: u32, piece: u32, #[serde(skip_serializing_if = "Option::is_none")] total_size: Option, } match self { UtMetadata::Request(piece) => { let message = Message { msg_type: 0, piece: *piece, total_size: None, }; bencode_serialize_to_writer(message, buf).unwrap() } UtMetadata::Data { piece, total_size, data, } => { let message = Message { msg_type: 1, piece: *piece, total_size: Some(*total_size), }; bencode_serialize_to_writer(message, buf).unwrap(); buf.write_all(data.as_ref()).unwrap(); } UtMetadata::Reject(piece) => { let message = Message { msg_type: 2, piece: *piece, total_size: None, }; bencode_serialize_to_writer(message, buf).unwrap(); } } } pub fn deserialize<'a>(buf: &'a [u8]) -> Result where ByteBuf: From<&'a [u8]>, { let mut de = BencodeDeserializer::new_from_buf(buf); #[derive(Deserialize)] struct Message { msg_type: u32, piece: u32, total_size: Option, } let message = Message::deserialize(&mut de).map_err(|e| MessageDeserializeError::Other(e.into()))?; let remaining = de.into_remaining(); match message.msg_type { // request 0 => { if !remaining.is_empty() { return Err(MessageDeserializeError::Other(anyhow::anyhow!( "trailing bytes when decoding UtMetadata" ))); } Ok(UtMetadata::Request(message.piece)) } // data 1 => { let total_size = message.total_size.ok_or_else(|| { MessageDeserializeError::Other(anyhow::anyhow!( "expected key total_size to be present in UtMetadata \"data\" message" )) })?; Ok(UtMetadata::Data { piece: message.piece, total_size, data: ByteBuf::from(remaining), }) } // reject 2 => { if !remaining.is_empty() { return Err(MessageDeserializeError::Other(anyhow::anyhow!( "trailing bytes when decoding UtMetadata" ))); } Ok(UtMetadata::Reject(message.piece)) } other => Err(MessageDeserializeError::Other(anyhow::anyhow!( "unrecognized ut_metadata message type {}", other ))), } } }