use goblin::Object;
use std::collections::HashMap;
#[derive(Debug, Clone)]
pub struct ScratchLeakHit {
pub func_va: u64,
pub alloc_calls: u32,
pub returns_none: bool,
}
fn va_to_file_offset(obj: &Object<'_>, va: u64) -> Option<usize> {
if let Object::PE(pe) = obj {
for sec in &pe.sections {
let svaddr = pe.image_base as u64 + sec.virtual_address as u64;
let sv = sec.virtual_size as u64;
if va >= svaddr && va < svaddr + sv {
let raddr = sec.pointer_to_raw_data as usize;
let rsize = sec.size_of_raw_data as usize;
let off_in_section = (va - svaddr) as usize;
if off_in_section < rsize {
return Some(raddr + off_in_section);
}
}
}
}
None
}
const ALLOCATOR_NAMES: &[&str] = &[
"PyMem_Malloc",
"PyMem_RawMalloc",
"PyMem_Calloc",
"PyObject_Malloc",
"_PyObject_New",
"_PyObject_NewVar",
"PyTuple_New",
"PyDict_New",
"PyList_New",
"PyBytes_FromStringAndSize",
"PyByteArray_FromStringAndSize",
"malloc",
"calloc",
"HeapAlloc",
"VirtualAlloc",
"LocalAlloc",
];
const NONE_STRUCT_NAMES: &[&str] = &["_Py_NoneStruct"];
pub fn check_function(
obj: &Object<'_>,
data: &[u8],
iat: &HashMap<u64, String>,
func_va: u64,
body_max: usize,
) -> Option<ScratchLeakHit> {
let off = va_to_file_offset(obj, func_va)?;
let scan_len = body_max.min(data.len() - off);
let body = &data[off..off + scan_len];
let mut alloc_calls = 0u32;
let mut none_via_iat = false;
let mut none_via_data = false;
let mut k = 0;
while k + 6 <= body.len() {
if body[k] == 0xff && body[k + 1] == 0x15 {
let d32 = i32::from_le_bytes([body[k + 2], body[k + 3], body[k + 4], body[k + 5]]);
let next_rip = func_va.wrapping_add((k + 6) as u64);
let target = next_rip.wrapping_add(d32 as i64 as u64);
if let Some(name) = iat.get(&target) {
if ALLOCATOR_NAMES.iter().any(|n| name == n) {
alloc_calls += 1;
}
}
k += 6;
continue;
}
if k + 7 <= body.len() && body[k] == 0x48 && body[k + 1] == 0x8b && body[k + 2] == 0x05 {
let d32 = i32::from_le_bytes([body[k + 3], body[k + 4], body[k + 5], body[k + 6]]);
let next_rip = func_va.wrapping_add((k + 7) as u64);
let target = next_rip.wrapping_add(d32 as i64 as u64);
if let Some(name) = iat.get(&target) {
if NONE_STRUCT_NAMES.iter().any(|n| name == n) {
none_via_iat = true;
}
}
k += 7;
continue;
}
if body[k] == 0xc3 {
k += 1;
continue;
}
if k + 2 <= body.len() && body[k] == 0xff && body[k + 1] == 0x00 {
none_via_data = true;
}
k += 1;
}
if alloc_calls >= 1 && (none_via_iat || none_via_data) {
Some(ScratchLeakHit {
func_va,
alloc_calls,
returns_none: none_via_iat || none_via_data,
})
} else {
None
}
}
pub fn scan_functions(
obj: &Object<'_>,
data: &[u8],
iat: &HashMap<u64, String>,
func_vas: &[u64],
body_max: usize,
) -> Vec<ScratchLeakHit> {
func_vas
.iter()
.filter_map(|&va| check_function(obj, data, iat, va, body_max))
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn render_includes_va_and_counts() {
let hits = vec![ScratchLeakHit {
func_va: 0x1800_1234,
alloc_calls: 2,
returns_none: true,
}];
let out = render(&hits);
assert_eq!(out.len(), 1);
assert!(out[0].contains("0x18001234"));
assert!(out[0].contains("alloc_calls=2"));
assert!(out[0].contains("returns_None=true"));
}
#[test]
fn allocator_list_includes_pyvmprotect_targets() {
for n in &[
"PyMem_Malloc",
"_PyObject_New",
"PyTuple_New",
"PyDict_New",
"PyList_New",
"HeapAlloc",
] {
assert!(
ALLOCATOR_NAMES.iter().any(|a| a == n),
"missing allocator: {}",
n
);
}
assert!(NONE_STRUCT_NAMES.iter().any(|a| *a == "_Py_NoneStruct"));
}
#[test]
fn empty_list_yields_no_output() {
let out = render(&[]);
assert!(out.is_empty());
}
}
pub fn render(hits: &[ScratchLeakHit]) -> Vec<String> {
hits.iter()
.map(|h| {
format!(
"{:#x}: alloc_calls={} returns_None={} — possible scratch leak",
h.func_va, h.alloc_calls, h.returns_none
)
})
.collect()
}