17use std::collections::{HashMap, HashSet};
18use std::io::{Cursor, Read};
19use std::path::{Path, PathBuf};
20use std::process::Command;
21use std::{env, fs, process};

23use object::write::{SectionId, StandardSection, Symbol, SymbolId, SymbolSection};
24use object::{Endianness, Object, ObjectSection, ObjectSymbol, SymbolFlags, SymbolKind, SymbolScope};
25use subsecond_types::{AddressMap, JumpTable};
26
27struct Args {
28    base: PathBuf,
29    objects: Vec<PathBuf>,
30    aslr_reference: u64,
31    out_dir: PathBuf,
32}
33
34fn usage() -> ! {
35    eprintln!(
36        "usage: hotpatch --base <exe> --objects <.o/.rlib>... --aslr-reference <addr> --out-dir <dir>"
37    );
38    process::exit(2);
39}
40
41fn parse_args() -> Args {
42    let mut base = None;
43    let mut objects = Vec::new();
44    let mut aslr_reference = None;
45    let mut out_dir = None;
46
47    let mut args = env::args().skip(1).peekable();
48    while let Some(arg) = args.next() {
49        match arg.as_str() {
50            "--base" => base = Some(PathBuf::from(args.next().unwrap_or_else(|| usage()))),
51            "--aslr-reference" => {
52                let raw = args.next().unwrap_or_else(|| usage());
53                aslr_reference = Some(parse_u64(&raw));
54            }
55            "--out-dir" => out_dir = Some(PathBuf::from(args.next().unwrap_or_else(|| usage()))),
56            "--objects" => {
57                while let Some(next) = args.peek() {
58                    if next.starts_with("--") {
59                        break;
60                    }
61                    objects.push(PathBuf::from(args.next().unwrap()));
62                }
63            }
64            _ => usage(),
65        }
66    }
67
68    if objects.is_empty() {
69        eprintln!("hotpatch: --objects needs at least one path");
70        usage();
71    }
72
73    Args {
74        base: base.unwrap_or_else(|| usage()),
75        objects,
76        aslr_reference: aslr_reference.unwrap_or_else(|| usage()),
77        out_dir: out_dir.unwrap_or_else(|| usage()),
78    }
79}
80
81fn parse_u64(s: &str) -> u64 {
82    if let Some(hex) = s.strip_prefix("0x").or_else(|| s.strip_prefix("0X")) {
83        u64::from_str_radix(hex, 16).unwrap_or_else(|e| {
84            eprintln!("hotpatch: bad --aslr-reference {s:?}: {e}");
85            process::exit(2);
86        })
87    } else {
88        s.parse().unwrap_or_else(|e| {
89            eprintln!("hotpatch: bad --aslr-reference {s:?}: {e}");
90            process::exit(2);
91        })
92    }
93}

98struct CachedSymbol {
99    address: u64,

104    size: u64,
105    kind: SymbolKind,
106    is_undefined: bool,
107    is_weak: bool,
108    flags: SymbolFlags<SectionId, SymbolId>,
109}

116fn build_symbol_cache(bytes: &[u8]) -> HashMap<String, CachedSymbol> {
117    let obj = object::read::File::parse(bytes).expect("parsing base binary as an object file");
118    obj.symbols()
119        .filter_map(|s| {
120            let flags = match s.flags() {
121                SymbolFlags::None => SymbolFlags::None,
122                SymbolFlags::Elf { st_info, st_other } => SymbolFlags::Elf { st_info, st_other },
123                _ => SymbolFlags::None,
124            };
125            Some((
126                s.name().ok()?.to_string(),
127                CachedSymbol {
128                    address: s.address(),
129                    size: s.size(),
130                    kind: s.kind(),
131                    is_undefined: s.is_undefined(),
132                    is_weak: s.is_weak(),
133                    flags,
134                },
135            ))
136        })
137        .collect()
138}

146fn tdata(bytes: &[u8]) -> Vec<u8> {
147    let obj = object::read::File::parse(bytes).expect("parsing base binary as an object file");
148    obj.section_by_name(".tdata").and_then(|s| s.data().ok()).unwrap_or(&[]).to_vec()
149}

160fn collect_symbols(
161    paths: &[PathBuf],
162    extract_dir: &Path,
163) -> (HashSet<String>, HashSet<String>, Vec<PathBuf>) {
164    let mut undefined = HashSet::new();
165    let mut defined = HashSet::new();
166    let mut link_objects = Vec::new();
167
168    for path in paths {
169        let bytes = fs::read(path).unwrap_or_else(|e| {
170            eprintln!("hotpatch: reading {}: {e}", path.display());
171            process::exit(1);
172        });
173
174        let is_archive = matches!(
175            path.extension().and_then(|e| e.to_str()),
176            Some("rlib" | "a")
177        );
178
179        if !is_archive {
180            collect_symbols_from_bytes(&bytes, path, &mut undefined, &mut defined);
181            link_objects.push(path.clone());
182            continue;
183        }
184
185        let stem = path
186            .file_stem()
187            .map(|s| s.to_string_lossy().into_owned())
188            .unwrap_or_else(|| "member".to_string());
189        let mut archive = ar::Archive::new(Cursor::new(bytes));
190        let mut member_index = 0usize;
191        while let Some(entry) = archive.next_entry() {
192            let mut entry = entry.unwrap_or_else(|e| {
193                eprintln!("hotpatch: reading archive member of {}: {e}", path.display());
194                process::exit(1);
195            });
196            let name = std::str::from_utf8(entry.header().identifier())
197                .unwrap_or_default()
198                .to_string();
199            if name.ends_with(".rmeta") || !(name.ends_with(".o") || name.ends_with(".obj")) {
200                continue;
201            }
202
203            let mut member_bytes = Vec::with_capacity(entry.header().size() as usize);
204            entry
205                .read_to_end(&mut member_bytes)
206                .unwrap_or_else(|e| panic!("reading archive member {name}: {e}"));
207            collect_symbols_from_bytes(&member_bytes, path, &mut undefined, &mut defined);
208
209            member_index += 1;
210            let out_path = extract_dir.join(format!("{stem}.{member_index}.o"));
211            fs::write(&out_path, &member_bytes)
212                .unwrap_or_else(|e| panic!("writing extracted member {}: {e}", out_path.display()));
213            link_objects.push(out_path);
214        }
215    }
216
217    (undefined, defined, link_objects)
218}

220fn collect_symbols_from_bytes(
221    bytes: &[u8],
222    path: &Path,
223    undefined: &mut HashSet<String>,
224    defined: &mut HashSet<String>,
225) {
226    let file = object::read::File::parse(bytes)
227        .unwrap_or_else(|e| panic!("parsing {} as an object file: {e}", path.display()));
228    for symbol in file.symbols() {
229        let Ok(name) = symbol.name() else { continue };
230        if symbol.is_undefined() {
231            undefined.insert(name.to_string());
232        } else if symbol.is_global() {
233            defined.insert(name.to_string());
234        }
235    }
236}

248fn create_undefined_symbol_stub(
249    cache: &HashMap<String, CachedSymbol>,
250    tdata: &[u8],
251    undefined_symbols: &[String],
252    aslr_offset: i64,
253) -> Vec<u8> {
254    let mut obj = object::write::Object::new(
255        object::BinaryFormat::Elf,
256        object::Architecture::X86_64,
257        Endianness::Little,
258    );
259    let text_section = obj.section_id(StandardSection::Text);
260
261    for name in undefined_symbols {
262        let Some(sym) = cache.get(name.as_str()) else {
263            continue;
264        };
265        if sym.is_undefined {
266            continue;
267        }
268
269        let abs_addr = (sym.address as i64 + aslr_offset) as u64;
270
271        match sym.kind {
272            SymbolKind::Text => {
273                // jmp [rip+0]; <8-byte absolute address> - patch.rs:1083-1091.
274                let mut code = vec![0xFF, 0x25, 0x00, 0x00, 0x00, 0x00];
275                code.extend_from_slice(&abs_addr.to_le_bytes());
276                let offset = obj.append_section_data(text_section, &code, 8);
277                obj.add_symbol(Symbol {
278                    name: name.as_bytes().to_vec(),
279                    value: offset,
280                    size: code.len() as u64,
281                    scope: SymbolScope::Linkage,
282                    kind: SymbolKind::Text,
283                    weak: false,
284                    section: SymbolSection::Section(text_section),
285                    flags: SymbolFlags::None,
286                });
287            }
288            SymbolKind::Tls => {
289                // patch.rs:1163-1226's ELF branch: `sym.address` is already the
290                // OFFSET into the TLS template (not a virtual address, unlike
291                // every other symbol kind here), so no `aslr_offset` applies.
292                // The new symbol gets its OWN fresh TLS slot in the patch,
293                // seeded with the base binary's init bytes - each patch gets
294                // its own copy, which is why a thread-local resets on a patch
295                // (research/subsecond-patch-build.md §5, same limitation as
296                // apps/hotdemo's `TICKS` static resetting on the tip crate's
297                // own recompile).
298                let start = sym.address as usize;
299                let end = start + sym.size as usize;
300                let init = tdata.get(start..end).map(<[u8]>::to_vec).unwrap_or_else(|| {
301                    // Beyond .tdata's bounds: a `.tbss` (zero-initialized)
302                    // thread-local, whose initial value genuinely is zero.
303                    vec![0u8; sym.size as usize]
304                });
305                let align = sym.size.max(1).min(8).next_power_of_two();
306                let tls_section = obj.section_id(StandardSection::Tls);
307                let sym_id = obj.add_symbol(Symbol {
308                    name: name.as_bytes().to_vec(),
309                    value: 0,
310                    size: 0,
311                    scope: SymbolScope::Linkage,
312                    kind: SymbolKind::Tls,
313                    weak: sym.is_weak,
314                    section: SymbolSection::Undefined,
315                    flags: SymbolFlags::None,
316                });
317                obj.add_symbol_data(sym_id, tls_section, &init, align);
318            }
319            _ => {
320                // Data, and anything Unknown (patch.rs treats Unknown as
321                // Data too - some platforms report statics that way).
322                let kind = match sym.kind {
323                    SymbolKind::Unknown => SymbolKind::Data,
324                    k => k,
325                };
326                obj.add_symbol(Symbol {
327                    name: name.as_bytes().to_vec(),
328                    value: abs_addr,
329                    size: 0,
330                    scope: SymbolScope::Linkage,
331                    kind,
332                    weak: sym.is_weak,
333                    section: SymbolSection::Absolute,
334                    flags: sym.flags,
335                });
336            }
337        }
338    }
339
340    obj.write().expect("writing the stub object")
341}

347fn link_patch(objects: &[PathBuf], stub: &Path, out: &Path) {
348    let cc = env::var("CC").unwrap_or_else(|_| "cc".to_string());
349    let status = Command::new(&cc)
350        .arg("-shared")
351        .arg("-Wl,--eh-frame-hdr")
352        .arg("-Wl,-z,noexecstack")
353        .arg("-Wl,-z,relro,-z,now")
354        .arg("-nodefaultlibs")
355        .arg("-Wl,-Bdynamic")
356        .args(objects)
357        .arg(stub)
358        .arg("-o")
359        .arg(out)
360        .status()
361        .unwrap_or_else(|e| panic!("running the linker ({cc}): {e}"));
362
363    if !status.success() {
364        eprintln!("hotpatch: link failed ({status})");
365        process::exit(1);
366    }
367}

375fn build_jump_table(
376    cache: &HashMap<String, CachedSymbol>,
377    patch_path: &Path,
378    patch_bytes: &[u8],
379) -> JumpTable {
380    let obj2 = object::read::File::parse(patch_bytes).expect("parsing the freshly linked patch");
381    let symbol_map = obj2.symbol_map();
382    let new_name_to_addr: HashMap<&str, u64> = symbol_map
383        .symbols()
384        .iter()
385        .map(|s| (s.name(), s.address()))
386        .collect();
387
388    let mut map = AddressMap::default();
389    for (new_name, new_addr) in &new_name_to_addr {
390        if let Some(old) = cache.get(*new_name) {
391            map.insert(old.address, *new_addr);
392        }
393    }
394
395    let new_base_address = *new_name_to_addr.get("main").unwrap_or_else(|| {
396        eprintln!("hotpatch: the patch has no 'main' symbol - built with debuginfo?");
397        process::exit(1);
398    });
399    let aslr_reference = cache
400        .get("main")
401        .unwrap_or_else(|| {
402            eprintln!("hotpatch: the base binary has no 'main' symbol");
403            process::exit(1);
404        })
405        .address;
406
407    JumpTable {
408        lib: patch_path
409            .canonicalize()
410            .unwrap_or_else(|e| panic!("canonicalizing {}: {e}", patch_path.display())),
411        map,
412        aslr_reference,
413        new_base_address,
414        ifunc_count: 0,
415    }
416}

418fn next_patch_number(out_dir: &Path) -> u32 {
419    let mut n = 0;
420    if let Ok(entries) = fs::read_dir(out_dir) {
421        for entry in entries.flatten() {
422            let name = entry.file_name();
423            let name = name.to_string_lossy();
424            if let Some(k) = name
425                .strip_prefix("patch-")
426                .and_then(|r| r.strip_suffix(".so"))
427                .and_then(|r| r.parse::<u32>().ok())
428            {
429                n = n.max(k);
430            }
431        }
432    }
433    n + 1
434}
435
436fn main() {
437    // CLI hygiene (code-rust.md): a piped `| head` shouldn't panic this
438    // process.
439    unsafe {
440        libc::signal(libc::SIGPIPE, libc::SIG_DFL);
441    }
442
443    let args = parse_args();
444    fs::create_dir_all(&args.out_dir)
445        .unwrap_or_else(|e| panic!("creating {}: {e}", args.out_dir.display()));
446    let extract_dir = args.out_dir.join("objs");
447    fs::create_dir_all(&extract_dir)
448        .unwrap_or_else(|e| panic!("creating {}: {e}", extract_dir.display()));
449
450    let base_bytes = fs::read(&args.base)
451        .unwrap_or_else(|e| panic!("reading base exe {}: {e}", args.base.display()));
452    let cache = build_symbol_cache(&base_bytes);
453    let tdata = tdata(&base_bytes);
454    let base_main = cache
455        .get("main")
456        .unwrap_or_else(|| {
457            eprintln!(
458                "hotpatch: {} has no 'main' symbol - built with -Cdebuginfo=1 and \
459                 -Wl,--export-dynamic-symbol,main?",
460                args.base.display()
461            );
462            process::exit(1);
463        })
464        .address;
465    let aslr_offset = args.aslr_reference as i64 - base_main as i64;
466
467    let (undefined, defined, link_objects) = collect_symbols(&args.objects, &extract_dir);
468    let undefined_symbols: Vec<String> = undefined.difference(&defined).cloned().collect();
469
470    let stub_bytes = create_undefined_symbol_stub(&cache, &tdata, &undefined_symbols, aslr_offset);
471
472    let n = next_patch_number(&args.out_dir);
473    let stub_path = args.out_dir.join(format!("patch-{n}.stub.o"));
474    fs::write(&stub_path, &stub_bytes)
475        .unwrap_or_else(|e| panic!("writing {}: {e}", stub_path.display()));
476
477    let so_path = args.out_dir.join(format!("patch-{n}.so"));
478    link_patch(&link_objects, &stub_path, &so_path);
479
480    let patch_bytes = fs::read(&so_path)
481        .unwrap_or_else(|e| panic!("reading freshly linked {}: {e}", so_path.display()));
482    let table = build_jump_table(&cache, &so_path, &patch_bytes);
483
484    let json_path = args.out_dir.join(format!("patch-{n}.json"));
485    let json = serde_json::to_vec_pretty(&table).expect("serializing the jump table");
486    fs::write(&json_path, &json)
487        .unwrap_or_else(|e| panic!("writing {}: {e}", json_path.display()));
488
489    println!(
490        "hotpatch: wrote {} and {} ({} undefined symbol(s) stubbed, {} address(es) mapped)",
491        so_path.display(),
492        json_path.display(),
493        undefined_symbols.len(),
494        table.map.len()
495    );
496}