librqbit_peer_protocol/extended/
mod.rs

1use bencode::bencode_serialize_to_writer;
2use bencode::from_bytes;
3use bencode::BencodeValue;
4use buffers::ByteBufT;
5use bytes::Bytes;
6use clone_to_owned::CloneToOwned;
7use serde::Deserialize;
8use ut_pex::UtPex;
9
10use crate::MY_EXTENDED_UT_PEX;
11
12use self::{handshake::ExtendedHandshake, ut_metadata::UtMetadata};
13
14use super::MessageDeserializeError;
15
16pub mod handshake;
17mod ip;
18
19pub use ip::{PeerIP, PeerIP4, PeerIP6, PeerIPAny};
20
21pub mod ut_metadata;
22pub mod ut_pex;
23
24use super::MY_EXTENDED_UT_METADATA;
25
26#[derive(Debug, Default)]
27pub struct PeerExtendedMessageIds {
28    pub ut_metadata: Option<u8>,
29    pub ut_pex: Option<u8>,
30}
31
32#[derive(Debug)]
33pub enum ExtendedMessage<ByteBuf: ByteBufT> {
34    Handshake(ExtendedHandshake<ByteBuf>),
35    UtMetadata(UtMetadata<ByteBuf>),
36    UtPex(UtPex<ByteBuf>),
37    Dyn(u8, BencodeValue<ByteBuf>),
38}
39
40impl<ByteBuf> CloneToOwned for ExtendedMessage<ByteBuf>
41where
42    ByteBuf: ByteBufT,
43    <ByteBuf as CloneToOwned>::Target: ByteBufT,
44{
45    type Target = ExtendedMessage<<ByteBuf as CloneToOwned>::Target>;
46
47    fn clone_to_owned(&self, within_buffer: Option<&Bytes>) -> Self::Target {
48        match self {
49            ExtendedMessage::Handshake(h) => {
50                ExtendedMessage::Handshake(h.clone_to_owned(within_buffer))
51            }
52            ExtendedMessage::Dyn(u, d) => ExtendedMessage::Dyn(*u, d.clone_to_owned(within_buffer)),
53            ExtendedMessage::UtMetadata(m) => {
54                ExtendedMessage::UtMetadata(m.clone_to_owned(within_buffer))
55            }
56            ExtendedMessage::UtPex(m) => ExtendedMessage::UtPex(m.clone_to_owned(within_buffer)),
57        }
58    }
59}
60
61impl<ByteBuf: ByteBufT> ExtendedMessage<ByteBuf> {
62    pub fn serialize(
63        &self,
64        out: &mut Vec<u8>,
65        extended_handshake_ut_metadata: &dyn Fn() -> PeerExtendedMessageIds,
66    ) -> anyhow::Result<()>
67    where
68        ByteBuf: AsRef<[u8]>,
69    {
70        match self {
71            ExtendedMessage::Dyn(msg_id, v) => {
72                out.push(*msg_id);
73                bencode_serialize_to_writer(v, out)?;
74            }
75            ExtendedMessage::Handshake(h) => {
76                out.push(0);
77                bencode_serialize_to_writer(h, out)?;
78            }
79            ExtendedMessage::UtMetadata(u) => {
80                let emsg_id = extended_handshake_ut_metadata()
81                    .ut_metadata
82                    .ok_or_else(|| {
83                        anyhow::anyhow!("need peer's handshake to serialize ut_metadata")
84                    })?;
85                out.push(emsg_id);
86                u.serialize(out);
87            }
88            ExtendedMessage::UtPex(m) => {
89                let emsg_id = extended_handshake_ut_metadata().ut_pex.ok_or_else(|| {
90                    anyhow::anyhow!(
91                        "need peer's handshake to serialize ut_pex, or peer does't support ut_pex"
92                    )
93                })?;
94                out.push(emsg_id);
95                bencode_serialize_to_writer(m, out)?;
96            }
97        }
98        Ok(())
99    }
100
101    pub fn deserialize<'a>(mut buf: &'a [u8]) -> Result<Self, MessageDeserializeError>
102    where
103        ByteBuf: Deserialize<'a> + From<&'a [u8]>,
104    {
105        let emsg_id = buf.first().copied().ok_or_else(|| {
106            MessageDeserializeError::Other(anyhow::anyhow!(
107                "cannot deserialize extended message: can't read first byte"
108            ))
109        })?;
110
111        buf = buf.get(1..).ok_or_else(|| {
112            MessageDeserializeError::Other(anyhow::anyhow!(
113                "cannot deserialize extended message: buffer empty"
114            ))
115        })?;
116
117        match emsg_id {
118            0 => Ok(ExtendedMessage::Handshake(from_bytes(buf)?)),
119            MY_EXTENDED_UT_METADATA => {
120                Ok(ExtendedMessage::UtMetadata(UtMetadata::deserialize(buf)?))
121            }
122            MY_EXTENDED_UT_PEX => Ok(ExtendedMessage::UtPex(from_bytes(buf)?)),
123            _ => Ok(ExtendedMessage::Dyn(emsg_id, from_bytes(buf)?)),
124        }
125    }
126}
librqbit_peer_protocol/extended/mod.rs

librqbit_peer_protocol/extended/
mod.rs
