From 74c087dcefe9e7c6893b0833417bcb61c39da413 Mon Sep 17 00:00:00 2001
From: Max Inden <mail@max-inden.de>
Date: Fri, 13 Oct 2023 14:17:10 +0200
Subject: [PATCH] fix(libp2p): add shortcut
 `with_tcp(...).with_bandwidth_logging()`

Add the shortcut method `.with_bandwidth_logging` to `SwarmBuilder<_, QuicPhase<_>>`, thus allowing `with_tcp(...).with_bandwidth_logging()`.

Pull-Request: #4626.
---
 libp2p/src/builder.rs            | 57 ++++++++++++++++++++++++++++++--
 libp2p/src/builder/phase/quic.rs | 20 ++++++++++-
 2 files changed, 73 insertions(+), 4 deletions(-)

diff --git a/libp2p/src/builder.rs b/libp2p/src/builder.rs
index 1b07689d..83541378 100644
--- a/libp2p/src/builder.rs
+++ b/libp2p/src/builder.rs
@@ -66,6 +66,8 @@ pub struct SwarmBuilder<Provider, Phase> {
 #[cfg(test)]
 mod tests {
     use crate::SwarmBuilder;
+    use libp2p_core::{muxing::StreamMuxerBox, transport::dummy::DummyTransport};
+    use libp2p_identity::PeerId;
     use libp2p_swarm::{NetworkBehaviour, Swarm};
 
     #[test]
@@ -228,9 +230,6 @@ mod tests {
     #[test]
     #[cfg(feature = "tokio")]
     fn other_transport() -> Result<(), Box<dyn std::error::Error>> {
-        use libp2p_core::{muxing::StreamMuxerBox, transport::dummy::DummyTransport};
-        use libp2p_identity::PeerId;
-
         let _ = SwarmBuilder::with_new_identity()
             .with_tokio()
             // Closure can either return a Transport directly.
@@ -319,4 +318,56 @@ mod tests {
             .unwrap()
             .build();
     }
+
+    #[test]
+    #[cfg(all(feature = "tokio", feature = "tcp", feature = "tls", feature = "yamux"))]
+    fn tcp_bandwidth_logging() -> Result<(), Box<dyn std::error::Error>> {
+        let (builder, _logging) = SwarmBuilder::with_new_identity()
+            .with_tokio()
+            .with_tcp(
+                Default::default(),
+                libp2p_tls::Config::new,
+                libp2p_yamux::Config::default,
+            )?
+            .with_bandwidth_logging();
+
+        builder
+            .with_behaviour(|_| libp2p_swarm::dummy::Behaviour)
+            .unwrap()
+            .build();
+
+        Ok(())
+    }
+
+    #[test]
+    #[cfg(all(feature = "tokio", feature = "quic"))]
+    fn quic_bandwidth_logging() -> Result<(), Box<dyn std::error::Error>> {
+        let (builder, _logging) = SwarmBuilder::with_new_identity()
+            .with_tokio()
+            .with_quic()
+            .with_bandwidth_logging();
+
+        builder
+            .with_behaviour(|_| libp2p_swarm::dummy::Behaviour)
+            .unwrap()
+            .build();
+
+        Ok(())
+    }
+
+    #[test]
+    #[cfg(feature = "tokio")]
+    fn other_transport_bandwidth_logging() -> Result<(), Box<dyn std::error::Error>> {
+        let (builder, _logging) = SwarmBuilder::with_new_identity()
+            .with_tokio()
+            .with_other_transport(|_| DummyTransport::<(PeerId, StreamMuxerBox)>::new())?
+            .with_bandwidth_logging();
+
+        builder
+            .with_behaviour(|_| libp2p_swarm::dummy::Behaviour)
+            .unwrap()
+            .build();
+
+        Ok(())
+    }
 }
diff --git a/libp2p/src/builder/phase/quic.rs b/libp2p/src/builder/phase/quic.rs
index 6ef056ae..e49b3b07 100644
--- a/libp2p/src/builder/phase/quic.rs
+++ b/libp2p/src/builder/phase/quic.rs
@@ -7,7 +7,7 @@ use libp2p_core::muxing::StreamMuxer;
     all(not(target_arch = "wasm32"), feature = "websocket")
 ))]
 use libp2p_core::{InboundUpgrade, Negotiated, OutboundUpgrade, UpgradeInfo};
-use std::marker::PhantomData;
+use std::{marker::PhantomData, sync::Arc};
 
 pub struct QuicPhase<T> {
     pub(crate) transport: T,
@@ -237,3 +237,21 @@ impl_quic_phase_with_websocket!(
     super::provider::Tokio,
     rw_stream_sink::RwStreamSink<libp2p_websocket::BytesConnection<libp2p_tcp::tokio::TcpStream>>
 );
+impl<Provider, T: AuthenticatedMultiplexedTransport> SwarmBuilder<Provider, QuicPhase<T>> {
+    pub fn with_bandwidth_logging(
+        self,
+    ) -> (
+        SwarmBuilder<
+            Provider,
+            BehaviourPhase<impl AuthenticatedMultiplexedTransport, NoRelayBehaviour>,
+        >,
+        Arc<crate::bandwidth::BandwidthSinks>,
+    ) {
+        self.without_quic()
+            .without_any_other_transports()
+            .without_dns()
+            .without_relay()
+            .without_websocket()
+            .with_bandwidth_logging()
+    }
+}