atm0s_sdn_identity/
node_addr.rs

1use serde::{Deserialize, Serialize};
2use std::{fmt::Display, str::FromStr};
3
4use crate::node_id::NodeId;
5pub use multiaddr::Protocol;
6
7#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
8pub struct NodeAddr(NodeId, multiaddr::Multiaddr);
9
10impl NodeAddr {
11    pub fn empty(node_id: NodeId) -> Self {
12        Self(node_id, multiaddr::Multiaddr::empty())
13    }
14
15    pub fn node_id(&self) -> NodeId {
16        self.0
17    }
18
19    pub fn multiaddr(&self) -> &multiaddr::Multiaddr {
20        &self.1
21    }
22
23    pub fn from_iter<'a>(node_id: NodeId, iter: impl IntoIterator<Item = Protocol<'a>>) -> Self {
24        Self(node_id, multiaddr::Multiaddr::from_iter(iter))
25    }
26
27    pub fn to_vec(&self) -> Vec<u8> {
28        let mut buf = self.0.to_be_bytes().to_vec();
29        buf.extend(self.1.to_vec());
30        buf
31    }
32
33    pub fn from_vec(buf: &[u8]) -> Option<Self> {
34        let node_id = NodeId::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
35        let multiaddr = multiaddr::Multiaddr::try_from(buf[4..].to_vec()).ok()?;
36        Some(Self(node_id, multiaddr))
37    }
38}
39
40impl FromStr for NodeAddr {
41    type Err = String;
42
43    fn from_str(s: &str) -> Result<Self, Self::Err> {
44        let mut split = s.split('@');
45        let node_id = split.next().ok_or("Missing NodeId".to_string())?.parse::<NodeId>().map_err(|e| e.to_string())?;
46        let multiaddr = split.next().unwrap_or("").parse::<multiaddr::Multiaddr>().map_err(|e| e.to_string())?;
47        Ok(Self(node_id, multiaddr))
48    }
49}
50
51impl Display for NodeAddr {
52    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53        if self.1.is_empty() {
54            write!(f, "{}", self.0)
55        } else {
56            write!(f, "{}@{}", self.0, self.1)
57        }
58    }
59}
60
61/// A builder for creating `NodeAddr` instances.
62pub struct NodeAddrBuilder {
63    node_id: NodeId,
64    addr: multiaddr::Multiaddr,
65}
66
67impl NodeAddrBuilder {
68    pub fn new(node_id: NodeId) -> Self {
69        Self {
70            node_id,
71            addr: multiaddr::Multiaddr::empty(),
72        }
73    }
74
75    pub fn node_id(&self) -> NodeId {
76        self.node_id
77    }
78
79    /// Adds a protocol to the node address.
80    pub fn add_protocol(&mut self, protocol: Protocol) {
81        self.addr.push(protocol);
82    }
83
84    /// Get the node address.
85    pub fn addr(&self) -> NodeAddr {
86        NodeAddr(self.node_id, self.addr.clone())
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use std::str::FromStr;
93
94    use multiaddr::Multiaddr;
95
96    #[test]
97    fn test_to_from_str() {
98        let addr = super::NodeAddr::from_str("1@/ip4/127.0.0.1").unwrap();
99        assert_eq!(addr, super::NodeAddr(1, "/ip4/127.0.0.1".parse().unwrap()));
100        assert_eq!(addr.to_string(), "1@/ip4/127.0.0.1");
101    }
102
103    #[test]
104    fn test_empty() {
105        let addr = super::NodeAddr::from_str("1").unwrap();
106        assert_eq!(addr, super::NodeAddr(1, Multiaddr::empty()));
107        assert_eq!(addr, super::NodeAddr::empty(1));
108        assert_eq!(addr.to_string(), "1");
109    }
110}