geiserx_ts_runtime 0.58.0

use std::{collections::HashMap, net::IpAddr, sync::Arc};

use kameo::{
    actor::ActorRef,
    message::{Context, Message},
};
use tokio::sync::mpsc;
use ts_dataplane::{AdvertisementTarget, DiscoAdvertisementState};
use ts_packet::PacketMut;
use ts_transport::{OverlayTransportId, PeerId, UnderlayTransportId};

use crate::{
    Error,
    env::Env,
    packetfilter::PacketFilterState,
    peer_tracker::PeerState,
    route_updater::{PeerRouteUpdate, SelfRouteUpdate},
    src_filter::SourceFilterState,
};

/// How long ago the application datapath last sent a packet towards the tailnet — the fork's
/// `tstun.Wrapper.IdleDuration`, the activity signal Go `wgengine/userspace.go` wires into
/// magicsock as `Conn.idleFunc`.
///
/// Go measures idleness at its single TUN device, which every path — the OS stack in TUN mode, the
/// userspace netstack in netstack mode, and forwarded exit/subnet flows — passes through. This tree
/// has no single device, but it does have a single *seam*: every overlay transport writes to the
/// dataplane through [`OverlayToDataplane::send`], and that is where the timestamp is taken. So the
/// application netstack, the TUN pump and the forwarder netstack all feed one clock, exactly as they
/// would all feed one `tstun.Wrapper` upstream.
///
/// Only the **send** direction is recorded, which is what the consumer needs: the periodic STUN
/// sweep exists to keep this node's NAT mapping open so peers can reach *us*, and it is our own
/// outbound traffic that both proves someone here still wants that mapping and refreshes it.
/// Inbound-only traffic keeps arriving over an already-working path.
///
/// The clock is monotonic and lock-free: an origin [`Instant`](std::time::Instant) fixed at
/// construction plus an [`AtomicU64`](std::sync::atomic::AtomicU64) of milliseconds since it. Before
/// the first send the counter is `0`, so [`idle_duration`](Self::idle_duration) reads as "idle since
/// the runtime started" — which is also how Go reads its zero `mono.Time`.
#[derive(Debug, Clone)]
pub struct DatapathActivity(Arc<DatapathActivityInner>);

#[derive(Debug)]
struct DatapathActivityInner {
    /// Fixed at construction; every recorded send is stored as a millisecond offset from it, so the
    /// whole clock fits in one atomic and can never go backwards.
    origin: std::time::Instant,
    /// Milliseconds from `origin` to the last datapath send, or `0` when there has been none.
    last_send_millis: std::sync::atomic::AtomicU64,
}

impl DatapathActivity {
    /// A clock with no send recorded yet, whose origin is now.
    pub fn new() -> Self {
        Self(Arc::new(DatapathActivityInner {
            origin: std::time::Instant::now(),
            last_send_millis: std::sync::atomic::AtomicU64::new(0),
        }))
    }

    /// Record that the datapath just sent towards the tailnet (Go `tstun.Wrapper.noteActivity`).
    ///
    /// `Relaxed` is the right ordering: the value is a coarse timestamp read by an unrelated task on
    /// a ~23s cadence, it is never used to publish other state, and a monotonic store is still
    /// monotonic without a fence.
    fn note_send(&self) {
        let millis = u64::try_from(self.0.origin.elapsed().as_millis()).unwrap_or(u64::MAX);
        self.0
            .last_send_millis
            .store(millis, std::sync::atomic::Ordering::Relaxed);
    }

    /// How long since the last datapath send, or since this clock's origin when there has been none.
    pub fn idle_duration(&self) -> std::time::Duration {
        let last = std::time::Duration::from_millis(
            self.0
                .last_send_millis
                .load(std::sync::atomic::Ordering::Relaxed),
        );
        // `saturating_sub`: a send recorded concurrently with this read can carry an offset a hair
        // past the `elapsed()` sampled just above, and "slightly negative idle" must read as zero,
        // not wrap.
        self.0.origin.elapsed().saturating_sub(last)
    }

    /// A clock that already reads as idle for `idle` — the backdated-origin equivalent of letting
    /// `idle` of wall-clock elapse, so the idle stop condition can be exercised without a sleep.
    /// [`note_send`](Self::note_send) still resets it, which is what makes the "traffic resumes"
    /// assertion go through the production path.
    #[cfg(test)]
    pub(crate) fn idle_for(idle: std::time::Duration) -> Self {
        let now = std::time::Instant::now();
        Self(Arc::new(DatapathActivityInner {
            origin: now.checked_sub(idle).unwrap_or(now),
            last_send_millis: std::sync::atomic::AtomicU64::new(0),
        }))
    }
}

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

/// Queue for packets sent from the overlay to the dataplane.
///
/// A newtype around the raw channel rather than an alias for it, because this is the one seam every
/// overlay transport's outbound packets cross: wrapping it makes recording datapath activity
/// (see [`DatapathActivity`]) structural — a future overlay transport gets it by construction rather
/// than by remembering to call something.
#[derive(Clone)]
pub struct OverlayToDataplane {
    tx: mpsc::UnboundedSender<Vec<PacketMut>>,
    activity: DatapathActivity,
}

impl OverlayToDataplane {
    /// Hand a batch of packets to the dataplane, noting the datapath activity.
    ///
    /// The timestamp is taken only when the send succeeds: a closed channel means the dataplane is
    /// gone (shutdown), and that is not the datapath carrying traffic. `Err` returns the batch, like
    /// the underlying channel's `send`.
    pub fn send(
        &self,
        packets: Vec<PacketMut>,
    ) -> Result<(), mpsc::error::SendError<Vec<PacketMut>>> {
        let sent = self.tx.send(packets);
        if sent.is_ok() {
            self.activity.note_send();
        }
        sent
    }

    /// A send half feeding `activity`, paired with its receiver, so a test can drive the real
    /// [`send`](Self::send) path (and thus the real activity recording) without standing up the
    /// dataplane actor.
    #[cfg(test)]
    pub(crate) fn for_test(activity: DatapathActivity) -> (Self, OverlayFromDataplane) {
        let (tx, rx) = mpsc::unbounded_channel();
        (Self { tx, activity }, rx)
    }
}

/// Queue for packets entering the overlay from the dataplane.
pub type OverlayFromDataplane = mpsc::UnboundedReceiver<Vec<PacketMut>>;

/// Which peers have had WireGuard traffic lately — the input to Go magicsock's per-endpoint
/// `lastSendExt`, which `endpoint.send` stamps for every packet WireGuard hands it.
///
/// Recorded where every underlay crosses the dataplane: [`UnderlayFromDataplane::recv`] for
/// packets leaving towards a peer and [`UnderlayToDataplane::send`] for packets arriving from one,
/// over DERP or a direct path alike. Upstream only stamps sends, but wireguard-go answers received
/// data with data or a keepalive of its own, so in Go received data keeps a session active too.
/// This tree's tunnel does not always answer, so received data is counted directly.
///
/// Only transport data with a payload counts. Go stamps `lastSendExt` for every packet
/// wireguard-go sends, handshakes included (wgengine/magicsock/endpoint.go:1072, v1.102.4), but
/// upstream (v1.102.4) sets no WireGuard persistent keepalive on any peer: wgcfg.Config carries no
/// peers (wgengine/wgcfg/config.go:16-25) and SetPeerConfigFunc supplies only allowed IPs
/// (wgengine/userspace.go:689), so an idle Go session exchanges nothing and never rekeys. This
/// fork sends a persistent keepalive to every peer with a session
/// ([`Env::persistent_keepalive_interval`]), and the initiator rekeys the keypair on the first
/// packet received after it is 165 s old (ts_tunnel `keep_key_fresh_receiving`), so counting
/// keepalives, or the handshakes they cause, would restart discovery for every idle session every
/// few minutes. The cost is that discovery for a fresh session starts on its first data packet,
/// one handshake round trip over DERP later than in Go.
///
/// The direct manager takes the peers with new traffic ([`next`](Self::next)) and marks them
/// active on the magicsock, which starts discovery for a peer that was idle. A peer is reported at
/// most once per [`PEER_ACTIVITY_REPORT_INTERVAL`], which bounds the hot-path cost to one lock and
/// one map lookup per batch; the activity window it feeds is 45s.
#[derive(Debug, Clone, Default)]
pub struct PeerActivity(Arc<PeerActivityInner>);

#[derive(Debug, Default)]
struct PeerActivityInner {
    state: std::sync::Mutex<PeerActivityState>,
    /// Wakes [`PeerActivity::next`]. `notify_one` stores a permit when nobody is waiting, so a
    /// report made between the consumer's check and its wait is not lost.
    notify: tokio::sync::Notify,
}

#[derive(Debug, Default)]
struct PeerActivityState {
    /// When each peer's traffic was last reported.
    reported: HashMap<PeerId, tokio::time::Instant>,
    /// Peers reported but not yet taken by [`PeerActivity::next`].
    pending: Vec<PeerId>,
}

/// The least time between two reports of the same peer's traffic.
const PEER_ACTIVITY_REPORT_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1);

impl PeerActivity {
    /// Record WireGuard traffic with `peer`. Ignored unless `packets` holds transport data with a
    /// payload: keepalives and handshakes are not traffic.
    pub fn note(&self, peer: PeerId, packets: &[PacketMut]) {
        if !packets.iter().any(|p| is_wireguard_data(p.as_ref())) {
            return;
        }
        let now = tokio::time::Instant::now();
        {
            let mut state = self.lock();
            if state
                .reported
                .get(&peer)
                .is_some_and(|t| now.saturating_duration_since(*t) < PEER_ACTIVITY_REPORT_INTERVAL)
            {
                return;
            }
            state.reported.insert(peer, now);
            if !state.pending.contains(&peer) {
                state.pending.push(peer);
            }
        }
        self.0.notify.notify_one();
    }

    /// Wait for traffic, then take every peer reported since the last call.
    pub async fn next(&self) -> Vec<PeerId> {
        loop {
            {
                let mut state = self.lock();
                if !state.pending.is_empty() {
                    return core::mem::take(&mut state.pending);
                }
            }
            self.0.notify.notified().await;
        }
    }

    /// Forget peers for which `live` is false, so a departed peer's entry does not linger.
    pub fn retain(&self, live: impl Fn(PeerId) -> bool) {
        let mut state = self.lock();
        state.reported.retain(|peer, _| live(*peer));
        state.pending.retain(|peer| live(*peer));
    }

    fn lock(&self) -> std::sync::MutexGuard<'_, PeerActivityState> {
        // The state is two independent collections with no invariant a panicking writer could
        // break, so a poisoned lock is safe to keep using.
        self.0
            .state
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
    }
}

/// Whether `packet` is WireGuard transport data with a payload: a transport-data message (type 4)
/// longer than a keepalive, which is its 16-byte header plus the 16-byte AEAD tag around an empty
/// payload (WireGuard §5.4.6 and §6.5). Handshake messages (types 1-3) are not data.
fn is_wireguard_data(packet: &[u8]) -> bool {
    packet.len() > 32 && packet[0] == 4
}

/// Queue for packets leaving the underlay to the dataplane.
///
/// A newtype around the raw channel, like [`OverlayToDataplane`], so recording per-peer traffic
/// ([`PeerActivity`]) is structural: every underlay transport records it by construction.
#[derive(Debug, Clone)]
pub struct UnderlayToDataplane {
    tx: mpsc::UnboundedSender<(PeerId, Vec<PacketMut>)>,
    activity: PeerActivity,
}

impl UnderlayToDataplane {
    /// Hand packets received from `peer` to the dataplane, noting the peer's traffic. As with
    /// [`OverlayToDataplane::send`], nothing is noted when the dataplane has gone away.
    pub fn send(
        &self,
        (peer, packets): (PeerId, Vec<PacketMut>),
    ) -> Result<(), mpsc::error::SendError<(PeerId, Vec<PacketMut>)>> {
        if self.tx.is_closed() {
            return Err(mpsc::error::SendError((peer, packets)));
        }
        self.activity.note(peer, &packets);
        self.tx.send((peer, packets))
    }

    /// Whether the dataplane end has gone away.
    pub fn is_closed(&self) -> bool {
        self.tx.is_closed()
    }

    /// A send half recording into `activity`, paired with its receiver, for tests.
    #[cfg(test)]
    pub(crate) fn for_test(
        activity: PeerActivity,
    ) -> (Self, mpsc::UnboundedReceiver<(PeerId, Vec<PacketMut>)>) {
        let (tx, rx) = mpsc::unbounded_channel();
        (Self { tx, activity }, rx)
    }
}

/// Queue for packets entering an underlay from the dataplane, noting per-peer traffic as each
/// batch is taken (see [`PeerActivity`]).
#[derive(Debug)]
pub struct UnderlayFromDataplane {
    rx: mpsc::UnboundedReceiver<(PeerId, Vec<PacketMut>)>,
    activity: PeerActivity,
}

impl UnderlayFromDataplane {
    /// The next batch of packets to send to a peer, or `None` once the dataplane has gone away.
    /// Cancel-safe, like the channel it wraps.
    pub async fn recv(&mut self) -> Option<(PeerId, Vec<PacketMut>)> {
        let (peer, packets) = self.rx.recv().await?;
        self.activity.note(peer, &packets);
        Some((peer, packets))
    }

    /// Whether the dataplane end has gone away.
    pub fn is_closed(&self) -> bool {
        self.rx.is_closed()
    }

    /// A receiver recording into `activity`, paired with its send half, for tests.
    #[cfg(test)]
    pub(crate) fn for_test(
        activity: PeerActivity,
    ) -> (Self, mpsc::UnboundedSender<(PeerId, Vec<PacketMut>)>) {
        let (tx, rx) = mpsc::unbounded_channel();
        (Self { rx, activity }, tx)
    }
}

/// A peer's disco public key, learned from a TSMP disco-key advertisement it sent us inside the
/// WireGuard tunnel (Go `packet.TSMPDiscoKeyAdvertisement`, upstream capability version 144).
///
/// Published on the bus by [`DataplaneActor`] and consumed by the peer tracker, which applies it to
/// the peer. This is the Rust shape of Go's `events.PeerDiscoKeyUpdate`, which `wgengine` turns
/// into a `magicsock.Conn.HandleDiscoKeyAdvertisement` call.
///
/// `peer` is the WireGuard peer whose session decrypted the advertisement — the dataplane's source
/// filter has already bound that peer's source addresses — so a peer can only ever advertise a key
/// for itself. Go arrives at the same peer by looking the advertisement's source IP up in the
/// netmap (`wgengine.userspaceEngine.peerForIP`).
#[derive(Debug, Clone, Copy)]
pub struct PeerDiscoKeyAdvertisement {
    /// The peer that advertised the key.
    pub peer: PeerId,
    /// The disco public key it advertised. Never the zero key: the dataplane drops those, as Go
    /// does, before they get this far.
    pub key: ts_keys::DiscoPublicKey,
}

pub struct DataplaneActor {
    dataplane: Arc<ts_dataplane::async_tokio::DataPlane>,
    task: tokio::task::JoinHandle<()>,
    /// Forwards TSMP-learned peer disco keys from the dataplane onto the bus.
    disco_key_task: tokio::task::JoinHandle<()>,
    /// Persistent-keepalive interval applied to every upserted peer (or `None` to disable). Snapshot
    /// of [`Env::persistent_keepalive_interval`] taken at actor start. See the peer-upsert handler.
    persistent_keepalive_interval: Option<std::time::Duration>,
    /// Working copy of the netmap state the dataplane needs to send a TSMP disco-key advertisement
    /// when a WireGuard session comes up (Go capability version 144). Two independent sources feed
    /// it — the self node (our own tailnet addresses) and the peer set (each peer's destination
    /// address) — so it is assembled here and pushed into the dataplane whole on every change.
    disco_advertisement: DiscoAdvertisementState,
    /// The shared datapath-idleness clock (the same handle as `Env::datapath_activity`), stamped
    /// into every [`OverlayToDataplane`] this actor hands out so each overlay transport's sends feed
    /// it. Read by the direct manager's periodic STUN gate.
    datapath_activity: DatapathActivity,
    /// The shared per-peer traffic record (the same handle as `Env::peer_activity`), stamped into
    /// every underlay channel this actor hands out. Read by the direct manager.
    peer_activity: PeerActivity,
}

impl Drop for DataplaneActor {
    fn drop(&mut self) {
        self.task.abort();
        self.disco_key_task.abort();
    }
}

#[kameo::messages]
impl DataplaneActor {
    /// Stop packet processing independently of this actor's mailbox or core-state lock.
    #[message]
    pub async fn packet_pump_abort_handle(&self) -> tokio::task::AbortHandle {
        self.task.abort_handle()
    }

    #[message]
    pub async fn new_overlay_transport(
        &self,
    ) -> (OverlayTransportId, OverlayToDataplane, OverlayFromDataplane) {
        let (id, tx, rx) = self.dataplane.new_overlay_transport().await;
        (
            id,
            OverlayToDataplane {
                tx,
                activity: self.datapath_activity.clone(),
            },
            rx,
        )
    }

    #[message]
    pub async fn new_underlay_transport(
        &self,
    ) -> (
        UnderlayTransportId,
        UnderlayFromDataplane,
        UnderlayToDataplane,
    ) {
        let (id, rx, tx) = self.dataplane.new_underlay_transport().await;
        (
            id,
            UnderlayFromDataplane {
                rx,
                activity: self.peer_activity.clone(),
            },
            UnderlayToDataplane {
                tx,
                activity: self.peer_activity.clone(),
            },
        )
    }

    /// Install (`Some`) or clear (`None`) the debug packet-capture hook on the running dataplane.
    /// `Some(hook)` begins teeing every plaintext packet crossing the datapath to `hook`; `None`
    /// stops capture. Mirrors Go `tstun.Wrapper.InstallCaptureHook` / `ClearCaptureSink`.
    #[message]
    pub async fn install_capture(&self, hook: Option<ts_dataplane::CaptureHook>) {
        let dp = &mut *self.dataplane.inner().await;
        dp.capture = hook;
    }
}

impl kameo::Actor for DataplaneActor {
    type Args = Env;
    type Error = Error;

    async fn on_start(env: Self::Args, slf: ActorRef<Self>) -> Result<Self, Self::Error> {
        let dataplane = Arc::new(ts_dataplane::async_tokio::DataPlane::new(
            // `.clone()`: `node_keys` is no longer `Copy` and `env` is a shared `Arc`.
            env.keys.node_keys.clone(),
        ));

        let persistent_keepalive_interval = env.persistent_keepalive_interval;
        let datapath_activity = env.datapath_activity.clone();
        let peer_activity = env.peer_activity.clone();

        // Go `tstun.Wrapper.PeerAPIPort`. The dataplane reads it for one thing: an inbound SYN to
        // the peerAPI port that the ACL drops must not be answered with a TSMP rejected-connection
        // message, because upstream's peerAPI carve-out has already admitted it by that point. It
        // is fixed for the life of the node (`Config::peerapi_port`), so it is set once here rather
        // than refreshed from the netmap.
        dataplane.inner().await.peerapi_port = env.peerapi_port;

        // Peers advertise their disco key over TSMP immediately after an eligible WireGuard
        // session is established, so the sink must be installed before the dataplane starts
        // running or the first advertisement is lost.
        let (disco_key_tx, mut disco_key_rx) = mpsc::unbounded_channel();
        dataplane.install_disco_key_sink(Some(disco_key_tx)).await;

        let disco_key_env = env.clone();
        let disco_key_task = tokio::task::spawn(async move {
            while let Some((peer, advert)) = disco_key_rx.recv().await {
                let key = ts_keys::DiscoPublicKey::from(advert.key);

                tracing::debug!(?peer, %key, "publishing TSMP-advertised peer disco key");

                if let Err(e) = disco_key_env
                    .publish(PeerDiscoKeyAdvertisement { peer, key })
                    .await
                {
                    tracing::error!(error = %e, "publishing TSMP disco-key advertisement");
                }
            }
        });

        env.subscribe::<PeerRouteUpdate>(&slf).await?;
        env.subscribe::<SelfRouteUpdate>(&slf).await?;
        env.subscribe::<PacketFilterState>(&slf).await?;
        env.subscribe::<SourceFilterState>(&slf).await?;
        env.subscribe::<Arc<PeerState>>(&slf).await?;
        // The self node, for the source address of our own TSMP disco-key advertisements. Peers
        // arrive via `PeerState`; this is the other half of `disco_advertisement`.
        env.subscribe::<Arc<ts_control::StateUpdate>>(&slf).await?;

        let task_dataplane = dataplane.clone();

        let task = tokio::task::spawn(async move {
            task_dataplane.run().await;
        });

        tracing::trace!("dataplane running");

        Ok(Self {
            dataplane,
            task,
            disco_key_task,
            persistent_keepalive_interval,
            datapath_activity,
            peer_activity,
            // Our disco key is generated once per run and never rotates (unlike the node key), so
            // it is snapshotted here rather than re-read per advertisement. Go reads it per call
            // (`Conn.DiscoPublicKey()`) because its disco key *can* be regenerated on rebind.
            disco_advertisement: DiscoAdvertisementState {
                disco_key: env.keys.disco_keys.public.to_bytes(),
                ..Default::default()
            },
        })
    }
}

impl Message<PeerRouteUpdate> for DataplaneActor {
    type Reply = ();

    async fn handle(&mut self, msg: PeerRouteUpdate, _ctx: &mut Context<Self, Self::Reply>) {
        tracing::trace!("applying peer route update");

        let dp = &mut *self.dataplane.inner().await;
        dp.or_out.swap(msg.inner.overlay_out_routes.clone());

        dp.ur_out.table = msg.inner.underlay_routes.clone();
    }
}

impl Message<SelfRouteUpdate> for DataplaneActor {
    type Reply = ();

    async fn handle(&mut self, msg: SelfRouteUpdate, _ctx: &mut Context<Self, Self::Reply>) {
        {
            let dp = &mut *self.dataplane.inner().await;
            dp.or_in.swap(msg.overlay_in_routes.as_ref().clone());
        }

        tracing::trace!("applied self route update");
    }
}

impl Message<PacketFilterState> for DataplaneActor {
    type Reply = ();

    async fn handle(&mut self, msg: PacketFilterState, _ctx: &mut Context<Self, Self::Reply>) {
        {
            let dp = &mut *self.dataplane.inner().await;
            dp.packet_filter = msg.0;
        }

        tracing::trace!("applied new packet filter");
    }
}

impl Message<SourceFilterState> for DataplaneActor {
    type Reply = ();

    async fn handle(&mut self, msg: SourceFilterState, _ctx: &mut Context<Self, Self::Reply>) {
        {
            let dp = &mut *self.dataplane.inner().await;
            dp.src_filter_in = msg.0;
        }

        tracing::trace!("applied new source filter");
    }
}

/// This node's own tailnet addresses, in the order Go's `self.Addresses()` presents them and
/// narrowed to the single-IP prefixes `magicsock.selfIPMatchingFamily` accepts.
///
/// Control lists a node's IPv4 address first, so that is the order kept here. The unspecified
/// placeholder [`ts_control::TailnetAddress`] synthesizes for a family the tailnet does not assign
/// (`0.0.0.0/32` / `::/128`) is not an address and is skipped — advertising it as our source would
/// name an address no peer can answer.
fn advertisement_self_addrs(node: &ts_control::Node) -> Vec<IpAddr> {
    let v4 = node.tailnet_address.ipv4.addr();
    let v6 = node.tailnet_address.ipv6.addr();

    let mut addrs = Vec::with_capacity(2);
    if !v4.is_unspecified() {
        addrs.push(IpAddr::from(v4));
    }
    if !v6.is_unspecified() {
        addrs.push(IpAddr::from(v6));
    }
    addrs
}

/// Where to send `node` a TSMP disco-key advertisement, or `None` if it has no usable tailnet
/// address at all.
///
/// The destination is Go's `endpoint.nodeAddr`, "the node's first tailscale address" — IPv4 on a
/// dual-stack tailnet, falling back to IPv6 for a tailnet that assigns no IPv4. `wireguard_only`
/// carries Go's `endpoint.isWireguardOnly` through to the refusal in
/// [`DiscoAdvertisementState::advertisement_for`].
fn advertisement_target(node: &ts_control::Node) -> Option<AdvertisementTarget> {
    let v4 = node.tailnet_address.ipv4.addr();
    let v6 = node.tailnet_address.ipv6.addr();

    let node_addr = if !v4.is_unspecified() {
        IpAddr::from(v4)
    } else if !v6.is_unspecified() {
        IpAddr::from(v6)
    } else {
        return None;
    };

    Some(AdvertisementTarget {
        node_addr,
        wireguard_only: node.is_wireguard_only,
    })
}

/// Rebuild the per-peer advertisement destinations from a peer snapshot.
///
/// Rebuilt wholesale rather than patched from the upsert/deletion sets: the snapshot is the
/// authoritative peer set, so a peer that left the netmap stops being advertised to without any
/// separate bookkeeping.
fn advertisement_targets(
    peers: &crate::peer_tracker::PeerDb,
) -> HashMap<PeerId, AdvertisementTarget> {
    peers
        .peers()
        .iter()
        .filter_map(|(&id, node)| Some((id, advertisement_target(node)?)))
        .collect()
}

impl Message<Arc<ts_control::StateUpdate>> for DataplaneActor {
    type Reply = ();

    async fn handle(
        &mut self,
        msg: Arc<ts_control::StateUpdate>,
        _ctx: &mut Context<Self, Self::Reply>,
    ) {
        // Only a response that carried a self node says anything about our own addresses; a
        // keep-alive or a peers-only update leaves the previous snapshot in place.
        let Some(node) = msg.node.as_ref() else {
            return;
        };

        let self_addrs = advertisement_self_addrs(node);
        if self_addrs == self.disco_advertisement.self_addrs {
            return;
        }
        self.disco_advertisement.self_addrs = self_addrs;

        let mut dp = self.dataplane.inner().await;
        dp.disco_advertisement = Some(Arc::new(self.disco_advertisement.clone()));

        tracing::trace!("applied self addresses for TSMP disco-key advertisement");
    }
}

impl Message<Arc<PeerState>> for DataplaneActor {
    type Reply = ();

    async fn handle(&mut self, msg: Arc<PeerState>, _ctx: &mut Context<Self, Self::Reply>) {
        self.disco_advertisement.peers = advertisement_targets(&msg.peers);

        {
            let mut dp = self.dataplane.inner().await;
            dp.disco_advertisement = Some(Arc::new(self.disco_advertisement.clone()));
            let wg = &mut dp.wireguard;

            for &upsert in &msg.upserts {
                let Some((_, node)) = msg.peers.get(&upsert) else {
                    tracing::error!(
                        ?upsert,
                        "dataplane: upsert id missing from peer snapshot; skipping"
                    );
                    continue;
                };

                apply_tunnel_peer(wg, upsert, node, self.persistent_keepalive_interval);
            }

            for delete in &msg.deletions {
                wg.remove_peer(ts_tunnel::PeerId(delete.0));
            }
        }

        tracing::trace!("applied new peer state");
    }
}

/// Keep expired peers in the control snapshot for status, but remove all tunnel authority.
fn apply_tunnel_peer(
    wg: &mut ts_tunnel::Endpoint,
    id: PeerId,
    node: &ts_control::Node,
    persistent_keepalive_interval: Option<std::time::Duration>,
) {
    let id = ts_tunnel::PeerId(id.0);
    if node.expired {
        wg.remove_peer(id);
        return;
    }
    wg.upsert_peer(
        id,
        ts_tunnel::PeerConfig {
            key: node.node_key,
            psk: [0u8; 32].into(),
            // Keep the peer's path warm; None is the embedder's opt-out.
            persistent_keepalive_interval,
        },
    );
}

#[cfg(test)]
mod tests {
    use ts_control::{Node, StableNodeId, TailnetAddress};

    use super::*;

    /// A node with the given tailnet addresses. `ipv4`/`ipv6` are prefix strings so a test can hand
    /// in the unspecified placeholders `TailnetAddress` synthesizes for a family the tailnet does
    /// not assign.
    fn node(stable: &str, ipv4: &str, ipv6: &str, is_wireguard_only: bool) -> Node {
        Node {
            id: 1,
            stable_id: StableNodeId(stable.to_string()),
            hostname: stable.to_string(),
            user_id: 0,
            tailnet: None,
            tags: vec![],
            // Go's `Node.Addresses` is what control actually assigned, and the wire -> domain
            // `From` impl copies that list verbatim: the unspecified placeholder is synthesized
            // onto `tailnet_address` alone, for a family the tailnet does not assign. Keep the
            // same split here, or a fixture credits a node with an address it was never given.
            addresses: [ipv4.parse().unwrap(), ipv6.parse().unwrap()]
                .into_iter()
                .filter(|net: &ipnet::IpNet| !net.addr().is_unspecified())
                .collect(),
            tailnet_address: TailnetAddress {
                ipv4: ipv4.parse().unwrap(),
                ipv6: ipv6.parse().unwrap(),
            },
            node_key: [0u8; 32].into(),
            node_key_expiry: None,
            expired: false,
            online: None,
            last_seen: None,
            key_signature: vec![],
            machine_key: None,
            disco_key: None,
            accepted_routes: vec![],
            underlay_addresses: vec![],
            derp_region: None,
            cap: Default::default(),
            cap_map: Default::default(),
            peerapi_port: None,
            peerapi_dns_proxy: false,
            is_wireguard_only,
            exit_node_dns_resolvers: vec![],
            peer_relay: false,
            ssh_host_keys: vec![],
            service_vips: Default::default(),
            unsigned_peer_api_only: false,
        }
    }

    /// Data in either direction counts as traffic with a peer; a WireGuard keepalive does not,
    /// since this fork keeps one going to every peer with a session, and neither does a handshake,
    /// since those keepalives cause periodic rekeys. A peer is reported once per report interval
    /// however many batches it sends.
    #[tokio::test(start_paused = true)]
    async fn peer_activity_counts_data_but_not_keepalives_or_handshakes() {
        let activity = PeerActivity::default();
        let peer = PeerId(7);
        let packet = |len: usize| PacketMut::from([vec![4u8, 0, 0, 0], vec![0; len - 4]].concat());
        let reported = async |activity: &PeerActivity| {
            tokio::time::timeout(std::time::Duration::from_millis(10), activity.next())
                .await
                .ok()
        };

        activity.note(peer, &[packet(32)]);
        assert_eq!(
            reported(&activity).await,
            None,
            "a keepalive is not traffic"
        );

        activity.note(peer, &[PacketMut::from(vec![1u8; 148])]);
        assert_eq!(
            reported(&activity).await,
            None,
            "a handshake initiation is not traffic"
        );

        activity.note(peer, &[PacketMut::from(vec![2u8; 92])]);
        assert_eq!(
            reported(&activity).await,
            None,
            "a handshake response is not traffic"
        );

        activity.note(peer, &[packet(32), packet(48)]);
        assert_eq!(reported(&activity).await, Some(vec![peer]), "data is");

        activity.note(peer, &[packet(48)]);
        assert_eq!(
            reported(&activity).await,
            None,
            "reported at most once per interval"
        );

        tokio::time::advance(PEER_ACTIVITY_REPORT_INTERVAL).await;
        activity.note(peer, &[packet(48)]);
        assert_eq!(reported(&activity).await, Some(vec![peer]));
    }

    #[test]
    fn expired_peer_update_removes_tunnel_and_recovery_recreates_it() {
        let mut peer = node("peer", "100.64.0.1/32", "fd7a::1/128", false);
        let mut wg = ts_tunnel::Endpoint::new(ts_tunnel::NodeKeyPair::new());
        let id = PeerId(1);
        let tunnel_id = ts_tunnel::PeerId(id.0);
        peer.node_key = ts_tunnel::NodeKeyPair::new().public;
        apply_tunnel_peer(&mut wg, id, &peer, None);
        assert_eq!(wg.peer_key(tunnel_id), Some(peer.node_key));
        // Stage a handshake so expiry must also remove its queue, ID and retransmit timer.
        assert!(
            !wg.send([(tunnel_id, vec![PacketMut::from(&b"staged"[..])])])
                .to_peers
                .is_empty()
        );
        assert!(wg.next_event().is_some());
        peer.expired = true;
        apply_tunnel_peer(&mut wg, id, &peer, None);
        assert_eq!(wg.peer_key(tunnel_id), None);
        assert!(wg.next_event().is_none());
        assert!(!wg.remove_peer(tunnel_id));
        apply_tunnel_peer(&mut wg, id, &peer, None); // Repeated expired snapshot is idempotent.
        assert_eq!(wg.peer_key(tunnel_id), None);
        peer.expired = false;
        apply_tunnel_peer(&mut wg, id, &peer, None);
        assert_eq!(wg.peer_key(tunnel_id), Some(peer.node_key));
        assert!(wg.next_event().is_none());
    }

    /// This node's advertisement source addresses (Go `self.Addresses()` as
    /// `selfIPMatchingFamily` walks it): IPv4 first, IPv6 second, and the unspecified placeholder
    /// for a family the tailnet does not assign left out entirely.
    #[test]
    fn self_advertisement_addrs_skip_the_unassigned_family() {
        let dual = node("self", "100.64.0.1/32", "fd7a:115c:a1e0::1/128", false);
        assert_eq!(
            advertisement_self_addrs(&dual),
            vec![
                IpAddr::from([100, 64, 0, 1]),
                "fd7a:115c:a1e0::1".parse::<IpAddr>().unwrap(),
            ],
            "IPv4 comes first, the order control sends"
        );

        let v4_only = node("self", "100.64.0.1/32", "::/128", false);
        assert_eq!(
            advertisement_self_addrs(&v4_only),
            vec![IpAddr::from([100, 64, 0, 1])],
            "the unspecified IPv6 placeholder is not an address"
        );

        let v6_only = node("self", "0.0.0.0/32", "fd7a:115c:a1e0::1/128", false);
        assert_eq!(
            advertisement_self_addrs(&v6_only),
            vec!["fd7a:115c:a1e0::1".parse::<IpAddr>().unwrap()],
            "an IPv6-only tailnet still has a source address"
        );

        assert!(
            advertisement_self_addrs(&node("self", "0.0.0.0/32", "::/128", false)).is_empty(),
            "a node with no addresses at all advertises from nowhere"
        );
    }

    /// A placeholder for a family the tailnet does not assign belongs to `tailnet_address` alone,
    /// never to `Node::addresses` (Go `Node.Addresses`, control's assigned list, which the wire ->
    /// domain `From` impl copies verbatim). `Node::is_router` asks "does this node route anything
    /// besides its own addresses" of exactly that list, so a placeholder left in it answers "the
    /// unassigned family is mine" and swallows a genuinely routed prefix.
    #[test]
    fn placeholder_families_are_not_assigned_addresses() {
        let v4_only = node("self", "100.64.0.1/32", "::/128", false);

        // A route to the family the tailnet never assigned is a routed prefix, not one of this
        // node's own — the placeholder must not answer for it.
        let mut routes_the_unassigned_family = v4_only.clone();
        routes_the_unassigned_family.accepted_routes = vec!["::/128".parse().unwrap()];
        assert!(
            routes_the_unassigned_family.is_router(),
            "a prefix this node was never assigned is a route, whatever family it is in"
        );

        // The same prefixes `advertisement_self_addrs` is willing to speak from, and no others:
        // the two views of "what did control assign us" cannot disagree.
        assert_eq!(
            v4_only
                .addresses
                .iter()
                .map(|net| net.addr())
                .collect::<Vec<_>>(),
            advertisement_self_addrs(&v4_only),
            "the unspecified IPv6 placeholder is not an assigned prefix"
        );

        let dual = node("self", "100.64.0.1/32", "fd7a:115c:a1e0::1/128", false);
        assert_eq!(
            dual.addresses.len(),
            2,
            "a dual-stack node keeps both assigned prefixes"
        );
    }

    /// A peer's advertisement destination is Go's `endpoint.nodeAddr`, "the node's first tailscale
    /// address" — IPv4 where the tailnet assigns one — and it carries `IsWireGuardOnly` through to
    /// the refusal in the dataplane.
    #[test]
    fn peer_advertisement_target_is_the_nodes_first_address() {
        let dual = node("peer", "100.64.0.2/32", "fd7a:115c:a1e0::2/128", false);
        assert_eq!(
            advertisement_target(&dual),
            Some(AdvertisementTarget {
                node_addr: IpAddr::from([100, 64, 0, 2]),
                wireguard_only: false,
            }),
        );

        let v6_only = node("peer", "0.0.0.0/32", "fd7a:115c:a1e0::2/128", false);
        assert_eq!(
            advertisement_target(&v6_only),
            Some(AdvertisementTarget {
                node_addr: "fd7a:115c:a1e0::2".parse().unwrap(),
                wireguard_only: false,
            }),
            "a peer with no IPv4 is addressed over IPv6",
        );

        assert_eq!(
            advertisement_target(&node("peer", "100.64.0.3/32", "::/128", true)),
            Some(AdvertisementTarget {
                node_addr: IpAddr::from([100, 64, 0, 3]),
                wireguard_only: true,
            }),
            "IsWireGuardOnly is carried, not filtered out here",
        );

        assert_eq!(
            advertisement_target(&node("peer", "0.0.0.0/32", "::/128", false)),
            None,
            "a peer with no address at all cannot be advertised to",
        );
    }

    /// The per-peer map is rebuilt from the snapshot, so a peer that left the netmap stops being
    /// advertised to without any separate deletion bookkeeping.
    #[test]
    fn advertisement_targets_follow_the_peer_snapshot() {
        let mut db = crate::peer_tracker::PeerDb::default();
        let a = db.upsert(&node("a", "100.64.0.2/32", "::/128", false));
        let b = db.upsert(&node("b", "100.64.0.3/32", "::/128", false));

        let targets = advertisement_targets(&db);
        assert_eq!(targets.len(), 2);
        assert_eq!(
            targets.get(&a).map(|t| t.node_addr),
            Some(IpAddr::from([100, 64, 0, 2]))
        );

        db.remove(&StableNodeId("b".to_string()));
        let targets = advertisement_targets(&db);
        assert_eq!(targets.len(), 1, "a departed peer is no longer a target");
        assert!(targets.contains_key(&a));
        assert!(!targets.contains_key(&b));
    }
}