scx_rustsched_loader/
kfunc_protos.rs

1// SPDX-License-Identifier: GPL-2.0-only
2// Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
3
4//! Match the return types of the object's kfunc prototypes to the running
5//! kernel's.
6//!
7//! libbpf refuses an extern kfunc whose prototype is not compatible with
8//! the kernel's, and a return type that is void on one side and a value on
9//! the other is not. Such changes happen while the cid-form interface
10//! settles: scx_bpf_kick_cid() returned s32 on kernels where it now returns
11//! nothing, and scx_bpf_cidperf_set() the other way around. The BPF side
12//! never reads what these return, so the object can be told whichever the
13//! kernel has.
14//!
15//! The object's BTF is rewritten before libbpf opens it. Each kfunc whose
16//! return type has to change gets a copy of its prototype with the kernel's
17//! choice, appended after the last type, since the linker may share one
18//! prototype between kfuncs that look alike. Type ids and string offsets
19//! stay where they were, so .BTF.ext still holds. The longer .BTF goes at
20//! the end of the image and its section header is pointed there.
21
22use std::collections::BTreeMap;
23
24use anyhow::{Context, bail, ensure};
25use libbpf_rs::{
26    Btf,
27    btf::{ReferencesType, types::Func},
28};
29
30use crate::Result;
31
32const BTF_MAGIC: u16 = 0xeb9f;
33const BTF_KIND_INT: u32 = 1;
34const BTF_KIND_FUNC: u32 = 12;
35const BTF_KIND_FUNC_PROTO: u32 = 13;
36const BTF_FUNC_EXTERN: u32 = 2;
37const BTF_INT_SIGNED: u32 = 1;
38
39fn u16_at(bytes: &[u8], at: usize) -> Result<u16> {
40    let word = bytes.get(at..at + 2).context("truncated BPF object")?;
41    Ok(u16::from_le_bytes(word.try_into()?))
42}
43
44fn u32_at(bytes: &[u8], at: usize) -> Result<u32> {
45    let word = bytes.get(at..at + 4).context("truncated BPF object")?;
46    Ok(u32::from_le_bytes(word.try_into()?))
47}
48
49fn u64_at(bytes: &[u8], at: usize) -> Result<u64> {
50    let word = bytes.get(at..at + 8).context("truncated BPF object")?;
51    Ok(u64::from_le_bytes(word.try_into()?))
52}
53
54fn put_u32(bytes: &mut [u8], at: usize, value: u32) {
55    bytes[at..at + 4].copy_from_slice(&value.to_le_bytes());
56}
57
58fn put_u64(bytes: &mut [u8], at: usize, value: u64) {
59    bytes[at..at + 8].copy_from_slice(&value.to_le_bytes());
60}
61
62/// The section header of `.BTF` in an ELF64 little-endian object: where
63/// the header is, and the section's offset and size.
64fn btf_section(object: &[u8]) -> Result<Option<(usize, usize, usize)>> {
65    ensure!(
66        object.get(..6) == Some(&[0x7f, b'E', b'L', b'F', 2, 1][..]),
67        "BPF object is not a little-endian ELF64 file"
68    );
69    let shoff = u64_at(object, 0x28)? as usize;
70    let shentsize = u16_at(object, 0x3a)? as usize;
71    let shnum = u16_at(object, 0x3c)? as usize;
72    let shstrndx = u16_at(object, 0x3e)? as usize;
73    let names = u64_at(object, shoff + shstrndx * shentsize + 24)? as usize;
74    for index in 0..shnum {
75        let header = shoff + index * shentsize;
76        let name = names + u32_at(object, header)? as usize;
77        if object.get(name..name + 5) == Some(&b".BTF\0"[..]) {
78            let offset = u64_at(object, header + 24)? as usize;
79            let size = u64_at(object, header + 32)? as usize;
80            return Ok(Some((header, offset, size)));
81        }
82    }
83    Ok(None)
84}
85
86/// The bytes a type takes after its 12-byte header.
87fn trailing(kind: u32, vlen: usize) -> usize {
88    match kind {
89        1 | 14 | 17 => 4,     // INT, VAR, DECL_TAG
90        3 => 12,              // ARRAY
91        4 | 5 => vlen * 12,   // STRUCT, UNION
92        6 | 13 => vlen * 8,   // ENUM, FUNC_PROTO
93        15 | 19 => vlen * 12, // DATASEC, ENUM64
94        _ => 0,
95    }
96}
97
98struct Type {
99    offset: usize,
100    kind: u32,
101    vlen: usize,
102    name: String,
103    size_or_type: u32,
104}
105
106/// Return `object` with the return types of its kfuncs matched to the
107/// running kernel's, and the names of the kfuncs that changed.
108pub fn adapt(object: &[u8]) -> Result<(Vec<u8>, Vec<String>)> {
109    let Some((header, offset, size)) = btf_section(object)? else {
110        return Ok((object.to_vec(), Vec::new()));
111    };
112    let btf = object
113        .get(offset..offset + size)
114        .context("truncated .BTF section")?;
115    ensure!(u16_at(btf, 0)? == BTF_MAGIC, "bad .BTF magic");
116    let hdr_len = u32_at(btf, 4)? as usize;
117    let type_off = hdr_len + u32_at(btf, 8)? as usize;
118    let type_len = u32_at(btf, 12)? as usize;
119    let str_off = hdr_len + u32_at(btf, 16)? as usize;
120    let str_len = u32_at(btf, 20)? as usize;
121    ensure!(
122        str_off == type_off + type_len && str_off + str_len <= size,
123        "unexpected .BTF layout"
124    );
125    let string = |at: u32| -> String {
126        let start = str_off + at as usize;
127        let end = btf[start..str_off + str_len]
128            .iter()
129            .position(|&byte| byte == 0)
130            .map_or(start, |length| start + length);
131        String::from_utf8_lossy(&btf[start..end]).into_owned()
132    };
133
134    // Every type, by id from 1.
135    let mut types = vec![];
136    let mut at = type_off;
137    while at < type_off + type_len {
138        let info = u32_at(btf, at + 4)?;
139        let (kind, vlen) = ((info >> 24) & 0x1f, (info & 0xffff) as usize);
140        types.push(Type {
141            offset: at,
142            kind,
143            vlen,
144            name: string(u32_at(btf, at)?),
145            size_or_type: u32_at(btf, at + 8)?,
146        });
147        at += 12 + trailing(kind, vlen);
148    }
149    ensure!(at == type_off + type_len, "unexpected .BTF type section");
150    let by_id = |id: u32| types.get((id as usize).wrapping_sub(1));
151
152    let kernel = Btf::from_vmlinux().context("read running kernel BTF")?;
153    // The kfuncs whose return type has to change: the FUNC's offset, its
154    // prototype's id, and whether the kernel returns a value.
155    let mut changes = BTreeMap::new();
156    for function in &types {
157        if function.kind != BTF_KIND_FUNC || function.vlen as u32 != BTF_FUNC_EXTERN {
158            continue;
159        }
160        let Some(prototype) = by_id(function.size_or_type) else {
161            continue;
162        };
163        if prototype.kind != BTF_KIND_FUNC_PROTO {
164            continue;
165        }
166        let Some(theirs) = kernel.type_by_name::<Func<'_>>(&function.name) else {
167            continue;
168        };
169        let theirs = theirs.referenced_type();
170        let Some(theirs) = libbpf_rs::btf::types::FuncProto::try_from(theirs).ok() else {
171            continue;
172        };
173        let kernel_returns = u32::from(theirs.referenced_type_id()) != 0;
174        if kernel_returns != (prototype.size_or_type != 0) {
175            changes.insert(
176                function.name.clone(),
177                (function.offset, function.size_or_type, kernel_returns),
178            );
179        }
180    }
181    if changes.is_empty() {
182        return Ok((object.to_vec(), Vec::new()));
183    }
184
185    // A 32-bit signed integer to return where the kernel returns a value:
186    // one the object has, or one appended with its name.
187    let mut strings = btf[str_off..str_off + str_len].to_vec();
188    let mut appended = vec![];
189    let mut next_id = types.len() as u32 + 1;
190    let int_id = match types.iter().position(|ty| {
191        ty.kind == BTF_KIND_INT
192            && ty.size_or_type == 4
193            && (u32_at(btf, ty.offset + 12).unwrap_or(0) >> 24) & BTF_INT_SIGNED != 0
194    }) {
195        Some(index) => index as u32 + 1,
196        None if changes.values().any(|&(_, _, returns)| returns) => {
197            let name = strings.len() as u32;
198            strings.extend_from_slice(b"int\0");
199            appended.extend_from_slice(&name.to_le_bytes());
200            appended.extend_from_slice(&(BTF_KIND_INT << 24).to_le_bytes());
201            appended.extend_from_slice(&4u32.to_le_bytes());
202            appended.extend_from_slice(&((BTF_INT_SIGNED << 24) | 32).to_le_bytes());
203            next_id += 1;
204            next_id - 1
205        }
206        None => 0,
207    };
208
209    let mut patched = object.to_vec();
210    let mut new_btf = btf[..str_off].to_vec();
211    for &(function, prototype, returns) in changes.values() {
212        let Some(original) = by_id(prototype) else {
213            bail!("kfunc prototype {prototype} out of range");
214        };
215        let start = original.offset;
216        let mut copy = btf[start..start + 12 + trailing(original.kind, original.vlen)].to_vec();
217        put_u32(&mut copy, 8, if returns { int_id } else { 0 });
218        appended.extend_from_slice(&copy);
219        // The FUNC, in the copy of the type section, now names the copy.
220        put_u32(&mut new_btf, function + 8, next_id);
221        next_id += 1;
222    }
223    new_btf.extend_from_slice(&appended);
224    new_btf.extend_from_slice(&strings);
225    put_u32(&mut new_btf, 12, (type_len + appended.len()) as u32);
226    put_u32(
227        &mut new_btf,
228        16,
229        (type_off - hdr_len + type_len + appended.len()) as u32,
230    );
231    put_u32(&mut new_btf, 20, strings.len() as u32);
232
233    while !patched.len().is_multiple_of(8) {
234        patched.push(0);
235    }
236    let new_offset = patched.len();
237    patched.extend_from_slice(&new_btf);
238    put_u64(&mut patched, header + 24, new_offset as u64);
239    put_u64(&mut patched, header + 32, new_btf.len() as u64);
240    Ok((patched, changes.into_keys().collect()))
241}
scx_rustsched_loader/kfunc_protos.rs

scx_rustsched_loader/
kfunc_protos.rs
