syd 3.57.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;

// 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) };
    static RECURSION_GUARD: RwLock<bool> = const { RwLock::new(false) };
}

pub(crate) fn get_tls_sink() -> Option<KcovId> {
    // Check recursion guard first.
    let guard = match RECURSION_GUARD.try_with(|g| *g.read().unwrap_or_else(|e| e.into_inner())) {
        Ok(g) => g,
        Err(_) => return None,
    };
    if guard {
        return None;
    }

    // Try TLS.
    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>,
}

// 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) {
    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);
    } else {
        bind.local = Some(id);
    }
    map.insert(tid, bind);
}

pub(crate) fn get_kcov_tid(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.or(b.remote))
}

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| b.remote)
}

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, b.remote))
    };
    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 (drop_local, drop_remote) = {
        let mut map = kcov_tid_map().write().unwrap_or_else(|e| e.into_inner());
        map.remove(&tid);
        let dl = local.filter(|id| {
            !map.values()
                .any(|b| b.local == Some(*id) || b.remote == Some(*id))
        });
        let dr = remote.filter(|id| {
            Some(*id) != local
                && !map
                    .values()
                    .any(|b| b.local == Some(*id) || b.remote == Some(*id))
        });
        (dl, dr)
    };
    if let Some(id) = drop_local {
        abi::kcov_mgr().close(id);
        abi::kcov_reg_remove(id);
    }
    if let Some(id) = drop_remote {
        abi::kcov_mgr().close(id);
        abi::kcov_reg_remove(id);
    }
}

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;
        }
        if bind.local.is_none() && bind.remote.is_none() {
            map.remove(&tid);
        }
    }
}

pub(crate) fn inherit_kcov_tid(parent_tid: Pid, child_tid: Pid) {
    let remote = {
        let map = kcov_tid_map().read().unwrap_or_else(|e| e.into_inner());
        map.get(&parent_tid).and_then(|b| b.remote)
    };
    if let Some(id) = remote {
        set_kcov_tid(child_tid, id, true);
    }
}

//
// 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, Arc<State>>>,
}

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

    pub(crate) fn open(&self, kcov_id: u64) -> 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, 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 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;
            }
        }
    }

    pub(crate) fn close(&self, id: KcovId) {
        let mut map = self.map.write().unwrap_or_else(|e| e.into_inner());
        map.remove(&id);
    }

    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).cloned().ok_or(Errno::EBADF)
    }
}

//
// 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());
    }
}