unb-server 2.0.3

use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

use unb_client::{Endpoint, EndpointSet};
use unb_runtime::WsError;

use crate::connection::{ConnectError, PeerConnection};
use crate::node::{Node, PeerLink};
use crate::session::{CandidateFailure, CandidateOutcome, CandidateSession};

/// Supplies a transport for one already-ordered endpoint candidate.
///
/// This specialist hook lets platform bindings retain the server's connection
/// lifecycle while constructing transports in their native runtime.
#[doc(hidden)]
pub trait EndpointDialer: Send + Sync {
    fn supports(&self, kind: unb_client::TransportKind) -> bool;

    fn dial(
        &self,
        endpoint: Endpoint,
    ) -> Pin<Box<dyn Future<Output = Result<unb_runtime::Pipe, WsError>> + Send + 'static>>;
}

pub(crate) struct ReconnectCandidate {
    pub(crate) identity: unb_core::NodeIdentity,
    pub(crate) selected: PeerLink,
    pub(crate) candidate_wire: Arc<unb_runtime::Wire>,
}

impl Node {
    /// Eagerly dial, verify the peer, and synchronize its routes.
    ///
    /// The returned logical connection is retained by the node and automatically
    /// maintains an unintentionally lost selected session.
    pub async fn connect(
        self: &Arc<Self>,
        endpoints: impl Into<EndpointSet>,
    ) -> Result<PeerConnection, ConnectError> {
        self.connect_using(endpoints.into(), None, None).await
    }

    #[cfg(feature = "hosting")]
    pub(crate) async fn connect_expected(
        self: &Arc<Self>,
        endpoints: impl Into<EndpointSet>,
        expected_peer: &str,
    ) -> Result<PeerConnection, ConnectError> {
        self.connect_using(endpoints.into(), None, Some(expected_peer))
            .await
    }

    /// Connect using a platform-owned transport constructor.
    ///
    /// Endpoint ordering, establishment, identity verification, route sync,
    /// logical-handle convergence, and reconnect coordination remain owned by
    /// the server. Only construction of each candidate transport is delegated.
    #[doc(hidden)]
    pub async fn connect_with_dialer(
        self: &Arc<Self>,
        endpoints: impl Into<EndpointSet>,
        dialer: Arc<dyn EndpointDialer>,
    ) -> Result<PeerConnection, ConnectError> {
        self.connect_using(endpoints.into(), Some(dialer), None)
            .await
    }

    async fn connect_using(
        self: &Arc<Self>,
        set: EndpointSet,
        dialer: Option<Arc<dyn EndpointDialer>>,
        expected_peer: Option<&str>,
    ) -> Result<PeerConnection, ConnectError> {
        let key = set.cache_key();
        let ordered = self
            .dial_policy
            .ordered_candidates(&key, &set)
            .into_iter()
            .filter(|endpoint| {
                dialer
                    .as_ref()
                    .is_none_or(|dialer| dialer.supports(endpoint.kind))
            })
            .collect::<Vec<_>>();
        if ordered.is_empty() {
            return Err(ConnectError::NoSupportedEndpoint);
        }
        let mut last_error = ConnectError::NoSupportedEndpoint;
        for endpoint in ordered {
            match self
                .try_candidate(&endpoint, expected_peer, dialer.as_ref())
                .await
            {
                Ok((candidate, outcome)) => {
                    let identity = match outcome {
                        CandidateOutcome::Promoted(identity)
                        | CandidateOutcome::Duplicate(identity) => identity,
                    };
                    let Some(link) = self.peer(&identity.node_id).await else {
                        candidate.wire.shutdown();
                        let _ = candidate.cleaned.await;
                        last_error = ConnectError::Establishment {
                            message: format!(
                                "verified peer {:?} has no selected live session",
                                identity.node_id
                            ),
                        };
                        continue;
                    };
                    if n0_future::time::timeout(
                        crate::session::ROUTE_SYNC_TIMEOUT,
                        self.wait_for_selected_route(&identity.node_id, &link.session_id),
                    )
                    .await
                    .is_err()
                    {
                        candidate.wire.shutdown();
                        let _ = candidate.cleaned.await;
                        last_error = ConnectError::Establishment {
                            message: format!(
                                "verified peer {:?} did not publish its synchronized node route",
                                identity.node_id
                            ),
                        };
                        continue;
                    }
                    let connection = {
                        let mut connections = self
                            .connections
                            .write()
                            .unwrap_or_else(|poisoned| poisoned.into_inner());
                        if let Some(connection) = connections
                            .get(&identity.node_id)
                            .filter(|connection| !connection.is_terminal())
                            .cloned()
                        {
                            if connection.bind(
                                identity.clone(),
                                link.session_id.clone(),
                                link.wire.clone(),
                            ) {
                                connection.replace_endpoints(set.clone());
                                connection.replace_dialer(dialer.clone());
                                connection
                            } else {
                                let connection = PeerConnection::new(
                                    Arc::downgrade(self),
                                    identity.clone(),
                                    set.clone(),
                                    link.session_id.clone(),
                                    link.wire.clone(),
                                    dialer.clone(),
                                );
                                connections.insert(identity.node_id, connection.clone());
                                connection
                            }
                        } else {
                            let connection = PeerConnection::new(
                                Arc::downgrade(self),
                                identity.clone(),
                                set.clone(),
                                link.session_id.clone(),
                                link.wire.clone(),
                                dialer.clone(),
                            );
                            connections.insert(identity.node_id, connection.clone());
                            connection
                        }
                    };
                    self.dial_policy.record_winner(&key, endpoint.kind);
                    return Ok(connection);
                }
                Err(error) => last_error = error,
            }
        }
        Err(last_error)
    }

    async fn wait_for_selected_route(&self, peer: &str, session_id: &str) {
        let mut changes = self.route_changes();
        loop {
            let selected_session = self.peer(peer).await.map(|selected| selected.session_id);
            let direct_route = matches!(
                self.snapshot.load().node_core.resolve(peer),
                unb_core::Resolution::Route(next_hop) if next_hop == peer
            );
            if selected_session.as_deref() == Some(session_id) && direct_route {
                return;
            }
            if changes.changed().await.is_err() {
                return;
            }
        }
    }

    /// Wait until a directly linked peer is both selectable and routable.
    ///
    /// A link's transport handshake can finish before its first route snapshot
    /// is applied. `link` promises immediate request readiness, so it must wait
    /// for this local publication rather than merely for the handshake ACK.
    async fn wait_for_direct_peer_route(&self, peer: &str) {
        let mut changes = self.route_changes();
        loop {
            let selected = self.peer(peer).await.is_some();
            let direct_route = matches!(
                self.snapshot.load().node_core.resolve(peer),
                unb_core::Resolution::Route(next_hop) if next_hop == peer
            );
            if selected && direct_route {
                return;
            }
            if changes.changed().await.is_err() {
                return;
            }
        }
    }

    async fn try_candidate(
        self: &Arc<Self>,
        endpoint: &Endpoint,
        expected_peer: Option<&str>,
        dialer: Option<&Arc<dyn EndpointDialer>>,
    ) -> Result<(CandidateSession, CandidateOutcome), ConnectError> {
        let deadline = self.dial_policy.attempt_timeout();
        let candidate_dial = async {
            match dialer {
                Some(dialer) => dialer.dial(endpoint.clone()).await,
                None => self.dial_policy.dial_candidate(endpoint).await,
            }
        };
        let pipe = match n0_future::time::timeout(deadline, candidate_dial).await {
            Ok(Ok(pipe)) => pipe,
            Ok(Err(error)) => {
                return Err(ConnectError::Dial {
                    transport: endpoint.kind,
                    message: error.to_string(),
                })
            }
            Err(_) => {
                return Err(ConnectError::DialTimedOut {
                    transport: endpoint.kind,
                })
            }
        };
        let candidate = self.establish(pipe, expected_peer.map(str::to_owned)).await;
        let outcome = candidate.observed_outcome().await;
        match outcome {
            Ok(outcome) => Ok((candidate, outcome)),
            Err(failure) => {
                candidate.wire.shutdown();
                let _ = candidate.cleaned.await;
                Err(match (expected_peer, failure) {
                    (
                        Some(expected),
                        CandidateFailure::Retired {
                            reason: unb_core::RetirementReason::UnexpectedPeer,
                            identity,
                        },
                    ) => ConnectError::IdentityMismatch {
                        expected: expected.to_string(),
                        actual: identity.map(|identity| identity.node_id),
                    },
                    (_, CandidateFailure::Session(error)) => ConnectError::Establishment {
                        message: error.to_string(),
                    },
                    (_, CandidateFailure::Retired { reason, .. }) => ConnectError::Establishment {
                        message: format!("session retired during establishment: {reason:?}"),
                    },
                    (_, CandidateFailure::MissingIdentity) => ConnectError::Establishment {
                        message: "session completed without an admitted identity".into(),
                    },
                })
            }
        }
    }

    pub(crate) async fn reconnect_peer(
        self: &Arc<Self>,
        peer: &str,
        set: &EndpointSet,
        dialer: Option<Arc<dyn EndpointDialer>>,
    ) -> Result<ReconnectCandidate, ConnectError> {
        let key = set.cache_key();
        let ordered = self
            .dial_policy
            .ordered_candidates(&key, set)
            .into_iter()
            .filter(|endpoint| {
                dialer
                    .as_ref()
                    .is_none_or(|dialer| dialer.supports(endpoint.kind))
            })
            .collect::<Vec<_>>();
        if ordered.is_empty() {
            return Err(ConnectError::NoSupportedEndpoint);
        }
        let mut last_error = ConnectError::NoSupportedEndpoint;
        for endpoint in ordered {
            match self
                .try_candidate(&endpoint, Some(peer), dialer.as_ref())
                .await
            {
                Ok((candidate, outcome)) => {
                    let identity = match outcome {
                        CandidateOutcome::Promoted(identity)
                        | CandidateOutcome::Duplicate(identity) => identity,
                    };
                    let Some(selected) = self.peer(&identity.node_id).await else {
                        candidate.wire.shutdown();
                        let _ = candidate.cleaned.await;
                        last_error = ConnectError::Establishment {
                            message: format!(
                                "verified peer {:?} has no selected live session",
                                identity.node_id
                            ),
                        };
                        continue;
                    };
                    if n0_future::time::timeout(
                        crate::session::ROUTE_SYNC_TIMEOUT,
                        self.wait_for_selected_route(&identity.node_id, &selected.session_id),
                    )
                    .await
                    .is_err()
                    {
                        candidate.wire.shutdown();
                        let _ = candidate.cleaned.await;
                        last_error = ConnectError::Establishment {
                            message: format!(
                                "verified peer {:?} did not publish its synchronized node route",
                                identity.node_id
                            ),
                        };
                        continue;
                    }
                    let Some(selected) = self.peer(&identity.node_id).await else {
                        candidate.wire.shutdown();
                        let _ = candidate.cleaned.await;
                        last_error = ConnectError::Establishment {
                            message: format!(
                                "verified peer {:?} lost its selected session after route synchronization",
                                identity.node_id
                            ),
                        };
                        continue;
                    };
                    self.dial_policy.record_winner(&key, endpoint.kind);
                    return Ok(ReconnectCandidate {
                        identity,
                        selected,
                        candidate_wire: candidate.wire,
                    });
                }
                Err(error) => last_error = error,
            }
        }
        Err(last_error)
    }

    pub async fn link(self: &Arc<Self>, other: &Arc<Node>) -> Result<(), WsError> {
        if Arc::ptr_eq(self, other) || self.identity.node_id == other.identity.node_id {
            return Err(WsError::Connect("a node cannot link to itself".into()));
        }
        let (dial_side, accept_side) = unb_client::pair();
        let left = self
            .establish(dial_side, Some(other.identity.node_id.clone()))
            .await;
        let right = other
            .establish(accept_side, Some(self.identity.node_id.clone()))
            .await;
        let result = match tokio::join!(
            left.outcome(&other.identity.node_id),
            right.outcome(&self.identity.node_id)
        ) {
            (Ok(CandidateOutcome::Promoted(_)), Ok(CandidateOutcome::Promoted(_))) => {
                n0_future::time::timeout(crate::session::ROUTE_SYNC_TIMEOUT, async {
                    left.wire.routes_acked().await?;
                    right.wire.routes_acked().await?;
                    tokio::join!(
                        self.wait_for_direct_peer_route(&other.identity.node_id),
                        other.wait_for_direct_peer_route(&self.identity.node_id),
                    );
                    Ok::<(), WsError>(())
                })
                .await
                .map_err(|_| {
                    WsError::Connect(
                        "linked peers did not publish their synchronized node routes".into(),
                    )
                })?
            }
            (Err(error), _) | (_, Err(error)) => Err(error),
            _ => Err(WsError::Connect("link closed during establishment".into())),
        };
        if result.is_err() {
            left.wire.shutdown();
            right.wire.shutdown();
            let _ = tokio::join!(left.cleaned, right.cleaned);
        }
        result
    }

    /// Attach one caller-provided outbound transport without imposing an expected peer name.
    /// Browser bindings use this after JavaScript completes the platform dial.
    pub async fn connect_transport_unchecked(
        self: &Arc<Self>,
        transport: unb_runtime::Pipe,
    ) -> Result<(), WsError> {
        let candidate = self.establish(transport, None).await;
        match candidate.outcome("candidate").await {
            Ok(CandidateOutcome::Promoted(_) | CandidateOutcome::Duplicate(_)) => {
                let synchronized = n0_future::time::timeout(
                    crate::session::ROUTE_SYNC_TIMEOUT,
                    candidate.wire.routes_acked(),
                )
                .await
                .map_err(|_| {
                    WsError::Connect(
                        "connected peer did not acknowledge its synchronized routes".into(),
                    )
                })
                .and_then(|result| result);
                if let Err(error) = synchronized {
                    candidate.wire.shutdown();
                    let _ = candidate.cleaned.await;
                    return Err(error);
                }
                Ok(())
            }
            Err(error) => {
                candidate.wire.shutdown();
                let _ = candidate.cleaned.await;
                Err(error)
            }
        }
    }
}