use anyhow::{Context, Result};
use std::fs::File;
use std::path::{Path, PathBuf};
#[cfg(coverage)]
use crate::vmm;
pub(crate) fn try_flush_profraw() {
#[cfg(coverage)]
{
if !vmm::guest_comms::is_guest() {
return;
}
{
use std::sync::atomic::{AtomicBool, Ordering};
static FLUSHED: AtomicBool = AtomicBool::new(false);
if FLUSHED.swap(true, Ordering::SeqCst) {
return;
}
}
unsafe extern "C" {
fn __llvm_profile_get_size_for_buffer() -> u64;
fn __llvm_profile_write_buffer(buf: *mut std::os::raw::c_char) -> std::os::raw::c_int;
}
let needed = unsafe { __llvm_profile_get_size_for_buffer() } as usize;
if needed == 0 {
vmm::guest_comms::send_dmesg(
b"ktstr coverage: __llvm_profile_get_size_for_buffer returned 0; no guest profile to flush\n",
);
return;
}
let mut buf: Vec<u8> = vec![0u8; needed];
if unsafe { __llvm_profile_write_buffer(buf.as_mut_ptr().cast::<std::os::raw::c_char>()) }
!= 0
{
vmm::guest_comms::send_dmesg(
b"ktstr coverage: __llvm_profile_write_buffer failed; guest coverage lost for this run\n",
);
return;
}
vmm::guest_comms::send_profraw(&buf);
}
}
pub(crate) fn find_symbol_vaddrs(elf: &goblin::elf::Elf<'_>, names: &[&str]) -> Vec<Option<u64>> {
let mut results = vec![None; names.len()];
let mut remaining = names.len();
for sym in elf.syms.iter() {
if remaining == 0 {
break;
}
let sym_name = match elf.strtab.get_at(sym.st_name) {
Some(n) => n,
None => continue,
};
for (i, name) in names.iter().enumerate() {
if results[i].is_none() && sym_name == *name {
results[i] = Some(sym.st_value);
remaining -= 1;
break;
}
}
}
results
}
static PROFRAW_COUNTER: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(0);
pub(crate) fn persist_guest_profraw(messages: &crate::vmm::host_comms::BulkDrainResult) {
use crate::vmm::wire::MsgType;
for entry in &messages.entries {
if MsgType::from_wire(entry.msg_type) == Some(MsgType::Profraw)
&& entry.crc_ok
&& !entry.payload.is_empty()
&& let Err(e) = write_profraw(&entry.payload)
{
eprintln!("ktstr_test: persist guest profraw: {e}");
}
}
}
pub(crate) fn write_profraw(data: &[u8]) -> Result<()> {
let target_dir = target_dir();
std::fs::create_dir_all(&target_dir)
.with_context(|| format!("create profraw dir: {}", target_dir.display()))?;
let id = PROFRAW_COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let path = target_dir.join(format!("ktstr-test-{}-{}.profraw", std::process::id(), id));
std::fs::write(&path, data).with_context(|| format!("write profraw: {}", path.display()))?;
Ok(())
}
pub fn target_dir() -> PathBuf {
if let Ok(d) = std::env::var("LLVM_COV_TARGET_DIR") {
return PathBuf::from(d);
}
if let Some(parent) = std::env::var("LLVM_PROFILE_FILE")
.ok()
.as_ref()
.and_then(|p| Path::new(p).parent())
.filter(|p| !p.as_os_str().is_empty())
{
return parent.to_path_buf();
}
let mut p = crate::resolve_current_exe().unwrap_or_else(|_| std::env::temp_dir());
p.pop(); p.push("llvm-cov-target");
p
}
fn redirect_pattern_for(
pid: libc::pid_t,
existing: Option<std::ffi::OsString>,
is_coverage_instrumented: bool,
target_dir: impl FnOnce() -> PathBuf,
) -> Option<PathBuf> {
if pid == 1 || existing.is_some() || !is_coverage_instrumented {
return None;
}
Some(target_dir().join("default-%p-%m.profraw"))
}
fn is_coverage_instrumented_binary() -> bool {
let exe_file = match File::open("/proc/self/exe") {
Ok(f) => f,
Err(_) => return false,
};
let mmap = match unsafe { memmap2::Mmap::map(&exe_file) } {
Ok(m) => m,
Err(_) => return false,
};
let elf = match goblin::elf::Elf::parse(&mmap) {
Ok(e) => e,
Err(_) => return false,
};
let vaddrs = find_symbol_vaddrs(
&elf,
&[
"__llvm_profile_write_buffer",
"__llvm_profile_get_size_for_buffer",
],
);
vaddrs.iter().any(|v| matches!(v, Some(va) if *va != 0))
}
#[doc(hidden)]
pub fn current_binary_is_coverage_instrumented() -> bool {
use std::sync::OnceLock;
static CACHE: OnceLock<bool> = OnceLock::new();
*CACHE.get_or_init(is_coverage_instrumented_binary)
}
ctor::declarative::ctor! {
#[ctor(unsafe, priority = 0)]
fn redirect_default_profraw_path() {
let pid = unsafe { libc::getpid() };
let existing = std::env::var_os("LLVM_PROFILE_FILE");
if pid == 1 || existing.is_some() {
return;
}
let instrumented = current_binary_is_coverage_instrumented();
if let Some(pattern) = redirect_pattern_for(pid, existing, instrumented, target_dir) {
unsafe {
std::env::set_var("LLVM_PROFILE_FILE", &pattern);
}
}
}
}
#[cfg(test)]
mod tests {
use super::super::test_helpers::{EnvVarGuard, lock_env};
use super::*;
#[test]
fn target_dir_with_env_var() {
let _lock = lock_env();
let _env = EnvVarGuard::set("LLVM_COV_TARGET_DIR", "/tmp/my-cov-dir");
let dir = target_dir();
assert_eq!(dir, PathBuf::from("/tmp/my-cov-dir"));
}
#[test]
fn target_dir_from_llvm_profile_file() {
let _lock = lock_env();
let _env_cov = EnvVarGuard::remove("LLVM_COV_TARGET_DIR");
let _env_prof =
EnvVarGuard::set("LLVM_PROFILE_FILE", "/tmp/cov-target/ktstr-%p-%m.profraw");
let dir = target_dir();
assert_eq!(dir, PathBuf::from("/tmp/cov-target"));
}
#[test]
fn target_dir_without_env_var() {
let _lock = lock_env();
let _env_cov = EnvVarGuard::remove("LLVM_COV_TARGET_DIR");
let _env_prof = EnvVarGuard::remove("LLVM_PROFILE_FILE");
let dir = target_dir();
assert!(
dir.ends_with("llvm-cov-target"),
"expected path ending in llvm-cov-target, got: {}",
dir.display()
);
}
#[test]
fn target_dir_bare_filename_llvm_profile_file_falls_through() {
let _lock = lock_env();
let _g_cov = EnvVarGuard::remove("LLVM_COV_TARGET_DIR");
let _g_prof = EnvVarGuard::set("LLVM_PROFILE_FILE", "default.profraw");
let dir = target_dir();
assert!(
!dir.as_os_str().is_empty(),
"bare-filename LLVM_PROFILE_FILE must fall through to the \
current_exe fallback, not return an empty PathBuf",
);
assert!(
dir.ends_with("llvm-cov-target"),
"fallback must land at the current_exe-relative llvm-cov-target \
dir, got: {}",
dir.display(),
);
}
#[test]
fn redirect_pattern_for_pid_1_returns_none() {
let pattern =
redirect_pattern_for(1, None, true, || PathBuf::from("/should/not/be/called"));
assert!(
pattern.is_none(),
"pid=1 (guest init) must skip env redirect"
);
}
#[test]
fn redirect_pattern_for_existing_env_returns_none() {
let pattern = redirect_pattern_for(
42,
Some(std::ffi::OsString::from("/operator/picked/path.profraw")),
true,
|| PathBuf::from("/should/not/be/called"),
);
assert!(
pattern.is_none(),
"existing LLVM_PROFILE_FILE must take precedence"
);
}
#[test]
fn redirect_pattern_for_empty_env_short_circuits() {
let pattern = redirect_pattern_for(42, Some(std::ffi::OsString::new()), true, || {
PathBuf::from("/should/not/be/called")
});
assert!(
pattern.is_none(),
"Some(\"\") in LLVM_PROFILE_FILE counts as set; redirect must defer"
);
}
#[test]
fn redirect_pattern_for_non_instrumented_binary_returns_none() {
let pattern =
redirect_pattern_for(42, None, false, || PathBuf::from("/should/not/be/called"));
assert!(
pattern.is_none(),
"non-coverage-instrumented binary must not pollute the env passed \
to children"
);
}
#[test]
fn redirect_pattern_for_host_unset_returns_target_pattern() {
let target = PathBuf::from("/synthetic/llvm-cov-target");
let pattern = redirect_pattern_for(42, None, true, || target.clone())
.expect("host pid + unset env + instrumented must produce a redirect pattern");
assert_eq!(
pattern,
PathBuf::from("/synthetic/llvm-cov-target/default-%p-%m.profraw"),
);
}
#[test]
fn redirect_pattern_for_filename_matches_cargo_ktstr_wrapper() {
let target = PathBuf::from("/x");
let pattern = redirect_pattern_for(42, None, true, || target.clone()).unwrap();
assert_eq!(
pattern.file_name().and_then(|n| n.to_str()),
Some("default-%p-%m.profraw"),
"filename suffix must match cargo-ktstr's profraw_inject_for",
);
}
#[test]
fn find_symbol_vaddrs_resolves_known_symbol() {
let exe = crate::resolve_current_exe().unwrap();
let data = std::fs::read(&exe).unwrap();
let elf = goblin::elf::Elf::parse(&data).unwrap();
let results = find_symbol_vaddrs(&elf, &["main"]);
assert_eq!(results.len(), 1);
assert!(
results[0].is_some(),
"main symbol should be resolved in test binary"
);
assert_ne!(results[0].unwrap(), 0, "main address should be nonzero");
}
#[test]
fn find_symbol_vaddrs_missing_symbol_returns_none() {
let exe = crate::resolve_current_exe().unwrap();
let data = std::fs::read(&exe).unwrap();
let elf = goblin::elf::Elf::parse(&data).unwrap();
let results = find_symbol_vaddrs(&elf, &["__nonexistent_symbol_xyz__"]);
assert_eq!(results.len(), 1);
assert!(results[0].is_none());
}
#[test]
fn find_symbol_vaddrs_mixed_results() {
let exe = crate::resolve_current_exe().unwrap();
let data = std::fs::read(&exe).unwrap();
let elf = goblin::elf::Elf::parse(&data).unwrap();
let results = find_symbol_vaddrs(&elf, &["main", "__nonexistent_symbol_xyz__"]);
assert_eq!(results.len(), 2);
assert!(results[0].is_some(), "main should resolve");
assert!(results[1].is_none(), "nonexistent should not resolve");
}
#[cfg(coverage)]
#[test]
fn write_buffer_symbol_retained_under_coverage() {
let exe = crate::resolve_current_exe().unwrap();
let data = std::fs::read(&exe).unwrap();
let elf = goblin::elf::Elf::parse(&data).unwrap();
let v = find_symbol_vaddrs(&elf, &["__llvm_profile_write_buffer"]);
assert!(
v[0].is_some(),
"__llvm_profile_write_buffer must survive --gc-sections under \
coverage; without it the guest flush silently no-ops",
);
}
#[test]
fn find_symbol_vaddrs_resolves_zero_size_symbol() {
let exe = crate::resolve_current_exe().unwrap();
let data = std::fs::read(&exe).unwrap();
let elf = goblin::elf::Elf::parse(&data).unwrap();
let zero_size_name = elf
.syms
.iter()
.filter(|s| s.st_size == 0)
.filter_map(|s| elf.strtab.get_at(s.st_name))
.find(|n| !n.is_empty())
.map(str::to_string)
.expect(
"test binary's .symtab should carry at least one named \
zero-size symbol (e.g. a linker marker like _edata / _end)",
);
let v = find_symbol_vaddrs(&elf, &[zero_size_name.as_str()]);
assert!(
v[0].is_some(),
"find_symbol_vaddrs must resolve zero-size symbol \
{zero_size_name:?}; the removed st_size==0 filter previously \
dropped such symbols, losing gc-sections'd coverage markers",
);
}
}