p2p/network/identify/
p2p_network_identify_protocol.rs

1use super::keys_proto;
2use crate::{
3    identity::PublicKey,
4    token::{self, StreamKind},
5};
6
7use malloc_size_of_derive::MallocSizeOf;
8use multiaddr::Multiaddr;
9use quick_protobuf::{BytesReader, MessageRead, MessageWrite, Writer};
10use serde::{Deserialize, Serialize};
11
12#[derive(Clone, Debug, Serialize, Deserialize, MallocSizeOf)]
13pub struct P2pNetworkIdentify {
14    pub protocol_version: Option<String>,
15    pub agent_version: Option<String>,
16    #[ignore_malloc_size_of = "doesn't allocate"]
17    pub public_key: Option<PublicKey>,
18    #[with_malloc_size_of_func = "measurement::multiaddr_vec"]
19    pub listen_addrs: Vec<Multiaddr>,
20    #[with_malloc_size_of_func = "measurement::multiaddr_opt"]
21    pub observed_addr: Option<Multiaddr>,
22    pub protocols: Vec<token::StreamKind>,
23}
24
25#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, thiserror::Error)]
26pub enum P2pNetworkIdentifyFromMessageError {
27    #[error("cant parse protocol: {0}")]
28    UnsupportedProtocol(String),
29    #[error("unsupported public key type: {0}")]
30    UnsupportedPubKeyType(String),
31    #[error("error parsing public key: {0}")]
32    ErrorParsingPubKey(String),
33    #[error("{0}")]
34    MultiaddrError(String),
35}
36
37impl TryFrom<super::pb::Identify> for P2pNetworkIdentify {
38    type Error = P2pNetworkIdentifyFromMessageError;
39
40    fn try_from(value: super::pb::Identify) -> Result<Self, Self::Error> {
41        let protocol_version = value.protocol_version;
42        let agent_version = value.agent_version;
43
44        let public_key = match value.public_key {
45            Some(pubkey) => Some(parse_public_key(&pubkey)?),
46            None => None,
47        };
48
49        let mut listen_addrs = Vec::new();
50
51        for addr in value.listen_addrs.iter().cloned() {
52            listen_addrs.push(
53                addr.try_into()
54                    .map_err(|err: multiaddr::Error| err.to_string())
55                    .map_err(P2pNetworkIdentifyFromMessageError::MultiaddrError)?,
56            )
57        }
58
59        let observed_addr = match value.observed_addr {
60            Some(addr) => Some(
61                addr.try_into()
62                    .map_err(|err: multiaddr::Error| err.to_string())
63                    .map_err(P2pNetworkIdentifyFromMessageError::MultiaddrError)?,
64            ),
65            None => None,
66        };
67
68        let mut protocols = Vec::new();
69
70        for proto in value.protocols.iter() {
71            protocols.push(parse_protocol(proto)?)
72        }
73
74        Ok(Self {
75            protocol_version,
76            agent_version,
77            public_key,
78            listen_addrs,
79            observed_addr,
80            protocols,
81        })
82    }
83}
84
85#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, thiserror::Error)]
86pub enum P2pNetworkMessageFromIdentifyError {
87    #[error("encoding error: {0}")]
88    EncodingError(String),
89}
90
91impl P2pNetworkIdentify {
92    pub fn to_proto_message(
93        &self,
94    ) -> Result<super::pb::Identify, P2pNetworkMessageFromIdentifyError> {
95        super::pb::Identify::try_from(self)
96    }
97}
98
99impl<'a> TryFrom<&'a P2pNetworkIdentify> for super::pb::Identify {
100    type Error = P2pNetworkMessageFromIdentifyError;
101
102    fn try_from(value: &'a P2pNetworkIdentify) -> Result<Self, Self::Error> {
103        let public_key = if let Some(key) = value.public_key.as_ref() {
104            let key_bytes = key.to_bytes();
105            let pubkey = keys_proto::PublicKey {
106                Type: crate::network::identify::KeyType::Ed25519,
107                Data: key_bytes.as_ref().into(),
108            };
109            let mut buf = Vec::with_capacity(pubkey.get_size());
110            let mut writer = Writer::new(&mut buf);
111
112            pubkey
113                .write_message(&mut writer)
114                .map_err(|e| P2pNetworkMessageFromIdentifyError::EncodingError(e.to_string()))?;
115
116            Some(buf)
117        } else {
118            None
119        };
120
121        Ok(Self {
122            protocol_version: value.protocol_version.as_ref().map(|v| v.into()),
123            agent_version: value.agent_version.as_ref().map(|v| v.into()),
124            public_key,
125            listen_addrs: value.listen_addrs.iter().map(|v| v.to_vec()).collect(),
126            observed_addr: value.observed_addr.as_ref().map(|v| v.to_vec()),
127            protocols: value
128                .protocols
129                .iter()
130                .map(|v| v.name_str().into())
131                .collect(),
132        })
133    }
134}
135
136pub fn parse_public_key(key_bytes: &[u8]) -> Result<PublicKey, P2pNetworkIdentifyFromMessageError> {
137    let mut reader = BytesReader::from_bytes(key_bytes);
138
139    keys_proto::PublicKey::from_reader(&mut reader, key_bytes).map_or_else(
140        |error| {
141            Err(P2pNetworkIdentifyFromMessageError::ErrorParsingPubKey(
142                error.to_string(),
143            ))
144        },
145        |pubkey| match pubkey {
146            keys_proto::PublicKey {
147                Type: keys_proto::KeyType::Ed25519,
148                Data: data,
149            } => {
150                let bytes = data[..].try_into().or(Err(
151                    P2pNetworkIdentifyFromMessageError::ErrorParsingPubKey(format!(
152                        "invalid size {}",
153                        data.len()
154                    )),
155                ))?;
156
157                PublicKey::from_bytes(bytes).map_err(|err| {
158                    P2pNetworkIdentifyFromMessageError::ErrorParsingPubKey(err.to_string())
159                })
160            }
161            _ => Err(P2pNetworkIdentifyFromMessageError::UnsupportedPubKeyType(
162                format!("{:?}", pubkey.Type),
163            )),
164        },
165    )
166}
167
168pub fn parse_protocol(name: &str) -> Result<StreamKind, P2pNetworkIdentifyFromMessageError> {
169    // buffer content should match one of tokens
170    for tok in token::Token::ALL.iter() {
171        if let token::Token::Protocol(token::Protocol::Stream(a)) = tok {
172            if a.name_str() == name {
173                return Ok(*a);
174            }
175        }
176    }
177
178    Err(P2pNetworkIdentifyFromMessageError::UnsupportedProtocol(
179        name.to_string(),
180    ))
181}
182
183#[derive(Clone, Debug, Serialize, Deserialize, thiserror::Error)]
184#[error("error decoding Multiaddr from bytes: {0}")]
185pub struct P2pNetworkIdentifyMultiaddrError(String);
186
187impl From<multiaddr::Error> for P2pNetworkIdentifyMultiaddrError {
188    fn from(value: multiaddr::Error) -> Self {
189        P2pNetworkIdentifyMultiaddrError(value.to_string())
190    }
191}
192
193mod measurement {
194    use std::mem;
195
196    use malloc_size_of::MallocSizeOfOps;
197
198    use super::Multiaddr;
199
200    pub fn multiaddr_vec(v: &Vec<Multiaddr>, _ops: &mut MallocSizeOfOps) -> usize {
201        v.capacity() * mem::size_of::<Multiaddr>() + v.iter().map(Multiaddr::len).sum::<usize>()
202    }
203
204    pub fn multiaddr_opt(v: &Option<Multiaddr>, _ops: &mut MallocSizeOfOps) -> usize {
205        v.as_ref().map_or(0, Multiaddr::len)
206    }
207}