142 lines
4.3 KiB
Rust
142 lines
4.3 KiB
Rust
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<ByteBuf> {
|
|
Request(u32),
|
|
Data {
|
|
piece: u32,
|
|
total_size: u32,
|
|
data: ByteBuf,
|
|
},
|
|
Reject(u32),
|
|
}
|
|
|
|
impl<ByteBuf: CloneToOwned> CloneToOwned for UtMetadata<ByteBuf> {
|
|
type Target = UtMetadata<<ByteBuf as CloneToOwned>::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<ByteBuf: ByteBufT> UtMetadata<ByteBuf> {
|
|
pub fn serialize(&self, buf: &mut Vec<u8>)
|
|
where
|
|
ByteBuf: AsRef<[u8]>,
|
|
{
|
|
#[derive(Serialize)]
|
|
struct Message {
|
|
msg_type: u32,
|
|
piece: u32,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
total_size: Option<u32>,
|
|
}
|
|
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<Self, MessageDeserializeError>
|
|
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<u32>,
|
|
}
|
|
|
|
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
|
|
))),
|
|
}
|
|
}
|
|
}
|