1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

use libp2p::{multiaddr::Protocol, Multiaddr, PeerId};

pub trait MultiaddrExt {
    /// Peer id
    fn peer_id(&self) -> Option<PeerId>;

    fn extract_peer_id(&mut self) -> Option<PeerId>;

    /// Relay peer id
    fn relay_peer_id(&self) -> Option<PeerId>;

    /// Address that only doesnt include peer protocols
    fn address(&self) -> Multiaddr;

    /// Determine if the address is a relay circuit
    fn is_relay(&self) -> bool;

    /// Determine if the address is being relayed to a peer
    fn is_relayed(&self) -> bool;

    /// Determine if address is loopback or local address
    fn is_loopback(&self) -> bool;

    /// Determine if address is private address
    fn is_private(&self) -> bool;
}

impl MultiaddrExt for Multiaddr {
    fn peer_id(&self) -> Option<PeerId> {
        if let Some(Protocol::P2p(peer)) = self.iter().last() {
            return Some(peer);
        }

        None
    }

    fn extract_peer_id(&mut self) -> Option<PeerId> {
        match self.pop() {
            Some(Protocol::P2p(peer)) => Some(peer),
            _ => None,
        }
    }

    fn relay_peer_id(&self) -> Option<PeerId> {
        if !self.is_relay() {
            return None;
        }

        while let Some(protocol) = self.iter().next() {
            //Find the first peer id and return it
            if let Protocol::P2p(peer) = protocol {
                return Some(peer);
            }
        }

        None
    }

    fn address(&self) -> Multiaddr {
        let mut addr = Multiaddr::empty();

        for proto in self.iter() {
            if matches!(
                proto,
                Protocol::Ip4(_)
                    | Protocol::Ip6(_)
                    | Protocol::Tcp(_)
                    | Protocol::Udp(_)
                    | Protocol::Quic
                    | Protocol::QuicV1
                    | Protocol::Dnsaddr(_)
            ) {
                addr.push(proto);
            }
        }

        addr
    }

    fn is_relay(&self) -> bool {
        self.iter()
            .any(|proto| matches!(proto, Protocol::P2pCircuit))
    }

    fn is_relayed(&self) -> bool {
        if !self.is_relay() {
            return false;
        }

        if self.peer_id().is_none() {
            return false;
        }

        true
    }

    fn is_loopback(&self) -> bool {
        self.iter().any(|proto| match proto {
            Protocol::Ip4(ip) => ip.is_loopback(),
            Protocol::Ip6(ip) => ip.is_loopback(),
            _ => false,
        })
    }

    fn is_private(&self) -> bool {
        self.iter().any(|proto| match proto {
            Protocol::Ip4(ip) => ip.is_private(),
            Protocol::Ip6(ip) => {
                (ip.segments()[0] & 0xffc0) != 0xfe80 && (ip.segments()[0] & 0xfe00) != 0xfc00
            }
            _ => false,
        })
    }
}

#[cfg(test)]
mod tests {
    use std::str::FromStr;

    use super::*;

    #[test]
    fn connection_targets() {
        let peer_id = PeerId::from_str("QmaCpDMGvV2BGHeYERUEnRQAwe3N8SzbUtfsmvsqQLuvuJ")
            .expect("Valid peer id");
        let multiaddr_wo_peer =
            Multiaddr::from_str("/ip4/104.131.131.82/tcp/4001").expect("Valid multiaddr");
        let multiaddr_with_peer =
            Multiaddr::from_str(&format!("{multiaddr_wo_peer}/p2p/{peer_id}"))
                .expect("valid multiaddr");
        let p2p_peer = Multiaddr::from_str(&format!("/p2p/{peer_id}")).expect("Valid multiaddr");

        assert!(multiaddr_wo_peer.peer_id().is_none());
        assert!(multiaddr_with_peer.peer_id().is_some());
        assert!(p2p_peer.peer_id().is_some());

        let peer_id_target = multiaddr_with_peer.peer_id().expect("Valid peer id");

        assert_eq!(peer_id_target, peer_id);
    }
}