feat(p2p): add upgrade to quic (#259)

* feat: add quic upgrade

* feat: add comment about peer addresses

* fix: linter

* fix: clean up tests

* fix: linter

* fix: review comments

* fix: review comments

Author: varex83

crates/p2p/examples/quic_upgrade.rs

@@ -0,0 +1,248 @@
+//! QUIC Upgrade Example
+//!
+//! This example demonstrates the QUIC upgrade behaviour that automatically
+//! upgrades TCP connections to QUIC when both peers support it.
+//!
+//! # Running the Example
+//!
+//! Run two instances in separate terminals:
+//!
+//! ```bash
+//! # Terminal 1: Start the first node (note the peer ID printed)
+//! cargo run -p pluto-p2p --example quic_upgrade -- --port 9000
+//! # Output: Started node with peer ID: 16Uiu2HAmXXX...
+//!
+//! # Terminal 2: Start the second node with the first node's peer ID as known peer
+//! cargo run -p pluto-p2p --example quic_upgrade -- --port 9001 \
+//!     --dial /ip4/127.0.0.1/tcp/9000 \
+//!     --peer 16Uiu2HAmXXX...
+//! ```
+//!
+//! # What Happens
+//!
+//! 1. Node 2 connects to Node 1 via TCP
+//! 2. Both nodes exchange identify information (including their QUIC addresses)
+//! 3. The peer addresses are stored in the peer store
+//! 4. The QUIC upgrade behaviour detects the TCP connection to a known peer
+//! 5. After ~1 minute, it attempts to dial the peer's QUIC address
+//! 6. On success: the redundant TCP connection is closed
+//! 7. On failure: exponential backoff is applied before retrying
+//!
+//! # Note
+//!
+//! The QUIC upgrade behaviour only upgrades connections to "known peers"
+//! (cluster members). In production, these are configured at startup.
+//! For this example, use `--peer` to specify the peer ID to upgrade.
+//!
+//! # Expected Output
+//!
+//! ```text
+//! [CONNECTED] 16Uiu2HAm... via TCP at /ip4/127.0.0.1/tcp/9000
+//! [IDENTIFY] Received from 16Uiu2HAm...
+//!   Agent: quic-upgrade-example/1.0.0
+//!   Addresses:
+//!     - [TCP] /ip4/127.0.0.1/tcp/9000
+//!     - [QUIC] /ip4/127.0.0.1/udp/9000/quic-v1
+//! ... (after ~1 minute) ...
+//! [CONNECTED] 16Uiu2HAm... via QUIC at /ip4/127.0.0.1/udp/9000/quic-v1
+//! [QUIC UPGRADE] Successfully upgraded connection to 16Uiu2HAm...!
+//! [DISCONNECTED] 16Uiu2HAm... TCP connection closed (remaining: 1)
+//! ```
+
+use std::str::FromStr;
+
+use anyhow::Result;
+use clap::Parser;
+use k256::elliptic_curve::rand_core::OsRng;
+use libp2p::{Multiaddr, PeerId, futures::StreamExt, relay, swarm::SwarmEvent};
+use pluto_p2p::{
+    behaviours::pluto::PlutoBehaviourEvent,
+    config::P2PConfig,
+    p2p::{Node, NodeType},
+    quic_upgrade::QuicUpgradeEvent,
+};
+use tokio::signal;
+
+/// Command line arguments.
+#[derive(Debug, Parser)]
+#[command(name = "quic_upgrade")]
+#[command(about = "Demonstrates QUIC upgrade behaviour")]
+pub struct Args {
+    /// The port to listen on (both TCP and UDP/QUIC).
+    #[arg(short, long, default_value = "9000")]
+    pub port: u16,
+
+    /// Address to dial (e.g., /ip4/127.0.0.1/tcp/9000).
+    #[arg(short, long)]
+    pub dial: Option<Multiaddr>,
+
+    /// Known peer ID(s) to attempt QUIC upgrade for.
+    /// The upgrade behaviour only upgrades connections to known peers.
+    #[arg(long)]
+    pub peer: Vec<String>,
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let args = Args::parse();
+    let key = k256::SecretKey::random(&mut OsRng);
+
+    // Create a config with the specified port
+    // Note: P2PConfig requires a specific IP (not 0.0.0.0), so we use 127.0.0.1 for
+    // local testing
+    let config = P2PConfig {
+        tcp_addrs: vec![format!("127.0.0.1:{}", args.port)],
+        udp_addrs: vec![format!("127.0.0.1:{}", args.port)],
+        ..Default::default()
+    };
+
+    // Parse known peer IDs from command line
+    let known_peers: Vec<PeerId> = args
+        .peer
+        .iter()
+        .filter_map(|s| PeerId::from_str(s).ok())
+        .collect();
+
+    if !known_peers.is_empty() {
+        println!("Known peers for QUIC upgrade: {:?}", known_peers);
+    }
+
+    let mut node = Node::new(
+        config,
+        key,
+        NodeType::QUIC, // Enable QUIC transport
+        false,          // Don't filter private addresses (for local testing)
+        known_peers,
+        |builder, _keypair, relay_client| {
+            builder
+                .with_user_agent("quic-upgrade-example/1.0.0")
+                .with_quic_enabled(true) // Enable QUIC upgrade behaviour
+                .with_inner(relay_client)
+        },
+    )?;
+
+    println!("Started node with peer ID: {}", node.local_peer_id());
+    println!("Listening on TCP and QUIC port: {}", args.port);
+
+    // Dial the remote peer if specified
+    if let Some(dial_addr) = &args.dial {
+        println!("Dialing remote peer via TCP: {dial_addr}");
+        node.dial(dial_addr.clone())?;
+    }
+
+    println!("\nWaiting for events... (Ctrl+C to quit)");
+    if args.peer.is_empty() {
+        println!("Note: No --peer specified. QUIC upgrade only works for known peers.");
+        println!("      Copy this node's peer ID and pass it to the other node with --peer\n");
+    } else {
+        println!("QUIC upgrade will be attempted ~1 minute after TCP connection is established.\n");
+    }
+
+    // Event loop
+    loop {
+        tokio::select! {
+            event = node.select_next_some() => {
+                handle_event(event);
+            }
+            _ = signal::ctrl_c() => {
+                println!("\nReceived Ctrl+C, shutting down...");
+                break;
+            }
+        }
+    }
+
+    Ok(())
+}
+
+fn handle_event(event: SwarmEvent<PlutoBehaviourEvent<relay::client::Behaviour>>) {
+    match event {
+        // New listen address
+        SwarmEvent::NewListenAddr { address, .. } => {
+            println!("[LISTEN] {address}");
+        }
+
+        // Connection established
+        SwarmEvent::ConnectionEstablished {
+            peer_id, endpoint, ..
+        } => {
+            let addr = match &endpoint {
+                libp2p::core::ConnectedPoint::Dialer { address, .. } => address,
+                libp2p::core::ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
+            };
+            let transport = if addr.to_string().contains("quic") {
+                "QUIC"
+            } else {
+                "TCP"
+            };
+            println!("[CONNECTED] {peer_id} via {transport} at {addr}");
+        }
+
+        // Connection closed
+        SwarmEvent::ConnectionClosed {
+            peer_id,
+            endpoint,
+            num_established,
+            ..
+        } => {
+            let addr = match &endpoint {
+                libp2p::core::ConnectedPoint::Dialer { address, .. } => address,
+                libp2p::core::ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
+            };
+            let transport = if addr.to_string().contains("quic") {
+                "QUIC"
+            } else {
+                "TCP"
+            };
+            println!(
+                "[DISCONNECTED] {peer_id} {transport} connection closed (remaining: {num_established})"
+            );
+        }
+
+        // QUIC upgrade events
+        SwarmEvent::Behaviour(PlutoBehaviourEvent::QuicUpgrade(event)) => match event {
+            QuicUpgradeEvent::Upgraded { peer } => {
+                println!("[QUIC UPGRADE] Successfully upgraded connection to {peer}!");
+            }
+            QuicUpgradeEvent::UpgradeFailed { peer, reason } => {
+                println!("[QUIC UPGRADE FAILED] {peer}: {reason}");
+            }
+        },
+
+        // Identify received - shows peer's addresses including QUIC
+        SwarmEvent::Behaviour(PlutoBehaviourEvent::Identify(
+            libp2p::identify::Event::Received { peer_id, info, .. },
+        )) => {
+            println!("[IDENTIFY] Received from {peer_id}");
+            println!("  Agent: {}", info.agent_version);
+            println!("  Addresses:");
+            for addr in &info.listen_addrs {
+                let transport = if addr.to_string().contains("quic") {
+                    "QUIC"
+                } else if addr.to_string().contains("tcp") {
+                    "TCP"
+                } else {
+                    "other"
+                };
+                println!("    - [{transport}] {addr}");
+            }
+        }
+
+        // Ping events
+        SwarmEvent::Behaviour(PlutoBehaviourEvent::Ping(event)) => {
+            if let Ok(rtt) = event.result {
+                println!("[PING] {} RTT: {:?}", event.peer, rtt);
+            }
+        }
+
+        // Connection errors
+        SwarmEvent::OutgoingConnectionError { peer_id, error, .. } => {
+            println!("[ERROR] Outgoing connection to {peer_id:?}: {error}");
+        }
+        SwarmEvent::IncomingConnectionError { error, .. } => {
+            println!("[ERROR] Incoming connection: {error}");
+        }
+
+        // Ignore other events
+        _ => {}
+    }
+}

crates/p2p/src/behaviours/pluto.rs

@@ -12,6 +12,7 @@ use crate::{
     conn_logger::{ConnectionLoggerBehaviour, DefaultConnectionLoggerMetrics},
     gater::ConnGater,
     p2p_context::P2PContext,
+    quic_upgrade::QuicUpgradeBehaviour,
 };
 
 pub use super::optional::OptionalBehaviour;
@@ -39,6 +40,7 @@ pub const DEFAULT_IDENTIFY_CACHE_SIZE: usize = 100;
 /// - **Identify**: Exchanges peer information and supported protocols
 /// - **Ping**: Measures latency and keeps connections alive
 /// - **AutoNAT**: Detects NAT status and public reachability
+/// - **QUIC upgrade**: Periodically upgrades TCP connections to QUIC
 #[derive(NetworkBehaviour)]
 pub struct PlutoBehaviour<B: NetworkBehaviour> {
     /// Connection logger behaviour - MUST be first so peer store is updated
@@ -52,6 +54,8 @@ pub struct PlutoBehaviour<B: NetworkBehaviour> {
     pub ping: ping::Behaviour,
     /// AutoNAT behaviour for NAT detection.
     pub autonat: autonat::Behaviour,
+    /// QUIC upgrade behaviour for upgrading TCP to QUIC connections.
+    pub quic_upgrade: QuicUpgradeBehaviour,
     /// Inner behaviour.
     pub inner: OptionalBehaviour<B>,
 }
@@ -70,6 +74,7 @@ impl<B: NetworkBehaviour> PlutoBehaviour<B> {
 /// - **Identify**: Protocol and agent identification
 /// - **Ping**: Latency measurement and keepalive
 /// - **AutoNAT**: NAT traversal detection
+/// - **QUIC upgrade**: Periodic TCP to QUIC connection upgrades
 #[derive(Debug, Clone)]
 pub struct PlutoBehaviourBuilder<B> {
     // Gater config
@@ -84,6 +89,9 @@ pub struct PlutoBehaviourBuilder<B> {
 
     p2p_context: P2PContext,
 
+    // QUIC upgrade config
+    quic_enabled: bool,
+
     // Inner behaviour
     inner: Option<B>,
 }
@@ -96,6 +104,7 @@ impl<B> Default for PlutoBehaviourBuilder<B> {
             user_agent: DEFAULT_USER_AGENT.clone(),
             autonat_config: autonat::Config::default(),
             p2p_context: P2PContext::default(),
+            quic_enabled: false,
             inner: None,
         }
     }
@@ -166,27 +175,43 @@ impl<B: NetworkBehaviour> PlutoBehaviourBuilder<B> {
         self
     }
 
+    /// Sets whether QUIC is enabled.
+    ///
+    /// When enabled, the behaviour will periodically attempt to upgrade
+    /// TCP connections to QUIC connections.
+    pub fn with_quic_enabled(mut self, enabled: bool) -> Self {
+        self.quic_enabled = enabled;
+        self
+    }
+
     /// Builds the [`PlutoBehaviour`] with the provided keypair.
     ///
     /// # Arguments
     ///
     /// * `key` - The keypair for this node, used for identify and autonat
     pub fn build(self, key: &Keypair) -> PlutoBehaviour<B> {
+        let local_peer_id = key.public().to_peer_id();
+
         let identify_config = identify::Config::new(self.identify_protocol, key.public())
             .with_agent_version(self.user_agent)
             .with_interval(DEFAULT_IDENTIFY_INTERVAL)
             .with_cache_size(DEFAULT_IDENTIFY_CACHE_SIZE);
 
         PlutoBehaviour {
-            conn_logger: ConnectionLoggerBehaviour::new(self.p2p_context),
+            conn_logger: ConnectionLoggerBehaviour::new(self.p2p_context.clone()),
             gater: self.gater.unwrap_or_else(ConnGater::new_open_gater),
             identify: identify::Behaviour::new(identify_config),
             ping: ping::Behaviour::new(
                 ping::Config::new()
                     .with_interval(DEFAULT_PING_INTERVAL)
                     .with_timeout(DEFAULT_PING_TIMEOUT),
             ),
-            autonat: autonat::Behaviour::new(key.public().to_peer_id(), self.autonat_config),
+            autonat: autonat::Behaviour::new(local_peer_id, self.autonat_config),
+            quic_upgrade: QuicUpgradeBehaviour::new(
+                self.p2p_context,
+                local_peer_id,
+                self.quic_enabled,
+            ),
             inner: self.inner.into(),
         }
     }

crates/p2p/src/conn_logger.rs

@@ -219,11 +219,19 @@ impl<M: ConnectionLoggerMetrics + 'static> NetworkBehaviour for ConnectionLogger
                     other_established = event.other_established,
                     "connection established"
                 );
+                // Extract remote address from the endpoint
+                let remote_addr = match &event.endpoint {
+                    libp2p::core::ConnectedPoint::Dialer { address, .. } => address.clone(),
+                    libp2p::core::ConnectedPoint::Listener { send_back_addr, .. } => {
+                        send_back_addr.clone()
+                    }
+                };
                 // Update peer store - this is done here so all other behaviours
                 // see the updated peer store immediately
                 self.p2p_context.peer_store_write_lock().add_peer(Peer {
                     id: event.peer_id,
                     connection_id: event.connection_id,
+                    remote_addr,
                 });
             }
             libp2p::swarm::FromSwarm::ConnectionClosed(event) => {
@@ -233,17 +241,21 @@ impl<M: ConnectionLoggerMetrics + 'static> NetworkBehaviour for ConnectionLogger
                     num_established = event.remaining_established,
                     "connection closed"
                 );
+                // Extract remote address from the endpoint
+                let addr = match &event.endpoint {
+                    libp2p::core::ConnectedPoint::Dialer { address, .. } => address.clone(),
+                    libp2p::core::ConnectedPoint::Listener { send_back_addr, .. } => {
+                        send_back_addr.clone()
+                    }
+                };
                 // Update peer store
                 self.p2p_context.peer_store_write_lock().remove_peer(Peer {
                     id: event.peer_id,
                     connection_id: event.connection_id,
+                    remote_addr: addr.clone(),
                 });
                 // Decrement the connection count based on the endpoint address
-                let addr = match &event.endpoint {
-                    libp2p::core::ConnectedPoint::Dialer { address, .. } => address,
-                    libp2p::core::ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
-                };
-                self.decrement_connection(event.peer_id, addr);
+                self.decrement_connection(event.peer_id, &addr);
             }
             _ => {}
         }

crates/p2p/src/lib.rs

@@ -40,3 +40,6 @@ pub mod conn_logger;
 
 /// Global context.
 pub mod p2p_context;
+
+/// QUIC connection upgrade behaviour.
+pub mod quic_upgrade;