|
| 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