From 4577c5a81da1ad4f7b310d7e82ecd0fc44d3e4d8 Mon Sep 17 00:00:00 2001
From: Elena Frank <elena.frank@protonmail.com>
Date: Tue, 4 Oct 2022 19:42:05 +0200
Subject: [PATCH] transports/tcp: Only translate tcp addresses (#2970)

Return `None` in in `<GenTcpTransport as Transport>::address_translation` if the
address is not a tcp address. Relevant if in case of something like
`OrTransport<TcpTransport, QuicTransport>`, where tcp would currently perform
the address translation for quic addresses.
---
 transports/tcp/CHANGELOG.md |  4 +++
 transports/tcp/src/lib.rs   | 60 +++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/transports/tcp/CHANGELOG.md b/transports/tcp/CHANGELOG.md
index d5582f43..83ad2d3c 100644
--- a/transports/tcp/CHANGELOG.md
+++ b/transports/tcp/CHANGELOG.md
@@ -7,8 +7,12 @@
 
 - Remove default features. If you previously depended on `async-std` you need to enable this explicitly now. See [PR 2918].
 
+- Return `None` in `GenTcpTransport::address_translation` if listen- or observed address are not tcp addresses.
+  See [PR 2970].
+
 [PR 2813]: https://github.com/libp2p/rust-libp2p/pull/2813
 [PR 2918]: https://github.com/libp2p/rust-libp2p/pull/2918
+[PR 2970]: https://github.com/libp2p/rust-libp2p/pull/2970
 
 # 0.36.0
 
diff --git a/transports/tcp/src/lib.rs b/transports/tcp/src/lib.rs
index f7b897c0..bf751079 100644
--- a/transports/tcp/src/lib.rs
+++ b/transports/tcp/src/lib.rs
@@ -520,6 +520,9 @@ where
     /// `None` is returned if one of the given addresses is not a TCP/IP
     /// address.
     fn address_translation(&self, listen: &Multiaddr, observed: &Multiaddr) -> Option<Multiaddr> {
+        if !is_tcp_addr(listen) || !is_tcp_addr(observed) {
+            return None;
+        }
         match &self.port_reuse {
             PortReuse::Disabled => address_translation(listen, observed),
             PortReuse::Enabled { .. } => Some(observed.clone()),
@@ -829,6 +832,23 @@ fn ip_to_multiaddr(ip: IpAddr, port: u16) -> Multiaddr {
     Multiaddr::empty().with(ip.into()).with(Protocol::Tcp(port))
 }
 
+fn is_tcp_addr(addr: &Multiaddr) -> bool {
+    use Protocol::*;
+
+    let mut iter = addr.iter();
+
+    let first = match iter.next() {
+        None => return false,
+        Some(p) => p,
+    };
+    let second = match iter.next() {
+        None => return false,
+        Some(p) => p,
+    };
+
+    matches!(first, Ip4(_) | Ip6(_) | Dns(_) | Dns4(_) | Dns6(_)) && matches!(second, Tcp(_))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -836,6 +856,7 @@ mod tests {
         channel::{mpsc, oneshot},
         future::poll_fn,
     };
+    use libp2p_core::PeerId;
 
     #[test]
     fn multiaddr_to_tcp_conversion() {
@@ -1240,4 +1261,43 @@ mod tests {
 
         test("/ip4/127.0.0.1/tcp/12345/tcp/12345".parse().unwrap());
     }
+
+    #[cfg(any(feature = "async-io", feature = "tcp"))]
+    #[test]
+    fn test_address_translation() {
+        #[cfg(feature = "async-io")]
+        let transport = TcpTransport::new(GenTcpConfig::new());
+        #[cfg(all(feature = "tokio", not(feature = "async-io")))]
+        let transport = TokioTcpTransport::new(GenTcpConfig::new());
+
+        let port = 42;
+        let tcp_listen_addr = Multiaddr::empty()
+            .with(Protocol::Ip4(Ipv4Addr::new(127, 0, 0, 1)))
+            .with(Protocol::Tcp(port));
+        let observed_ip = Ipv4Addr::new(123, 45, 67, 8);
+        let tcp_observed_addr = Multiaddr::empty()
+            .with(Protocol::Ip4(observed_ip))
+            .with(Protocol::Tcp(1))
+            .with(Protocol::P2p(PeerId::random().into()));
+
+        let translated = transport
+            .address_translation(&tcp_listen_addr, &tcp_observed_addr)
+            .unwrap();
+        let mut iter = translated.iter();
+        assert_eq!(iter.next(), Some(Protocol::Ip4(observed_ip)));
+        assert_eq!(iter.next(), Some(Protocol::Tcp(port)));
+        assert_eq!(iter.next(), None);
+
+        let quic_addr = Multiaddr::empty()
+            .with(Protocol::Ip4(Ipv4Addr::new(87, 65, 43, 21)))
+            .with(Protocol::Udp(1))
+            .with(Protocol::Quic);
+
+        assert!(transport
+            .address_translation(&tcp_listen_addr, &quic_addr)
+            .is_none());
+        assert!(transport
+            .address_translation(&quic_addr, &tcp_observed_addr)
+            .is_none());
+    }
 }