fips-core 0.3.57

use super::*;

/// End-to-end FSP session storage keyed by remote node address.
#[derive(Default)]
pub(in crate::node) struct SessionRegistry {
    sessions: HashMap<NodeAddr, SessionEntry>,
    worker_registrations: DecryptSessionRegistrations,
}

impl SessionRegistry {
    pub(in crate::node) fn insert(
        &mut self,
        node_addr: NodeAddr,
        entry: SessionEntry,
    ) -> Option<SessionEntry> {
        self.sessions.insert(node_addr, entry)
    }

    pub(in crate::node) fn remove(&mut self, node_addr: &NodeAddr) -> Option<SessionEntry> {
        self.sessions.remove(node_addr)
    }

    pub(in crate::node) fn get(&self, node_addr: &NodeAddr) -> Option<&SessionEntry> {
        self.sessions.get(node_addr)
    }

    pub(in crate::node) fn get_mut(&mut self, node_addr: &NodeAddr) -> Option<&mut SessionEntry> {
        self.sessions.get_mut(node_addr)
    }

    #[cfg(test)]
    pub(in crate::node) fn contains_key(&self, node_addr: &NodeAddr) -> bool {
        self.sessions.contains_key(node_addr)
    }

    pub(in crate::node) fn is_empty(&self) -> bool {
        self.sessions.is_empty()
    }

    pub(in crate::node) fn len(&self) -> usize {
        self.sessions.len()
    }

    pub(in crate::node) fn iter(&self) -> impl Iterator<Item = (&NodeAddr, &SessionEntry)> {
        self.sessions.iter()
    }

    pub(in crate::node) fn iter_mut(
        &mut self,
    ) -> impl Iterator<Item = (&NodeAddr, &mut SessionEntry)> {
        self.sessions.iter_mut()
    }

    pub(in crate::node) fn values(&self) -> impl Iterator<Item = &SessionEntry> {
        self.sessions.values()
    }

    pub(in crate::node) fn record_fsp_send_bookkeeping(
        &mut self,
        node_addr: &NodeAddr,
        input: FspSendBookkeepingInput,
    ) -> Option<FspSendBookkeeping> {
        let entry = self.sessions.get_mut(node_addr)?;
        let mut result = FspSendBookkeeping {
            data_recorded: false,
            mmp_recorded: false,
            touched: false,
            next_hop_recorded: false,
        };

        if let Some(next_hop) = input.next_hop {
            entry.record_outbound_next_hop(next_hop);
            result.next_hop_recorded = true;
        }
        if let Some(data_bytes) = input.data_bytes {
            entry.record_sent(data_bytes);
            result.data_recorded = true;
        }
        if let Some(mmp) = entry.mmp_mut() {
            mmp.sender
                .record_sent(input.counter, input.timestamp, input.frame_bytes);
            result.mmp_recorded = true;
        }
        if let Some(touch_ms) = input.touch_ms {
            entry.touch(touch_ms);
            result.touched = true;
        }

        Some(result)
    }

    pub(in crate::node) fn record_fsp_send_bookkeeping_batch<I>(
        &mut self,
        node_addr: &NodeAddr,
        inputs: I,
    ) -> Option<usize>
    where
        I: IntoIterator<Item = FspSendBookkeepingInput>,
    {
        let entry = self.sessions.get_mut(node_addr)?;
        let mut data_packets = 0usize;
        let mut data_bytes = 0usize;
        let mut last_touch_ms = None;
        let mut last_next_hop = None;

        for input in inputs {
            if let Some(next_hop) = input.next_hop {
                last_next_hop = Some(next_hop);
            }
            if let Some(bytes) = input.data_bytes {
                data_packets += 1;
                data_bytes += bytes;
            }
            if let Some(mmp) = entry.mmp_mut() {
                mmp.sender
                    .record_sent(input.counter, input.timestamp, input.frame_bytes);
            }
            if input.touch_ms.is_some() {
                last_touch_ms = input.touch_ms;
            }
        }

        if let Some(next_hop) = last_next_hop {
            entry.record_outbound_next_hop(next_hop);
        }
        if data_packets > 0 {
            entry.record_sent_batch(data_packets, data_bytes);
        }
        if let Some(touch_ms) = last_touch_ms {
            entry.touch(touch_ms);
        }

        Some(data_packets)
    }

    pub(in crate::node) fn direct_endpoint_data_can_send(&self, node_addr: &NodeAddr) -> bool {
        self.sessions
            .get(node_addr)
            .is_some_and(|entry| entry.is_established())
    }

    pub(in crate::node) fn record_direct_endpoint_data_send(
        &mut self,
        node_addr: &NodeAddr,
        payload_len: usize,
        now_ms: u64,
        next_hop: NodeAddr,
    ) -> Option<()> {
        let entry = self.sessions.get_mut(node_addr)?;
        if !entry.is_established() {
            return None;
        }
        entry.record_outbound_next_hop(next_hop);
        entry.record_sent(payload_len);
        entry.touch(now_ms);
        Some(())
    }

    pub(in crate::node) fn record_direct_endpoint_data_send_batch<I>(
        &mut self,
        node_addr: &NodeAddr,
        inputs: I,
    ) -> Option<usize>
    where
        I: IntoIterator<Item = (usize, u64, NodeAddr)>,
    {
        let entry = self.sessions.get_mut(node_addr)?;
        if !entry.is_established() {
            return None;
        }

        let mut packets = 0usize;
        let mut bytes = 0usize;
        let mut last_touch_ms = None;
        let mut last_next_hop = None;
        for (payload_len, now_ms, next_hop) in inputs {
            packets = packets.saturating_add(1);
            bytes = bytes.saturating_add(payload_len);
            last_touch_ms = Some(now_ms);
            last_next_hop = Some(next_hop);
        }

        if packets == 0 {
            return Some(0);
        }
        if let Some(next_hop) = last_next_hop {
            entry.record_outbound_next_hop(next_hop);
        }
        entry.record_sent_batch(packets, bytes);
        if let Some(now_ms) = last_touch_ms {
            entry.touch(now_ms);
        }
        Some(packets)
    }

    pub(in crate::node) fn record_direct_endpoint_data_receive(
        &mut self,
        node_addr: &NodeAddr,
        payload_len: usize,
        now_ms: u64,
    ) -> bool {
        let Some(entry) = self.sessions.get_mut(node_addr) else {
            return false;
        };
        if !entry.is_established() {
            return false;
        }
        entry.record_recv(payload_len);
        entry.touch(now_ms);
        // Direct-FMP endpoint data is accepted only when the authenticated
        // FMP peer is the session peer. Treat it as session liveness so an
        // active direct fast path does not look outbound-only stale.
        entry.touch_inbound_frame(now_ms);
        true
    }

    pub(in crate::node) fn record_direct_endpoint_data_receive_batch(
        &mut self,
        node_addr: &NodeAddr,
        packets: usize,
        bytes: usize,
        now_ms: u64,
    ) -> bool {
        let Some(entry) = self.sessions.get_mut(node_addr) else {
            return false;
        };
        if !entry.is_established() {
            return false;
        }
        if packets > 0 {
            entry.record_recv_batch(packets, bytes);
            entry.touch(now_ms);
            entry.touch_inbound_frame(now_ms);
        }
        true
    }

    #[cfg(unix)]
    pub(in crate::node) fn reserve_endpoint_data_fsp_worker_send(
        &mut self,
        node_addr: &NodeAddr,
        input: FspWorkerSendReservationInput,
    ) -> Result<Option<FspSendReservation>, FspWorkerSendReservationError> {
        let entry = self
            .sessions
            .get_mut(node_addr)
            .ok_or(FspWorkerSendReservationError::MissingSession)?;
        if let Some(mmp) = entry.mmp_mut() {
            mmp.path_mtu.seed_source_mtu(input.path_mtu);
        }
        if !entry.is_established() {
            return Err(FspWorkerSendReservationError::NotEstablished);
        }
        entry
            .reserve_fsp_worker_send(input.flags, input.payload_len)
            .map_err(|_| FspWorkerSendReservationError::CounterReservationFailed)
    }

    pub(in crate::node) fn record_worker_registration(
        &mut self,
        session_key: DecryptSessionKey,
        accepted: bool,
    ) -> bool {
        self.worker_registrations
            .record_worker_registration(session_key, accepted)
    }

    pub(in crate::node) fn unregister_worker_session_if_registered(
        &mut self,
        session_key: &DecryptSessionKey,
    ) -> bool {
        self.worker_registrations
            .unregister_if_registered(session_key)
    }

    pub(in crate::node) fn is_worker_registered(&self, session_key: &DecryptSessionKey) -> bool {
        self.worker_registrations.is_registered(session_key)
    }

    #[cfg(test)]
    pub(in crate::node) fn worker_registration_is_empty(&self) -> bool {
        self.worker_registrations.is_empty()
    }
}

impl<'a> IntoIterator for &'a SessionRegistry {
    type Item = (&'a NodeAddr, &'a SessionEntry);
    type IntoIter = std::collections::hash_map::Iter<'a, NodeAddr, SessionEntry>;

    fn into_iter(self) -> Self::IntoIter {
        self.sessions.iter()
    }
}

/// Rx-loop mirror of sessions accepted by decrypt-worker shards.
#[derive(Debug, Default)]
pub(in crate::node) struct DecryptSessionRegistrations {
    sessions: HashSet<DecryptSessionKey>,
}

impl DecryptSessionRegistrations {
    pub(in crate::node) fn record_worker_registration(
        &mut self,
        session_key: DecryptSessionKey,
        accepted: bool,
    ) -> bool {
        if !accepted {
            return false;
        }
        self.sessions.insert(session_key);
        true
    }

    pub(in crate::node) fn unregister_if_registered(
        &mut self,
        session_key: &DecryptSessionKey,
    ) -> bool {
        self.sessions.remove(session_key)
    }

    pub(in crate::node) fn is_registered(&self, session_key: &DecryptSessionKey) -> bool {
        self.sessions.contains(session_key)
    }

    #[cfg(test)]
    pub(in crate::node) fn is_empty(&self) -> bool {
        self.sessions.is_empty()
    }
}

/// Send-scheduling policy derived from the configured peer roster.
#[derive(Debug, Default)]
pub(in crate::node) struct ConfiguredPeerSendWeights {
    entries: HashMap<NodeAddr, u8>,
}

impl ConfiguredPeerSendWeights {
    pub(in crate::node) fn from_config(config: &Config) -> Self {
        let entries = config
            .peers()
            .iter()
            .filter_map(|peer| {
                PeerIdentity::from_npub(&peer.npub).ok().map(|identity| {
                    (
                        *identity.node_addr(),
                        encrypt_worker::EXPLICIT_PEER_SEND_WEIGHT,
                    )
                })
            })
            .collect();
        Self { entries }
    }

    pub(in crate::node) fn weight_for(&self, peer_addr: &NodeAddr) -> u8 {
        self.entries
            .get(peer_addr)
            .copied()
            .unwrap_or(encrypt_worker::DEFAULT_SEND_WEIGHT)
    }

    #[cfg(any(target_os = "linux", target_os = "macos"))]
    pub(in crate::node) fn contains(&self, peer_addr: &NodeAddr) -> bool {
        self.entries.contains_key(peer_addr)
    }

    #[cfg(test)]
    pub(in crate::node) fn len(&self) -> usize {
        self.entries.len()
    }
}

/// Pending outbound FMP handshakes keyed by `(transport_id, our_index)`.
#[derive(Debug, Default)]
pub(in crate::node) struct PendingOutboundHandshakes {
    entries: HashMap<(TransportId, u32), LinkId>,
}

impl PendingOutboundHandshakes {
    pub(in crate::node) fn insert(
        &mut self,
        key: (TransportId, u32),
        link_id: LinkId,
    ) -> Option<LinkId> {
        self.entries.insert(key, link_id)
    }

    pub(in crate::node) fn remove(&mut self, key: &(TransportId, u32)) -> Option<LinkId> {
        self.entries.remove(key)
    }

    #[cfg(test)]
    pub(in crate::node) fn get(&self, key: &(TransportId, u32)) -> Option<&LinkId> {
        self.entries.get(key)
    }

    #[cfg(test)]
    pub(in crate::node) fn contains_key(&self, key: &(TransportId, u32)) -> bool {
        self.entries.contains_key(key)
    }

    #[cfg(test)]
    pub(in crate::node) fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    #[cfg(test)]
    pub(in crate::node) fn retain<F>(&mut self, f: F)
    where
        F: FnMut(&(TransportId, u32), &mut LinkId) -> bool,
    {
        self.entries.retain(f);
    }

    pub(in crate::node) fn match_msg2(
        &self,
        transport_id: TransportId,
        receiver_idx: u32,
    ) -> Option<((TransportId, u32), LinkId)> {
        let exact_key = (transport_id, receiver_idx);
        if let Some(link_id) = self.entries.get(&exact_key).copied() {
            return Some((exact_key, link_id));
        }

        let mut matches = self
            .entries
            .iter()
            .filter(|((_, idx), _)| *idx == receiver_idx);
        match (matches.next(), matches.next()) {
            (Some((fallback_key, link_id)), None) => Some((*fallback_key, *link_id)),
            _ => None,
        }
    }
}