use std::path::Path;
#[derive(Debug)]
pub(super) enum Probe {
Literal(String),
Unauditable { file: String },
}
pub(super) fn collect_probes(dir: &Path, probes: &mut Vec<Probe>) -> Result<(), String> {
let read = std::fs::read_dir(dir).map_err(|e| format!("cannot read {}: {e}", dir.display()))?;
let mut paths = Vec::new();
for entry in read {
let entry =
entry.map_err(|e| format!("cannot read a dir entry under {}: {e}", dir.display()))?;
let file_type = entry
.file_type()
.map_err(|e| format!("cannot stat {}: {e}", entry.path().display()))?;
paths.push((file_type.is_dir(), entry.path()));
}
paths.sort();
for (is_dir, path) in paths {
if is_dir {
collect_probes(&path, probes)?;
} else if path.extension().and_then(|e| e.to_str()) == Some("rs") {
let source = std::fs::read_to_string(&path)
.map_err(|e| format!("cannot read source {}: {e}", path.display()))?;
scan_source(&source, &path.display().to_string(), probes);
}
}
Ok(())
}
fn skip_block_comment(b: &[u8], mut i: usize) -> usize {
let mut depth = 1usize;
i += 2; while i + 1 < b.len() && depth > 0 {
if b[i] == b'/' && b[i + 1] == b'*' {
depth += 1;
i += 2;
} else if b[i] == b'*' && b[i + 1] == b'/' {
depth -= 1;
i += 2;
} else {
i += 1;
}
}
if depth > 0 { b.len() } else { i }
}
pub(super) fn scan_source(source: &str, file: &str, probes: &mut Vec<Probe>) {
let b = source.as_bytes();
let mut i = 0;
while i < b.len() {
if let Some(next) = skip_literal_or_comment(b, i) {
i = next;
continue;
}
let left_boundary = i == 0 || !is_ident_byte(b[i - 1]);
if left_boundary {
if let Some(rest) = match_probe_marker(b, i) {
let (probe, next) = capture_probe(b, rest, file);
if let Some(probe) = probe {
probes.push(probe);
}
i = next;
continue;
}
}
if b[i] == b'!' {
let mut name_end = i;
while name_end > 0 && b[name_end - 1].is_ascii_whitespace() {
name_end -= 1;
}
let mut name_start = name_end;
while name_start > 0 && is_ident_byte(b[name_start - 1]) {
name_start -= 1;
}
let is_raw_ident = name_start >= 2
&& b[name_start - 1] == b'#'
&& b[name_start - 2] == b'r'
&& (name_start == 2 || !is_ident_byte(b[name_start - 3]));
if name_start < name_end && (is_raw_ident || !is_rust_keyword(&b[name_start..name_end]))
{
if let Some(end) = foreign_macro_body_end(b, i) {
i = end;
continue;
}
}
}
i += 1;
}
}
fn skip_literal_or_comment(b: &[u8], i: usize) -> Option<usize> {
if b[i] == b'/' && i + 1 < b.len() && b[i + 1] == b'/' {
let mut j = i;
while j < b.len() && b[j] != b'\n' {
j += 1;
}
return Some(j);
}
if b[i] == b'/' && i + 1 < b.len() && b[i + 1] == b'*' {
return Some(skip_block_comment(b, i));
}
if let Some(end) = raw_or_byte_string_end(b, i) {
return Some(end);
}
if b[i] == b'"' {
let mut j = i + 1;
while j < b.len() && b[j] != b'"' {
if b[j] == b'\\' {
j += 1;
}
j += 1;
}
return Some((j + 1).min(b.len()));
}
if b[i] == b'\'' {
let is_char =
(i + 1 < b.len() && b[i + 1] == b'\\') || (i + 2 < b.len() && b[i + 2] == b'\'');
if is_char {
let mut j = i + 1;
while j < b.len() && b[j] != b'\'' {
if b[j] == b'\\' {
j += 1;
}
j += 1;
}
return Some((j + 1).min(b.len()));
}
}
None
}
fn preceding_ident_is(b: &[u8], end: usize, target: &[u8]) -> bool {
let mut start = end;
while start > 0 && is_ident_byte(b[start - 1]) {
start -= 1;
}
&b[start..end] == target
}
fn foreign_macro_body_end(b: &[u8], bang: usize) -> Option<usize> {
let mut i = skip_trivia(b, bang + 1);
let mut name_end = bang;
while name_end > 0 && b[name_end - 1].is_ascii_whitespace() {
name_end -= 1;
}
if preceding_ident_is(b, name_end, b"macro_rules") {
let name_start = i;
while i < b.len() && (is_ident_byte(b[i]) || b[i] == b'#') {
i += 1;
}
if i == name_start {
return None; }
i = skip_trivia(b, i);
}
if !matches!(b.get(i), Some(b'{') | Some(b'(') | Some(b'[')) {
return None;
}
let mut depth = 0usize;
while i < b.len() {
if let Some(next) = skip_literal_or_comment(b, i) {
i = next;
continue;
}
match b[i] {
b'{' | b'(' | b'[' => depth += 1,
b'}' | b')' | b']' => {
depth = depth.saturating_sub(1);
if depth == 0 {
return Some(i + 1);
}
}
_ => {}
}
i += 1;
}
Some(b.len())
}
fn raw_or_byte_string_end(b: &[u8], i: usize) -> Option<usize> {
let mut j = i;
let byte = j < b.len() && b[j] == b'b';
if byte {
j += 1;
}
let raw = j < b.len() && b[j] == b'r';
if raw {
j += 1;
let mut hashes = 0;
while j < b.len() && b[j] == b'#' {
hashes += 1;
j += 1;
}
if j >= b.len() || b[j] != b'"' {
return None;
}
j += 1;
while j < b.len() {
if b[j] == b'"' {
let mut k = j + 1;
let mut h = 0;
while k < b.len() && h < hashes && b[k] == b'#' {
k += 1;
h += 1;
}
if h == hashes {
return Some(k);
}
}
j += 1;
}
return Some(b.len());
}
if byte && j < b.len() && b[j] == b'"' {
j += 1;
while j < b.len() && b[j] != b'"' {
if b[j] == b'\\' {
j += 1;
}
j += 1;
}
return Some((j + 1).min(b.len()));
}
None
}
fn match_probe_marker(b: &[u8], i: usize) -> Option<usize> {
const NAME: &[u8] = b"assert_boundary";
if i + NAME.len() > b.len() || &b[i..i + NAME.len()] != NAME {
return None;
}
let after_name = i + NAME.len();
if b.get(after_name).is_some_and(|&c| is_ident_byte(c)) {
return None;
}
let bang = skip_trivia(b, after_name);
if b.get(bang) != Some(&b'!') {
return None;
}
Some(bang + 1)
}
fn is_ident_byte(byte: u8) -> bool {
byte.is_ascii_alphanumeric() || byte == b'_' || byte >= 0x80
}
fn is_rust_keyword(word: &[u8]) -> bool {
let Ok(word) = std::str::from_utf8(word) else {
return false;
};
matches!(
word,
"as" | "break"
| "const"
| "continue"
| "crate"
| "dyn"
| "else"
| "enum"
| "extern"
| "false"
| "fn"
| "for"
| "if"
| "impl"
| "in"
| "let"
| "loop"
| "match"
| "mod"
| "move"
| "mut"
| "pub"
| "ref"
| "return"
| "self"
| "Self"
| "static"
| "struct"
| "super"
| "trait"
| "true"
| "type"
| "unsafe"
| "use"
| "where"
| "while"
| "async"
| "await"
| "abstract"
| "become"
| "box"
| "do"
| "final"
| "macro"
| "override"
| "priv"
| "typeof"
| "unsized"
| "virtual"
| "yield"
| "try"
| "gen"
)
}
fn skip_trivia(b: &[u8], mut i: usize) -> usize {
loop {
while i < b.len() && b[i].is_ascii_whitespace() {
i += 1;
}
if b.get(i) == Some(&b'/') && b.get(i + 1) == Some(&b'/') {
while i < b.len() && b[i] != b'\n' {
i += 1;
}
continue;
}
if b.get(i) == Some(&b'/') && b.get(i + 1) == Some(&b'*') {
i = skip_block_comment(b, i);
continue;
}
return i;
}
}
fn capture_probe(b: &[u8], i: usize, file: &str) -> (Option<Probe>, usize) {
let i = skip_trivia(b, i);
if !matches!(b.get(i), Some(&b'(') | Some(&b'{') | Some(&b'[')) {
return (None, i);
}
let i = skip_trivia(b, i + 1);
if i >= b.len() {
return (None, i);
}
if b[i] == b'r' && matches!(b.get(i + 1), Some(b'"') | Some(b'#')) {
if let Some((seam, next)) = raw_string_value(b, i) {
return (Some(Probe::Literal(seam)), next);
}
return (
Some(Probe::Unauditable {
file: file.to_string(),
}),
i,
);
}
if b[i] == b'"' {
let mut j = i + 1;
let start = j;
while j < b.len() && b[j] != b'"' {
if b[j] == b'\\' {
j += 1;
}
j += 1;
}
if j >= b.len() {
return (None, j);
}
return match decode_str_escapes(&b[start..j]) {
Some(seam) => (Some(Probe::Literal(seam)), j + 1),
None => (
Some(Probe::Unauditable {
file: file.to_string(),
}),
j + 1,
),
};
}
(
Some(Probe::Unauditable {
file: file.to_string(),
}),
i,
)
}
fn raw_string_value(b: &[u8], i: usize) -> Option<(String, usize)> {
let mut j = i + 1; let mut hashes = 0;
while b.get(j) == Some(&b'#') {
hashes += 1;
j += 1;
}
if b.get(j) != Some(&b'"') {
return None;
}
j += 1;
let start = j;
while j < b.len() {
if b[j] == b'"' {
let mut k = j + 1;
let mut h = 0;
while h < hashes && b.get(k) == Some(&b'#') {
k += 1;
h += 1;
}
if h == hashes {
return Some((String::from_utf8_lossy(&b[start..j]).into_owned(), k));
}
}
j += 1;
}
None
}
fn decode_str_escapes(inner: &[u8]) -> Option<String> {
let s = std::str::from_utf8(inner).ok()?;
let mut out = String::with_capacity(s.len());
let mut chars = s.chars();
while let Some(c) = chars.next() {
if c != '\\' {
out.push(c);
continue;
}
match chars.next()? {
'n' => out.push('\n'),
'r' => out.push('\r'),
't' => out.push('\t'),
'\\' => out.push('\\'),
'0' => out.push('\0'),
'\'' => out.push('\''),
'"' => out.push('"'),
'x' => {
let hi = chars.next()?.to_digit(16)?;
let lo = chars.next()?.to_digit(16)?;
let v = hi * 16 + lo;
if v > 0x7F {
return None;
}
out.push(char::from_u32(v)?);
}
'u' => {
if chars.next()? != '{' {
return None;
}
let mut value: u32 = 0;
let mut digits = 0;
loop {
match chars.next()? {
'}' => break,
'_' if digits == 0 => return None,
'_' => continue,
d => {
let hd = d.to_digit(16)?;
digits += 1;
if digits > 6 {
return None;
}
value = value * 16 + hd;
}
}
}
if digits == 0 {
return None;
}
out.push(char::from_u32(value)?);
}
_ => return None,
}
}
Some(out)
}