use std::collections::HashMap;
use anyhow::{Context, Result};
use tldr_core::Language;
use super::types::FunctionId;
pub fn build_cfg_for_function(
file_contents: &str,
function_id: &FunctionId,
language: Language,
) -> Result<tldr_core::CfgInfo> {
tldr_core::get_cfg_context(file_contents, &function_id.qualified_name, language)
.map_err(|e| anyhow::anyhow!(e))
.with_context(|| format!("CFG construction for {}", function_id))
}
pub fn build_dfg_for_function(
file_contents: &str,
function_id: &FunctionId,
language: Language,
) -> Result<tldr_core::DfgInfo> {
tldr_core::get_dfg_context(file_contents, &function_id.qualified_name, language)
.map_err(|e| anyhow::anyhow!(e))
.with_context(|| format!("DFG construction for {}", function_id))
}
pub fn build_ssa_for_function(
file_contents: &str,
function_id: &FunctionId,
language: Language,
) -> Result<tldr_core::ssa::SsaFunction> {
tldr_core::ssa::construct_ssa(
file_contents,
&function_id.qualified_name,
language,
tldr_core::ssa::SsaType::Minimal,
)
.map_err(|e| anyhow::anyhow!(e))
.with_context(|| format!("SSA construction for {}", function_id))
}
pub fn build_taint_for_function(
file_contents: &str,
function_id: &FunctionId,
language: Language,
) -> Result<tldr_core::TaintInfo> {
let cfg = build_cfg_for_function(file_contents, function_id, language)?;
let dfg = build_dfg_for_function(file_contents, function_id, language)?;
let statements: HashMap<u32, String> = file_contents
.lines()
.enumerate()
.map(|(i, line)| ((i + 1) as u32, line.to_string()))
.collect();
let tree = tldr_core::ast::parser::parse(file_contents, language)
.map_err(|e| anyhow::anyhow!(e))
.ok();
tldr_core::compute_taint_with_tree(
&cfg,
&dfg.refs,
&statements,
tree.as_ref(),
Some(file_contents.as_bytes()),
language,
)
.map_err(|e| anyhow::anyhow!(e))
.with_context(|| format!("Taint analysis for {}", function_id))
}
pub fn build_abstract_interp(
file_contents: &str,
function_id: &FunctionId,
language: Language,
) -> Result<tldr_core::AbstractInterpInfo> {
let cfg = build_cfg_for_function(file_contents, function_id, language)?;
let dfg = build_dfg_for_function(file_contents, function_id, language)?;
let source_lines: Vec<&str> = file_contents.lines().collect();
let lang_str = match language {
Language::Python => "python",
Language::Rust => "rust",
Language::JavaScript | Language::TypeScript => "javascript",
Language::Go => "go",
Language::Java => "java",
_ => "unknown",
};
tldr_core::compute_abstract_interp(&cfg, &dfg, Some(&source_lines), lang_str)
.map_err(|e| anyhow::anyhow!(e))
.with_context(|| format!("Abstract interpretation for {}", function_id))
}
#[cfg(test)]
mod tests {
use super::*;
const PYTHON_SOURCE: &str = "def do_work(x):\n y = x + 1\n return y\n";
const RUST_SOURCE: &str = "fn do_work() {\n let x = 1;\n let y = x + 1;\n y\n}\n";
#[test]
fn test_build_cfg_for_python_function() {
let fid = FunctionId::new("src/example.py", "do_work", 1);
let result = build_cfg_for_function(PYTHON_SOURCE, &fid, Language::Python);
assert!(
result.is_ok(),
"CFG construction should succeed for valid Python, got: {:?}",
result.err()
);
let cfg = result.unwrap();
assert_eq!(cfg.function, "do_work");
assert!(!cfg.blocks.is_empty(), "CFG should have at least one block");
}
#[test]
fn test_build_cfg_for_rust_function() {
let fid = FunctionId::new("src/example.rs", "do_work", 1);
let result = build_cfg_for_function(RUST_SOURCE, &fid, Language::Rust);
assert!(
result.is_ok(),
"CFG construction should succeed for valid Rust, got: {:?}",
result.err()
);
let cfg = result.unwrap();
assert_eq!(cfg.function, "do_work");
}
#[test]
fn test_build_cfg_nonexistent_function_returns_empty() {
let fid = FunctionId::new("src/example.py", "nonexistent", 1);
let result = build_cfg_for_function(PYTHON_SOURCE, &fid, Language::Python);
assert!(result.is_ok(), "should return Ok with empty CFG");
let cfg = result.unwrap();
assert!(cfg.blocks.is_empty(), "CFG should have no blocks for nonexistent function");
}
#[test]
fn test_build_dfg_for_python_function() {
let fid = FunctionId::new("src/example.py", "do_work", 1);
let result = build_dfg_for_function(PYTHON_SOURCE, &fid, Language::Python);
assert!(
result.is_ok(),
"DFG construction should succeed for valid Python, got: {:?}",
result.err()
);
let dfg = result.unwrap();
assert_eq!(dfg.function, "do_work");
assert!(!dfg.refs.is_empty(), "DFG should have variable references");
}
#[test]
fn test_build_dfg_nonexistent_function_returns_error() {
let fid = FunctionId::new("src/example.py", "missing", 1);
let result = build_dfg_for_function(PYTHON_SOURCE, &fid, Language::Python);
assert!(
result.is_err(),
"DFG should fail for nonexistent function"
);
}
#[test]
fn test_build_ssa_for_python_function() {
let fid = FunctionId::new("src/example.py", "do_work", 1);
let result = build_ssa_for_function(PYTHON_SOURCE, &fid, Language::Python);
assert!(
result.is_ok(),
"SSA construction should succeed for valid Python, got: {:?}",
result.err()
);
let ssa = result.unwrap();
assert_eq!(ssa.function, "do_work");
}
#[test]
fn test_build_taint_for_python_function() {
let fid = FunctionId::new("src/example.py", "do_work", 1);
let result = build_taint_for_function(PYTHON_SOURCE, &fid, Language::Python);
assert!(
result.is_ok(),
"Taint analysis should succeed for valid Python, got: {:?}",
result.err()
);
let taint = result.unwrap();
assert_eq!(taint.function_name, "do_work");
}
#[test]
fn test_build_abstract_interp_for_python_function() {
let fid = FunctionId::new("src/example.py", "do_work", 1);
let result = build_abstract_interp(PYTHON_SOURCE, &fid, Language::Python);
assert!(
result.is_ok(),
"Abstract interp should succeed for valid Python, got: {:?}",
result.err()
);
let ai = result.unwrap();
assert_eq!(ai.function_name, "do_work");
}
#[test]
fn test_error_context_contains_function_id() {
let fid = FunctionId::new("src/special.py", "Foo::bar_method", 10);
let result = build_dfg_for_function(PYTHON_SOURCE, &fid, Language::Python);
assert!(result.is_err());
let err_msg = format!("{:#}", result.unwrap_err());
assert!(
err_msg.contains("Foo::bar_method"),
"error should contain function name, got: {}",
err_msg
);
}
}