fips-core 0.3.55

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)
    }

    #[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,
        }
    }
}