zakura-network 6.0.0

//! Native discovery service (stream kind 4) on the Zakura transport.
//!
//! Discovery is a single long-lived ordered stream per peer. Each side runs a
//! [`DiscoverySink`] (the reader, which imports peer records and answers
//! `GetPeers`) and a [`DiscoverySource`] (the writer, which periodically gossips
//! the local self-record and asks for more peers). The wire format is the
//! [`DiscoveryMessage`] payload carried inside a generic transport [`Frame`]
//! (`message_type = DISCOVERY_FRAME_MESSAGE_TYPE`, `flags = 0`), identical to the
//! original native-discovery wire so peers interoperate.

use std::{
    collections::HashMap,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc, Mutex as StdMutex,
    },
    time::Duration,
};

use iroh::NodeId;
use tokio::sync::Notify;
use tokio_util::sync::CancellationToken;

use crate::zakura::{
    handle_pipe_exit, spawn_supervised_peer_task, spawn_supervised_pipe, BlockSyncHandle,
    CloseCause, Flow, Frame, FramedRecv, FramedSend, HeaderSyncEvent, HeaderSyncHandle,
    OrderedSendError, OrderedSessionDemand, OrderedStreamOpening, OrderedStreamPolicy, Peer,
    PeerStreamSession, Pipe, Service, ServiceAdmissionDecision, ServicePeerDirection, SinkReject,
    Stream, StreamMode, ZakuraConnId, ZakuraPeerId, LOCAL_MAX_CONTROL_FRAME_BYTES,
    ZAKURA_CAP_DISCOVERY,
};

#[cfg(test)]
use super::pipe::decode_discovery_frame;
use super::pipe::{discovery_pipe, DsEnv, DsLocal, DISCOVERY_FRAME_MESSAGE_TYPE};
use super::protocol::{
    BlockSyncServiceSummary, DiscoveryBookError, DiscoveryMessage, DiscoveryRecordError,
    GetServices, HeaderSyncServiceSummary, ServiceSummaryEnvelope, Services, ZakuraDiscoveryHandle,
    ZakuraNodeRecord, ZakuraServiceId, DEFAULT_LIVE_SERVICE_SUMMARY_TTL,
    MAX_DISCOVERY_RECORDS_PER_RESPONSE, ZAKURA_DISCOVERY_STREAM_VERSION, ZAKURA_STREAM_DISCOVERY,
};

/// Maximum time discovery waits for first-party exchange responses before releasing the session.
const DISCOVERY_INITIAL_EXCHANGE_TIMEOUT: Duration = Duration::from_secs(2);

const DISCOVERY_SERVICE_STREAMS: [Stream; 1] = [Stream {
    kind: ZAKURA_STREAM_DISCOVERY,
    version: ZAKURA_DISCOVERY_STREAM_VERSION,
    // Advisory until the transport wires Stream::frame_cap end-to-end; the
    // authoritative inbound cap is app_frame_cap_for_stream_kind.
    frame_cap: LOCAL_MAX_CONTROL_FRAME_BYTES,
    capability: ZAKURA_CAP_DISCOVERY,
    mode: StreamMode::Ordered,
}];

/// Service-declared streams for native discovery.
pub(crate) fn discovery_streams() -> &'static [Stream] {
    &DISCOVERY_SERVICE_STREAMS
}

/// Cloneable typed sender for one native discovery ordered stream.
#[derive(Clone, Debug)]
pub struct DiscoveryPeerSession {
    peer_id: ZakuraPeerId,
    direction: ServicePeerDirection,
    send: FramedSend,
    cancel: CancellationToken,
}

impl DiscoveryPeerSession {
    fn new(session: &PeerStreamSession, direction: ServicePeerDirection) -> Self {
        Self {
            peer_id: session.peer_id().clone(),
            direction,
            send: session.sender(),
            cancel: session.cancel_token(),
        }
    }

    /// Authenticated peer identity for this discovery stream.
    pub fn peer_id(&self) -> &ZakuraPeerId {
        &self.peer_id
    }

    /// Direction of the underlying Zakura connection.
    pub fn direction(&self) -> ServicePeerDirection {
        self.direction
    }

    /// Peer disconnect/local shutdown cancellation token.
    pub fn cancel_token(&self) -> CancellationToken {
        self.cancel.clone()
    }

    /// Send this node's signed self-record.
    pub fn try_send_hello(&self, record: ZakuraNodeRecord) -> Result<(), OrderedSendError> {
        self.try_send_message(DiscoveryMessage::Hello { record })
    }

    /// Ask this peer for more peer records.
    pub fn try_send_get_peers(
        &self,
        limit: u16,
        wanted_services: Vec<ZakuraServiceId>,
        exclude_node_ids: Vec<NodeId>,
    ) -> Result<(), OrderedSendError> {
        self.try_send_message(DiscoveryMessage::GetPeers {
            limit,
            wanted_services,
            exclude_node_ids,
        })
    }

    /// Send peer records to this peer.
    pub fn try_send_peers(&self, records: Vec<ZakuraNodeRecord>) -> Result<(), OrderedSendError> {
        self.try_send_message(DiscoveryMessage::Peers { records })
    }

    /// Ask this peer for its own live service summaries.
    pub fn try_send_get_services(
        &self,
        wanted_services: Vec<ZakuraServiceId>,
    ) -> Result<(), OrderedSendError> {
        self.try_send_message(DiscoveryMessage::GetServices(GetServices {
            wanted_services,
        }))
    }

    /// Send this node's first-party live service summaries.
    pub fn try_send_services(&self, services: Services) -> Result<(), OrderedSendError> {
        self.try_send_message(DiscoveryMessage::Services(services))
    }

    fn try_send_message(&self, message: DiscoveryMessage) -> Result<(), OrderedSendError> {
        let payload = message
            .encode()
            .map_err(|error| OrderedSendError::Encode(Box::new(error)))?;
        match self.send.try_send(Frame {
            message_type: DISCOVERY_FRAME_MESSAGE_TYPE,
            flags: 0,
            payload,
        }) {
            Ok(()) => Ok(()),
            Err(tokio::sync::mpsc::error::TrySendError::Full(_frame)) => {
                Err(OrderedSendError::Full)
            }
            Err(tokio::sync::mpsc::error::TrySendError::Closed(_frame)) => {
                Err(OrderedSendError::Closed)
            }
        }
    }
}

/// Native discovery service backed by a [`ZakuraDiscoveryHandle`] runtime.
#[derive(Clone, Debug)]
pub struct DiscoveryService {
    handle: ZakuraDiscoveryHandle,
    header_sync: Option<HeaderSyncHandle>,
    block_sync: Option<BlockSyncHandle>,
    session_states: Arc<StdMutex<SessionStateMap>>,
    connection_owners: Arc<StdMutex<Vec<Arc<dyn Service>>>>,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum DiscoverySessionState {
    Active,
    Retired,
}

/// The session-state slot for one `(peer, connection)`, tagged with the stream
/// session that owns it so a finished exchange cannot retire a newer session
/// admitted on the same connection.
#[derive(Clone, Copy, Debug)]
struct DiscoverySessionRecord {
    session_id: u64,
    state: DiscoverySessionState,
    /// Consecutive sessions on this connection that ended without a successful
    /// exchange. Reaching `MAX_DISCOVERY_SESSION_FAILURES` retires the record so
    /// the transport stops reopening a stream the peer keeps breaking.
    failed_attempts: u32,
}

/// Consecutive failed sessions on one `(peer, connection)` before discovery
/// retires it. Generous: an honest peer completes the exchange on the first
/// session, and a fresh connection always starts a fresh record.
const MAX_DISCOVERY_SESSION_FAILURES: u32 = 3;

type SessionStateMap = HashMap<(ZakuraPeerId, ZakuraConnId), DiscoverySessionRecord>;

fn retire_discovery_session(
    session_states: &StdMutex<SessionStateMap>,
    peer: &ZakuraPeerId,
    conn_id: ZakuraConnId,
    session_id: u64,
) -> bool {
    let mut session_states = session_states
        .lock()
        .expect("discovery session-state mutex is never poisoned");
    let Some(record) = session_states.get_mut(&(peer.clone(), conn_id)) else {
        return false;
    };
    if record.session_id != session_id {
        return false;
    }
    record.state = DiscoverySessionState::Retired;
    true
}

/// Charge one failed session to the owning record; retires it at the failure
/// bound. Session-scoped like `retire_discovery_session` so a stale task cannot
/// charge a newer session admitted on the same connection. Returns whether the
/// record was retired by this failure.
fn record_discovery_session_failure(
    session_states: &StdMutex<SessionStateMap>,
    peer: &ZakuraPeerId,
    conn_id: ZakuraConnId,
    session_id: u64,
) -> bool {
    let mut session_states = session_states
        .lock()
        .expect("discovery session-state mutex is never poisoned");
    let Some(record) = session_states.get_mut(&(peer.clone(), conn_id)) else {
        return false;
    };
    if record.session_id != session_id || record.state != DiscoverySessionState::Active {
        return false;
    }
    record.failed_attempts = record.failed_attempts.saturating_add(1);
    if record.failed_attempts >= MAX_DISCOVERY_SESSION_FAILURES {
        record.state = DiscoverySessionState::Retired;
        tracing::info!(
            ?peer,
            conn_id,
            attempts = record.failed_attempts,
            "retiring Zakura discovery sessions on this connection after repeated failed exchanges"
        );
        return true;
    }
    false
}

/// Reset the owning record's failure count after a successful exchange, so the
/// intended shared-connection refresh loop never accumulates failures.
fn record_discovery_session_success(
    session_states: &StdMutex<SessionStateMap>,
    peer: &ZakuraPeerId,
    conn_id: ZakuraConnId,
    session_id: u64,
) {
    let mut session_states = session_states
        .lock()
        .expect("discovery session-state mutex is never poisoned");
    let Some(record) = session_states.get_mut(&(peer.clone(), conn_id)) else {
        return;
    };
    if record.session_id != session_id {
        return;
    }
    record.failed_attempts = 0;
}

impl DiscoveryService {
    /// Builds a discovery service driven by `handle`.
    pub fn new(handle: ZakuraDiscoveryHandle) -> Self {
        Self {
            handle,
            header_sync: None,
            block_sync: None,
            session_states: Arc::new(StdMutex::new(HashMap::new())),
            connection_owners: Arc::new(StdMutex::new(Vec::new())),
        }
    }

    /// Builds a discovery service with header-sync and block-sync summary providers.
    pub(crate) fn with_sync_services(
        handle: ZakuraDiscoveryHandle,
        header_sync: HeaderSyncHandle,
        block_sync: Option<BlockSyncHandle>,
    ) -> Self {
        Self {
            handle,
            header_sync: Some(header_sync),
            block_sync,
            session_states: Arc::new(StdMutex::new(HashMap::new())),
            connection_owners: Arc::new(StdMutex::new(Vec::new())),
        }
    }

    pub(crate) fn set_connection_owners(&self, owners: Vec<Arc<dyn Service>>) {
        *self
            .connection_owners
            .lock()
            .expect("discovery connection-owner mutex is never poisoned") = owners;
    }

    /// Returns the underlying discovery runtime handle.
    pub fn handle(&self) -> &ZakuraDiscoveryHandle {
        &self.handle
    }
}

impl Service for DiscoveryService {
    fn name(&self) -> &'static str {
        "discovery"
    }

    fn streams(&self) -> &[Stream] {
        discovery_streams()
    }

    fn ordered_stream_policy(&self, _kind: u16) -> OrderedStreamPolicy {
        OrderedStreamPolicy {
            opening: OrderedStreamOpening::InitiatorOnly,
            reopen: true,
        }
    }

    fn ordered_session_demand(
        &self,
        conn_id: ZakuraConnId,
        peer: &ZakuraPeerId,
        _negotiated: u64,
        direction: ServicePeerDirection,
    ) -> OrderedSessionDemand {
        if self
            .session_states
            .lock()
            .expect("discovery session-state mutex is never poisoned")
            .get(&(peer.clone(), conn_id))
            .is_some_and(|record| record.state == DiscoverySessionState::Retired)
        {
            return OrderedSessionDemand::Retire;
        }

        let mut peers = self.handle.subscribe_peer_snapshot();
        let snapshot = *peers.borrow_and_update();
        let slots_free = match direction {
            ServicePeerDirection::Inbound => snapshot.inbound_slots_free,
            ServicePeerDirection::Outbound => snapshot.outbound_slots_free,
        };
        if slots_free == 0 {
            return OrderedSessionDemand::WaitForChange(Box::pin(async move {
                if peers.changed().await.is_err() {
                    std::future::pending::<()>().await;
                }
            }));
        }

        OrderedSessionDemand::OpenNow
    }

    fn wants_peer(
        &self,
        _peer: &ZakuraPeerId,
        _negotiated: u64,
        direction: ServicePeerDirection,
    ) -> bool {
        // Discovery escalation only checks this reactor's local room; live
        // summaries are first-party advisory data imported by the runtime.
        let snapshot = self.handle.peer_snapshot();
        match direction {
            ServicePeerDirection::Inbound => snapshot.inbound_slots_free > 0,
            ServicePeerDirection::Outbound => snapshot.outbound_slots_free > 0,
        }
    }

    fn add_peer(&self, mut peer: Peer) {
        let Some((session_id, recv, send)) =
            peer.take_stream_with_session_id(ZAKURA_STREAM_DISCOVERY)
        else {
            return;
        };
        let Some(peer_node_id) = node_id_from_peer_id(&peer.id) else {
            // A peer id that is not a 32-byte node id cannot be a discovery
            // author; drop the stream without registering an exchange.
            return;
        };
        let session = PeerStreamSession::new(
            peer.id.clone(),
            ZAKURA_STREAM_DISCOVERY,
            recv,
            send,
            peer.service_cancel_token(),
        );
        let discovery_session = DiscoveryPeerSession::new(&session, peer.direction);
        let conn_id = peer.conn_id;
        {
            let mut session_states = self
                .session_states
                .lock()
                .expect("discovery session-state mutex is never poisoned");
            let previous = session_states.get(&(peer.id.clone(), conn_id));
            if previous.is_some_and(|record| record.state == DiscoverySessionState::Retired) {
                // A retired record must not be resurrected by a new stream on the
                // same connection — the remote initiates streams on inbound
                // connections and could otherwise churn sessions at its own pace.
                return;
            }
            // Preserve the failure count across reopens so consecutive broken
            // sessions on this connection stay bounded.
            let failed_attempts = previous.map_or(0, |record| record.failed_attempts);
            session_states.insert(
                (peer.id.clone(), conn_id),
                DiscoverySessionRecord {
                    session_id,
                    state: DiscoverySessionState::Active,
                    failed_attempts,
                },
            );
        }
        let service_cancel = discovery_session.cancel_token();
        let connection_cancel = peer.cancel_token();
        let close_cause = peer.close_cause();
        let (_peer_id, _stream_kind, recv, _send, _session_cancel) = session.into_parts();

        let handle = self.handle.clone();
        let header_sync = self.header_sync.clone();
        let block_sync = self.block_sync.clone();
        let session_states = self.session_states.clone();
        let connection_owners = self
            .connection_owners
            .lock()
            .expect("discovery connection-owner mutex is never poisoned")
            .clone();
        // SR-1: a panic in the admission task (before it hands off to the
        // exchange) must still disconnect this one peer and cancel its discovery
        // session instead of leaving admitted state behind a half-live
        // connection. Normal/parked exits cancel `service_cancel` inline below;
        // `on_panic` covers the unwind path only.
        let admit_peer_id = discovery_session.peer_id().clone();
        let panic_service_cancel = service_cancel.clone();
        let panic_connection_cancel = connection_cancel.clone();
        let panic_close_cause = close_cause.clone();
        spawn_supervised_peer_task(
            admit_peer_id,
            || {},
            move || {
                panic_close_cause.record("service_panic");
                panic_service_cancel.cancel();
                panic_connection_cancel.cancel();
            },
            async move {
                let decision = handle
                    .admit_peer_session(
                        conn_id,
                        session_id,
                        discovery_session.peer_id().clone(),
                        discovery_session.direction(),
                    )
                    .await;
                if decision != ServiceAdmissionDecision::Admit {
                    metrics::counter!("zakura.discovery.peer.parked").increment(1);
                    tracing::info!(
                        peer = ?discovery_session.peer_id(),
                        direction = ?discovery_session.direction(),
                        ?decision,
                        "locally parking Zakura discovery service session"
                    );
                    // A parked session ends without an exchange; charge it so
                    // reopen-then-park cycles stay bounded on this connection.
                    record_discovery_session_failure(
                        &session_states,
                        discovery_session.peer_id(),
                        conn_id,
                        session_id,
                    );
                    service_cancel.cancel();
                    return;
                }

                spawn_discovery_exchange(DiscoveryExchangeStart {
                    handle,
                    header_sync,
                    block_sync,
                    connection_owners,
                    peer_node_id,
                    discovery_session,
                    conn_id,
                    session_id,
                    recv,
                    service_cancel,
                    connection_cancel,
                    close_cause,
                    session_states,
                });
            },
        );
    }

    fn remove_peer(&self, peer: &ZakuraPeerId, conn_id: ZakuraConnId) {
        self.session_states
            .lock()
            .expect("discovery session-state mutex is never poisoned")
            .remove(&(peer.clone(), conn_id));
        let handle = self.handle.clone();
        let peer = peer.clone();
        tokio::spawn(async move {
            handle.remove_peer(&peer, conn_id).await;
        });
    }
}

struct DiscoveryExchangeStart {
    handle: ZakuraDiscoveryHandle,
    header_sync: Option<HeaderSyncHandle>,
    block_sync: Option<BlockSyncHandle>,
    connection_owners: Vec<Arc<dyn Service>>,
    peer_node_id: NodeId,
    discovery_session: DiscoveryPeerSession,
    conn_id: ZakuraConnId,
    session_id: u64,
    recv: FramedRecv,
    service_cancel: CancellationToken,
    connection_cancel: CancellationToken,
    close_cause: CloseCause,
    session_states: Arc<StdMutex<SessionStateMap>>,
}

fn spawn_discovery_exchange(start: DiscoveryExchangeStart) {
    let DiscoveryExchangeStart {
        handle,
        header_sync,
        block_sync,
        connection_owners,
        peer_node_id,
        discovery_session,
        conn_id,
        session_id,
        recv,
        service_cancel,
        connection_cancel,
        close_cause,
        session_states,
    } = start;
    let peer_id = discovery_session.peer_id().clone();
    let progress = Arc::new(DiscoveryExchangeProgress::default());
    let sink = DiscoverySink {
        handle: handle.clone(),
        header_sync,
        block_sync,
        peer_node_id,
        session: discovery_session.clone(),
        conn_id,
        session_id,
        progress: progress.clone(),
    };
    let sink_service_cancel = service_cancel.clone();
    let reject_connection_cancel = connection_cancel.clone();
    let panic_connection_cancel = connection_cancel.clone();
    let reject_close_cause = close_cause.clone();
    let panic_close_cause = close_cause.clone();
    let sink_peer_id = peer_id.clone();
    // A protocol reject is fatal to the connection; normal/parked exits leave it
    // for the source task to tear down once it knows no other service owns the
    // peer (below). Panic teardown is in `on_panic`.
    let pipe = async move {
        let mut pipe = discovery_pipe(sink_peer_id);
        handle_pipe_exit(
            "discovery",
            &reject_connection_cancel,
            &reject_close_cause,
            run_discovery_pipe(&mut pipe, recv, sink).await,
        );
    };
    let on_panic = move || {
        panic_close_cause.record("service_panic");
        panic_connection_cancel.cancel();
    };
    // Let the returned handle drop to detach the supervised reader task; the
    // `PipeTeardown` still runs on every exit path.
    spawn_supervised_pipe(peer_id.clone(), sink_service_cancel, || {}, on_panic, pipe);

    let source = DiscoverySource {
        handle: handle.clone(),
        session: discovery_session,
        conn_id,
        session_id,
        progress,
    };
    // SR-1: a panic in the source task skips its `service_cancel.cancel()`,
    // `handle.remove_session()`, and discovery-only connection cancellation,
    // leaving admitted discovery state behind a half-live connection. On the
    // unwind path, disconnect this one peer; the connection teardown then drives
    // the async session removal through the registry. Normal exits run the inline
    // cleanup below, so `on_panic` is the panic-only path.
    let source_task_peer_id = peer_id.clone();
    let panic_source_service_cancel = service_cancel.clone();
    let panic_source_connection_cancel = connection_cancel.clone();
    let panic_source_close_cause = close_cause.clone();
    let source_close_cause = close_cause.clone();
    spawn_supervised_peer_task(
        source_task_peer_id,
        || {},
        move || {
            panic_source_close_cause.record("service_panic");
            panic_source_service_cancel.cancel();
            panic_source_connection_cancel.cancel();
        },
        async move {
            let exchanged = source.run_initial_exchange().await;
            if exchanged {
                record_discovery_session_success(&session_states, &peer_id, conn_id, session_id);
            } else {
                // The stream broke before the exchange completed; charge the
                // session so an attacker cannot get endless reopens by breaking
                // the stream pre-exchange.
                record_discovery_session_failure(&session_states, &peer_id, conn_id, session_id);
            }
            let mut other_service_owner =
                exchanged && peer_has_other_service_owner(&connection_owners, &peer_id, conn_id);
            while other_service_owner && source.refresh_after_interval().await.is_ok() {
                other_service_owner =
                    peer_has_other_service_owner(&connection_owners, &peer_id, conn_id);
            }
            let closes_discovery_only_connection = exchanged
                && !other_service_owner
                && handle
                    .is_current_session(&peer_id, conn_id, session_id)
                    .await;
            if closes_discovery_only_connection {
                handle.mark_short_lived_exchange(&peer_node_id).await;
                retire_discovery_session(&session_states, &peer_id, conn_id, session_id);
            }
            service_cancel.cancel();
            handle.remove_session(&peer_id, conn_id, session_id).await;
            if closes_discovery_only_connection {
                source_close_cause.record("discovery_exchange_complete");
                connection_cancel.cancel();
            }
        },
    );
}

/// Reader half of the discovery stream: imports peer records and answers queries.
struct DiscoverySink {
    handle: ZakuraDiscoveryHandle,
    header_sync: Option<HeaderSyncHandle>,
    block_sync: Option<BlockSyncHandle>,
    peer_node_id: NodeId,
    session: DiscoveryPeerSession,
    conn_id: ZakuraConnId,
    session_id: u64,
    progress: Arc<DiscoveryExchangeProgress>,
}

async fn run_discovery_pipe(
    pipe: &mut Pipe<DsLocal, DsEnv>,
    mut recv: FramedRecv,
    sink: DiscoverySink,
) -> Result<(), SinkReject> {
    let cancel = sink.session.cancel_token();
    loop {
        let frame = tokio::select! {
            biased;
            _ = cancel.cancelled() => return Ok(()),
            frame = recv.recv() => frame,
        };
        let Some(frame) = frame else {
            return Ok(());
        };

        match pipe.run_one(frame) {
            Flow::Continue(()) | Flow::Done => {}
            Flow::Reject(reject) => return Err(reject),
        }

        let Some(message) = pipe.local_mut().take_decoded() else {
            continue;
        };
        sink.handle_message(message).await?;
    }
}

impl DiscoverySink {
    async fn handle_message(&self, message: DiscoveryMessage) -> Result<(), SinkReject> {
        if !self
            .handle
            .is_current_session(self.session.peer_id(), self.conn_id, self.session_id)
            .await
        {
            return Ok(());
        }

        match message {
            DiscoveryMessage::Hello { record } => self.handle_hello(record).await,
            DiscoveryMessage::GetPeers {
                limit,
                wanted_services,
                exclude_node_ids,
            } => {
                let records = self
                    .handle
                    .sample_peers(
                        self.peer_node_id,
                        usize::from(limit),
                        &wanted_services,
                        &exclude_node_ids,
                    )
                    .await;
                self.send_peers(records)
            }
            DiscoveryMessage::Peers { records } => {
                self.handle
                    .import_peer_records(records, Some(self.peer_node_id))
                    .await;
                self.progress.mark_peers();
                Ok(())
            }
            DiscoveryMessage::GetServices(query) => {
                let services = self.local_services_response(query).await?;
                self.send_services(services)
            }
            DiscoveryMessage::Services(services) => self.handle_services(services).await,
        }
    }

    async fn local_services_response(&self, query: GetServices) -> Result<Services, SinkReject> {
        let mut summaries = Vec::new();

        if service_wanted(&query.wanted_services, &ZakuraServiceId::header_sync()) {
            if let Some(header_sync) = &self.header_sync {
                let (best_height, best_hash) = header_sync.best_header_tip();
                let summary = HeaderSyncServiceSummary::from_snapshot(
                    best_height,
                    best_hash,
                    None,
                    true,
                    header_sync.peer_snapshot(),
                );
                summaries.push(
                    ServiceSummaryEnvelope::header_sync(&summary).map_err(SinkReject::local)?,
                );
            }
        }

        if service_wanted(&query.wanted_services, &ZakuraServiceId::discovery()) {
            let summary = self.handle.local_discovery_summary().await;
            summaries.push(ServiceSummaryEnvelope::discovery(&summary).map_err(SinkReject::local)?);
        }

        if service_wanted(&query.wanted_services, &ZakuraServiceId::block_sync()) {
            if let Some(block_sync) = &self.block_sync {
                let summary = BlockSyncServiceSummary::from_status_and_snapshot(
                    block_sync.local_status(),
                    block_sync.peer_snapshot(),
                );
                summaries
                    .push(ServiceSummaryEnvelope::block_sync(&summary).map_err(SinkReject::local)?);
            }
        }

        Ok(self.handle.local_services_response(summaries))
    }

    async fn handle_services(&self, services: Services) -> Result<(), SinkReject> {
        if services.node_id != self.peer_node_id {
            return Err(SinkReject::protocol(
                "Zakura discovery SERVICES authored by a different node id",
            ));
        }

        let header_summaries =
            decode_header_sync_summaries(&services).map_err(SinkReject::protocol)?;
        self.handle
            .import_connected_peer_services(services, self.peer_node_id)
            .await
            .map_err(SinkReject::protocol)?;
        self.progress.mark_services();

        if let Some(header_sync) = &self.header_sync {
            for summary in header_summaries {
                if let Err(error) = header_sync
                    .send(HeaderSyncEvent::AdvisoryHeaderSummary {
                        peer: self.session.peer_id().clone(),
                        summary,
                    })
                    .await
                {
                    tracing::debug!(
                        ?error,
                        peer = ?self.session.peer_id(),
                        "failed to queue first-party Zakura header-sync advisory summary"
                    );
                    break;
                }
            }
        }

        Ok(())
    }

    async fn handle_hello(&self, record: ZakuraNodeRecord) -> Result<(), SinkReject> {
        if record.body.node_id != self.peer_node_id {
            return Err(SinkReject::protocol(
                "Zakura discovery hello authored by a different node id",
            ));
        }
        match self
            .handle
            .import_connected_peer_record(record, self.peer_node_id)
            .await
        {
            Ok(_) => Ok(()),
            Err(error) if is_advisory_self_record_import_error(&error) => {
                tracing::debug!(?error, "ignoring advisory discovery hello import error");
                Ok(())
            }
            Err(error) => Err(SinkReject::protocol(error)),
        }?;
        self.progress.mark_hello();
        Ok(())
    }

    fn send_peers(&self, records: Vec<ZakuraNodeRecord>) -> Result<(), SinkReject> {
        match self.session.try_send_peers(records) {
            Ok(()) | Err(OrderedSendError::Full) => Ok(()),
            Err(OrderedSendError::Closed) => {
                Err(SinkReject::local("Zakura discovery send channel closed"))
            }
            Err(OrderedSendError::Encode(error)) => Err(SinkReject::local(error)),
        }
    }

    fn send_services(&self, services: Services) -> Result<(), SinkReject> {
        match self.session.try_send_services(services) {
            Ok(()) | Err(OrderedSendError::Full) => Ok(()),
            Err(OrderedSendError::Closed) => {
                Err(SinkReject::local("Zakura discovery send channel closed"))
            }
            Err(OrderedSendError::Encode(error)) => Err(SinkReject::local(error)),
        }
    }
}

fn service_wanted(wanted_services: &[ZakuraServiceId], service_id: &ZakuraServiceId) -> bool {
    wanted_services.is_empty() || wanted_services.iter().any(|wanted| wanted == service_id)
}

fn decode_header_sync_summaries(
    services: &Services,
) -> Result<Vec<HeaderSyncServiceSummary>, crate::BoxError> {
    let mut summaries = Vec::new();
    for envelope in &services.summaries {
        if let Some(summary) = envelope.decode_header_sync()? {
            summaries.push(summary);
        }
    }
    Ok(summaries)
}

/// Writer half of the discovery stream: periodic self-record gossip + peer asks.
struct DiscoverySource {
    handle: ZakuraDiscoveryHandle,
    session: DiscoveryPeerSession,
    conn_id: ZakuraConnId,
    session_id: u64,
    progress: Arc<DiscoveryExchangeProgress>,
}

impl DiscoverySource {
    async fn run_initial_exchange(&self) -> bool {
        if self.exchange().await.is_err() {
            return false;
        }
        let cancel = self.session.cancel_token();
        let completed = tokio::select! {
            biased;
            _ = cancel.cancelled() => false,
            _ = self.progress.wait_complete() => true,
            _ = tokio::time::sleep(DISCOVERY_INITIAL_EXCHANGE_TIMEOUT) => false,
        };
        if !completed {
            return false;
        }

        self.handle
            .is_current_session(self.session.peer_id(), self.conn_id, self.session_id)
            .await
    }

    async fn refresh_after_interval(&self) -> Result<(), ()> {
        let cancel = self.session.cancel_token();
        let refresh_interval = discovery_exchange_interval(self.handle.refresh_interval().await);
        tokio::select! {
            biased;
            _ = cancel.cancelled() => Err(()),
            _ = tokio::time::sleep(refresh_interval) => self.exchange().await,
        }
    }

    /// Gossips the current self-record and asks the peer for records and services.
    ///
    /// Returns `Err(())` once the stream's send side is gone, so the caller
    /// stops the periodic loop.
    async fn exchange(&self) -> Result<(), ()> {
        if !self
            .handle
            .is_current_session(self.session.peer_id(), self.conn_id, self.session_id)
            .await
        {
            return Err(());
        }

        let record = self
            .handle
            .current_self_record_for_gossip()
            .await
            .map_err(|error| {
                tracing::debug!(
                    ?error,
                    peer = ?self.session.peer_id(),
                    "failed to refresh Zakura discovery self-record"
                );
            })?;
        let record = (*record).clone();
        self.handle_send_result(self.session.try_send_hello(record))?;

        let limit = self
            .handle
            .peer_sample_limit()
            .await
            .min(MAX_DISCOVERY_RECORDS_PER_RESPONSE);
        // `peer_sample_limit` is bounded by MAX_DISCOVERY_RECORDS_PER_RESPONSE
        // (<= u16::MAX), so the cast cannot truncate.
        let exclude_node_ids = self.handle.peer_sample_exclusions().await;
        self.handle_send_result(self.session.try_send_get_peers(
            limit as u16,
            Vec::new(),
            exclude_node_ids,
        ))?;

        self.handle_send_result(self.session.try_send_get_services(Vec::new()))
    }

    fn handle_send_result(&self, result: Result<(), OrderedSendError>) -> Result<(), ()> {
        match result {
            Ok(()) | Err(OrderedSendError::Full) => Ok(()),
            Err(OrderedSendError::Closed) => Err(()),
            Err(OrderedSendError::Encode(error)) => {
                tracing::debug!(
                    ?error,
                    peer = ?self.session.peer_id(),
                    "failed to encode Zakura discovery message"
                );
                Ok(())
            }
        }
    }
}

#[derive(Default)]
struct DiscoveryExchangeProgress {
    hello: AtomicBool,
    peers: AtomicBool,
    services: AtomicBool,
    notify: Notify,
}

impl DiscoveryExchangeProgress {
    fn mark_hello(&self) {
        self.hello.store(true, Ordering::Relaxed);
        self.notify.notify_waiters();
    }

    fn mark_peers(&self) {
        self.peers.store(true, Ordering::Relaxed);
        self.notify.notify_waiters();
    }

    fn mark_services(&self) {
        self.services.store(true, Ordering::Relaxed);
        self.notify.notify_waiters();
    }

    fn complete(&self) -> bool {
        self.hello.load(Ordering::Relaxed)
            && self.peers.load(Ordering::Relaxed)
            && self.services.load(Ordering::Relaxed)
    }

    async fn wait_complete(&self) {
        while !self.complete() {
            self.notify.notified().await;
        }
    }
}

fn peer_has_other_service_owner(
    connection_owners: &[Arc<dyn Service>],
    peer_id: &ZakuraPeerId,
    conn_id: ZakuraConnId,
) -> bool {
    connection_owners
        .iter()
        .any(|owner| owner.owns_connection_for_peer(peer_id, conn_id))
}

fn discovery_exchange_interval(record_refresh_interval: Duration) -> Duration {
    record_refresh_interval.min(DEFAULT_LIVE_SERVICE_SUMMARY_TTL / 2)
}

/// Returns the iroh node id encoded by a discovery peer id, if it is a 32-byte
/// node id.
fn node_id_from_peer_id(peer_id: &ZakuraPeerId) -> Option<NodeId> {
    let bytes: [u8; 32] = peer_id.as_bytes().try_into().ok()?;
    NodeId::from_bytes(&bytes).ok()
}

/// A peer-hello import error that should be logged and ignored rather than
/// closing the live connection. These mean the peer's record is not locally
/// dialable or has drifted out of the freshness window, neither of which is the
/// connected peer's fault.
fn is_advisory_self_record_import_error(error: &DiscoveryBookError) -> bool {
    matches!(
        error,
        DiscoveryBookError::NoUsableDirectAddress
            | DiscoveryBookError::NonDialableDirectAddress { .. }
            | DiscoveryBookError::Record(DiscoveryRecordError::Expired)
            | DiscoveryBookError::Record(DiscoveryRecordError::FarFutureExpiry)
    )
}

#[cfg(test)]
mod tests {
    use std::{
        collections::HashMap,
        net::{IpAddr, Ipv4Addr, SocketAddr},
        time::{Duration, SystemTime, UNIX_EPOCH},
    };

    use iroh::SecretKey;
    use tokio::{sync::watch, task::JoinHandle};

    use super::*;
    use crate::zakura::discovery::protocol::{
        DiscoveryServiceSummary, ZakuraLiveServiceSummary, ZakuraNodeRecordBody,
        SUMMARY_TAG_HEADER_SYNC_RETIRED,
    };
    use crate::zakura::{
        framed_channel, spawn_block_sync_reactor, spawn_header_sync_reactor, BlockSyncFrontiers,
        BlockSyncService, BlockSyncStartup, HeaderSyncAction, HeaderSyncFrontiers,
        HeaderSyncMessage, HeaderSyncPeerSession, HeaderSyncService, HeaderSyncStartup,
        HeaderSyncStatus, HeaderSyncWireRequestIdentity, ServicePeerLimits, ZakuraBlockSyncConfig,
        ZakuraDiscoveryConfig, ZakuraDiscoveryLocalConfig, ZakuraHandshakeConfig,
        ZakuraHeaderSyncConfig, LOCAL_MAX_MESSAGE_BYTES, MAX_BS_RESPONSE_BYTES,
        ZAKURA_CAP_BLOCK_SYNC, ZAKURA_CAP_DISCOVERY, ZAKURA_CAP_HEADER_SYNC,
        ZAKURA_CAP_LEGACY_GOSSIP, ZAKURA_STREAM_HEADER_SYNC,
    };
    use zakura_chain::{block, parameters::Network};

    struct HeaderAdvisoryFixture {
        discovery_handle: ZakuraDiscoveryHandle,
        header_sync: HeaderSyncHandle,
        header_actions: tokio::sync::mpsc::Receiver<HeaderSyncAction>,
        header_task: JoinHandle<()>,
        peer_node_id: NodeId,
        peer_id: ZakuraPeerId,
        peer_send: FramedSend,
        _peer_recv: FramedRecv,
    }

    #[derive(Debug)]
    struct TestConnectionOwner {
        peer: ZakuraPeerId,
        conn_id: ZakuraConnId,
    }

    impl Service for TestConnectionOwner {
        fn name(&self) -> &'static str {
            "test-connection-owner"
        }

        fn streams(&self) -> &[Stream] {
            &[]
        }

        fn owns_connection_for_peer(&self, peer: &ZakuraPeerId, conn_id: ZakuraConnId) -> bool {
            peer == &self.peer && conn_id == self.conn_id
        }

        fn add_peer(&self, _peer: Peer) {}

        fn remove_peer(&self, _peer: &ZakuraPeerId, _conn_id: ZakuraConnId) {}
    }

    impl Drop for HeaderAdvisoryFixture {
        fn drop(&mut self) {
            self.header_task.abort();
        }
    }

    fn current_test_unix_secs() -> u64 {
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("system clock is after Unix epoch")
            .as_secs()
    }

    #[test]
    fn periodic_exchange_refreshes_before_live_service_summaries_expire() {
        assert_eq!(
            discovery_exchange_interval(Duration::from_secs(10 * 60)),
            Duration::from_secs(15)
        );
        assert_eq!(
            discovery_exchange_interval(Duration::from_secs(10)),
            Duration::from_secs(10)
        );
    }

    #[test]
    fn connection_ownership_is_scoped_to_the_exact_connection() {
        let peer = ZakuraPeerId::new(vec![37; 32]).expect("test peer id is within bounds");
        let owners: Vec<Arc<dyn Service>> = vec![Arc::new(TestConnectionOwner {
            peer: peer.clone(),
            conn_id: 2,
        })];

        assert!(!peer_has_other_service_owner(&owners, &peer, 1));
        assert!(peer_has_other_service_owner(&owners, &peer, 2));
    }

    fn header_summary(best_height: block::Height) -> HeaderSyncServiceSummary {
        HeaderSyncServiceSummary {
            best_height,
            best_hash: block::Hash([7; 32]),
            finalized_height: None,
            serving_headers: true,
            inbound_slots_free: 1,
            inbound_slots_max: 1,
            outbound_slots_free: 1,
            outbound_slots_max: 1,
        }
    }

    fn spawn_test_header_sync() -> Result<
        (
            HeaderSyncHandle,
            tokio::sync::mpsc::Receiver<HeaderSyncAction>,
            JoinHandle<()>,
        ),
        crate::BoxError,
    > {
        let network = Network::new_regtest(Default::default());
        let anchor = (block::Height(0), network.genesis_hash());
        let mut startup = HeaderSyncStartup::new(
            network,
            anchor,
            HeaderSyncFrontiers {
                finalized_height: anchor.0,
                verified_block_tip: anchor.0,
                verified_block_hash: anchor.1,
            },
            Some(anchor),
            ZakuraHeaderSyncConfig::default(),
            LOCAL_MAX_MESSAGE_BYTES,
        );
        startup.range_state_actions_enabled = true;
        spawn_header_sync_reactor(startup).map_err(Into::into)
    }

    fn signed_header_sync_record(
        secret_key: &SecretKey,
        handshake: &ZakuraHandshakeConfig,
    ) -> Result<ZakuraNodeRecord, crate::BoxError> {
        let body = ZakuraNodeRecordBody {
            node_id: secret_key.public(),
            direct_addrs: vec![SocketAddr::new(
                IpAddr::V4(Ipv4Addr::new(45, 33, 30, 44)),
                8233,
            )],
            services: vec![ZakuraServiceId::header_sync()],
            zakura_protocol_min: handshake.zakura_protocol_min,
            zakura_protocol_max: handshake.zakura_protocol_max,
            network_id: handshake.network_id,
            chain_id: handshake.chain_id,
            sequence: 1,
            expires_at_unix_secs: current_test_unix_secs().saturating_add(60),
        };
        Ok(ZakuraNodeRecord::sign(body, secret_key)?)
    }

    fn spawn_header_advisory_fixture(
        peer_seed: u8,
    ) -> Result<HeaderAdvisoryFixture, crate::BoxError> {
        let (connected_tx, connected_rx) = watch::channel(Vec::new());
        let handshake = ZakuraHandshakeConfig::for_network(&Network::Mainnet);
        let local_secret = SecretKey::from_bytes(&[31u8; 32]);
        let discovery_handle = ZakuraDiscoveryHandle::new(
            ZakuraDiscoveryLocalConfig {
                secret_key: local_secret,
                direct_addrs: Vec::new(),
                services: vec![ZakuraServiceId::discovery()],
                zakura_protocol_min: handshake.zakura_protocol_min,
                zakura_protocol_max: handshake.zakura_protocol_max,
                network_id: handshake.network_id,
                chain_id: handshake.chain_id,
                last_authored_sequence: None,
            },
            ZakuraDiscoveryConfig::default(),
            connected_rx,
        )?;
        let (header_sync, header_actions, header_task) = spawn_test_header_sync()?;
        let service = DiscoveryService::with_sync_services(
            discovery_handle.clone(),
            header_sync.clone(),
            None,
        );
        let peer_node_id = SecretKey::from_bytes(&[peer_seed; 32]).public();
        let peer_id = ZakuraPeerId::new(peer_node_id.as_bytes().to_vec())?;
        connected_tx.send_replace(vec![peer_id.clone()]);

        let (peer_send, service_recv) = framed_channel(8);
        let (service_send, peer_recv) = framed_channel(8);
        let streams = HashMap::from([(ZAKURA_STREAM_DISCOVERY, (service_recv, service_send))]);

        service.add_peer(Peer::new(
            peer_id.clone(),
            None,
            ZAKURA_CAP_DISCOVERY,
            streams,
            CancellationToken::new(),
        ));

        Ok(HeaderAdvisoryFixture {
            discovery_handle,
            header_sync,
            header_actions,
            header_task,
            peer_node_id,
            peer_id,
            peer_send,
            _peer_recv: peer_recv,
        })
    }

    async fn send_discovery_message(
        fixture: &HeaderAdvisoryFixture,
        message: DiscoveryMessage,
    ) -> Result<(), crate::BoxError> {
        fixture
            .peer_send
            .send(Frame {
                message_type: DISCOVERY_FRAME_MESSAGE_TYPE,
                flags: 0,
                payload: message.encode()?,
            })
            .await?;
        Ok(())
    }

    fn discovery_frame(message: DiscoveryMessage) -> Result<Frame, crate::BoxError> {
        Ok(Frame {
            message_type: DISCOVERY_FRAME_MESSAGE_TYPE,
            flags: 0,
            payload: message.encode()?,
        })
    }

    fn signed_discovery_record(
        secret_key: &SecretKey,
        handshake: &ZakuraHandshakeConfig,
    ) -> Result<ZakuraNodeRecord, crate::BoxError> {
        let body = ZakuraNodeRecordBody {
            node_id: secret_key.public(),
            direct_addrs: vec![SocketAddr::new(
                IpAddr::V4(Ipv4Addr::new(45, 33, 30, 45)),
                8233,
            )],
            services: vec![ZakuraServiceId::discovery()],
            zakura_protocol_min: handshake.zakura_protocol_min,
            zakura_protocol_max: handshake.zakura_protocol_max,
            network_id: handshake.network_id,
            chain_id: handshake.chain_id,
            sequence: 1,
            expires_at_unix_secs: current_test_unix_secs().saturating_add(60),
        };
        Ok(ZakuraNodeRecord::sign(body, secret_key)?)
    }

    async fn complete_peer_side_discovery_exchange(
        peer_send: &FramedSend,
        peer_recv: &mut FramedRecv,
        peer_secret: &SecretKey,
        handshake: &ZakuraHandshakeConfig,
    ) -> Result<(), crate::BoxError> {
        let mut saw_hello = false;
        let mut saw_get_peers = false;
        let mut saw_get_services = false;
        while !(saw_hello && saw_get_peers && saw_get_services) {
            let frame = tokio::time::timeout(Duration::from_secs(2), peer_recv.recv())
                .await?
                .expect("discovery source sends exchange frames");
            match decode_discovery_frame(&frame)? {
                DiscoveryMessage::Hello { .. } => saw_hello = true,
                DiscoveryMessage::GetPeers { .. } => saw_get_peers = true,
                DiscoveryMessage::GetServices(_) => saw_get_services = true,
                DiscoveryMessage::Peers { .. } | DiscoveryMessage::Services(_) => {}
            }
        }

        peer_send
            .send(discovery_frame(DiscoveryMessage::Hello {
                record: signed_discovery_record(peer_secret, handshake)?,
            })?)
            .await?;
        peer_send
            .send(discovery_frame(DiscoveryMessage::Peers {
                records: Vec::new(),
            })?)
            .await?;
        let summary = DiscoveryServiceSummary {
            peer_exchange_slots_free: 1,
            max_records_per_response: 1,
            expected_disconnect_after_exchange: true,
        };
        peer_send
            .send(discovery_frame(DiscoveryMessage::Services(Services {
                node_id: peer_secret.public(),
                expires_at_unix_secs: u64::MAX,
                summaries: vec![ServiceSummaryEnvelope::discovery(&summary)?],
            }))?)
            .await?;

        Ok(())
    }

    async fn wait_for_next_discovery_hello(
        peer_recv: &mut FramedRecv,
    ) -> Result<ZakuraNodeRecord, crate::BoxError> {
        tokio::time::timeout(Duration::from_secs(2), async {
            loop {
                let frame = peer_recv
                    .recv()
                    .await
                    .expect("discovery refresh stream remains open");
                if let DiscoveryMessage::Hello { record } = decode_discovery_frame(&frame)? {
                    return Ok(record);
                }
            }
        })
        .await?
    }

    async fn wait_for_discovery_inbound_peers(handle: &ZakuraDiscoveryHandle, expected: usize) {
        tokio::time::timeout(Duration::from_secs(2), async {
            loop {
                if handle.peer_snapshot().inbound_peers == expected {
                    return;
                }
                tokio::time::sleep(Duration::from_millis(10)).await;
            }
        })
        .await
        .expect("discovery peer snapshot reaches expected inbound count");
    }

    async fn advisory_backoff_after_empty_headers(
        fixture: &mut HeaderAdvisoryFixture,
    ) -> Result<bool, crate::BoxError> {
        let (send, _recv) = framed_channel(32);
        let session = HeaderSyncPeerSession::from_parts_with_direction(
            fixture.peer_id.clone(),
            ServicePeerDirection::Inbound,
            send,
            CancellationToken::new(),
        );
        fixture
            .header_sync
            .send(HeaderSyncEvent::PeerConnected(session))
            .await?;
        fixture
            .header_sync
            .send(HeaderSyncEvent::WireMessage {
                peer: fixture.peer_id.clone(),
                msg: HeaderSyncMessage::Status(HeaderSyncStatus {
                    tip_height: block::Height(1),
                    tip_hash: block::Hash([9; 32]),
                    anchor_height: block::Height(0),
                    max_headers_per_response: 1,
                    max_inflight_requests: 1,
                }),
            })
            .await?;

        let request_id = tokio::time::timeout(Duration::from_secs(2), async {
            loop {
                if let Some(HeaderSyncAction::SendMessage {
                    peer,
                    request_id,
                    msg: HeaderSyncMessage::GetHeaders { .. },
                }) = fixture.header_actions.recv().await
                {
                    if peer == fixture.peer_id {
                        return request_id
                            .expect("an outbound GetHeaders always carries a request ID");
                    }
                }
            }
        })
        .await
        .expect("header sync schedules a request before empty response");

        fixture
            .header_sync
            .send(HeaderSyncEvent::WireHeaders {
                wire_request: HeaderSyncWireRequestIdentity {
                    peer: fixture.peer_id.clone(),
                    session_id: 0,
                    request_id,
                },
                entries: Vec::new(),
            })
            .await?;
        tokio::time::sleep(Duration::from_millis(20)).await;

        Ok(fixture
            .header_sync
            .candidate_state()
            .backed_off_node_ids
            .contains(&fixture.peer_node_id))
    }

    #[tokio::test]
    async fn get_services_returns_local_first_party_discovery_summary(
    ) -> Result<(), crate::BoxError> {
        let (_connected_tx, connected_rx) = watch::channel(Vec::new());
        let handshake = ZakuraHandshakeConfig::for_network(&Network::Mainnet);
        let local_secret = SecretKey::from_bytes(&[21u8; 32]);
        let handle = ZakuraDiscoveryHandle::new(
            ZakuraDiscoveryLocalConfig {
                secret_key: local_secret.clone(),
                direct_addrs: Vec::new(),
                services: vec![ZakuraServiceId::discovery()],
                zakura_protocol_min: handshake.zakura_protocol_min,
                zakura_protocol_max: handshake.zakura_protocol_max,
                network_id: handshake.network_id,
                chain_id: handshake.chain_id,
                last_authored_sequence: None,
            },
            ZakuraDiscoveryConfig {
                peer_limits: ServicePeerLimits {
                    max_inbound_peers: 4,
                    ..ServicePeerLimits::default()
                },
                ..ZakuraDiscoveryConfig::default()
            },
            connected_rx,
        )?;
        let service = DiscoveryService::new(handle.clone());
        let peer_node_id = SecretKey::from_bytes(&[22u8; 32]).public();
        let peer_id = ZakuraPeerId::new(peer_node_id.as_bytes().to_vec())?;
        let (peer_send, service_recv) = framed_channel(8);
        let (service_send, mut peer_recv) = framed_channel(8);
        let streams = HashMap::from([(ZAKURA_STREAM_DISCOVERY, (service_recv, service_send))]);

        service.add_peer(Peer::new(
            peer_id,
            None,
            ZAKURA_CAP_DISCOVERY,
            streams,
            CancellationToken::new(),
        ));

        peer_send
            .send(Frame {
                message_type: DISCOVERY_FRAME_MESSAGE_TYPE,
                flags: 0,
                payload: DiscoveryMessage::GetServices(GetServices {
                    wanted_services: vec![ZakuraServiceId::discovery()],
                })
                .encode()?,
            })
            .await?;

        let services = tokio::time::timeout(Duration::from_secs(2), async {
            loop {
                let frame = peer_recv.recv().await.expect("discovery stream stays open");
                let message = decode_discovery_frame(&frame).expect("outbound frame decodes");
                if let DiscoveryMessage::Services(services) = message {
                    return services;
                }
            }
        })
        .await
        .expect("service response is sent");

        assert_eq!(services.node_id, local_secret.public());
        assert_eq!(services.summaries.len(), 1);
        assert_eq!(
            services.summaries[0].service_id,
            ZakuraServiceId::discovery()
        );
        let summary = services.summaries[0]
            .decode_discovery()?
            .expect("discovery summary tag decodes");
        assert_eq!(summary.peer_exchange_slots_free, 3);
        assert!(summary.expected_disconnect_after_exchange);
        assert_eq!(
            summary.max_records_per_response,
            u16::try_from(MAX_DISCOVERY_RECORDS_PER_RESPONSE)
                .expect("record response cap fits in u16")
        );

        Ok(())
    }

    #[tokio::test]
    async fn get_services_returns_local_first_party_block_sync_summary(
    ) -> Result<(), crate::BoxError> {
        let (_connected_tx, connected_rx) = watch::channel(Vec::new());
        let handshake = ZakuraHandshakeConfig::for_network(&Network::Mainnet);
        let local_secret = SecretKey::from_bytes(&[24u8; 32]);
        let discovery_handle = ZakuraDiscoveryHandle::new(
            ZakuraDiscoveryLocalConfig {
                secret_key: local_secret.clone(),
                direct_addrs: Vec::new(),
                services: vec![ZakuraServiceId::discovery(), ZakuraServiceId::block_sync()],
                zakura_protocol_min: handshake.zakura_protocol_min,
                zakura_protocol_max: handshake.zakura_protocol_max,
                network_id: handshake.network_id,
                chain_id: handshake.chain_id,
                last_authored_sequence: None,
            },
            ZakuraDiscoveryConfig::default(),
            connected_rx,
        )?;
        let (header_sync, _header_actions, header_task) = spawn_test_header_sync()?;
        let (tip_tx, tip_rx) = watch::channel((block::Height(5), block::Hash([5; 32])));
        drop(tip_tx);
        let (block_sync, _block_actions, block_task) =
            spawn_block_sync_reactor(BlockSyncStartup::new(
                BlockSyncFrontiers {
                    finalized_height: block::Height(0),
                    verified_block_tip: block::Height(5),
                    verified_block_hash: block::Hash([5; 32]),
                },
                (block::Height(5), block::Hash([5; 32])),
                tip_rx,
                ZakuraBlockSyncConfig::default(),
            ));
        let service = DiscoveryService::with_sync_services(
            discovery_handle,
            header_sync,
            Some(block_sync.clone()),
        );
        let peer_node_id = SecretKey::from_bytes(&[25u8; 32]).public();
        let peer_id = ZakuraPeerId::new(peer_node_id.as_bytes().to_vec())?;
        let (peer_send, service_recv) = framed_channel(8);
        let (service_send, mut peer_recv) = framed_channel(8);
        let streams = HashMap::from([(ZAKURA_STREAM_DISCOVERY, (service_recv, service_send))]);

        service.add_peer(Peer::new(
            peer_id,
            None,
            ZAKURA_CAP_DISCOVERY | ZAKURA_CAP_BLOCK_SYNC,
            streams,
            CancellationToken::new(),
        ));

        peer_send
            .send(Frame {
                message_type: DISCOVERY_FRAME_MESSAGE_TYPE,
                flags: 0,
                payload: DiscoveryMessage::GetServices(GetServices {
                    wanted_services: vec![ZakuraServiceId::block_sync()],
                })
                .encode()?,
            })
            .await?;

        let services = tokio::time::timeout(Duration::from_secs(2), async {
            loop {
                let frame = peer_recv.recv().await.expect("discovery stream stays open");
                let message = decode_discovery_frame(&frame).expect("outbound frame decodes");
                if let DiscoveryMessage::Services(services) = message {
                    return services;
                }
            }
        })
        .await
        .expect("service response is sent");

        assert_eq!(services.node_id, local_secret.public());
        assert_eq!(services.summaries.len(), 1);
        assert_eq!(
            services.summaries[0].service_id,
            ZakuraServiceId::block_sync()
        );
        let summary = services.summaries[0]
            .decode_block_sync()?
            .expect("block summary tag decodes");
        assert_eq!(summary.servable_low, block::Height(0));
        assert_eq!(summary.servable_high, block::Height(5));
        assert_eq!(summary.tip_hash, block::Hash([5; 32]));
        assert_eq!(
            usize::from(summary.free_slots),
            block_sync.peer_snapshot().inbound_slots_free
        );
        assert_eq!(
            summary.max_blocks_per_response,
            ZakuraBlockSyncConfig::default().advertised_max_blocks_per_response()
        );
        assert_eq!(summary.max_response_bytes, MAX_BS_RESPONSE_BYTES);

        header_task.abort();
        block_task.abort();
        Ok(())
    }

    #[tokio::test]
    async fn inbound_services_updates_first_party_live_summary_cache() -> Result<(), crate::BoxError>
    {
        let (connected_tx, connected_rx) = watch::channel(Vec::new());
        let handshake = ZakuraHandshakeConfig::for_network(&Network::Mainnet);
        let local_secret = SecretKey::from_bytes(&[23u8; 32]);
        let handle = ZakuraDiscoveryHandle::new(
            ZakuraDiscoveryLocalConfig {
                secret_key: local_secret,
                direct_addrs: Vec::new(),
                services: vec![ZakuraServiceId::discovery()],
                zakura_protocol_min: handshake.zakura_protocol_min,
                zakura_protocol_max: handshake.zakura_protocol_max,
                network_id: handshake.network_id,
                chain_id: handshake.chain_id,
                last_authored_sequence: None,
            },
            ZakuraDiscoveryConfig::default(),
            connected_rx,
        )?;
        let service = DiscoveryService::new(handle.clone());
        let peer_node_id = SecretKey::from_bytes(&[24u8; 32]).public();
        let peer_id = ZakuraPeerId::new(peer_node_id.as_bytes().to_vec())?;
        connected_tx.send_replace(vec![peer_id.clone()]);

        let (peer_send, service_recv) = framed_channel(8);
        let (service_send, _peer_recv) = framed_channel(8);
        let streams = HashMap::from([(ZAKURA_STREAM_DISCOVERY, (service_recv, service_send))]);

        service.add_peer(Peer::new(
            peer_id,
            None,
            ZAKURA_CAP_DISCOVERY,
            streams,
            CancellationToken::new(),
        ));

        let summary = DiscoveryServiceSummary {
            peer_exchange_slots_free: 7,
            max_records_per_response: 11,
            expected_disconnect_after_exchange: false,
        };
        peer_send
            .send(Frame {
                message_type: DISCOVERY_FRAME_MESSAGE_TYPE,
                flags: 0,
                payload: DiscoveryMessage::Services(Services {
                    node_id: peer_node_id,
                    expires_at_unix_secs: u64::MAX,
                    summaries: vec![ServiceSummaryEnvelope::discovery(&summary)?],
                })
                .encode()?,
            })
            .await?;

        let cached = tokio::time::timeout(Duration::from_secs(2), async {
            loop {
                if let Some(cached) = handle.live_service_summaries(peer_node_id).await {
                    if !cached.is_empty() {
                        return cached;
                    }
                }
                tokio::time::sleep(Duration::from_millis(10)).await;
            }
        })
        .await
        .expect("inbound SERVICES is imported");

        assert_eq!(cached.len(), 1);
        assert_eq!(
            cached[0].summary,
            ZakuraLiveServiceSummary::Discovery(summary)
        );

        Ok(())
    }

    #[tokio::test]
    async fn first_party_header_services_emit_header_sync_advisory() -> Result<(), crate::BoxError>
    {
        let mut fixture = spawn_header_advisory_fixture(25)?;
        let summary = header_summary(block::Height(10));

        send_discovery_message(
            &fixture,
            DiscoveryMessage::Services(Services {
                node_id: fixture.peer_node_id,
                expires_at_unix_secs: u64::MAX,
                summaries: vec![ServiceSummaryEnvelope::header_sync(&summary)?],
            }),
        )
        .await?;

        tokio::time::timeout(Duration::from_secs(2), async {
            loop {
                if let Some(cached) = fixture
                    .discovery_handle
                    .live_service_summaries(fixture.peer_node_id)
                    .await
                {
                    if cached.iter().any(|cached_summary| {
                        cached_summary.summary == ZakuraLiveServiceSummary::HeaderSync(summary)
                    }) {
                        return;
                    }
                }
                tokio::time::sleep(Duration::from_millis(10)).await;
            }
        })
        .await
        .expect("first-party header summary is cached");

        assert!(
            advisory_backoff_after_empty_headers(&mut fixture).await?,
            "first-party header SERVICES should emit a header-sync advisory event"
        );

        Ok(())
    }

    #[tokio::test]
    async fn first_party_retired_header_services_emit_header_sync_advisory(
    ) -> Result<(), crate::BoxError> {
        let mut fixture = spawn_header_advisory_fixture(35)?;
        let summary = header_summary(block::Height(10));
        let mut legacy_envelope = ServiceSummaryEnvelope::header_sync(&summary)?;
        legacy_envelope.service_id = ZakuraServiceId::header_sync_retired();
        legacy_envelope.summary_tag = SUMMARY_TAG_HEADER_SYNC_RETIRED;

        send_discovery_message(
            &fixture,
            DiscoveryMessage::Services(Services {
                node_id: fixture.peer_node_id,
                expires_at_unix_secs: u64::MAX,
                summaries: vec![legacy_envelope],
            }),
        )
        .await?;

        tokio::time::timeout(Duration::from_secs(2), async {
            loop {
                if let Some(cached) = fixture
                    .discovery_handle
                    .live_service_summaries(fixture.peer_node_id)
                    .await
                {
                    if cached.iter().any(|cached_summary| {
                        cached_summary.service_id == ZakuraServiceId::header_sync_retired()
                            && cached_summary.summary
                                == ZakuraLiveServiceSummary::HeaderSync(summary)
                    }) {
                        return;
                    }
                }
                tokio::time::sleep(Duration::from_millis(10)).await;
            }
        })
        .await
        .expect("retired first-party header summary is cached");

        assert!(
            advisory_backoff_after_empty_headers(&mut fixture).await?,
            "retired first-party header SERVICES should emit a header-sync advisory event"
        );

        Ok(())
    }

    #[tokio::test]
    async fn mismatched_services_node_id_does_not_emit_header_sync_advisory(
    ) -> Result<(), crate::BoxError> {
        let mut fixture = spawn_header_advisory_fixture(26)?;
        let claimed_node_id = SecretKey::from_bytes(&[27u8; 32]).public();
        let summary = header_summary(block::Height(10));

        send_discovery_message(
            &fixture,
            DiscoveryMessage::Services(Services {
                node_id: claimed_node_id,
                expires_at_unix_secs: u64::MAX,
                summaries: vec![ServiceSummaryEnvelope::header_sync(&summary)?],
            }),
        )
        .await?;
        tokio::time::sleep(Duration::from_millis(20)).await;

        assert_eq!(
            fixture
                .discovery_handle
                .live_service_summaries(fixture.peer_node_id)
                .await,
            None
        );
        assert_eq!(
            fixture
                .discovery_handle
                .live_service_summaries(claimed_node_id)
                .await,
            None
        );
        assert!(
            !advisory_backoff_after_empty_headers(&mut fixture).await?,
            "mismatched SERVICES node id must not emit a header-sync advisory event"
        );

        Ok(())
    }

    #[tokio::test]
    async fn peers_response_does_not_emit_header_sync_advisory() -> Result<(), crate::BoxError> {
        let mut fixture = spawn_header_advisory_fixture(28)?;
        let handshake = ZakuraHandshakeConfig::for_network(&Network::Mainnet);
        let record_secret = SecretKey::from_bytes(&[29u8; 32]);
        let record = signed_header_sync_record(&record_secret, &handshake)?;

        send_discovery_message(
            &fixture,
            DiscoveryMessage::Peers {
                records: vec![record.clone()],
            },
        )
        .await?;
        tokio::time::sleep(Duration::from_millis(20)).await;

        assert_eq!(
            fixture
                .discovery_handle
                .live_service_summaries(record.body.node_id)
                .await,
            None
        );
        assert!(
            !advisory_backoff_after_empty_headers(&mut fixture).await?,
            "PEERS/gossiped records must not emit live header-sync advisory events"
        );

        Ok(())
    }

    #[tokio::test]
    async fn stale_discovery_source_cannot_send_after_stream_replacement(
    ) -> Result<(), crate::BoxError> {
        let (_connected_tx, connected_rx) = watch::channel(Vec::new());
        let handshake = ZakuraHandshakeConfig::for_network(&Network::Mainnet);
        let handle = ZakuraDiscoveryHandle::new(
            ZakuraDiscoveryLocalConfig {
                secret_key: SecretKey::from_bytes(&[38u8; 32]),
                direct_addrs: Vec::new(),
                services: vec![ZakuraServiceId::discovery()],
                zakura_protocol_min: handshake.zakura_protocol_min,
                zakura_protocol_max: handshake.zakura_protocol_max,
                network_id: handshake.network_id,
                chain_id: handshake.chain_id,
                last_authored_sequence: None,
            },
            ZakuraDiscoveryConfig::default(),
            connected_rx,
        )?;
        let peer_node_id = SecretKey::from_bytes(&[39u8; 32]).public();
        let peer_id = ZakuraPeerId::new(peer_node_id.as_bytes().to_vec())?;
        assert_eq!(
            handle
                .admit_peer_session(2, 1, peer_id.clone(), ServicePeerDirection::Inbound)
                .await,
            ServiceAdmissionDecision::Admit
        );

        let (send, mut recv) = framed_channel(4);
        let source = DiscoverySource {
            handle: handle.clone(),
            session: DiscoveryPeerSession {
                peer_id: peer_id.clone(),
                direction: ServicePeerDirection::Inbound,
                send,
                cancel: CancellationToken::new(),
            },
            conn_id: 2,
            session_id: 1,
            progress: Arc::new(DiscoveryExchangeProgress::default()),
        };
        assert_eq!(
            handle
                .admit_peer_session(2, 2, peer_id, ServicePeerDirection::Inbound)
                .await,
            ServiceAdmissionDecision::Admit
        );

        assert_eq!(source.exchange().await, Err(()));
        assert!(
            tokio::time::timeout(Duration::from_millis(20), recv.recv())
                .await
                .is_err(),
            "a replaced discovery stream must not emit another exchange"
        );
        Ok(())
    }

    #[tokio::test]
    async fn discovery_source_replaced_during_settle_does_not_complete_exchange(
    ) -> Result<(), crate::BoxError> {
        let (_connected_tx, connected_rx) = watch::channel(Vec::new());
        let handshake = ZakuraHandshakeConfig::for_network(&Network::Mainnet);
        let handle = ZakuraDiscoveryHandle::new(
            ZakuraDiscoveryLocalConfig {
                secret_key: SecretKey::from_bytes(&[48u8; 32]),
                direct_addrs: Vec::new(),
                services: vec![ZakuraServiceId::discovery()],
                zakura_protocol_min: handshake.zakura_protocol_min,
                zakura_protocol_max: handshake.zakura_protocol_max,
                network_id: handshake.network_id,
                chain_id: handshake.chain_id,
                last_authored_sequence: None,
            },
            ZakuraDiscoveryConfig::default(),
            connected_rx,
        )?;
        let peer_node_id = SecretKey::from_bytes(&[49u8; 32]).public();
        let peer_id = ZakuraPeerId::new(peer_node_id.as_bytes().to_vec())?;
        assert_eq!(
            handle
                .admit_peer_session(2, 1, peer_id.clone(), ServicePeerDirection::Inbound)
                .await,
            ServiceAdmissionDecision::Admit
        );

        let (send, mut recv) = framed_channel(4);
        let progress = Arc::new(DiscoveryExchangeProgress::default());
        let source = DiscoverySource {
            handle: handle.clone(),
            session: DiscoveryPeerSession {
                peer_id: peer_id.clone(),
                direction: ServicePeerDirection::Inbound,
                send,
                cancel: CancellationToken::new(),
            },
            conn_id: 2,
            session_id: 1,
            progress: progress.clone(),
        };
        let source_task = tokio::spawn(async move { source.run_initial_exchange().await });
        for _ in 0..3 {
            recv.recv()
                .await
                .expect("initial discovery request is sent");
        }

        assert_eq!(
            handle
                .admit_peer_session(2, 2, peer_id, ServicePeerDirection::Inbound)
                .await,
            ServiceAdmissionDecision::Admit
        );
        progress.mark_hello();
        progress.mark_peers();
        progress.mark_services();

        assert!(!source_task.await?);
        Ok(())
    }

    #[tokio::test(start_paused = true)]
    async fn silent_discovery_source_does_not_complete_exchange() -> Result<(), crate::BoxError> {
        let (_connected_tx, connected_rx) = watch::channel(Vec::new());
        let handshake = ZakuraHandshakeConfig::for_network(&Network::Mainnet);
        let handle = ZakuraDiscoveryHandle::new(
            ZakuraDiscoveryLocalConfig {
                secret_key: SecretKey::from_bytes(&[50u8; 32]),
                direct_addrs: Vec::new(),
                services: vec![ZakuraServiceId::discovery()],
                zakura_protocol_min: handshake.zakura_protocol_min,
                zakura_protocol_max: handshake.zakura_protocol_max,
                network_id: handshake.network_id,
                chain_id: handshake.chain_id,
                last_authored_sequence: None,
            },
            ZakuraDiscoveryConfig::default(),
            connected_rx,
        )?;
        let peer_node_id = SecretKey::from_bytes(&[51u8; 32]).public();
        let peer_id = ZakuraPeerId::new(peer_node_id.as_bytes().to_vec())?;
        assert_eq!(
            handle
                .admit_peer_session(2, 1, peer_id.clone(), ServicePeerDirection::Inbound)
                .await,
            ServiceAdmissionDecision::Admit
        );

        let (send, _recv) = framed_channel(4);
        let source = DiscoverySource {
            handle,
            session: DiscoveryPeerSession {
                peer_id,
                direction: ServicePeerDirection::Inbound,
                send,
                cancel: CancellationToken::new(),
            },
            conn_id: 2,
            session_id: 1,
            progress: Arc::new(DiscoveryExchangeProgress::default()),
        };

        assert!(
            !source.run_initial_exchange().await,
            "a peer that sends no discovery responses must fail the initial exchange"
        );
        Ok(())
    }

    #[tokio::test]
    async fn discovery_exchange_closes_connection_when_block_sync_is_not_admitted(
    ) -> Result<(), crate::BoxError> {
        let (connected_tx, connected_rx) = watch::channel(Vec::new());
        let handshake = ZakuraHandshakeConfig::for_network(&Network::Mainnet);
        let local_secret = SecretKey::from_bytes(&[40u8; 32]);
        let handle = ZakuraDiscoveryHandle::new(
            ZakuraDiscoveryLocalConfig {
                secret_key: local_secret,
                direct_addrs: Vec::new(),
                services: vec![ZakuraServiceId::discovery()],
                zakura_protocol_min: handshake.zakura_protocol_min,
                zakura_protocol_max: handshake.zakura_protocol_max,
                network_id: handshake.network_id,
                chain_id: handshake.chain_id,
                last_authored_sequence: None,
            },
            ZakuraDiscoveryConfig::default(),
            connected_rx,
        )?;
        let service = DiscoveryService::new(handle.clone());
        let peer_secret = SecretKey::from_bytes(&[41u8; 32]);
        let peer_node_id = peer_secret.public();
        let peer_id = ZakuraPeerId::new(peer_node_id.as_bytes().to_vec())?;
        connected_tx.send_replace(vec![peer_id.clone()]);

        let connection_cancel = CancellationToken::new();
        let (peer_send, service_recv) = framed_channel(16);
        let (service_send, mut peer_recv) = framed_channel(16);
        let streams = HashMap::from([(ZAKURA_STREAM_DISCOVERY, (service_recv, service_send))]);

        service.add_peer(Peer::new(
            peer_id.clone(),
            None,
            ZAKURA_CAP_DISCOVERY | ZAKURA_CAP_BLOCK_SYNC,
            streams,
            connection_cancel.clone(),
        ));

        wait_for_discovery_inbound_peers(&handle, 1).await;
        complete_peer_side_discovery_exchange(&peer_send, &mut peer_recv, &peer_secret, &handshake)
            .await?;
        tokio::time::timeout(Duration::from_secs(2), connection_cancel.cancelled())
            .await
            .expect("unadmitted block sync does not keep the shared connection alive");
        // A finished discovery-only exchange retires the session so the handler
        // cannot reopen the stream while the connection tears down.
        assert!(matches!(
            service.ordered_session_demand(
                0,
                &peer_id,
                ZAKURA_CAP_DISCOVERY | ZAKURA_CAP_BLOCK_SYNC,
                ServicePeerDirection::Inbound,
            ),
            OrderedSessionDemand::Retire,
        ));
        wait_for_discovery_inbound_peers(&handle, 0).await;

        connected_tx.send_replace(Vec::new());
        assert!(handle
            .dial_candidates(&[ZakuraServiceId::discovery()], &[])
            .await
            .is_empty());

        Ok(())
    }

    #[tokio::test]
    async fn admitted_legacy_gossip_keeps_discovery_connection_alive() -> Result<(), crate::BoxError>
    {
        let (connected_tx, connected_rx) = watch::channel(Vec::new());
        let handshake = ZakuraHandshakeConfig::for_network(&Network::Mainnet);
        let handle = ZakuraDiscoveryHandle::new(
            ZakuraDiscoveryLocalConfig {
                secret_key: SecretKey::from_bytes(&[46u8; 32]),
                direct_addrs: Vec::new(),
                services: vec![ZakuraServiceId::discovery()],
                zakura_protocol_min: handshake.zakura_protocol_min,
                zakura_protocol_max: handshake.zakura_protocol_max,
                network_id: handshake.network_id,
                chain_id: handshake.chain_id,
                last_authored_sequence: None,
            },
            ZakuraDiscoveryConfig {
                refresh_interval: Duration::from_millis(20),
                ..ZakuraDiscoveryConfig::default()
            },
            connected_rx,
        )?;
        let service = DiscoveryService::new(handle.clone());
        let peer_secret = SecretKey::from_bytes(&[47u8; 32]);
        let peer_node_id = peer_secret.public();
        let peer_id = ZakuraPeerId::new(peer_node_id.as_bytes().to_vec())?;
        service.set_connection_owners(vec![Arc::new(TestConnectionOwner {
            peer: peer_id.clone(),
            conn_id: 0,
        })]);
        connected_tx.send_replace(vec![peer_id.clone()]);

        let connection_cancel = CancellationToken::new();
        let (peer_send, service_recv) = framed_channel(16);
        let (service_send, mut peer_recv) = framed_channel(16);
        service.add_peer(Peer::new(
            peer_id,
            None,
            ZAKURA_CAP_DISCOVERY | ZAKURA_CAP_LEGACY_GOSSIP,
            HashMap::from([(ZAKURA_STREAM_DISCOVERY, (service_recv, service_send))]),
            connection_cancel.clone(),
        ));

        wait_for_discovery_inbound_peers(&handle, 1).await;
        complete_peer_side_discovery_exchange(&peer_send, &mut peer_recv, &peer_secret, &handshake)
            .await?;
        wait_for_next_discovery_hello(&mut peer_recv).await?;
        assert!(
            tokio::time::timeout(Duration::from_millis(100), connection_cancel.cancelled())
                .await
                .is_err(),
            "an admitted legacy-gossip service owns the shared connection"
        );
        assert_eq!(handle.peer_snapshot().inbound_peers, 1);

        connection_cancel.cancel();
        wait_for_discovery_inbound_peers(&handle, 0).await;
        Ok(())
    }

    #[tokio::test]
    async fn admitted_block_sync_keeps_discovery_connection_alive() -> Result<(), crate::BoxError> {
        let (connected_tx, connected_rx) = watch::channel(Vec::new());
        let handshake = ZakuraHandshakeConfig::for_network(&Network::Mainnet);
        let handle = ZakuraDiscoveryHandle::new(
            ZakuraDiscoveryLocalConfig {
                secret_key: SecretKey::from_bytes(&[44u8; 32]),
                direct_addrs: Vec::new(),
                services: vec![ZakuraServiceId::discovery(), ZakuraServiceId::block_sync()],
                zakura_protocol_min: handshake.zakura_protocol_min,
                zakura_protocol_max: handshake.zakura_protocol_max,
                network_id: handshake.network_id,
                chain_id: handshake.chain_id,
                last_authored_sequence: None,
            },
            ZakuraDiscoveryConfig {
                refresh_interval: Duration::from_millis(20),
                ..ZakuraDiscoveryConfig::default()
            },
            connected_rx,
        )?;
        let service = DiscoveryService::new(handle.clone());
        let (block_sync, _block_events) =
            BlockSyncService::new_for_test(ZakuraBlockSyncConfig::default());
        let block_sync = Arc::new(block_sync);
        service.set_connection_owners(vec![block_sync.clone()]);

        let peer_secret = SecretKey::from_bytes(&[45u8; 32]);
        let peer_node_id = peer_secret.public();
        let peer_id = ZakuraPeerId::new(peer_node_id.as_bytes().to_vec())?;
        connected_tx.send_replace(vec![peer_id.clone()]);
        let connection_cancel = CancellationToken::new();

        let (_peer_block_send, service_block_recv) = framed_channel(4);
        let (service_block_send, _peer_block_recv) = framed_channel(4);
        block_sync.add_peer(Peer::new(
            peer_id.clone(),
            None,
            ZAKURA_CAP_BLOCK_SYNC,
            HashMap::from([(
                crate::zakura::ZAKURA_STREAM_BLOCK_SYNC,
                (service_block_recv, service_block_send),
            )]),
            connection_cancel.clone(),
        ));

        let (peer_send, service_recv) = framed_channel(16);
        let (service_send, mut peer_recv) = framed_channel(16);
        service.add_peer(Peer::new(
            peer_id,
            None,
            ZAKURA_CAP_DISCOVERY | ZAKURA_CAP_BLOCK_SYNC,
            HashMap::from([(ZAKURA_STREAM_DISCOVERY, (service_recv, service_send))]),
            connection_cancel.clone(),
        ));

        wait_for_discovery_inbound_peers(&handle, 1).await;
        complete_peer_side_discovery_exchange(&peer_send, &mut peer_recv, &peer_secret, &handshake)
            .await?;
        wait_for_next_discovery_hello(&mut peer_recv).await?;
        assert!(
            tokio::time::timeout(Duration::from_millis(100), connection_cancel.cancelled())
                .await
                .is_err(),
            "an admitted block-sync service owns the shared connection"
        );
        assert_eq!(handle.peer_snapshot().inbound_peers, 1);

        connection_cancel.cancel();
        wait_for_discovery_inbound_peers(&handle, 0).await;
        Ok(())
    }

    #[tokio::test]
    async fn discovery_refreshes_while_header_sync_owns_connection() -> Result<(), crate::BoxError>
    {
        let (connected_tx, connected_rx) = watch::channel(Vec::new());
        let handshake = ZakuraHandshakeConfig::for_network(&Network::Mainnet);
        let local_secret = SecretKey::from_bytes(&[42u8; 32]);
        let discovery_handle = ZakuraDiscoveryHandle::new(
            ZakuraDiscoveryLocalConfig {
                secret_key: local_secret,
                direct_addrs: Vec::new(),
                services: vec![ZakuraServiceId::discovery()],
                zakura_protocol_min: handshake.zakura_protocol_min,
                zakura_protocol_max: handshake.zakura_protocol_max,
                network_id: handshake.network_id,
                chain_id: handshake.chain_id,
                last_authored_sequence: None,
            },
            ZakuraDiscoveryConfig {
                refresh_interval: Duration::from_millis(20),
                ..ZakuraDiscoveryConfig::default()
            },
            connected_rx,
        )?;
        let (header_sync, _header_actions, header_task) = spawn_test_header_sync()?;
        let service = DiscoveryService::with_sync_services(
            discovery_handle.clone(),
            header_sync.clone(),
            None,
        );
        let header_service = Arc::new(HeaderSyncService::new(header_sync.clone()));
        service.set_connection_owners(vec![header_service.clone()]);
        let peer_secret = SecretKey::from_bytes(&[43u8; 32]);
        let peer_node_id = peer_secret.public();
        let peer_id = ZakuraPeerId::new(peer_node_id.as_bytes().to_vec())?;
        connected_tx.send_replace(vec![peer_id.clone()]);

        let connection_cancel = CancellationToken::new();
        let (_peer_header_send, service_header_recv) = framed_channel(8);
        let (service_header_send, _peer_header_recv) = framed_channel(8);
        header_service.add_peer(Peer::new(
            peer_id.clone(),
            None,
            ZAKURA_CAP_HEADER_SYNC,
            HashMap::from([(
                ZAKURA_STREAM_HEADER_SYNC,
                (service_header_recv, service_header_send),
            )]),
            connection_cancel.clone(),
        ));
        tokio::time::timeout(Duration::from_secs(2), async {
            loop {
                if header_sync
                    .candidate_state()
                    .admitted_node_ids
                    .contains(&peer_node_id)
                {
                    return;
                }
                tokio::time::sleep(Duration::from_millis(10)).await;
            }
        })
        .await
        .expect("header sync admits the peer");

        let (peer_send, service_recv) = framed_channel(16);
        let (service_send, mut peer_recv) = framed_channel(16);
        let streams = HashMap::from([(ZAKURA_STREAM_DISCOVERY, (service_recv, service_send))]);
        service.add_peer(Peer::new(
            peer_id.clone(),
            None,
            ZAKURA_CAP_DISCOVERY | ZAKURA_CAP_HEADER_SYNC,
            streams,
            connection_cancel.clone(),
        ));

        wait_for_discovery_inbound_peers(&discovery_handle, 1).await;
        complete_peer_side_discovery_exchange(&peer_send, &mut peer_recv, &peer_secret, &handshake)
            .await?;
        let refreshed_record = wait_for_next_discovery_hello(&mut peer_recv).await?;
        assert_eq!(
            refreshed_record.body.node_id,
            discovery_handle.local_node_id()
        );
        assert_eq!(discovery_handle.peer_snapshot().inbound_peers, 1);
        assert_eq!(header_sync.peer_snapshot().inbound_peers, 1);
        assert!(
            tokio::time::timeout(Duration::from_millis(100), connection_cancel.cancelled())
                .await
                .is_err(),
            "discovery releases only its own session while header sync owns the connection"
        );
        // The refresh loop still owns this session, so the handler must keep
        // the stream eligible instead of retiring it for the connection.
        assert!(!matches!(
            service.ordered_session_demand(
                0,
                &peer_id,
                ZAKURA_CAP_DISCOVERY | ZAKURA_CAP_HEADER_SYNC,
                ServicePeerDirection::Inbound,
            ),
            OrderedSessionDemand::Retire,
        ));

        connection_cancel.cancel();
        wait_for_discovery_inbound_peers(&discovery_handle, 0).await;
        header_task.abort();
        Ok(())
    }

    #[test]
    fn teardown_before_retirement_does_not_recreate_discovery_session_state() {
        let peer = ZakuraPeerId::new(vec![71; 32]).expect("test peer id is within bounds");
        let conn_id = 9;
        let states = StdMutex::new(HashMap::from([(
            (peer.clone(), conn_id),
            DiscoverySessionRecord {
                session_id: 1,
                state: DiscoverySessionState::Active,
                failed_attempts: 0,
            },
        )]));

        states
            .lock()
            .expect("test session-state mutex is never poisoned")
            .remove(&(peer.clone(), conn_id));

        assert!(!retire_discovery_session(&states, &peer, conn_id, 1));
        assert!(states
            .lock()
            .expect("test session-state mutex is never poisoned")
            .is_empty());
    }

    #[test]
    fn stale_exchange_cannot_retire_a_newer_session_on_the_same_connection() {
        let peer = ZakuraPeerId::new(vec![72; 32]).expect("test peer id is within bounds");
        let conn_id = 9;
        let states = StdMutex::new(HashMap::from([(
            (peer.clone(), conn_id),
            DiscoverySessionRecord {
                session_id: 1,
                state: DiscoverySessionState::Active,
                failed_attempts: 0,
            },
        )]));

        // A reopened stream admits session 2 on the same connection while
        // session 1's exchange is still finishing.
        states
            .lock()
            .expect("test session-state mutex is never poisoned")
            .insert(
                (peer.clone(), conn_id),
                DiscoverySessionRecord {
                    session_id: 2,
                    state: DiscoverySessionState::Active,
                    failed_attempts: 0,
                },
            );

        assert!(
            !retire_discovery_session(&states, &peer, conn_id, 1),
            "session 1's finished exchange must not retire session 2's slot"
        );
        assert_eq!(
            states
                .lock()
                .expect("test session-state mutex is never poisoned")
                .get(&(peer.clone(), conn_id))
                .map(|record| record.state),
            Some(DiscoverySessionState::Active),
        );

        assert!(retire_discovery_session(&states, &peer, conn_id, 2));
        assert_eq!(
            states
                .lock()
                .expect("test session-state mutex is never poisoned")
                .get(&(peer.clone(), conn_id))
                .map(|record| record.state),
            Some(DiscoverySessionState::Retired),
        );
    }

    fn record_for(
        states: &StdMutex<SessionStateMap>,
        peer: &ZakuraPeerId,
        conn_id: ZakuraConnId,
    ) -> DiscoverySessionRecord {
        *states
            .lock()
            .expect("test session-state mutex is never poisoned")
            .get(&(peer.clone(), conn_id))
            .expect("test session record exists")
    }

    #[test]
    fn failed_sessions_retire_discovery_after_bounded_attempts() {
        let peer = ZakuraPeerId::new(vec![73; 32]).expect("test peer id is within bounds");
        let conn_id = 9;
        let states = StdMutex::new(HashMap::from([(
            (peer.clone(), conn_id),
            DiscoverySessionRecord {
                session_id: 1,
                state: DiscoverySessionState::Active,
                failed_attempts: 0,
            },
        )]));

        for attempt in 1..MAX_DISCOVERY_SESSION_FAILURES {
            assert!(
                !record_discovery_session_failure(&states, &peer, conn_id, 1),
                "attempt {attempt} must not yet retire the record"
            );
            let record = record_for(&states, &peer, conn_id);
            assert_eq!(record.state, DiscoverySessionState::Active);
            assert_eq!(record.failed_attempts, attempt);
        }

        assert!(
            record_discovery_session_failure(&states, &peer, conn_id, 1),
            "the failure bound must retire the record"
        );
        assert_eq!(
            record_for(&states, &peer, conn_id).state,
            DiscoverySessionState::Retired,
        );
    }

    #[test]
    fn successful_exchange_resets_discovery_failure_count() {
        let peer = ZakuraPeerId::new(vec![74; 32]).expect("test peer id is within bounds");
        let conn_id = 9;
        let states = StdMutex::new(HashMap::from([(
            (peer.clone(), conn_id),
            DiscoverySessionRecord {
                session_id: 1,
                state: DiscoverySessionState::Active,
                failed_attempts: 0,
            },
        )]));

        for _ in 1..MAX_DISCOVERY_SESSION_FAILURES {
            assert!(!record_discovery_session_failure(
                &states, &peer, conn_id, 1
            ));
        }
        record_discovery_session_success(&states, &peer, conn_id, 1);
        assert_eq!(record_for(&states, &peer, conn_id).failed_attempts, 0);

        for _ in 1..MAX_DISCOVERY_SESSION_FAILURES {
            assert!(
                !record_discovery_session_failure(&states, &peer, conn_id, 1),
                "the reset counter must grant a full fresh failure budget"
            );
        }
        assert_eq!(
            record_for(&states, &peer, conn_id).state,
            DiscoverySessionState::Active,
        );
    }

    #[test]
    fn stale_session_failure_cannot_charge_a_newer_session() {
        let peer = ZakuraPeerId::new(vec![75; 32]).expect("test peer id is within bounds");
        let conn_id = 9;
        let states = StdMutex::new(HashMap::from([(
            (peer.clone(), conn_id),
            DiscoverySessionRecord {
                session_id: 2,
                state: DiscoverySessionState::Active,
                failed_attempts: 0,
            },
        )]));

        assert!(
            !record_discovery_session_failure(&states, &peer, conn_id, 1),
            "session 1's failure must not charge session 2's slot"
        );
        assert_eq!(record_for(&states, &peer, conn_id).failed_attempts, 0);
    }

    #[tokio::test]
    async fn add_peer_refuses_a_retired_discovery_record() -> Result<(), crate::BoxError> {
        let (connected_tx, connected_rx) = watch::channel(Vec::new());
        let handshake = ZakuraHandshakeConfig::for_network(&Network::Mainnet);
        let handle = ZakuraDiscoveryHandle::new(
            ZakuraDiscoveryLocalConfig {
                secret_key: SecretKey::from_bytes(&[48u8; 32]),
                direct_addrs: Vec::new(),
                services: vec![ZakuraServiceId::discovery()],
                zakura_protocol_min: handshake.zakura_protocol_min,
                zakura_protocol_max: handshake.zakura_protocol_max,
                network_id: handshake.network_id,
                chain_id: handshake.chain_id,
                last_authored_sequence: None,
            },
            ZakuraDiscoveryConfig::default(),
            connected_rx,
        )?;
        let service = DiscoveryService::new(handle.clone());
        let peer_secret = SecretKey::from_bytes(&[49u8; 32]);
        let peer_id = ZakuraPeerId::new(peer_secret.public().as_bytes().to_vec())?;
        connected_tx.send_replace(vec![peer_id.clone()]);

        service
            .session_states
            .lock()
            .expect("discovery session-state mutex is never poisoned")
            .insert(
                (peer_id.clone(), 0),
                DiscoverySessionRecord {
                    session_id: 7,
                    state: DiscoverySessionState::Retired,
                    failed_attempts: MAX_DISCOVERY_SESSION_FAILURES,
                },
            );

        let connection_cancel = CancellationToken::new();
        let (_peer_send, service_recv) = framed_channel(16);
        let (service_send, _peer_recv) = framed_channel(16);
        let streams = HashMap::from([(ZAKURA_STREAM_DISCOVERY, (service_recv, service_send))]);
        service.add_peer(Peer::new(
            peer_id.clone(),
            None,
            ZAKURA_CAP_DISCOVERY,
            streams,
            connection_cancel.clone(),
        ));

        // Asserted synchronously (current-thread runtime): a refused add_peer
        // must not have replaced the retired record or spawned an exchange.
        let record = record_for(&service.session_states, &peer_id, 0);
        assert_eq!(
            record.session_id, 7,
            "a retired record must not be resurrected by a new stream"
        );
        assert_eq!(record.state, DiscoverySessionState::Retired);
        assert!(matches!(
            service.ordered_session_demand(
                0,
                &peer_id,
                ZAKURA_CAP_DISCOVERY,
                ServicePeerDirection::Inbound,
            ),
            OrderedSessionDemand::Retire,
        ));
        Ok(())
    }

    #[tokio::test]
    async fn add_peer_preserves_failure_count_across_reopens() -> Result<(), crate::BoxError> {
        let (connected_tx, connected_rx) = watch::channel(Vec::new());
        let handshake = ZakuraHandshakeConfig::for_network(&Network::Mainnet);
        let handle = ZakuraDiscoveryHandle::new(
            ZakuraDiscoveryLocalConfig {
                secret_key: SecretKey::from_bytes(&[50u8; 32]),
                direct_addrs: Vec::new(),
                services: vec![ZakuraServiceId::discovery()],
                zakura_protocol_min: handshake.zakura_protocol_min,
                zakura_protocol_max: handshake.zakura_protocol_max,
                network_id: handshake.network_id,
                chain_id: handshake.chain_id,
                last_authored_sequence: None,
            },
            ZakuraDiscoveryConfig::default(),
            connected_rx,
        )?;
        let service = DiscoveryService::new(handle.clone());
        let peer_secret = SecretKey::from_bytes(&[51u8; 32]);
        let peer_id = ZakuraPeerId::new(peer_secret.public().as_bytes().to_vec())?;
        connected_tx.send_replace(vec![peer_id.clone()]);

        service
            .session_states
            .lock()
            .expect("discovery session-state mutex is never poisoned")
            .insert(
                (peer_id.clone(), 0),
                DiscoverySessionRecord {
                    session_id: 7,
                    state: DiscoverySessionState::Active,
                    failed_attempts: MAX_DISCOVERY_SESSION_FAILURES - 1,
                },
            );

        let connection_cancel = CancellationToken::new();
        let (_peer_send, service_recv) = framed_channel(16);
        let (service_send, _peer_recv) = framed_channel(16);
        let streams = HashMap::from([(ZAKURA_STREAM_DISCOVERY, (service_recv, service_send))]);
        service.add_peer(Peer::new(
            peer_id.clone(),
            None,
            ZAKURA_CAP_DISCOVERY,
            streams,
            connection_cancel.clone(),
        ));

        let record = record_for(&service.session_states, &peer_id, 0);
        assert_ne!(
            record.session_id, 7,
            "a live reopen must install the new stream session"
        );
        assert_eq!(record.state, DiscoverySessionState::Active);
        assert_eq!(
            record.failed_attempts,
            MAX_DISCOVERY_SESSION_FAILURES - 1,
            "consecutive failures must stay counted across reopens"
        );
        Ok(())
    }
}