syd 3.58.0

//
// Syd: rock-solid application kernel
// src/kcov/mod.rs: KCOV userspace ABI shim for syzkaller
//
// Copyright (c) 2025, 2026 Ali Polatel <alip@chesswob.org>
// SPDX-License-Identifier: GPL-3.0

use std::{
    fmt,
    sync::{Arc, OnceLock, RwLock},
};

use nix::{errno::Errno, unistd::Pid};
use serde::{Serialize, Serializer};

use crate::{hash::SydHashMap, proc::util::proc_task_nr};

// KCOV ABI handlers
pub(crate) mod abi;

// KCOV API utilities
pub(crate) mod api;

// Thread-local sink describing where the live writer should send records.
#[derive(Clone, Copy, Debug)]
pub(crate) struct TlsSink {
    pub(crate) id: KcovId,
}

thread_local! {
    static TLS_SINK: RwLock<Option<TlsSink>> = const { RwLock::new(None) };
}

pub(crate) fn get_tls_sink() -> Option<KcovId> {
    if let Some(id) = TLS_SINK
        .try_with(|s| {
            s.read()
                .unwrap_or_else(|e| e.into_inner())
                .map(|sink| sink.id)
        })
        .ok()
        .flatten()
    {
        return Some(id);
    }

    None
}

pub(crate) fn set_tls_sink(id: KcovId) {
    let _ =
        TLS_SINK.try_with(|s| *s.write().unwrap_or_else(|e| e.into_inner()) = Some(TlsSink { id }));
}

pub(crate) fn clear_tls_sink() {
    let _ = TLS_SINK.try_with(|s| *s.write().unwrap_or_else(|e| e.into_inner()) = None);
}

#[derive(Clone, Copy, Debug, Default)]
struct KcovBind {
    local: Option<KcovId>,
    remote: Option<KcovId>,
    remote_owner: bool,
    handle: u64,
}

// Global TID map: TID -> KcovBind
static KCOV_TID_MAP: OnceLock<RwLock<SydHashMap<Pid, KcovBind>>> = OnceLock::new();

fn kcov_tid_map() -> &'static RwLock<SydHashMap<Pid, KcovBind>> {
    KCOV_TID_MAP.get_or_init(|| RwLock::new(SydHashMap::default()))
}

pub(crate) fn set_kcov_tid(tid: Pid, id: KcovId, is_remote: bool, common_handle: u64) {
    let mut map = kcov_tid_map().write().unwrap_or_else(|e| e.into_inner());
    let mut bind = map.get(&tid).copied().unwrap_or_default();
    if is_remote {
        bind.remote = Some(id);
        bind.remote_owner = true;
        if common_handle != 0 {
            bind.handle = common_handle;
        }
    } else {
        bind.local = Some(id);
    }
    map.insert(tid, bind);
}

pub(crate) fn clear_kcov_handle(tid: Pid) {
    let mut map = kcov_tid_map().write().unwrap_or_else(|e| e.into_inner());
    if let Some(bind) = map.get_mut(&tid) {
        bind.handle = 0;
        if bind.local.is_none() && bind.remote.is_none() {
            map.remove(&tid);
        }
    }
}

pub(crate) fn get_kcov_tid(tid: Pid) -> Option<KcovId> {
    let bind = {
        let map = kcov_tid_map().read().unwrap_or_else(|e| e.into_inner());
        map.get(&tid).copied()
    }?;
    if let Some(id) = bind.local.or(bind.remote) {
        return Some(id);
    }
    if bind.handle != 0 {
        return abi::kcov_handle_lookup(bind.handle);
    }
    None
}

pub(crate) fn get_kcov_tid_local(tid: Pid) -> Option<KcovId> {
    let map = kcov_tid_map().read().unwrap_or_else(|e| e.into_inner());
    map.get(&tid).and_then(|b| b.local)
}

pub(crate) fn get_kcov_tid_remote(tid: Pid) -> Option<KcovId> {
    let map = kcov_tid_map().read().unwrap_or_else(|e| e.into_inner());
    map.get(&tid)
        .and_then(|b| if b.remote_owner { b.remote } else { None })
}

fn find_kcov_binder(id: KcovId) -> Option<Pid> {
    let map = kcov_tid_map().read().unwrap_or_else(|e| e.into_inner());
    map.iter()
        .find_map(|(t, b)| (b.local == Some(id) || b.remote == Some(id)).then_some(*t))
}

pub(crate) fn remove_kcov_tid(tid: Pid) {
    let (local, remote) = {
        let map = kcov_tid_map().read().unwrap_or_else(|e| e.into_inner());
        map.get(&tid).map_or((None, None), |b| {
            (b.local, if b.remote_owner { b.remote } else { None })
        })
    };
    if let Some(id) = local {
        abi::kcov_mgr().exit_reset(id);
        abi::kcov_clear_mode(id);
    }
    if let Some(id) = remote {
        abi::kcov_mgr().exit_reset(id);
        abi::kcov_clear_mode(id);
        abi::kcov_handle_drop(id);
    }

    let mut map = kcov_tid_map().write().unwrap_or_else(|e| e.into_inner());
    map.remove(&tid);
    drop(map);

    let rebind = |id: KcovId, tgid: Pid| {
        find_kcov_binder(id).or_else(|| match proc_task_nr(tgid) {
            Ok(n) if n > 1 || tgid != tid => Some(tgid),
            _ => None,
        })
    };
    abi::kcov_remove_tid(tid, rebind);
    abi::kcov_mgr().remove_tid(tid, rebind);
}

pub(crate) fn unbind_kcov_tid(tid: Pid, id: KcovId) {
    let mut map = kcov_tid_map().write().unwrap_or_else(|e| e.into_inner());
    if let Some(bind) = map.get_mut(&tid) {
        if bind.local == Some(id) {
            bind.local = None;
        }
        if bind.remote == Some(id) {
            bind.remote = None;
            bind.remote_owner = false;
            bind.handle = 0;
        }
        if bind.local.is_none() && bind.remote.is_none() && bind.handle == 0 {
            map.remove(&tid);
        }
    }
}

pub(crate) fn inherit_kcov_tid(parent_tid: Pid, child_tid: Pid) {
    let mut map = kcov_tid_map().write().unwrap_or_else(|e| e.into_inner());
    let handle = map.get(&parent_tid).map_or(0, |b| b.handle);
    if handle != 0 {
        let mut bind = map.get(&child_tid).copied().unwrap_or_default();
        bind.handle = handle;
        map.insert(child_tid, bind);
    }
}

//
// Public API
//

// KCOV modes (pc/cmp).
#[derive(Copy, Clone, Eq, PartialEq, Debug)]
pub(crate) enum TraceMode {
    Pc,
    Cmp,
}

impl fmt::Display for TraceMode {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Pc => write!(f, "pc"),
            Self::Cmp => write!(f, "cmp"),
        }
    }
}

impl Serialize for TraceMode {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.serialize_str(&self.to_string())
    }
}

// /sys/kernel/debug/kcov handle.
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
pub(crate) struct KcovId(u64);

impl KcovId {
    // Create a new KcovId.
    pub(crate) const fn new(id: u64) -> Self {
        Self(id)
    }
}

pub(crate) struct Kcov {
    map: RwLock<SydHashMap<KcovId, (Pid, Pid, Arc<State>)>>,
}

impl Kcov {
    pub(crate) fn new() -> Self {
        Self {
            map: RwLock::new(SydHashMap::default()),
        }
    }

    pub(crate) fn open(&self, kcov_id: u64, tid: Pid, tgid: Pid) -> Result<(), Errno> {
        let kcov_id = KcovId(kcov_id);
        let state_arc = Arc::new(State::new());

        let mut map = self.map.write().unwrap_or_else(|e| e.into_inner());
        map.insert(kcov_id, (tid, tgid, state_arc));

        Ok(())
    }

    pub(crate) fn init_trace(
        &self,
        kcov_id: KcovId,
        words: u64,
        wordsize: u8,
    ) -> Result<(), Errno> {
        self.get(kcov_id)?.init_trace(words, wordsize)
    }

    pub(crate) fn reset(&self, id: KcovId) {
        if let Ok(st) = self.get(id) {
            let mut core = st.core.write().unwrap_or_else(|e| e.into_inner());
            core.phase = Phase::Disabled;
            core.mode = None;
        }
    }

    pub(crate) fn can_enable(&self, id: KcovId) -> Result<(), Errno> {
        let st = self.get(id)?;
        let core = st.core.read().unwrap_or_else(|e| e.into_inner());
        if core.phase != Phase::Init {
            return Err(Errno::EINVAL);
        }
        Ok(())
    }

    pub(crate) fn enable(&self, id: KcovId, mode: TraceMode) -> Result<(), Errno> {
        let st = self.get(id)?;
        st.enable(mode)?;

        set_tls_sink(id);

        Ok(())
    }

    pub(crate) fn disable(&self, id: KcovId) -> Result<(), Errno> {
        let st = self.get(id)?;
        st.disable()?;

        clear_tls_sink();

        Ok(())
    }

    pub(crate) fn exit_reset(&self, id: KcovId) {
        if let Ok(st) = self.get(id) {
            let mut core = st.core.write().unwrap_or_else(|e| e.into_inner());
            if core.phase == Phase::Enabled {
                core.phase = Phase::Init;
                core.mode = None;
            }
        }
    }

    fn get(&self, kcov_id: KcovId) -> Result<Arc<State>, Errno> {
        let read_guard = self.map.read().unwrap_or_else(|e| e.into_inner());
        read_guard
            .get(&kcov_id)
            .map(|(_, _, st)| Arc::clone(st))
            .ok_or(Errno::ENOTTY)
    }

    pub(crate) fn remove_tid<F: Fn(KcovId, Pid) -> Option<Pid>>(&self, tid: Pid, rebind: F) {
        let mut map = self.map.write().unwrap_or_else(|e| e.into_inner());
        map.retain(|id, (t, tgid, _)| {
            if *t != tid {
                return true;
            }
            if let Some(new_tid) = rebind(*id, *tgid) {
                *t = new_tid;
                true
            } else {
                false
            }
        });
    }
}

//
// Internals
//

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
enum Phase {
    Disabled,
    Init,
    Enabled,
}

impl fmt::Display for Phase {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let s = match self {
            Self::Disabled => "disabled",
            Self::Init => "init",
            Self::Enabled => "enabled",
        };
        f.write_str(s)
    }
}

impl Serialize for Phase {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.serialize_str(&self.to_string())
    }
}

struct State {
    core: RwLock<Core>,
}

struct Core {
    mode: Option<TraceMode>,
    phase: Phase,
}

impl State {
    fn new() -> Self {
        Self {
            core: RwLock::new(Core {
                mode: None,
                phase: Phase::Disabled,
            }),
        }
    }

    fn init_trace(&self, words: u64, wordsize: u8) -> Result<(), Errno> {
        if words < 2 || words > (i32::MAX as u64) / u64::from(wordsize) {
            return Err(Errno::EINVAL);
        }

        let mut core = self.core.write().unwrap_or_else(|e| e.into_inner());
        if core.phase != Phase::Disabled {
            return Err(Errno::EBUSY);
        }

        core.mode = None;
        core.phase = Phase::Init;

        Ok(())
    }

    fn enable(&self, mode: TraceMode) -> Result<(), Errno> {
        let mut core = self.core.write().unwrap_or_else(|e| e.into_inner());

        match core.phase {
            Phase::Init => {
                core.mode = Some(mode);
                core.phase = Phase::Enabled;
                Ok(())
            }

            Phase::Enabled | Phase::Disabled => Err(Errno::EINVAL),
        }
    }

    fn disable(&self) -> Result<(), Errno> {
        let mut core = self.core.write().unwrap_or_else(|e| e.into_inner());

        if core.phase != Phase::Enabled {
            return Err(Errno::EINVAL);
        }

        core.phase = Phase::Init;
        core.mode = None;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_kcov_id_new_1() {
        let id = KcovId::new(42);
        assert_eq!(id, KcovId(42));
    }

    #[test]
    fn test_kcov_id_eq_1() {
        assert_eq!(KcovId::new(1), KcovId::new(1));
    }

    #[test]
    fn test_kcov_id_ne_1() {
        assert_ne!(KcovId::new(1), KcovId::new(2));
    }

    #[test]
    fn test_tls_sink_none_by_default_1() {
        clear_tls_sink();
        assert!(get_tls_sink().is_none());
    }

    #[test]
    fn test_tls_sink_set_get_1() {
        let id = KcovId::new(99);
        set_tls_sink(id);
        assert_eq!(get_tls_sink(), Some(id));
        clear_tls_sink();
    }

    #[test]
    fn test_tls_sink_clear_1() {
        let id = KcovId::new(77);
        set_tls_sink(id);
        clear_tls_sink();
        assert!(get_tls_sink().is_none());
    }
}