Skip to content

Commit 42089bb

Browse files
committed
Split out elf and stack crates from xtask
I intend to do some rework of the max-stack analysis code, and would like to break out the max-stack analysis into its own crate to make testing this in isolation easier. As this code also depends on the `xtask::elf` module, I broke that out to avoid circular deps. This is the minimum change necessary to split things out before any refactoring.
1 parent 2f318dc commit 42089bb

11 files changed

Lines changed: 337 additions & 280 deletions

File tree

Cargo.lock

Lines changed: 23 additions & 3 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

build/elf/Cargo.toml

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
[package]
2+
name = "elf"
3+
version = "0.1.0"
4+
edition = "2024"
5+
6+
[dependencies]
7+
anyhow = { workspace = true }
8+
goblin = { workspace = true }
9+
scroll = { workspace = true }
10+
11+
[lints]
12+
workspace = true

build/stack/Cargo.toml

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
[package]
2+
name = "stack"
3+
version = "0.1.0"
4+
edition = "2024"
5+
6+
[dependencies]
7+
anyhow = { workspace = true }
8+
goblin = { workspace = true }
9+
capstone = { workspace = true }
10+
rustc-demangle = { workspace = true }
11+
leb128 = { workspace = true }
12+
13+
elf = { path = "../elf" }
14+
15+
[lints]
16+
workspace = true

build/stack/src/lib.rs

Lines changed: 269 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,269 @@
1+
// This Source Code Form is subject to the terms of the Mozilla Public
2+
// License, v. 2.0. If a copy of the MPL was not distributed with this
3+
// file, You can obtain one at https://mozilla.org/MPL/2.0/.
4+
5+
use std::{
6+
collections::{BTreeMap, BTreeSet},
7+
path::Path,
8+
};
9+
10+
use anyhow::{Context, Result, anyhow, bail};
11+
12+
/// Estimates the maximum stack size for the given task
13+
///
14+
/// This does not take dynamic function calls into account, which could cause
15+
/// underestimation. Overestimation is less likely, but still may happen if
16+
/// there are logically impossible call trees (e.g. `A -> B` and `B -> C`, but
17+
/// `B` never calls `C` if called by `A`).
18+
pub fn get_max_stack(
19+
elf: &Path,
20+
task_name: &str,
21+
verbose: bool,
22+
) -> Result<Vec<(u64, String)>> {
23+
// Open the statically-linked ELF file
24+
let data = std::fs::read(elf).context("could not open ELF file")?;
25+
let elf = goblin::elf::Elf::parse(&data)?;
26+
27+
// Read the .stack_sizes section, which is an array of
28+
// `(address: u32, stack size: unsigned leb128)` tuples
29+
let sizes = elf::get_section_by_name(&elf, ".stack_sizes")
30+
.context("could not get .stack_sizes")?;
31+
let mut sizes = &data[sizes.sh_offset as usize..][..sizes.sh_size as usize];
32+
let mut addr_to_frame_size = BTreeMap::new();
33+
while !sizes.is_empty() {
34+
let (addr, rest) = sizes.split_at(4);
35+
let addr = u32::from_le_bytes(addr.try_into().unwrap());
36+
sizes = rest;
37+
let size = leb128::read::unsigned(&mut sizes)?;
38+
addr_to_frame_size.insert(addr, size);
39+
}
40+
41+
// There are `$t` and `$d` symbols which indicate the beginning of text
42+
// versus data in the `.text` region. We collect them into a `BTreeMap`
43+
// here so that we can avoid trying to decode inline data words.
44+
let mut text_regions = BTreeMap::new();
45+
for sym in elf.syms.iter() {
46+
if sym.st_name == 0
47+
|| sym.st_size != 0
48+
|| sym.st_type() != goblin::elf::sym::STT_NOTYPE
49+
{
50+
continue;
51+
}
52+
53+
let addr = sym.st_value as u32;
54+
let is_text = match elf.strtab.get_at(sym.st_name) {
55+
Some("$t") => true,
56+
Some("$d") => false,
57+
Some(_) => continue,
58+
None => {
59+
bail!("bad symbol in {task_name}: {}", sym.st_name);
60+
}
61+
};
62+
text_regions.insert(addr, is_text);
63+
}
64+
let is_code = |addr| {
65+
let mut iter = text_regions.range(..=addr);
66+
*iter.next_back().unwrap().1
67+
};
68+
69+
// We'll be packing everything into this data structure
70+
#[derive(Debug)]
71+
struct FunctionData {
72+
name: String,
73+
short_name: String,
74+
frame_size: Option<u64>,
75+
calls: BTreeSet<u32>,
76+
}
77+
78+
let text = elf::get_section_by_name(&elf, ".text")
79+
.context("could not get .text")?;
80+
81+
use capstone::{
82+
Capstone, InsnGroupId, InsnGroupType,
83+
arch::{ArchOperand, BuildsCapstone, BuildsCapstoneExtraMode, arm},
84+
};
85+
let cs = Capstone::new()
86+
.arm()
87+
.mode(arm::ArchMode::Thumb)
88+
.extra_mode(std::iter::once(arm::ArchExtraMode::MClass))
89+
.detail(true)
90+
.build()
91+
.map_err(|e| anyhow!("failed to initialize disassembler: {e:?}"))?;
92+
93+
// Disassemble each function, building a map of its call sites
94+
let mut fns = BTreeMap::new();
95+
for sym in elf.syms.iter() {
96+
// We only care about named function symbols here
97+
if sym.st_name == 0 || !sym.is_function() || sym.st_size == 0 {
98+
continue;
99+
}
100+
101+
let Some(name) = elf.strtab.get_at(sym.st_name) else {
102+
bail!("bad symbol in {task_name}: {}", sym.st_name);
103+
};
104+
105+
// Clear the lowest bit, which indicates that the function contains
106+
// thumb instructions (always true for our systems!)
107+
let val = sym.st_value & !1;
108+
let base_addr = val as u32;
109+
110+
// Get the text region for this function
111+
let offset = (val - text.sh_addr + text.sh_offset) as usize;
112+
let text = &data[offset..][..sym.st_size as usize];
113+
114+
// Split the text region into instruction-only chunks
115+
let mut chunks = vec![];
116+
let mut chunk = None;
117+
for (i, b) in text.iter().enumerate() {
118+
let addr = base_addr + i as u32;
119+
if is_code(addr) {
120+
chunk.get_or_insert((addr, vec![])).1.push(*b);
121+
} else {
122+
chunks.extend(chunk.take());
123+
}
124+
}
125+
chunks.extend(chunk); // don't forget the trailing chunk!
126+
127+
let frame_size = addr_to_frame_size.get(&base_addr).copied();
128+
let mut calls = BTreeSet::new();
129+
for (addr, chunk) in chunks {
130+
let instrs = cs
131+
.disasm_all(&chunk, addr.into())
132+
.map_err(|e| anyhow!("disassembly failed: {e:?}"))?;
133+
for (i, instr) in instrs.iter().enumerate() {
134+
let detail = cs.insn_detail(instr).map_err(|e| {
135+
anyhow!("could not get instruction details: {e}")
136+
})?;
137+
138+
// Detect tail calls, which are jumps at the final instruction
139+
// when the function itself has no stack frame.
140+
let can_tail = frame_size == Some(0) && i == instrs.len() - 1;
141+
if detail.groups().iter().any(|g| {
142+
g == &InsnGroupId(InsnGroupType::CS_GRP_CALL as u8)
143+
|| (g == &InsnGroupId(InsnGroupType::CS_GRP_JUMP as u8)
144+
&& can_tail)
145+
}) {
146+
let arch = detail.arch_detail();
147+
let ops = arch.operands();
148+
let op = ops.last().unwrap_or_else(|| {
149+
panic!("missing operand!");
150+
});
151+
152+
let ArchOperand::ArmOperand(op) = op else {
153+
panic!("bad operand type: {op:?}");
154+
};
155+
// We can't resolve indirect calls, alas
156+
let arm::ArmOperandType::Imm(target) = op.op_type else {
157+
continue;
158+
};
159+
let target = u32::try_from(target).unwrap();
160+
161+
// Avoid recursive calls into the same function (or midway
162+
// into the function, which is a thing we've seen before!
163+
// it's weird!)
164+
if !(base_addr..base_addr + sym.st_size as u32)
165+
.contains(&target)
166+
{
167+
calls.insert(target);
168+
}
169+
}
170+
}
171+
}
172+
173+
let name = rustc_demangle::demangle(name).to_string();
174+
175+
// Strip the trailing hash from the name for ease of printing
176+
let short_name = if let Some(i) = name.rfind("::") {
177+
&name[..i]
178+
} else {
179+
&name
180+
}
181+
.to_owned();
182+
183+
fns.insert(
184+
base_addr,
185+
FunctionData {
186+
name,
187+
short_name,
188+
frame_size,
189+
calls,
190+
},
191+
);
192+
}
193+
194+
fn recurse(
195+
call_stack: &mut Vec<u32>,
196+
recurse_depth: usize,
197+
mut stack_depth: u64,
198+
fns: &BTreeMap<u32, FunctionData>,
199+
deepest: &mut Option<(u64, Vec<u32>)>,
200+
verbose: bool,
201+
) {
202+
let addr = *call_stack.last().unwrap();
203+
let Some(f) = fns.get(&addr) else {
204+
panic!("found jump to unknown function at {call_stack:08x?}");
205+
};
206+
let frame_size = f.frame_size.unwrap_or(0);
207+
stack_depth += frame_size;
208+
if verbose {
209+
let indent = recurse_depth * 2;
210+
println!(
211+
" {:indent$}{addr:08x}: {} [+{frame_size} => {stack_depth}]",
212+
"",
213+
f.short_name,
214+
indent = indent
215+
);
216+
}
217+
218+
if deepest
219+
.as_ref()
220+
.map(|(max_depth, _)| stack_depth > *max_depth)
221+
.unwrap_or(true)
222+
{
223+
*deepest = Some((stack_depth, call_stack.to_owned()));
224+
}
225+
for j in &f.calls {
226+
if call_stack.contains(j) {
227+
// Skip recursive / mutually recursive calls, because we can't
228+
// reason about them.
229+
continue;
230+
} else {
231+
call_stack.push(*j);
232+
recurse(
233+
call_stack,
234+
recurse_depth + 1,
235+
stack_depth,
236+
fns,
237+
deepest,
238+
verbose,
239+
);
240+
call_stack.pop();
241+
}
242+
}
243+
}
244+
245+
// Find stack sizes by traversing the graph
246+
if verbose {
247+
println!("finding stack sizes for {task_name}");
248+
}
249+
let start_addr = fns
250+
.iter()
251+
.find(|(_addr, v)| v.name.as_str() == "_start")
252+
.map(|(addr, _v)| *addr)
253+
.ok_or_else(|| anyhow!("could not find _start"))?;
254+
let mut deepest = None;
255+
recurse(&mut vec![start_addr], 0, 0, &fns, &mut deepest, verbose);
256+
257+
// Check against our configured task stack size
258+
let Some((_max_depth, max_stack)) = deepest else {
259+
unreachable!("must have at least one call stack");
260+
};
261+
262+
let mut out = vec![];
263+
for m in max_stack {
264+
let f = fns.get(&m).unwrap();
265+
let name = &f.short_name;
266+
out.push((f.frame_size.unwrap_or(0), name.clone()));
267+
}
268+
Ok(out)
269+
}

0 commit comments

Comments
 (0)