jsdet_core/

taint.rs

//! Taint tracking for the jsdet core.
//!
//! This module provides real taint tracking in the Rust core, complementing
//! the JavaScript-level taint tracking in the QuickJS engine.
//!
//! # Architecture
//!
//! The taint system has two layers:
//!
//! 1. **JS Engine Layer** (QuickJS/WASM): Taint is stored on JS string objects
//!    via the `__jsdet_set_taint`/`__jsdet_get_taint` intrinsics. This tracks
//!    taint through JS operations like concat, slice, replace.
//!
//! 2. **Rust Core Layer** (this module): Taint is stored on [`Value`] objects
//!    via [`TaintLabel`]. This tracks taint as values pass through the bridge
//!    between JS and Rust.
//!
//! # Usage
//!
//! ```
//! use jsdet_core::taint::{TaintLabel, TaintTracker};
//! use jsdet_core::observation::Value;
//!
//! // Create a tracker
//! let mut tracker = TaintTracker::new();
//!
//! // Mark a source as tainted
//! let tainted_value = Value::tainted_string("attacker_input", TaintLabel::new(1));
//!
//! // Check if value reaches a sink
//! if let Some(flow) = Value::check_taint_at_sink("eval", &[tainted_value]) {
//!     println!("Taint flow detected: {} -> {}", flow.sink, flow.label.0);
//! }
//! ```

use crate::observation::{TaintFlow, TaintLabel, Value};
use std::collections::HashMap;

/// A taint source registration.
#[derive(Debug, Clone)]
pub struct Source {
    /// The API name that produces tainted data (e.g., "chrome.runtime.onMessage").
    pub api: String,
    /// The taint label to assign.
    pub label: TaintLabel,
    /// Human-readable description.
    pub description: String,
}

/// A taint sink registration.
#[derive(Debug, Clone)]
pub struct Sink {
    /// The API name that is dangerous (e.g., "eval", "chrome.tabs.executeScript").
    pub api: String,
    /// Which argument positions are dangerous (0-indexed).
    pub dangerous_args: Vec<usize>,
    /// Severity if tainted data reaches here.
    pub severity: Severity,
    /// CWE identifier.
    pub cwe: String,
}

/// Severity levels for taint flows.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Severity {
    Critical,
    High,
    Medium,
    Low,
    Info,
}

impl std::fmt::Display for Severity {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Critical => write!(f, "critical"),
            Self::High => write!(f, "high"),
            Self::Medium => write!(f, "medium"),
            Self::Low => write!(f, "low"),
            Self::Info => write!(f, "info"),
        }
    }
}

/// Tracks taint through a single execution session.
///
/// This is the main interface for cross-function taint tracking in the Rust core.
/// It maintains the set of sources and sinks, and records confirmed taint flows.
#[derive(Debug, Default)]
pub struct TaintTracker {
    /// Registered taint sources.
    sources: HashMap<String, Source>,
    /// Registered taint sinks.
    sinks: HashMap<String, Sink>,
    /// Confirmed taint flows (tainted data reached a sink).
    flows: Vec<TaintFlow>,
    /// Next taint label ID to assign.
    next_label: u32,
}

impl TaintTracker {
    /// Create a new empty taint tracker.
    pub fn new() -> Self {
        Self {
            sources: HashMap::new(),
            sinks: HashMap::new(),
            flows: Vec::new(),
            next_label: 1, // Start at 1 (0 = clean)
        }
    }

    /// Register a new taint source.
    ///
    /// Returns the assigned taint label for this source.
    /// CRITICAL FIX: Uses saturating arithmetic to prevent overflow.
    pub fn register_source(
        &mut self,
        api: impl Into<String>,
        description: impl Into<String>,
    ) -> TaintLabel {
        let label = TaintLabel::new(self.next_label);
        // CRITICAL FIX: Use saturating_add to prevent overflow panic
        self.next_label = self.next_label.saturating_add(1);
        // Ensure we never return label 0 (CLEAN) even after overflow
        if label.0 == 0 {
            return TaintLabel::new(1);
        }

        let source = Source {
            api: api.into(),
            label,
            description: description.into(),
        };

        self.sources.insert(source.api.clone(), source);
        label
    }

    /// Register a new taint sink.
    pub fn register_sink(
        &mut self,
        api: impl Into<String>,
        dangerous_args: Vec<usize>,
        severity: Severity,
        cwe: impl Into<String>,
    ) {
        let sink = Sink {
            api: api.into(),
            dangerous_args,
            severity,
            cwe: cwe.into(),
        };
        self.sinks.insert(sink.api.clone(), sink);
    }

    /// Check if an API is a registered source.
    pub fn is_source(&self, api: &str) -> Option<&Source> {
        self.sources.get(api)
    }

    /// Check if an API is a registered sink.
    pub fn is_sink(&self, api: &str) -> Option<&Sink> {
        self.sinks.get(api)
    }

    /// Apply taint to a value returned from a source API.
    ///
    /// If the API is a registered source, the value is marked with the
    /// corresponding taint label. Otherwise, the value is returned unchanged.
    pub fn apply_source_taint(&self, api: &str, value: Value) -> Value {
        if let Some(source) = self.is_source(api) {
            value.with_taint(source.label)
        } else {
            value
        }
    }

    /// Check for taint flows at a sink API call.
    ///
    /// If any of the dangerous arguments are tainted, records a taint flow
    /// and returns it. Returns None if no tainted data reached the sink.
    pub fn check_sink(&mut self, api: &str, args: &[Value]) -> Option<TaintFlow> {
        let sink = self.is_sink(api)?;

        // Check only the dangerous argument positions
        let dangerous_values: Vec<(usize, &Value)> = args
            .iter()
            .enumerate()
            .filter(|(idx, _)| sink.dangerous_args.contains(idx))
            .collect();

        if let Some(flow) = Value::check_taint_at_sink(
            api,
            &dangerous_values
                .iter()
                .map(|(_, v)| (*v).clone())
                .collect::<Vec<_>>(),
        ) {
            self.flows.push(flow.clone());
            Some(flow)
        } else {
            None
        }
    }

    /// Get all recorded taint flows.
    pub fn flows(&self) -> &[TaintFlow] {
        &self.flows
    }

    /// Take all recorded taint flows (clears internal list).
    pub fn take_flows(&mut self) -> Vec<TaintFlow> {
        std::mem::take(&mut self.flows)
    }

    /// Returns true if any taint flows were recorded.
    pub fn has_flows(&self) -> bool {
        !self.flows.is_empty()
    }

    /// Count of confirmed taint flows.
    pub fn flow_count(&self) -> usize {
        self.flows.len()
    }

    /// Get all registered sources.
    pub fn sources(&self) -> &HashMap<String, Source> {
        &self.sources
    }

    /// Get all registered sinks.
    pub fn sinks(&self) -> &HashMap<String, Sink> {
        &self.sinks
    }
}

/// Propagate taint through a string concatenation operation.
///
/// Takes multiple values and returns a new string Value with the combined
/// taint labels. If any input is tainted, the result is tainted.
///
/// # Example
///
/// ```
/// use jsdet_core::taint::propagate_concat;
/// use jsdet_core::observation::{Value, TaintLabel};
///
/// let a = Value::string("hello ");
/// let b = Value::tainted_string("world", TaintLabel::new(1));
///
/// let result = propagate_concat(&[a, b]).unwrap();
/// assert!(result.is_tainted());
/// assert_eq!(result.as_str(), Some("hello world"));
/// ```
pub fn propagate_concat(values: &[Value]) -> Option<Value> {
    if values.is_empty() {
        return Some(Value::string(""));
    }

    // Build the concatenated string
    let mut result = String::new();
    let mut combined_label = TaintLabel::CLEAN;

    for value in values {
        match value {
            Value::String(s, label) => {
                result.push_str(s);
                if combined_label.is_clean() && label.is_tainted() {
                    combined_label = *label;
                }
            }
            _ => return None, // Non-string in concat
        }
    }

    Some(Value::String(result, combined_label))
}

/// Propagate taint through a string slice operation.
///
/// The result carries the same taint label as the source.
pub fn propagate_slice(value: &Value, start: usize, end: usize) -> Option<Value> {
    value.slice(start, end)
}

/// Propagate taint through a string replace operation.
///
/// The result carries the taint label of the source string.
pub fn propagate_replace(value: &Value, from: &str, to: &str) -> Option<Value> {
    value.replace(from, to)
}

/// Propagate taint through JSON.parse.
///
/// If the input JSON string is tainted, the resulting Json value
/// is also tainted.
pub fn propagate_json_parse(value: &Value) -> Option<Value> {
    match value {
        Value::String(s, label) => Some(Value::Json(s.clone(), *label)),
        _ => None,
    }
}

/// Propagate taint through JSON.stringify.
///
/// If the input value is tainted, the resulting string is also tainted.
pub fn propagate_json_stringify(value: &Value) -> Option<Value> {
    match value {
        Value::Json(s, label) => Some(Value::String(s.clone(), *label)),
        _ => None,
    }
}

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

    #[test]
    fn taint_label_basic() {
        let clean = TaintLabel::CLEAN;
        assert!(!clean.is_tainted());
        assert!(clean.is_clean());

        let tainted = TaintLabel::new(1);
        assert!(tainted.is_tainted());
        assert!(!tainted.is_clean());
    }

    #[test]
    fn taint_label_combine() {
        let clean = TaintLabel::CLEAN;
        let t1 = TaintLabel::new(1);
        let t2 = TaintLabel::new(2);

        assert_eq!(clean.combine(clean), clean);
        assert_eq!(clean.combine(t1), t1);
        assert_eq!(t1.combine(clean), t1);
        assert_eq!(t1.combine(t2), t1); // First wins when both tainted
    }

    #[test]
    fn value_taint_tracking() {
        let clean = Value::string("clean");
        let tainted = Value::tainted_string("tainted", TaintLabel::new(1));

        assert!(!clean.is_tainted());
        assert!(tainted.is_tainted());

        assert_eq!(clean.taint_label(), TaintLabel::CLEAN);
        assert_eq!(tainted.taint_label(), TaintLabel::new(1));
    }

    #[test]
    fn value_with_taint() {
        let clean = Value::string("test");
        let tainted = clean.with_taint(TaintLabel::new(5));

        assert!(tainted.is_tainted());
        assert_eq!(tainted.taint_label(), TaintLabel::new(5));
    }

    #[test]
    fn check_taint_at_sink_detects_tainted() {
        let args = vec![
            Value::string("safe"),
            Value::tainted_string("dangerous", TaintLabel::new(1)),
        ];

        let flow = Value::check_taint_at_sink("eval", &args);
        assert!(flow.is_some());

        let flow = flow.unwrap();
        assert_eq!(flow.sink, "eval");
        assert_eq!(flow.label, TaintLabel::new(1));
        assert_eq!(flow.tainted_args, vec![1]);
    }

    #[test]
    fn check_taint_at_sink_returns_none_for_clean() {
        let args = vec![Value::string("safe1"), Value::string("safe2")];

        let flow = Value::check_taint_at_sink("eval", &args);
        assert!(flow.is_none());
    }

    #[test]
    fn value_concat_propagates_taint() {
        let a = Value::string("hello ");
        let b = Value::tainted_string("world", TaintLabel::new(1));

        let result = a.concat(&b).unwrap();

        assert!(result.is_tainted());
        assert_eq!(result.taint_label(), TaintLabel::new(1));
        assert_eq!(result.as_str(), Some("hello world"));
    }

    #[test]
    fn value_slice_preserves_taint() {
        let s = Value::tainted_string("abcdef", TaintLabel::new(2));

        let result = s.slice(1, 4).unwrap();

        assert!(result.is_tainted());
        assert_eq!(result.taint_label(), TaintLabel::new(2));
        assert_eq!(result.as_str(), Some("bcd"));
    }

    #[test]
    fn value_replace_preserves_taint() {
        let s = Value::tainted_string("hello world", TaintLabel::new(3));

        let result = s.replace("world", "universe").unwrap();

        assert!(result.is_tainted());
        assert_eq!(result.taint_label(), TaintLabel::new(3));
        assert_eq!(result.as_str(), Some("hello universe"));
    }

    #[test]
    fn value_equality_ignores_taint() {
        let a = Value::string("test");
        let b = Value::tainted_string("test", TaintLabel::new(1));

        // Equality ignores taint - only the value matters
        assert_eq!(a, b);

        // But taint status is different
        assert!(!a.is_tainted());
        assert!(b.is_tainted());
    }

    #[test]
    fn taint_tracker_registration() {
        let mut tracker = TaintTracker::new();

        let label = tracker.register_source("chrome.runtime.onMessage", "Message from extension");
        assert_eq!(label, TaintLabel::new(1));

        tracker.register_sink("eval", vec![0], Severity::Critical, "CWE-95");

        assert!(tracker.is_source("chrome.runtime.onMessage").is_some());
        assert!(tracker.is_sink("eval").is_some());
        assert!(tracker.is_source("fetch").is_none());
    }

    #[test]
    fn taint_tracker_detects_flow() {
        let mut tracker = TaintTracker::new();

        tracker.register_source("source.api", "Test source");
        tracker.register_sink("sink.api", vec![0], Severity::High, "CWE-79");

        // Apply source taint (should apply the label)
        let tainted = tracker.apply_source_taint("source.api", Value::string("evil"));
        assert!(tainted.is_tainted());

        // Check at sink
        let flow = tracker.check_sink("sink.api", &[tainted]);
        assert!(flow.is_some());
        assert_eq!(tracker.flow_count(), 1);
    }

    #[test]
    fn propagate_concat_multiple() {
        let values = vec![
            Value::string("a"),
            Value::tainted_string("b", TaintLabel::new(1)),
            Value::string("c"),
        ];

        let result = propagate_concat(&values).unwrap();
        assert!(result.is_tainted());
        assert_eq!(result.as_str(), Some("abc"));
    }

    #[test]
    fn propagate_concat_empty() {
        let result = propagate_concat(&[]).unwrap();
        assert_eq!(result.as_str(), Some(""));
        assert!(!result.is_tainted());
    }

    #[test]
    fn propagate_json_parse_stringifies() {
        let json_str = Value::tainted_string(r#"{"key":"value"}"#, TaintLabel::new(1));

        let parsed = propagate_json_parse(&json_str).unwrap();
        assert!(matches!(parsed, Value::Json(_, _)));
        assert!(parsed.is_tainted());
    }

    #[test]
    fn propagate_json_stringify_preserves_taint() {
        let json = Value::tainted_json(r#"{"key":"value"}"#, TaintLabel::new(2));

        let stringified = propagate_json_stringify(&json).unwrap();
        assert!(matches!(stringified, Value::String(_, _)));
        assert!(stringified.is_tainted());
        assert_eq!(stringified.taint_label(), TaintLabel::new(2));
    }
}