Skip to main content

sbpf_coverage/
util.rs

1use std::{
2    env::current_dir,
3    ffi::OsStr,
4    path::{Path, PathBuf},
5    process::{Command, Stdio},
6};
7
8use crate::Loader;
9use crate::addr2line::gimli::{self, DW_AT_language, DW_AT_producer, DwAt, DwTag, LittleEndian};
10use crate::{Object, ObjectSection};
11use anyhow::anyhow;
12use sha2::{Digest, Sha256};
13
14pub trait StripCurrentDir {
15    fn strip_current_dir(&self) -> &Self;
16}
17
18impl StripCurrentDir for Path {
19    fn strip_current_dir(&self) -> &Self {
20        let Ok(current_dir) = current_dir() else {
21            return self;
22        };
23        self.strip_prefix(current_dir).unwrap_or(self)
24    }
25}
26
27pub fn find_files_with_extension(dirs: &[PathBuf], extension: &str) -> Vec<PathBuf> {
28    let mut so_files = Vec::new();
29
30    for dir in dirs {
31        if dir.is_dir()
32            && let Ok(entries) = std::fs::read_dir(dir)
33        {
34            for entry in entries.flatten() {
35                let path = entry.path();
36                if path.is_file() && path.extension().is_some_and(|ext| ext == extension) {
37                    so_files.push(path);
38                }
39            }
40        }
41    }
42
43    so_files
44}
45
46pub fn compute_hash(slice: &[u8]) -> String {
47    hex::encode(Sha256::digest(slice).as_slice())
48}
49
50pub fn get_section_start_address(loader: &Loader, section: &str) -> anyhow::Result<u64> {
51    Ok(loader
52        .get_section_range(section.as_bytes())
53        .ok_or(anyhow!("Can't get {} section begin address", section))?
54        .begin)
55}
56
57pub fn get_dwarf_attribute(
58    object: &object::File,
59    tag: DwTag,
60    attribute: DwAt,
61) -> anyhow::Result<String> {
62    let load_section = |id: gimli::SectionId| -> Result<_, LittleEndian> {
63        let data = object
64            .section_by_name(id.name())
65            .map(|s| s.data().unwrap_or(&[]))
66            .unwrap_or(&[]);
67        Ok(gimli::EndianSlice::new(data, LittleEndian))
68    };
69
70    let dwarf = addr2line::gimli::Dwarf::load(&load_section)
71        .map_err(|_| anyhow!("Failed to load DWARF sections"))?;
72    let mut iter = dwarf.units();
73    while let Ok(Some(header)) = iter.next() {
74        let Ok(unit) = dwarf.unit(header) else {
75            continue;
76        };
77        let mut entries = unit.entries();
78        while let Ok(Some(entry)) = entries.next_dfs() {
79            if let Some(val) = entry.attr_value(attribute)
80                && entry.tag() == tag
81            {
82                match attribute {
83                    a if a == DW_AT_producer => {
84                        if let Ok(s) = dwarf.attr_string(&unit, val) {
85                            return Ok(s.to_string_lossy().to_string());
86                        }
87                    }
88                    a if a == DW_AT_language => {
89                        if let gimli::AttributeValue::Language(lang) = val {
90                            return Ok(lang.to_string());
91                        }
92                    }
93                    _ => continue,
94                }
95            }
96        }
97    }
98    Err(anyhow!(
99        "No DWARF entry found for {:?} with attribute {:?}",
100        tag,
101        attribute
102    ))
103}
104
105pub fn execute_cmd<I, S>(program: &Path, args: I) -> Option<String>
106where
107    I: IntoIterator<Item = S>,
108    S: AsRef<OsStr>,
109{
110    let child = Command::new(program)
111        .args(args)
112        .stdout(Stdio::piped())
113        .spawn()
114        .map_err(|e| {
115            eprintln!("Failed to execute {}: {}", program.display(), e);
116        })
117        .ok()?;
118
119    let output = child
120        .wait_with_output()
121        .map_err(|e| eprintln!("failed to wait on child: {}", e))
122        .ok()?;
123    Some(
124        output
125            .stdout
126            .as_slice()
127            .iter()
128            .map(|&c| c as char)
129            .collect::<String>()
130            .trim()
131            .into(),
132    )
133}
134
135/// Returns the nth line from the given string, or empty if out of bounds.
136pub fn read_nth_line(file_content: &str, line_number: usize) -> String {
137    file_content
138        .lines()
139        .nth(line_number)
140        .unwrap_or("")
141        .to_string()
142}