librqbit_peer_protocol/extended/
handshake.rs

1use std::{collections::HashMap, net::IpAddr};
2
3use buffers::{ByteBuf, ByteBufT};
4use bytes::Bytes;
5use clone_to_owned::CloneToOwned;
6use serde::{Deserialize, Serialize};
7
8use crate::{
9    EXTENDED_UT_METADATA_KEY, EXTENDED_UT_PEX_KEY, MY_EXTENDED_UT_METADATA, MY_EXTENDED_UT_PEX,
10};
11
12use super::{PeerExtendedMessageIds, PeerIP4, PeerIP6, PeerIPAny};
13
14#[derive(Deserialize, Serialize, Debug, Default)]
15pub struct ExtendedHandshake<ByteBuf: ByteBufT> {
16    #[serde(bound(deserialize = "ByteBuf: From<&'de [u8]>"))]
17    pub m: HashMap<ByteBuf, u8>,
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub p: Option<u32>,
20    #[serde(skip_serializing_if = "Option::is_none")]
21    pub v: Option<ByteBuf>,
22    #[serde(skip_serializing_if = "Option::is_none")]
23    pub yourip: Option<PeerIPAny>,
24    #[serde(skip_serializing_if = "Option::is_none")]
25    pub ipv6: Option<PeerIP6>,
26    #[serde(skip_serializing_if = "Option::is_none")]
27    pub ipv4: Option<PeerIP4>,
28    #[serde(skip_serializing_if = "Option::is_none")]
29    pub reqq: Option<u32>,
30    #[serde(skip_serializing_if = "Option::is_none")]
31    pub metadata_size: Option<u32>,
32    #[serde(skip_serializing_if = "Option::is_none")]
33    pub complete_ago: Option<i32>,
34    #[serde(skip_serializing_if = "Option::is_none")]
35    pub upload_only: Option<u32>,
36}
37
38impl ExtendedHandshake<ByteBuf<'static>> {
39    pub fn new() -> Self {
40        let mut features = HashMap::new();
41        features.insert(ByteBuf(EXTENDED_UT_METADATA_KEY), MY_EXTENDED_UT_METADATA);
42        features.insert(ByteBuf(EXTENDED_UT_PEX_KEY), MY_EXTENDED_UT_PEX);
43        Self {
44            m: features,
45            ..Default::default()
46        }
47    }
48}
49
50impl<ByteBuf> ExtendedHandshake<ByteBuf>
51where
52    ByteBuf: ByteBufT,
53{
54    fn get_msgid(&self, msg_type: &[u8]) -> Option<u8> {
55        self.m.get(msg_type).copied()
56    }
57
58    pub fn ut_metadata(&self) -> Option<u8> {
59        self.get_msgid(EXTENDED_UT_METADATA_KEY)
60    }
61
62    pub fn ut_pex(&self) -> Option<u8> {
63        self.get_msgid(EXTENDED_UT_PEX_KEY)
64    }
65
66    pub fn peer_extended_messages(&self) -> PeerExtendedMessageIds {
67        PeerExtendedMessageIds {
68            ut_metadata: self.ut_metadata(),
69            ut_pex: self.ut_pex(),
70        }
71    }
72
73    pub fn ip_addr(&self) -> Option<IpAddr> {
74        if let Some(ref b) = self.ipv4 {
75            return Some(b.0.into());
76        }
77        if let Some(ref b) = self.ipv6 {
78            return Some(b.0.into());
79        }
80        None
81    }
82
83    pub fn port(&self) -> Option<u16> {
84        self.p.and_then(|p| u16::try_from(p).ok())
85    }
86}
87
88impl<ByteBuf> CloneToOwned for ExtendedHandshake<ByteBuf>
89where
90    ByteBuf: ByteBufT,
91    <ByteBuf as CloneToOwned>::Target: ByteBufT,
92{
93    type Target = ExtendedHandshake<<ByteBuf as CloneToOwned>::Target>;
94
95    fn clone_to_owned(&self, within_buffer: Option<&Bytes>) -> Self::Target {
96        ExtendedHandshake {
97            m: self.m.clone_to_owned(within_buffer),
98            p: self.p,
99            v: self.v.clone_to_owned(within_buffer),
100            yourip: self.yourip,
101            ipv6: self.ipv6,
102            ipv4: self.ipv4,
103            reqq: self.reqq,
104            metadata_size: self.metadata_size,
105            complete_ago: self.complete_ago,
106            upload_only: self.upload_only,
107        }
108    }
109}