refactor(swap): deduplicate upgrading TorBackend to Transport

Author: nabijaczleweli

swap/src/asb/network.rs

@@ -1,4 +1,4 @@
-use crate::common::tor::TorBackend;
+use crate::common::tor::{TorBackend, TorBackendSwap};
 use crate::network::rendezvous::XmrBtcNamespace;
 use crate::network::swap_setup::alice;
 use crate::network::transport::authenticate_and_multiplex;
@@ -18,10 +18,7 @@ use swap_feed::LatestRate;
 
 pub mod transport {
     use arti_client::config::onion_service::OnionServiceConfigBuilder;
-    use libp2p::{
-        core::transport::{OptionalTransport, OrTransport},
-        dns, identity, tcp, Transport,
-    };
+    use libp2p::{dns, identity, tcp, Transport};
     use libp2p_tor::AddressConversion;
 
     use super::*;
@@ -48,15 +45,8 @@ pub mod transport {
         num_intro_points: u8,
     ) -> Result<OnionTransportWithAddresses> {
         let mut onion_addresses = vec![];
-        let maybe_tor_transport = match maybe_tor_client {
-            TorBackend::Socks(universal_config) => OrTransport::new(
-                OptionalTransport::none(),
-                OptionalTransport::some(universal_config.transport()),
-            ),
-            TorBackend::Arti(tor_client) => {
-                let mut tor_transport =
-                    libp2p_tor::TorTransport::from_client(tor_client, AddressConversion::DnsOnly);
-
+        let maybe_tor_transport =
+            maybe_tor_client.into_transport(AddressConversion::DnsOnly, |arti_tor_transport| {
                 if register_hidden_service {
                     let onion_service_config = OnionServiceConfigBuilder::default()
                         .nickname(
@@ -68,7 +58,7 @@ pub mod transport {
                         .build()
                         .expect("We specified a valid nickname");
 
-                    match tor_transport
+                    match arti_tor_transport
                         .add_onion_service(onion_service_config, ASB_ONION_SERVICE_PORT)
                     {
                         Ok(addr) => {
@@ -83,16 +73,7 @@ pub mod transport {
                         }
                     }
                 }
-
-                OrTransport::new(
-                    OptionalTransport::some(tor_transport),
-                    OptionalTransport::none(),
-                )
-            }
-            TorBackend::None => {
-                OrTransport::new(OptionalTransport::none(), OptionalTransport::none())
-            }
-        };
+            });
 
         let tcp = maybe_tor_transport
             .or_transport(tcp::tokio::Transport::new(tcp::Config::new().nodelay(true)));

swap/src/cli/transport.rs

@@ -1,8 +1,8 @@
-use crate::common::tor::TorBackend;
+use crate::common::tor::{TorBackend, TorBackendSwap};
 use crate::network::transport::authenticate_and_multiplex;
 use anyhow::Result;
 use libp2p::core::muxing::StreamMuxerBox;
-use libp2p::core::transport::{Boxed, OptionalTransport, OrTransport};
+use libp2p::core::transport::Boxed;
 use libp2p::dns;
 use libp2p::tcp;
 use libp2p::{identity, PeerId, Transport};
@@ -23,20 +23,7 @@ pub fn new(
     let tcp = tcp::tokio::Transport::new(tcp::Config::new().nodelay(true));
     let tcp_with_dns = dns::tokio::Transport::system(tcp)?;
 
-    let maybe_tor_transport = match maybe_tor_client {
-        TorBackend::Socks(universal_config) => OrTransport::new(
-            OptionalTransport::none(),
-            OptionalTransport::some(universal_config.transport()),
-        ),
-        TorBackend::Arti(client) => OrTransport::new(
-            OptionalTransport::some(libp2p_tor::TorTransport::from_client(
-                client,
-                AddressConversion::IpAndDns,
-            )),
-            OptionalTransport::none(),
-        ),
-        TorBackend::None => OrTransport::new(OptionalTransport::none(), OptionalTransport::none()),
-    };
+    let maybe_tor_transport = maybe_tor_client.into_transport(AddressConversion::IpAndDns, |_| {});
 
     let transport = maybe_tor_transport.or_transport(tcp_with_dns).boxed();
 

swap/src/common/tor.rs

@@ -5,7 +5,9 @@ use crate::cli::api::tauri_bindings::{
 };
 use arti_client::{config::TorClientConfigBuilder, status::BootstrapStatus, Error, TorClient};
 use futures::StreamExt;
+use libp2p::core::transport::{OptionalTransport, OrTransport};
 use swap_env::env::{is_whonix, may_init_tor};
+use libp2p_tor::{AddressConversion, TorTransport};
 use swap_tor::*;
 use tor_rtcompat::tokio::TokioRustlsRuntime;
 
@@ -45,7 +47,13 @@ pub async fn create_tor_client(data_dir: &Path, tor: bool) -> Result<TorBackend,
 pub trait TorBackendSwap {
     async fn bootstrap(&self, tauri_handle: Option<TauriHandle>) -> anyhow::Result<()>;
     fn clone_for_monero_rpc(&self, enable_monero_tor: bool) -> TorBackend;
+    fn into_transport(
+        self,
+        arti_address_conversion: AddressConversion,
+        arti_transport_hook: impl FnOnce(&mut TorTransport),
+    ) -> IntoTransportT;
 }
+type IntoTransportT = OrTransport<OptionalTransport<TorTransport>, OptionalTransport<Socks5Transport>>;
 impl TorBackendSwap for TorBackend {
     async fn bootstrap(&self, tauri_handle: Option<TauriHandle>) -> anyhow::Result<()> {
         match self {
@@ -66,6 +74,31 @@ impl TorBackendSwap for TorBackend {
             TorBackend::Socks(..) | TorBackend::None => self.clone(),
         }
     }
+
+    fn into_transport(
+        self,
+        arti_address_conversion: AddressConversion,
+        arti_transport_hook: impl FnOnce(&mut TorTransport),
+    ) -> IntoTransportT {
+        match self {
+            TorBackend::Socks(universal_config) => OrTransport::new(
+                OptionalTransport::none(),
+                OptionalTransport::some(universal_config.transport()),
+            ),
+            TorBackend::Arti(tor_client) => {
+                let mut tor_transport =
+                    libp2p_tor::TorTransport::from_client(tor_client, arti_address_conversion);
+                arti_transport_hook(&mut tor_transport);
+                OrTransport::new(
+                    OptionalTransport::some(tor_transport),
+                    OptionalTransport::none(),
+                )
+            }
+            TorBackend::None => {
+                OrTransport::new(OptionalTransport::none(), OptionalTransport::none())
+            }
+        }
+    }
 }
 
 const TOR_CONNECT_TIMEOUT: Duration = Duration::from_secs(30);