stoffel-vm 0.1.3

// src/net/p2p.rs
//! # Peer-to-Peer Networking for StoffelVM
//!
//! This module provides the networking capabilities for StoffelVM, enabling
//! secure communication between distributed parties for multiparty computation.
//!
//! The networking layer is built on the QUIC protocol, which offers:
//! - Encrypted connections using TLS 1.3
//! - Low latency with 0-RTT connection establishment
//! - Stream multiplexing for concurrent data transfers
//! - Connection migration for network changes
//!
//! The module defines two primary abstractions:
//! - `PeerConnection`: Represents a connection to a single peer
//! - `NetworkManager`: Manages multiple peer connections
//!
//! The current implementation uses the Quinn library for QUIC support.

use ark_ff::Field;
use async_trait::async_trait;
use quinn::{ClientConfig, Connection, Endpoint, IdleTimeout, ServerConfig, TransportConfig};
use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::future::Future;
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::Arc;
use stoffelnet::network_utils::{
    ClientId, Message, Network, NetworkError, Node, PartyId, VerifiedOrdering,
};
use tokio::sync::Mutex;
use tracing::debug;
use uuid::Uuid;

type PeerFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, String>> + Send + 'a>>;
type BoxedPeerConnection = Box<dyn PeerConnection>;
type SharedPeerConnection = Arc<Mutex<BoxedPeerConnection>>;
type PeerConnectionMap = Arc<Mutex<HashMap<PartyId, SharedPeerConnection>>>;

const DEFAULT_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
const DEFAULT_KEEP_ALIVE_INTERVAL: std::time::Duration = std::time::Duration::from_secs(5);

fn client_endpoint_bind_addr() -> SocketAddr {
    SocketAddr::from(([0, 0, 0, 0], 0))
}

fn party_id_from_uuid(uuid: Uuid) -> PartyId {
    let mut id = PartyId::default();
    for (index, byte) in uuid.as_bytes().iter().enumerate() {
        let shift = (index % std::mem::size_of::<PartyId>()) * 8;
        id ^= PartyId::from(*byte) << shift;
    }
    id
}

fn random_party_id() -> PartyId {
    party_id_from_uuid(Uuid::new_v4())
}

fn party_id_to_u128(id: PartyId) -> u128 {
    let mut bytes = [0u8; 16];
    let id_bytes = id.to_le_bytes();
    bytes[..id_bytes.len()].copy_from_slice(&id_bytes);
    u128::from_le_bytes(bytes)
}

/// Represents a connection to a peer (re-export of stoffelnet traits is not used directly here)
pub trait PeerConnection: Send + Sync {
    /// Sends data to the peer on the default stream
    ///
    /// This is a convenience method that sends data on stream ID 0.
    /// For more control, use `send_on_stream`.
    ///
    /// # Arguments
    /// * `data` - The data to send
    ///
    /// # Returns
    /// * `Ok(())` - If the data was sent successfully
    /// * `Err(String)` - If there was an error sending the data
    fn send<'a>(&'a mut self, data: &'a [u8]) -> PeerFuture<'a, ()>;

    /// Receives data from the peer on the default stream
    ///
    /// This is a convenience method that receives data from stream ID 0.
    /// For more control, use `receive_from_stream`.
    ///
    /// # Returns
    /// * `Ok(Vec<u8>)` - The received data
    /// * `Err(String)` - If there was an error receiving data
    fn receive<'a>(&'a mut self) -> PeerFuture<'a, Vec<u8>>;

    /// Sends data on a specific stream
    ///
    /// This method allows sending data on a specific stream ID, enabling
    /// multiplexed communication with the peer.
    ///
    /// # Arguments
    /// * `stream_id` - The ID of the stream to send on
    /// * `data` - The data to send
    ///
    /// # Returns
    /// * `Ok(())` - If the data was sent successfully
    /// * `Err(String)` - If there was an error sending the data
    fn send_on_stream<'a>(&'a mut self, stream_id: u64, data: &'a [u8]) -> PeerFuture<'a, ()>;

    /// Receives data from a specific stream
    ///
    /// This method allows receiving data from a specific stream ID, enabling
    /// multiplexed communication with the peer.
    ///
    /// # Arguments
    /// * `stream_id` - The ID of the stream to receive from
    ///
    /// # Returns
    /// * `Ok(Vec<u8>)` - The received data
    /// * `Err(String)` - If there was an error receiving data
    fn receive_from_stream<'a>(&'a mut self, stream_id: u64) -> PeerFuture<'a, Vec<u8>>;

    /// Returns the address of the remote peer
    ///
    /// This method provides the network address of the connected peer,
    /// which can be useful for logging, debugging, or identity verification.
    fn remote_address(&self) -> SocketAddr;

    /// Closes the connection
    ///
    /// This method gracefully terminates the connection with the peer.
    /// After calling this method, no more data can be sent or received.
    ///
    /// # Returns
    /// * `Ok(())` - If the connection was closed successfully
    /// * `Err(String)` - If there was an error closing the connection
    fn close<'a>(&'a mut self) -> PeerFuture<'a, ()>;
}

/// Manages network connections for the VM
pub trait NetworkManager: Send + Sync {
    /// Establishes a connection to a new peer
    ///
    /// This method initiates an outgoing connection to a peer at the specified address.
    /// It handles the connection establishment process, including any necessary
    /// handshaking, encryption setup, and protocol negotiation.
    ///
    /// # Arguments
    /// * `address` - The network address of the peer to connect to
    ///
    /// # Returns
    /// * `Ok(Box<dyn PeerConnection>)` - A connection to the peer
    /// * `Err(String)` - If the connection could not be established
    fn connect<'a>(&'a mut self, address: SocketAddr) -> PeerFuture<'a, BoxedPeerConnection>;

    /// Accepts an incoming connection
    ///
    /// This method accepts a pending incoming connection from a peer.
    /// It should be called after `listen()` has been called to set up
    /// the listening endpoint.
    ///
    /// This method will block until a connection is available or an error occurs.
    ///
    /// # Returns
    /// * `Ok(Box<dyn PeerConnection>)` - A connection to the peer
    /// * `Err(String)` - If no connection could be accepted
    fn accept<'a>(&'a mut self) -> PeerFuture<'a, BoxedPeerConnection>;

    /// Listens for incoming connections
    ///
    /// This method sets up a network endpoint to listen for incoming connections
    /// at the specified address. After calling this method, `accept()` can be
    /// called to accept incoming connections.
    ///
    /// # Arguments
    /// * `bind_address` - The local address to bind to for listening
    ///
    /// # Returns
    /// * `Ok(())` - If the listening endpoint was set up successfully
    /// * `Err(String)` - If the listening endpoint could not be set up
    fn listen<'a>(&'a mut self, bind_address: SocketAddr) -> PeerFuture<'a, ()>;
}

/// QUIC-based implementation of PeerConnection
///
/// This struct implements the PeerConnection trait using the QUIC protocol
/// via the Quinn library. It manages a QUIC connection to a remote peer and
/// provides methods for sending and receiving data over that connection.
///
/// QUIC provides several benefits for secure multiparty computation:
/// - Built-in encryption and authentication
/// - Reliable, ordered delivery of data
/// - Stream multiplexing for concurrent operations
/// - Connection migration for network changes
pub struct QuicPeerConnection {
    /// The underlying QUIC connection
    connection: Connection,
    /// The remote peer's address
    remote_addr: SocketAddr,
    /// Map of stream IDs to send/receive stream pairs
    streams: Arc<Mutex<HashMap<u64, (quinn::SendStream, quinn::RecvStream)>>>,
    /// Whether this connection is on the server side
    is_server: bool,
}

impl QuicPeerConnection {
    /// Creates a new QUIC peer connection
    ///
    /// # Arguments
    /// * `connection` - The underlying QUIC connection
    /// * `is_server` - Whether this connection is on the server side
    ///
    /// The `is_server` parameter determines the behavior when creating new streams:
    /// - Server connections accept incoming streams
    /// - Client connections open new streams
    pub fn new(connection: Connection, is_server: bool) -> Self {
        let remote_addr = connection.remote_address();
        Self {
            connection,
            remote_addr,
            streams: Arc::new(Mutex::new(HashMap::new())),
            is_server,
        }
    }

    /// Gets or creates a bidirectional stream with the given ID
    ///
    /// This method manages the lifecycle of QUIC streams:
    /// 1. If a stream with the given ID already exists, it is reused
    /// 2. Otherwise, a new stream is created:
    ///    - For server connections, by accepting an incoming stream
    ///    - For client connections, by opening a new stream
    ///
    /// # Arguments
    /// * `stream_id` - The ID of the stream to get or create
    ///
    /// # Returns
    /// * `Ok((SendStream, RecvStream))` - The send and receive halves of the stream
    /// * `Err(String)` - If the stream could not be created
    async fn get_or_create_stream(
        &mut self,
        stream_id: u64,
    ) -> Result<(quinn::SendStream, quinn::RecvStream), String> {
        let mut streams = self.streams.lock().await;
        if let Some((send, recv)) = streams.remove(&stream_id) {
            // Reuse existing stream
            Ok((send, recv))
        } else {
            drop(streams); // Release the lock before async operations
            if self.is_server {
                // Server should accept incoming streams
                let (send, recv) = self
                    .connection
                    .accept_bi()
                    .await
                    .map_err(|e| format!("Failed to accept bidirectional stream: {}", e))?;
                Ok((send, recv))
            } else {
                // Client should create new streams
                let (send, recv) = self
                    .connection
                    .open_bi()
                    .await
                    .map_err(|e| format!("Failed to open bidirectional stream: {}", e))?;
                Ok((send, recv))
            }
        }
    }
}

impl PeerConnection for QuicPeerConnection {
    fn send<'a>(&'a mut self, data: &'a [u8]) -> PeerFuture<'a, ()> {
        Box::pin(async move { self.send_on_stream(0, data).await })
    }

    fn receive<'a>(&'a mut self) -> PeerFuture<'a, Vec<u8>> {
        Box::pin(async move { self.receive_from_stream(0).await })
    }

    fn send_on_stream<'a>(&'a mut self, stream_id: u64, data: &'a [u8]) -> PeerFuture<'a, ()> {
        Box::pin(async move {
            let (mut send, recv) = self.get_or_create_stream(stream_id).await?;

            send.write_all(data)
                .await
                .map_err(|e| format!("Failed to send data: {}", e))?;

            // Store the stream back for reuse
            let mut streams = self.streams.lock().await;
            streams.insert(stream_id, (send, recv));

            Ok(())
        })
    }

    fn receive_from_stream<'a>(&'a mut self, stream_id: u64) -> PeerFuture<'a, Vec<u8>> {
        Box::pin(async move {
            let (send, mut recv) = self.get_or_create_stream(stream_id).await?;

            // Read a chunk of data (up to 65536 bytes)
            let mut buf = vec![0u8; 65536];
            match recv.read(&mut buf).await {
                Ok(Some(n)) => {
                    buf.truncate(n);

                    // Store the stream back for reuse
                    let mut streams = self.streams.lock().await;
                    streams.insert(stream_id, (send, recv));

                    Ok(buf)
                }
                Ok(None) => Err("Connection closed by peer".to_string()),
                Err(e) => Err(format!("Failed to receive data: {}", e)),
            }
        })
    }

    fn remote_address(&self) -> SocketAddr {
        self.remote_addr
    }

    fn close<'a>(&'a mut self) -> PeerFuture<'a, ()> {
        Box::pin(async move {
            self.connection.close(0u32.into(), b"Connection closed");
            Ok(())
        })
    }
}

/// A node in the QUIC network
///
/// This struct represents a participant in the secure multiparty computation
/// network. It implements the Node trait from stoffelmpc-network.
#[derive(Debug, Clone)]
pub struct QuicNode {
    /// The UUID of this node
    uuid: Uuid,
    /// Explicit MPC/network party identifier.
    party_id: PartyId,
    /// The network address of this node
    address: SocketAddr,
}

impl QuicNode {
    /// Creates a new node with a random UUID
    ///
    /// # Arguments
    /// * `address` - The network address of the node
    pub fn new_with_random_id(address: SocketAddr) -> Self {
        let uuid = Uuid::new_v4();
        Self {
            uuid,
            party_id: party_id_from_uuid(uuid),
            address,
        }
    }

    /// Creates a new node with a specific UUID
    ///
    /// # Arguments
    /// * `uuid` - The UUID of the node
    /// * `address` - The network address of the node
    pub fn new(uuid: Uuid, address: SocketAddr) -> Self {
        Self {
            uuid,
            party_id: party_id_from_uuid(uuid),
            address,
        }
    }

    /// Creates a new node with a specific ID
    ///
    /// # Arguments
    /// * `id` - The explicit party ID of the node
    /// * `address` - The network address of the node
    pub fn from_party_id(id: PartyId, address: SocketAddr) -> Self {
        let uuid = Uuid::from_u128(party_id_to_u128(id));
        Self {
            uuid,
            party_id: id,
            address,
        }
    }

    /// Returns the network address of this node
    pub fn address(&self) -> SocketAddr {
        self.address
    }

    /// Returns the UUID of this node
    pub fn uuid(&self) -> Uuid {
        self.uuid
    }
}

impl Node for QuicNode {
    fn id(&self) -> PartyId {
        self.party_id
    }

    fn scalar_id<F: Field>(&self) -> F {
        F::from(party_id_to_u128(self.party_id))
    }
}

/// Configuration for the QUIC network
///
/// This struct contains configuration parameters for the QUIC network,
/// such as timeout values, retry settings, and other network-specific options.
#[derive(Debug, Clone)]
pub struct QuicNetworkConfig {
    /// Timeout for network operations in milliseconds
    pub timeout_ms: u64,
    /// Maximum number of retry attempts for network operations
    pub max_retries: u32,
    /// Whether to use secure connections (TLS)
    pub use_tls: bool,
}

impl Default for QuicNetworkConfig {
    fn default() -> Self {
        Self {
            timeout_ms: 30000, // 30 seconds
            max_retries: 3,
            use_tls: true,
        }
    }
}

/// A message type for QUIC-based communication
///
/// This struct implements the Message trait from stoffelmpc-network,
/// providing a standard way to serialize and deserialize messages
/// for secure multiparty computation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QuicMessage {
    /// The ID of the sender of this message
    sender_id: PartyId,
    /// The actual message content
    content: Vec<u8>,
}

impl QuicMessage {
    /// Creates a new message
    ///
    /// # Arguments
    /// * `sender_id` - The ID of the sender
    /// * `content` - The content of the message
    pub fn new(sender_id: PartyId, content: Vec<u8>) -> Self {
        Self { sender_id, content }
    }

    /// Returns the content of the message
    pub fn content(&self) -> &[u8] {
        &self.content
    }
}

impl Message for QuicMessage {
    fn sender_id(&self) -> PartyId {
        self.sender_id
    }

    fn bytes(&self) -> &[u8] {
        &self.content
    }
}

/// QUIC-based implementation of NetworkManager
///
/// This struct implements the NetworkManager trait using the QUIC protocol
/// via the Quinn library. It manages QUIC endpoints for both client and server
/// roles, and provides methods for establishing connections and accepting
/// incoming connections.
///
/// The implementation uses self-signed certificates for TLS, which is suitable
/// for development but should be replaced with proper certificate management
/// in production.
pub struct QuicNetworkManager {
    /// The QUIC endpoint for sending and receiving connections
    endpoint: Option<Endpoint>,
    /// Configuration for the server role
    #[allow(dead_code)]
    server_config: Option<ServerConfig>,
    /// Configuration for the client role
    #[allow(dead_code)]
    client_config: Option<ClientConfig>,
    /// The nodes in the network
    nodes: Vec<QuicNode>,
    /// The ID of this node in the network
    node_id: PartyId,
    /// Network configuration
    network_config: QuicNetworkConfig,
    /// Active connections to other server nodes and clients are handled in stoffelnet impl.
    /// Each connection has its own lock to avoid holding the HashMap lock across await points.
    connections: PeerConnectionMap,
}

impl Default for QuicNetworkManager {
    fn default() -> Self {
        Self::new()
    }
}

impl QuicNetworkManager {
    /// Creates a new QUIC network manager
    ///
    /// This initializes a network manager with no active endpoints or configurations.
    /// Before using the manager, you must call either `connect()` or `listen()`
    /// to set up the appropriate endpoint.
    pub fn new() -> Self {
        Self {
            endpoint: None,
            server_config: None,
            client_config: None,
            nodes: Vec::new(),
            node_id: random_party_id(),
            network_config: QuicNetworkConfig::default(),
            connections: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Creates a new QUIC network manager with the specified node ID
    ///
    /// # Arguments
    /// * `node_id` - The ID of this node in the network
    pub fn with_node_id(node_id: PartyId) -> Self {
        let mut manager = Self::new();
        manager.node_id = node_id;
        manager
    }

    /// Creates a new QUIC network manager with a random UUID-based node ID
    pub fn with_random_id() -> Self {
        Self::new() // new() already generates a random UUID-based ID
    }

    /// Creates a new QUIC network manager with the specified configuration
    ///
    /// # Arguments
    /// * `config` - The network configuration
    pub fn with_config(config: QuicNetworkConfig) -> Self {
        let mut manager = Self::new();
        manager.network_config = config;
        manager
    }

    fn default_transport_config() -> Result<TransportConfig, String> {
        let idle_timeout = IdleTimeout::try_from(DEFAULT_IDLE_TIMEOUT)
            .map_err(|e| format!("Invalid QUIC idle timeout: {e}"))?;
        let mut transport = TransportConfig::default();
        transport.max_concurrent_uni_streams(0u32.into());
        transport.max_idle_timeout(Some(idle_timeout));
        transport.keep_alive_interval(Some(DEFAULT_KEEP_ALIVE_INTERVAL));
        Ok(transport)
    }

    /// Return the local address currently bound by this manager's endpoint.
    pub fn local_addr(&self) -> Result<SocketAddr, String> {
        self.endpoint
            .as_ref()
            .ok_or_else(|| {
                "Endpoint not initialized. Call listen() or connect() first.".to_string()
            })?
            .local_addr()
            .map_err(|e| format!("Failed to read endpoint local address: {e}"))
    }

    /// Adds a node to the network
    ///
    /// # Arguments
    /// * `node` - The node to add
    pub fn add_node(&mut self, node: QuicNode) {
        self.nodes.push(node);
    }

    /// Adds a node with a random UUID to the network
    ///
    /// # Arguments
    /// * `address` - The network address of the node
    pub fn add_node_with_random_id(&mut self, address: SocketAddr) {
        let node = QuicNode::new_with_random_id(address);
        self.nodes.push(node);
    }

    /// Adds a node with a specific UUID to the network
    ///
    /// # Arguments
    /// * `uuid` - The UUID of the node
    /// * `address` - The network address of the node
    pub fn add_node_with_uuid(&mut self, uuid: Uuid, address: SocketAddr) {
        let node = QuicNode::new(uuid, address);
        self.nodes.push(node);
    }

    /// Adds a node with a specific party ID to the network
    ///
    /// # Arguments
    /// * `id` - The ID of the node
    /// * `address` - The network address of the node
    pub fn add_node_with_party_id(&mut self, id: PartyId, address: SocketAddr) {
        let node = QuicNode::from_party_id(id, address);
        self.nodes.push(node);
    }

    /// Creates an insecure client configuration for QUIC
    ///
    /// This method creates a client configuration that:
    /// 1. Skips server certificate verification (insecure, but useful for development)
    /// 2. Sets up ALPN protocols for protocol negotiation
    /// 3. Configures transport parameters
    ///
    /// # Warning
    /// This configuration is insecure and should only be used for development.
    /// In production, proper certificate verification should be implemented.
    ///
    /// # Returns
    /// * `Ok(ClientConfig)` - The client configuration
    /// * `Err(String)` - If the configuration could not be created
    fn create_insecure_client_config() -> Result<ClientConfig, String> {
        // Create a client crypto configuration that skips certificate verification
        let mut crypto = rustls::ClientConfig::builder()
            .dangerous()
            .with_custom_certificate_verifier(Arc::new(SkipServerVerification::new()))
            .with_no_client_auth();

        // Set ALPN protocol to match the server
        crypto.alpn_protocols = vec![b"quic-example".to_vec()];

        // Create a QUIC client configuration with the crypto configuration
        let mut config = ClientConfig::new(Arc::new(
            quinn::crypto::rustls::QuicClientConfig::try_from(crypto)
                .map_err(|e| format!("Failed to create QUIC client config: {}", e))?,
        ));

        // Set transport config with reasonable timeouts
        config.transport_config(Arc::new(Self::default_transport_config()?));

        Ok(config)
    }

    /// Creates a self-signed server configuration for QUIC
    ///
    /// This method creates a server configuration that:
    /// 1. Generates a self-signed certificate for TLS
    /// 2. Sets up ALPN protocols for protocol negotiation
    /// 3. Configures transport parameters
    ///
    /// # Warning
    /// This configuration uses a self-signed certificate, which is suitable for
    /// development but not for production. In production, proper certificates
    /// should be used.
    ///
    /// # Returns
    /// * `Ok(ServerConfig)` - The server configuration
    /// * `Err(String)` - If the configuration could not be created
    fn create_self_signed_server_config() -> Result<ServerConfig, String> {
        // Generate self-signed certificate
        let cert = rcgen::generate_simple_self_signed(vec!["localhost".into()])
            .map_err(|e| format!("Failed to generate certificate: {}", e))?;

        // Convert the certificate and key to DER format
        let cert_der = CertificateDer::from(cert.cert.der().to_vec());
        let key_der =
            PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(cert.signing_key.serialize_der()));

        // Create a server crypto configuration with the certificate
        let mut server_crypto = rustls::ServerConfig::builder()
            .with_no_client_auth()
            .with_single_cert(vec![cert_der], key_der)
            .map_err(|e| format!("Failed to create server crypto config: {}", e))?;

        // Set ALPN protocol
        server_crypto.alpn_protocols = vec![b"quic-example".to_vec()];

        // Create a QUIC server configuration with the crypto configuration
        let mut server_config = ServerConfig::with_crypto(Arc::new(
            quinn::crypto::rustls::QuicServerConfig::try_from(server_crypto)
                .map_err(|e| format!("Failed to create QUIC server config: {}", e))?,
        ));

        // Configure transport parameters with reasonable timeouts
        server_config.transport = Arc::new(Self::default_transport_config()?);

        Ok(server_config)
    }
}

impl NetworkManager for QuicNetworkManager {
    fn connect<'a>(&'a mut self, address: SocketAddr) -> PeerFuture<'a, BoxedPeerConnection> {
        Box::pin(async move {
            // Create client config for outgoing connections
            let client_config = Self::create_insecure_client_config()?;

            // Determine which endpoint to use for the connection
            let connection = if let Some(endpoint) = self.endpoint.as_mut() {
                // Try to use existing server endpoint with client config
                endpoint.set_default_client_config(client_config.clone());
                eprintln!("[quic] Using existing endpoint to connect to {}", address);

                match endpoint.connect(address, "localhost") {
                    Ok(connecting) => {
                        eprintln!("[quic] Awaiting connection handshake to {}", address);
                        connecting.await.map_err(|e| {
                            format!("Failed to establish connection to {}: {}", address, e)
                        })?
                    }
                    Err(e) => {
                        // Server endpoint failed, create a dedicated client endpoint
                        eprintln!(
                            "[quic] Server endpoint connect failed ({}), creating client endpoint for {}",
                            e, address
                        );
                        let mut client_endpoint = Endpoint::client(client_endpoint_bind_addr())
                            .map_err(|e| format!("Failed to create client endpoint: {}", e))?;
                        client_endpoint.set_default_client_config(client_config);

                        let connecting =
                            client_endpoint.connect(address, "localhost").map_err(|e| {
                                format!("Failed to initiate connection to {}: {}", address, e)
                            })?;

                        eprintln!(
                            "[quic] Awaiting connection handshake to {} (client endpoint)",
                            address
                        );
                        connecting.await.map_err(|e| {
                            format!("Failed to establish connection to {}: {}", address, e)
                        })?
                    }
                }
            } else {
                // No endpoint exists, create a client endpoint
                eprintln!(
                    "[quic] Creating new client endpoint to connect to {}",
                    address
                );
                let mut endpoint = Endpoint::client(client_endpoint_bind_addr())
                    .map_err(|e| format!("Failed to create client endpoint: {}", e))?;
                endpoint.set_default_client_config(client_config);
                let endpoint = self.endpoint.insert(endpoint);

                let connecting = endpoint
                    .connect(address, "localhost")
                    .map_err(|e| format!("Failed to initiate connection to {}: {}", address, e))?;

                eprintln!("[quic] Awaiting connection handshake to {}", address);
                connecting
                    .await
                    .map_err(|e| format!("Failed to establish connection to {}: {}", address, e))?
            };

            eprintln!("[quic] Connection established to {}", address);

            // Send identification handshake as SERVER with our node_id for stoffelnet compatibility
            if let Ok((mut send, _recv)) = connection.open_bi().await {
                let handshake = format!("ROLE:SERVER:{}\n", self.node_id);
                let _ = send.write_all(handshake.as_bytes()).await;
            }

            // Find the node ID for this address or generate a new one
            let node_id = self
                .nodes
                .iter()
                .find(|node| node.address() == address)
                .map(|node| node.id())
                .unwrap_or_else(|| {
                    // If we don't have a node for this address, create one
                    let node = QuicNode::new_with_random_id(address);
                    let id = node.id();
                    self.nodes.push(node);
                    id
                });

            // Store a clone of the connection in the connections hashmap
            let mut connections = self.connections.lock().await;
            connections.insert(
                node_id,
                Arc::new(Mutex::new(
                    Box::new(QuicPeerConnection::new(connection.clone(), false))
                        as Box<dyn PeerConnection>,
                )),
            );

            // Return the original connection
            Ok(Box::new(QuicPeerConnection::new(connection, false)) as Box<dyn PeerConnection>)
        })
    }

    fn accept<'a>(&'a mut self) -> PeerFuture<'a, BoxedPeerConnection> {
        Box::pin(async move {
            let endpoint = self
                .endpoint
                .as_ref()
                .ok_or_else(|| "Endpoint not initialized. Call listen() first.".to_string())?;

            let incoming = endpoint
                .accept()
                .await
                .ok_or_else(|| "No incoming connections".to_string())?;

            let connection = incoming
                .await
                .map_err(|e| format!("Failed to accept connection: {}", e))?;

            // Get the remote address of the connection
            let remote_addr = connection.remote_address();

            // Try to read a role identification handshake
            let mut parsed_id: Option<PartyId> = None;
            if let Ok((mut _send, mut recv)) = connection.accept_bi().await {
                let mut buf = vec![0u8; 256];
                if let Ok(Some(n)) = recv.read(&mut buf).await {
                    let line = String::from_utf8_lossy(&buf[..n])
                        .lines()
                        .next()
                        .unwrap_or("")
                        .to_string();
                    if let Some(rest) = line.strip_prefix("ROLE:SERVER:") {
                        if let Ok(id) = rest.trim().parse::<usize>() {
                            parsed_id = Some(id);
                        }
                    }
                }
            }
            let node_id = parsed_id.unwrap_or_else(|| {
                let node = QuicNode::new_with_random_id(remote_addr);
                let id = node.id();
                self.nodes.push(node);
                id
            });
            let mut connections = self.connections.lock().await;
            connections.insert(
                node_id,
                Arc::new(Mutex::new(
                    Box::new(QuicPeerConnection::new(connection.clone(), true))
                        as Box<dyn PeerConnection>,
                )),
            );

            // Return the original connection
            Ok(Box::new(QuicPeerConnection::new(connection, true)) as Box<dyn PeerConnection>)
        })
    }

    fn listen<'a>(&'a mut self, bind_address: SocketAddr) -> PeerFuture<'a, ()> {
        Box::pin(async move {
            let server_config = Self::create_self_signed_server_config()?;
            let endpoint = Endpoint::server(server_config, bind_address)
                .map_err(|e| format!("Failed to create server endpoint: {}", e))?;

            self.endpoint = Some(endpoint);
            Ok(())
        })
    }
}

/// Implementation of the Network trait for QuicNetworkManager
///
/// This implementation uses the QUIC protocol for communication between nodes.
#[async_trait]
impl Network for QuicNetworkManager {
    type NodeType = QuicNode;
    type NetworkConfig = QuicNetworkConfig;

    async fn send(&self, recipient: PartyId, message: &[u8]) -> Result<usize, NetworkError> {
        // Get the connection Arc while holding the HashMap lock briefly
        let connection_arc = {
            let connections = self.connections.lock().await;
            connections.get(&recipient).cloned()
        };
        // HashMap lock is now released

        let connection_arc = connection_arc.ok_or(NetworkError::PartyNotFound(recipient))?;

        // Log debug info about the outgoing message
        let preview_len = message.len().min(64);
        let hex_preview: String = message[..preview_len]
            .iter()
            .map(|b| format!("{:02x}", b))
            .collect();
        debug!(
            "[MSG:SEND] from node {} to node {} ({} bytes, hex[0..{}]={})",
            self.node_id,
            recipient,
            message.len(),
            preview_len,
            hex_preview
        );

        // Lock the individual connection and send (HashMap lock not held)
        let mut connection = connection_arc.lock().await;
        match connection.send(message).await {
            Ok(_) => {
                debug!(
                    "[MSG:SENT] from node {} to node {} ({} bytes)",
                    self.node_id,
                    recipient,
                    message.len()
                );
                Ok(message.len())
            }
            Err(e) => {
                debug!(
                    "[MSG:SEND-FAIL] from node {} to node {}: {}",
                    self.node_id, recipient, e
                );
                Err(NetworkError::SendError)
            }
        }
    }

    async fn broadcast(&self, message: &[u8]) -> Result<usize, NetworkError> {
        let mut total_bytes = 0;

        // Prepare a hex preview
        let preview_len = message.len().min(64);
        let hex_preview: String = message[..preview_len]
            .iter()
            .map(|b| format!("{:02x}", b))
            .collect();
        debug!(
            "[MSG:BROADCAST] from node {} to all ({} bytes, hex[0..{}]={})",
            self.node_id,
            message.len(),
            preview_len,
            hex_preview
        );

        // Collect connection Arcs while holding the HashMap lock briefly
        let targets: Vec<(PartyId, Arc<Mutex<Box<dyn PeerConnection>>>)> = {
            let connections = self.connections.lock().await;
            self.nodes
                .iter()
                .filter(|node| node.id() != self.node_id)
                .filter_map(|node| {
                    connections
                        .get(&node.id())
                        .map(|conn| (node.id(), conn.clone()))
                })
                .collect()
        };
        // HashMap lock is now released

        // Log skipped nodes
        for node in &self.nodes {
            if node.id() != self.node_id && !targets.iter().any(|(id, _)| *id == node.id()) {
                debug!(
                    "[MSG:BROADCAST-SKIP] from node {} -> node {}: no connection",
                    self.node_id,
                    node.id()
                );
            }
        }

        // Send to each target without holding the HashMap lock
        for (node_id, connection_arc) in targets {
            let mut connection = connection_arc.lock().await;
            match connection.send(message).await {
                Ok(_) => {
                    debug!(
                        "[MSG:BROADCAST-SENT] from node {} -> node {} ({} bytes)",
                        self.node_id,
                        node_id,
                        message.len()
                    );
                    total_bytes += message.len();
                }
                Err(e) => {
                    debug!(
                        "[MSG:BROADCAST-FAIL] from node {} -> node {}: {}",
                        self.node_id, node_id, e
                    );
                    // Continue with other nodes even if one fails
                }
            }
        }

        if total_bytes > 0 {
            Ok(total_bytes)
        } else {
            // If we didn't send any messages, return an error
            Err(NetworkError::SendError)
        }
    }

    fn parties(&self) -> Vec<&Self::NodeType> {
        self.nodes.iter().collect()
    }

    fn parties_mut(&mut self) -> Vec<&mut Self::NodeType> {
        self.nodes.iter_mut().collect()
    }

    fn config(&self) -> &Self::NetworkConfig {
        &self.network_config
    }

    fn node(&self, id: PartyId) -> Option<&Self::NodeType> {
        self.nodes.iter().find(|node| node.id() == id)
    }

    fn node_mut(&mut self, id: PartyId) -> Option<&mut Self::NodeType> {
        self.nodes.iter_mut().find(|node| node.id() == id)
    }

    async fn send_to_client(
        &self,
        client: ClientId,
        _message: &[u8],
    ) -> Result<usize, NetworkError> {
        // Not used in current VM path
        Err(NetworkError::ClientNotFound(client))
    }

    fn clients(&self) -> Vec<ClientId> {
        Vec::new()
    }

    fn is_client_connected(&self, _client: ClientId) -> bool {
        false
    }

    // --- party identification ---

    fn local_party_id(&self) -> PartyId {
        self.node_id
    }

    fn party_count(&self) -> usize {
        self.nodes.len()
    }

    fn verified_ordering(&self) -> Option<VerifiedOrdering> {
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn loopback_addr() -> SocketAddr {
        SocketAddr::from(([127, 0, 0, 1], 0))
    }

    #[test]
    fn client_endpoint_bind_address_is_unspecified_ephemeral_ipv4() {
        let addr = client_endpoint_bind_addr();
        assert!(addr.ip().is_unspecified());
        assert_eq!(addr.port(), 0);
    }

    #[test]
    fn default_transport_config_is_constructible_without_panics() {
        QuicNetworkManager::default_transport_config()
            .expect("valid default QUIC transport config");
    }

    #[test]
    fn quic_node_from_party_id_preserves_explicit_party_id() {
        let node = QuicNode::from_party_id(PartyId::MAX, loopback_addr());

        assert_eq!(node.id(), PartyId::MAX);
    }

    #[test]
    fn quic_node_uuid_constructor_uses_explicit_uuid_fold() {
        let uuid = Uuid::from_u128(party_id_to_u128(PartyId::MAX) + 1);
        let node = QuicNode::new(uuid, loopback_addr());

        assert_eq!(node.uuid(), uuid);
        assert_eq!(node.id(), party_id_from_uuid(uuid));
    }

    #[test]
    fn quic_node_scalar_id_uses_party_id() {
        let node = QuicNode::from_party_id(42, loopback_addr());

        assert_eq!(
            node.scalar_id::<ark_bn254::Fr>(),
            ark_bn254::Fr::from(42u128)
        );
    }

    #[test]
    fn local_addr_reports_uninitialized_endpoint() {
        let manager = QuicNetworkManager::new();

        let err = manager
            .local_addr()
            .expect_err("endpoint is not initialized");
        assert!(err.contains("Endpoint not initialized"));
    }
}

/// Certificate verifier that accepts any server certificate
///
/// This is a dummy implementation of the ServerCertVerifier trait that
/// accepts any server certificate without verification. It is used for
/// development and testing purposes only.
///
/// # Security Warning
///
/// This implementation is **extremely insecure** and vulnerable to
/// man-in-the-middle attacks. It should never be used in production.
/// In a production environment, proper certificate verification should
/// be implemented, typically using a trusted certificate authority.
#[derive(Debug)]
struct SkipServerVerification;

impl SkipServerVerification {
    /// Creates a new SkipServerVerification instance
    ///
    /// This is a simple constructor that returns a new instance of
    /// the SkipServerVerification struct.
    fn new() -> Self {
        Self
    }
}

impl rustls::client::danger::ServerCertVerifier for SkipServerVerification {
    fn verify_server_cert(
        &self,
        _end_entity: &CertificateDer<'_>,
        _intermediates: &[CertificateDer<'_>],
        _server_name: &rustls::pki_types::ServerName<'_>,
        _ocsp: &[u8],
        _now: rustls::pki_types::UnixTime,
    ) -> Result<rustls::client::danger::ServerCertVerified, rustls::Error> {
        Ok(rustls::client::danger::ServerCertVerified::assertion())
    }

    fn verify_tls12_signature(
        &self,
        message: &[u8],
        cert: &CertificateDer<'_>,
        dss: &rustls::DigitallySignedStruct,
    ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
        rustls::crypto::verify_tls12_signature(
            message,
            cert,
            dss,
            &rustls::crypto::ring::default_provider().signature_verification_algorithms,
        )
    }

    fn verify_tls13_signature(
        &self,
        message: &[u8],
        cert: &CertificateDer<'_>,
        dss: &rustls::DigitallySignedStruct,
    ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
        rustls::crypto::verify_tls13_signature(
            message,
            cert,
            dss,
            &rustls::crypto::ring::default_provider().signature_verification_algorithms,
        )
    }

    fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
        rustls::crypto::ring::default_provider()
            .signature_verification_algorithms
            .supported_schemes()
    }
}