Add BTI landing pads for range extension thunks

A range extension thunk jumps to its destination with `br x16`, which
faults under BTI if the destination doesn't begin with a landing pad, as
is the case for a static function that is only called directly. For such
a function, we now put `bti c; b <function>` at the beginning of the first
thunk after the function's section, which is always within direct branch
range, and other thunks jump there instead, as lld's landing pad thunks do.

src/passes.rs

@@ -2255,12 +2255,13 @@ pub fn compute_section_sizes<E: Target>(ctx: &mut Context<E>) {
     };
 
     // create_range_extension_thunks is not thread-safe
+    let landing_pads = crate::thunks::find_landing_pads(ctx);
     for i in 0..ctx.chunks.len() {
         let id = ctx.chunks[i];
         if let ChunkId::Output(osec) = id
             && needs_thunks(ctx, id)
         {
-            crate::thunks::create_range_extension_thunks(ctx, osec);
+            crate::thunks::create_range_extension_thunks(ctx, osec, &landing_pads);
         }
     }
 

src/target/arm64.rs

@@ -25,9 +25,9 @@ use crate::context::Context;
 use crate::elf::*;
 use crate::input_sections::NonAllocReloc;
 use crate::input_sections::{InputSection, check_tlsle, scan_absrel, scan_pcrel, scan_tlsdesc};
-use crate::symbol::{NEEDS_GOT, NEEDS_GOTTP, NEEDS_PLT, NEEDS_TLSGD, Symbol};
+use crate::symbol::{NEEDS_GOT, NEEDS_GOTTP, NEEDS_PLT, NEEDS_TLSGD, Symbol, SymbolId};
 use crate::target::{Family, Target, ThunkLayout};
-use crate::thunks::Thunk;
+use crate::thunks::{self, LANDING_PAD_SIZE, Thunk};
 use crate::util::endian::{read_ul32, write_ul32};
 use crate::util::{bits, is_int};
 use crate::{error, fatal};
@@ -119,6 +119,11 @@ impl<const LE: bool> Arm64Target<LE> {
                 || sym.aux(&ctx.symbols).is_some_and(|aux| !aux.thunk_addrs.is_empty()))
     }
 
+    /// The address that a thunk entry for `sym` jumps to.
+    fn thunk_dest(ctx: &Context<Self>, sym: SymbolId) -> u64 {
+        thunks::landing_pad_addr(ctx, sym).unwrap_or_else(|| ctx.symbols[sym].addr(ctx))
+    }
+
     /// Whether the ADRP+ADD pair at relocation `i` can become NOP+ADR,
     /// which the psABI allows when the target is within ±1 MiB.
     fn relaxes_adrp_add(
@@ -810,10 +815,10 @@ impl<const LE: bool> Target for Arm64Target<LE> {
         // page(0x1100) – page(0xfff) is 0x1000, even though the latter
         // distance is shorter than the former.
         let is_small = |prel: i64| is_int(prel + 0x1000, 33) && is_int(prel - 0x1000, 33);
-        let mut offsets = vec![0];
-        let mut off = 0;
+        let mut off = thunk.landing_pads.len() as u64 * LANDING_PAD_SIZE;
+        let mut offsets = vec![off];
         for &sym in &thunk.symbols {
-            let s = ctx.symbols[sym].addr(ctx);
+            let s = Self::thunk_dest(ctx, sym);
             let p = addr + off;
             let prel = page(s).wrapping_sub(page(p)) as i64;
             off += if is_small(prel) { 12 } else { 24 };
@@ -838,9 +843,25 @@ impl<const LE: bool> Target for Arm64Target<LE> {
             0x8b11_0210, // add  x16, x16, x17
             0xd61f_0200, // br   x16
         ];
+        // Landing pad for a function that doesn't start with one
+        const LANDING_PAD: [u32; 2] = [
+            0xd503_245f, // bti  c
+            0x1400_0000, // b    0
+        ];
+
+        for (i, &sym) in thunk.landing_pads.iter().enumerate() {
+            let s = ctx.symbols[sym].addr(ctx);
+            let off = i * LANDING_PAD_SIZE as usize;
+            let disp = s.wrapping_sub(addr + off as u64 + 4);
+            debug_assert!(is_int(disp as i64, 28));
+            for (j, &v) in LANDING_PAD.iter().enumerate() {
+                write_insn(&mut buf[off + j * 4..], v);
+            }
+            or_insn(&mut buf[off + 4..], bits(disp, 27, 2) as u32);
+        }
 
         for (i, &sym) in thunk.symbols.iter().enumerate() {
-            let s = ctx.symbols[sym].addr(ctx);
+            let s = Self::thunk_dest(ctx, sym);
             let p = addr + thunk.offsets[i];
             let entry = &mut buf[thunk.offsets[i] as usize..thunk.offsets[i + 1] as usize];
             if entry.len() == 12 {

src/thunks.rs

@@ -21,28 +21,38 @@
 //! removed. Sections only shrink in the second pass, so no existing
 //! reference to a thunk goes out of range because of it.
 
+use std::collections::HashMap;
+
 use rayon::prelude::*;
 
+use crate::chunks::note_property;
 use crate::chunks::output_section::OutputSection;
 use crate::chunks::{ChunkId, OutputSectionId};
 use crate::context::Context;
 use crate::elf::*;
-use crate::input_sections::InputSection;
+use crate::input_sections::{InputSection, InputSectionId};
 use crate::symbol::{AddrFlags, Symbol, SymbolId};
 use crate::target::{Family, Target};
 use crate::util::align_to;
+use crate::util::endian::read_ul32;
 
 /// A block of branch stubs placed between input sections.
 #[derive(Debug)]
 pub struct Thunk {
     /// Offset within the output section.
     pub offset: u64,
+    /// Functions whose landing pads are at the beginning of this thunk. See
+    /// find_landing_pads().
+    pub landing_pads: Vec<SymbolId>,
     pub symbols: Vec<SymbolId>,
     /// Offset of each stub within the thunk; the last entry is the size.
     pub offsets: Vec<u64>,
     pub name: String,
 }
 
+/// The size of a landing pad, which consists of `bti c` and `b <function>`.
+pub const LANDING_PAD_SIZE: u64 = 8;
+
 impl Thunk {
     pub fn size(&self) -> u64 {
         self.offsets.last().copied().unwrap_or(0)
@@ -55,9 +65,8 @@ impl Thunk {
     /// The entry offsets of a thunk with fixed-size entries.
     pub fn fixed_offsets<E: Target>(&self) -> Vec<u64> {
         let layout = E::THUNK.expect("target without thunks");
-        (0..=self.symbols.len())
-            .map(|i| layout.header_size + i as u64 * layout.entry_size)
-            .collect()
+        let base = self.landing_pads.len() as u64 * LANDING_PAD_SIZE + layout.header_size;
+        (0..=self.symbols.len()).map(|i| base + i as u64 * layout.entry_size).collect()
     }
 }
 
@@ -150,6 +159,104 @@ fn executable_sections<E: Target>(ctx: &Context<E>) -> Vec<OutputSectionId> {
         .collect()
 }
 
+/// The key to sort symbols in a thunk for deterministic output.
+fn sort_key<E: Target>(ctx: &Context<E>, id: SymbolId) -> (u32, u32) {
+    let sym = &ctx.symbols[id];
+    (sym.file().map_or(0, |f| ctx.file(f).priority), sym.sym_idx())
+}
+
+/// The input section containing the code of `sym`. A section folded by ICF
+/// has no code of its own, so this returns the section it was folded into.
+fn code_section<E: Target>(ctx: &Context<E>, sym: &Symbol) -> Option<InputSectionId> {
+    let id = sym.input_section()?;
+    match ctx.input_section(id).icf_leader() {
+        Some(leader) => ctx.objs[leader.file.index()].section_id(leader.shndx as usize),
+        None => Some(id),
+    }
+}
+
+/// Whether `sym` begins with an instruction that an indirect branch may land
+/// on if BTI is enabled, i.e., `bti c`, `bti j`, `bti jc`, `paciasp` or
+/// `pacibsp`. If `sym` is not within `isec`'s contents, we can't tell, so
+/// we assume it is one as lld does.
+fn is_landing_pad<E: Target>(isec: &InputSection<E>, sym: &Symbol) -> bool {
+    let off = sym.value as usize;
+    let Some(loc) = isec.contents().get(off..off + 4) else {
+        return true;
+    };
+    matches!(read_ul32(loc), 0xd503_245f | 0xd503_249f | 0xd503_24df | 0xd503_233f | 0xd503_237f)
+}
+
+/// If ARM64 BTI is enabled, an indirect branch must land on a landing pad
+/// instruction such as `bti c`, and a thunk jumps to its destination with
+/// an indirect branch. However, the compiler doesn't emit a landing pad at
+/// the beginning of a function that is only called directly, so such a
+/// function can't be a destination of a thunk as is.
+///
+/// For such a function, we create a landing pad consisting of `bti c` and
+/// `b <function>` at the beginning of the first thunk after the function's
+/// section, and other thunks jump to the landing pad instead. The thunk is
+/// always within direct branch range of the function because it's
+/// reachable from all sections in its batch, and the batch starts at or
+/// before the function's section.
+///
+/// This function returns functions that may need a landing pad, grouped by
+/// the sections containing them. We don't know yet which calls are out of
+/// range, so any function called from another section is a candidate.
+pub fn find_landing_pads<E: Target>(ctx: &Context<E>) -> HashMap<InputSectionId, Vec<SymbolId>> {
+    let mut map: HashMap<InputSectionId, Vec<SymbolId>> = HashMap::new();
+    if ctx.args.relocatable || !note_property::is_bti(ctx) {
+        return map;
+    }
+
+    let members: Vec<InputSectionId> = executable_sections(ctx)
+        .into_iter()
+        .flat_map(|id| ctx.output_sections[id.index()].members.iter().copied())
+        .collect();
+
+    let mut syms: Vec<(InputSectionId, SymbolId)> = members
+        .par_iter()
+        .flat_map_iter(|&member| {
+            let isec = ctx.input_section(member);
+            let file = &ctx.objs[isec.file.index()];
+            isec.rels(file).iter().filter_map(move |rel| {
+                let id = file.base.symbols[rel.r_sym() as usize];
+                let sym = &ctx.symbols[id];
+                if !rel.is_func_call::<E>() || sym.has_plt(&ctx.symbols) {
+                    return None;
+                }
+                let target = code_section(ctx, sym)?;
+                if target == member || is_landing_pad(ctx.input_section(target), sym) {
+                    return None;
+                }
+                Some((target, id))
+            })
+        })
+        .collect();
+
+    syms.par_sort_by_key(|&(target, id)| (target, sort_key(ctx, id)));
+    syms.dedup();
+    for (target, id) in syms {
+        map.entry(target).or_default().push(id);
+    }
+    map
+}
+
+/// Returns the address of the landing pad of `id` if it has one.
+pub fn landing_pad_addr<E: Target>(ctx: &Context<E>, id: SymbolId) -> Option<u64> {
+    if !note_property::is_bti(ctx) {
+        return None;
+    }
+    let isec = ctx.input_section(code_section(ctx, &ctx.symbols[id])?);
+    let osec = &ctx.output_sections[isec.output_section?.index()];
+
+    // Find the first thunk after the section.
+    let thunk = osec.thunks.get(osec.thunks.partition_point(|t| t.offset <= isec.offset()))?;
+    let key = sort_key(ctx, id);
+    let i = thunk.landing_pads.binary_search_by_key(&key, |&x| sort_key(ctx, x)).ok()?;
+    Some(thunk.addr(osec) + i as u64 * LANDING_PAD_SIZE)
+}
+
 /// We create thunks from the beginning of the section to the end.
 /// We manage progress using four offsets which increase monotonically.
 /// The locations they point to are always A <= B <= C <= D.
@@ -169,7 +276,13 @@ fn executable_sections<E: Target>(ctx: &Context<E>) -> Vec<OutputSectionId> {
 ///     <-------->       Smaller than BRANCH_DISTANCE
 ///          <-------->  Smaller than BRANCH_DISTANCE
 ///     <------------->  Reachable from the current batch
-pub fn create_range_extension_thunks<E: Target>(ctx: &mut Context<E>, id: OutputSectionId) {
+///
+/// `landing_pads` is the result of find_landing_pads().
+pub fn create_range_extension_thunks<E: Target>(
+    ctx: &mut Context<E>,
+    id: OutputSectionId,
+    landing_pads: &HashMap<InputSectionId, Vec<SymbolId>>,
+) {
     let members = std::mem::take(&mut ctx.output_sections[id.index()].members);
     if members.is_empty() {
         return;
@@ -190,6 +303,8 @@ pub fn create_range_extension_thunks<E: Target>(ctx: &mut Context<E>, id: Output
     let mut offset = 0u64;
     // The smallest thunk index that is reachable from the current batch.
     let mut t = 0usize;
+    // Sections before this index precede an already created thunk.
+    let mut prev_d = 0usize;
 
     while b < n {
         // Move D foward as far as we can jump from B to a thunk at D.
@@ -261,8 +376,21 @@ pub fn create_range_extension_thunks<E: Target>(ctx: &mut Context<E>, id: Output
                     symbols
                 })
         };
-        // Add symbols to the thunk
-        let mut thunk = Thunk { offset, symbols, offsets: Vec::new(), name: String::new() };
+        // Add symbols to the thunk. Functions in the sections placed since
+        // the previous thunk get their landing pads in this thunk.
+        let mut thunk = Thunk {
+            offset,
+            landing_pads: members[prev_d..d]
+                .iter()
+                .filter_map(|member| landing_pads.get(member))
+                .flatten()
+                .copied()
+                .collect(),
+            symbols,
+            offsets: Vec::new(),
+            name: String::new(),
+        };
+        prev_d = d;
         // Now that we know the number of symbols in the thunk, we can compute
         // the thunk's size.
         thunk.offsets = thunk.fixed_offsets::<E>();
@@ -285,10 +413,8 @@ pub fn create_range_extension_thunks<E: Target>(ctx: &mut Context<E>, id: Output
     {
         let ctx: &Context<E> = ctx;
         thunks.par_iter_mut().for_each(|thunk| {
-            thunk.symbols.sort_by_key(|&id| {
-                let sym = &ctx.symbols[id];
-                (sym.file().map_or(0, |f| ctx.file(f).priority), sym.sym_idx())
-            });
+            thunk.landing_pads.sort_by_key(|&id| sort_key(ctx, id));
+            thunk.symbols.sort_by_key(|&id| sort_key(ctx, id));
         });
     }
 
@@ -336,17 +462,31 @@ pub fn remove_redundant_thunks<E: Target>(ctx: &mut Context<E>) {
         }
     }
 
+    // Remove symbols from thunks if they don't actually need range extension
+    // thunks. The same goes for landing pads.
     for &id in &sections {
-        let mut thunks = std::mem::take(&mut ctx.output_sections[id.index()].thunks);
-        {
+        let symbols = &ctx.symbols;
+        ctx.output_sections[id.index()].thunks.par_iter_mut().for_each(|thunk| {
+            thunk.symbols.retain(|&sym| symbols[sym].is_marked());
+            thunk.landing_pads.retain(|&sym| symbols[sym].is_marked());
+        });
+    }
+
+    for &id in &sections {
+        // A thunk's size depends on the distances to its destinations, which
+        // may be landing pads in other thunks, so compute sizes while all
+        // thunks are in place.
+        let offsets: Vec<Vec<u64>> = {
             let ctx: &Context<E> = ctx;
             let osec = &ctx.output_sections[id.index()];
-            // Remove symbols from thunks if they don't actually need range
-            // extension thunks
-            thunks.par_iter_mut().for_each(|thunk| {
-                thunk.symbols.retain(|&sym| ctx.symbols[sym].is_marked());
-                thunk.offsets = E::thunk_offsets(ctx, thunk, thunk.addr(osec));
-            });
+            osec.thunks
+                .par_iter()
+                .map(|thunk| E::thunk_offsets(ctx, thunk, thunk.addr(osec)))
+                .collect()
+        };
+        let mut thunks = std::mem::take(&mut ctx.output_sections[id.index()].thunks);
+        for (thunk, offsets) in thunks.iter_mut().zip(offsets) {
+            thunk.offsets = offsets;
         }
 
         // Recompute section sizes

tests/arch-aarch64-thunk-landing-pad.sh

@@ -0,0 +1,62 @@
+#!/usr/bin/env bash
+. $(dirname $0)/common.inc
+
+test_cflags -mbranch-protection=bti || skip
+
+# A static function that is only called directly doesn't begin with a BTI
+# landing pad, so a range extension thunk has to jump to it via a landing
+# pad created by the linker.
+cat <<EOF > $t/a.c
+__attribute__((section(".high"), noinline)) static int high() { return 42; }
+
+__attribute__((section(".low"))) void _start() {
+  register long x0 __asm__("x0") = high() == 42 ? 0 : 1;
+  register long x8 __asm__("x8") = 93; // exit
+  __asm__ volatile("svc 0" :: "r"(x0), "r"(x8));
+}
+EOF
+
+$CC -mbranch-protection=bti -O2 -c -o $t/a.o $t/a.c
+$OBJDUMP -d $t/a.o | grep -A1 '<high>:' | not grep -E 'bti\s+c'
+
+$CC -B. -static -nostdlib -o $t/exe1 $t/a.o \
+  -Wl,--section-start=.low=0x10000000,--section-start=.high=0x20000000
+readelf -n $t/exe1 | grep 'AArch64 feature: BTI'
+$OBJDUMP -d -j .high $t/exe1 | grep -A1 -E 'bti\s+c' | grep -E 'b\s+20000000 <high>'
+$QEMU $t/exe1
+
+$CC -mbranch-protection=none -O2 -c -o $t/b.o $t/a.c
+$CC -B. -static -nostdlib -o $t/exe2 $t/b.o \
+  -Wl,--section-start=.low=0x10000000,--section-start=.high=0x20000000
+$OBJDUMP -d -j .high $t/exe2 | not grep -E 'bti\s+c'
+$QEMU $t/exe2
+
+# .text.high2 is folded into .text.high1 by ICF, so the landing pad for
+# .text.high2 has to be placed next to .text.high1.
+cat <<EOF | $CC -c -o $t/c.o -xassembler -
+  .section .text.high1, "ax"
+high1:
+  mov w0, #42
+  ret
+
+  .section .text.high2, "ax"
+high2:
+  mov w0, #42
+  ret
+
+  .section .low, "ax"
+  .globl _start
+_start:
+  bl high1
+  mov w19, w0
+  bl high2
+  add w0, w0, w19
+  cmp w0, #84
+  cset x0, ne
+  mov x8, #93 // exit
+  svc #0
+EOF
+
+$CC -B. -static -nostdlib -o $t/exe3 $t/c.o -Wl,-z,force-bti,--icf=all \
+  -Wl,--section-start=.low=0x10000000,--section-start=.text=0x20000000
+$QEMU $t/exe3