Allow section override when using patchable-function-entries

Sometimes it is necessary to group patchable function entrypoint
records in distinct linker sections. This is the case for some bpf
functions within the linux kernel which shouldn't be visible to
ftrace.

Extend `-Zpatchable-function-entry` to accept an argument of the form
`prefix_nops,total_nops,record_section`, which places all entry
record into a user specified section.

Likewise, extend the `patchable_function_entry` attribute to accept an
optional `section="name"` option to place a function into a specific
section.

This is made possible by llvm attribute `patchable-function-entry-section`
added in llvm 21.

Author: pmur

compiler/rustc_attr_parsing/src/attributes/codegen_attrs.rs

@@ -9,8 +9,8 @@ use super::prelude::*;
 use crate::attributes::AttributeSafety;
 use crate::session_diagnostics::{
     EmptyExportName, NakedFunctionIncompatibleAttribute, NullOnExport, NullOnObjcClass,
-    NullOnObjcSelector, ObjcClassExpectedStringLiteral, ObjcSelectorExpectedStringLiteral,
-    SanitizeInvalidStatic,
+    NullOnObjcSelector, NullOnSection, ObjcClassExpectedStringLiteral,
+    ObjcSelectorExpectedStringLiteral, SanitizeInvalidStatic,
 };
 use crate::target_checking::Policy::AllowSilent;
 
@@ -781,14 +781,16 @@ pub(crate) struct PatchableFunctionEntryParser;
 impl SingleAttributeParser for PatchableFunctionEntryParser {
     const PATH: &[Symbol] = &[sym::patchable_function_entry];
     const ALLOWED_TARGETS: AllowedTargets<'_> = AllowedTargets::AllowList(&[Allow(Target::Fn)]);
-    const TEMPLATE: AttributeTemplate = template!(List: &["prefix_nops = m, entry_nops = n"]);
+    const TEMPLATE: AttributeTemplate =
+        template!(List: &["prefix_nops = m, entry_nops = n, section = \"section\""]);
     const STABILITY: AttributeStability = unstable!(patchable_function_entry);
 
     fn convert(cx: &mut AcceptContext<'_, '_>, args: &ArgParser) -> Option<AttributeKind> {
         let meta_item_list = cx.expect_list(args, cx.attr_span)?;
 
         let mut prefix = None;
         let mut entry = None;
+        let mut section = None;
 
         if meta_item_list.len() == 0 {
             cx.adcx().expected_at_least_one_argument(meta_item_list.span);
@@ -821,6 +823,27 @@ impl SingleAttributeParser for PatchableFunctionEntryParser {
                     }
                     &mut entry
                 }
+                sym::section => {
+                    // Duplicate entries are not allowed
+                    if section.is_some() {
+                        errored = true;
+                        cx.adcx().duplicate_key(ident.span, sym::section);
+                        continue;
+                    }
+                    // Only a string type value is allowed.
+                    let Some(value_str) = value.value_as_str() else {
+                        cx.adcx().expect_string_literal(value);
+                        continue;
+                    };
+                    // The section name does not allow null characters.
+                    if value_str.as_str().contains('\0') {
+                        errored = true;
+                        cx.emit_err(NullOnSection { span: value.value_span });
+                    }
+                    section = Some(value_str);
+                    // Integer parsing is not needed, process next item.
+                    continue;
+                }
                 _ => {
                     errored = true;
                     cx.adcx().expected_specific_argument(
@@ -853,10 +876,7 @@ impl SingleAttributeParser for PatchableFunctionEntryParser {
         if errored {
             None
         } else {
-            Some(AttributeKind::PatchableFunctionEntry {
-                prefix: prefix.unwrap_or(0),
-                entry: entry.unwrap_or(0),
-            })
+            Some(AttributeKind::PatchableFunctionEntry { prefix, entry, section })
         }
     }
 }

compiler/rustc_attr_parsing/src/session_diagnostics.rs

@@ -300,6 +300,13 @@ pub(crate) struct NullOnObjcSelector {
     pub span: Span,
 }
 
+#[derive(Diagnostic)]
+#[diag("`section` may not contain null characters", code = E0648)]
+pub(crate) struct NullOnSection {
+    #[primary_span]
+    pub span: Span,
+}
+
 #[derive(Diagnostic)]
 #[diag("`objc::class!` expected a string literal")]
 pub(crate) struct ObjcClassExpectedStringLiteral {

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -84,11 +84,25 @@ fn patchable_function_entry_attrs<'ll>(
     attr: Option<PatchableFunctionEntry>,
 ) -> SmallVec<[&'ll Attribute; 2]> {
     let mut attrs = SmallVec::new();
-    let patchable_spec = attr.unwrap_or_else(|| {
-        PatchableFunctionEntry::from_config(sess.opts.unstable_opts.patchable_function_entry)
-    });
-    let entry = patchable_spec.entry();
-    let prefix = patchable_spec.prefix();
+
+    let mut entry = sess.opts.unstable_opts.patchable_function_entry.entry();
+    let mut prefix = sess.opts.unstable_opts.patchable_function_entry.prefix();
+    let mut section = sess.opts.unstable_opts.patchable_function_entry.section();
+    let section_sym;
+
+    // Apply attribute specified overrides, if any.
+    if let Some(patchable_spec) = attr {
+        if let Some(sym) = patchable_spec.section() {
+            section_sym = sym;
+            section = Some(section_sym.as_str());
+        }
+        // Override the cmdline specified options if any nop count is specified.
+        if patchable_spec.entry().is_some() || patchable_spec.prefix().is_some() {
+            entry = patchable_spec.entry().unwrap_or(0);
+            prefix = patchable_spec.prefix().unwrap_or(0);
+        }
+    }
+
     if entry > 0 {
         attrs.push(llvm::CreateAttrStringValue(
             cx.llcx,
@@ -103,6 +117,13 @@ fn patchable_function_entry_attrs<'ll>(
             &format!("{}", prefix),
         ));
     }
+    if let Some(section) = section {
+        attrs.push(llvm::CreateAttrStringValue(
+            cx.llcx,
+            "patchable-function-entry-section",
+            section,
+        ));
+    }
     attrs
 }
 

compiler/rustc_codegen_llvm/src/context.rs

@@ -14,7 +14,6 @@ use rustc_data_structures::base_n::{ALPHANUMERIC_ONLY, ToBaseN};
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::small_c_str::SmallCStr;
 use rustc_hir::def_id::DefId;
-use rustc_middle::middle::codegen_fn_attrs::PatchableFunctionEntry;
 use rustc_middle::mono::CodegenUnit;
 use rustc_middle::ty::layout::{
     FnAbiError, FnAbiOfHelpers, FnAbiRequest, HasTypingEnv, LayoutError, LayoutOfHelpers,
@@ -343,14 +342,13 @@ pub(crate) unsafe fn create_module<'ll>(
 
         // Add "kcfi-offset" module flag with -Z patchable-function-entry (See
         // https://reviews.llvm.org/D141172).
-        let pfe =
-            PatchableFunctionEntry::from_config(sess.opts.unstable_opts.patchable_function_entry);
-        if pfe.prefix() > 0 {
+        let patchable_prefix_nops = sess.opts.unstable_opts.patchable_function_entry.prefix();
+        if patchable_prefix_nops > 0 {
             llvm::add_module_flag_u32(
                 llmod,
                 llvm::ModuleFlagMergeBehavior::Override,
                 "kcfi-offset",
-                pfe.prefix().into(),
+                patchable_prefix_nops.into(),
             );
         }
 

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -290,9 +290,11 @@ fn process_builtin_attrs(
             AttributeKind::RustcOffloadKernel => {
                 codegen_fn_attrs.flags |= CodegenFnAttrFlags::OFFLOAD_KERNEL
             }
-            AttributeKind::PatchableFunctionEntry { prefix, entry } => {
+            AttributeKind::PatchableFunctionEntry { prefix, entry, section } => {
                 codegen_fn_attrs.patchable_function_entry =
-                    Some(PatchableFunctionEntry::from_prefix_and_entry(*prefix, *entry));
+                    Some(PatchableFunctionEntry::from_prefix_entry_and_section(
+                        *prefix, *entry, *section,
+                    ));
             }
             AttributeKind::InstrumentFn(instrument_fn) => {
                 codegen_fn_attrs.instrument_fn = match instrument_fn {

compiler/rustc_hir/src/attrs/data_structures.rs

@@ -1224,8 +1224,9 @@ pub enum AttributeKind {
 
     /// Represents `#[patchable_function_entry]`
     PatchableFunctionEntry {
-        prefix: u8,
-        entry: u8,
+        prefix: Option<u8>,
+        entry: Option<u8>,
+        section: Option<Symbol>,
     },
 
     /// Represents `#[path]`

compiler/rustc_interface/src/tests.rs

@@ -866,7 +866,7 @@ fn test_unstable_options_tracking_hash() {
     tracked!(panic_in_drop, PanicStrategy::Abort);
     tracked!(
         patchable_function_entry,
-        PatchableFunctionEntry::from_total_and_prefix_nops(10, 5)
+        PatchableFunctionEntry::from_parts(10, 5, None)
             .expect("total must be greater than or equal to prefix")
     );
     tracked!(plt, Some(true));

compiler/rustc_middle/src/middle/codegen_fn_attrs.rs

@@ -114,7 +114,7 @@ pub struct CodegenFnAttrs {
     // FIXME(#82232, #143834): temporarily renamed to mitigate `#[align]` nameres ambiguity
     pub alignment: Option<Align>,
     /// The `#[patchable_function_entry(...)]` attribute. Indicates how many nops should be around
-    /// the function entry.
+    /// the function entry, or override default section to record entry location.
     pub patchable_function_entry: Option<PatchableFunctionEntry>,
     /// The `#[rustc_objc_class = "..."]` attribute.
     pub objc_class: Option<Symbol>,
@@ -162,24 +162,33 @@ pub struct TargetFeature {
 #[derive(Copy, Clone, Debug, TyEncodable, TyDecodable, StableHash)]
 pub struct PatchableFunctionEntry {
     /// Nops to prepend to the function
-    prefix: u8,
+    prefix: Option<u8>,
     /// Nops after entry, but before body
-    entry: u8,
+    entry: Option<u8>,
+    /// Optional, specific section to record entry location in
+    section: Option<Symbol>,
 }
 
 impl PatchableFunctionEntry {
-    pub fn from_config(config: rustc_session::config::PatchableFunctionEntry) -> Self {
-        Self { prefix: config.prefix(), entry: config.entry() }
+    pub fn from_prefix_entry_and_section(
+        prefix: Option<u8>,
+        entry: Option<u8>,
+        section: Option<Symbol>,
+    ) -> Self {
+        Self { prefix, entry, section }
     }
     pub fn from_prefix_and_entry(prefix: u8, entry: u8) -> Self {
-        Self { prefix, entry }
+        Self { prefix: Some(prefix), entry: Some(entry), section: None }
     }
-    pub fn prefix(&self) -> u8 {
+    pub fn prefix(&self) -> Option<u8> {
         self.prefix
     }
-    pub fn entry(&self) -> u8 {
+    pub fn entry(&self) -> Option<u8> {
         self.entry
     }
+    pub fn section(&self) -> Option<Symbol> {
+        self.section
+    }
 }
 
 #[derive(Clone, Copy, PartialEq, Eq, TyEncodable, TyDecodable, StableHash)]

compiler/rustc_session/src/config.rs

@@ -3330,23 +3330,29 @@ impl DumpMonoStatsFormat {
 
 /// `-Z patchable-function-entry` representation - how many nops to put before and after function
 /// entry.
-#[derive(Clone, Copy, PartialEq, Hash, Debug, Default)]
+#[derive(Clone, PartialEq, Hash, Debug, Default)]
 pub struct PatchableFunctionEntry {
     /// Nops before the entry
     prefix: u8,
     /// Nops after the entry
     entry: u8,
+    /// An optional section name to record the entry location
+    section: Option<String>,
 }
 
 impl PatchableFunctionEntry {
-    pub fn from_total_and_prefix_nops(
+    pub fn from_parts(
         total_nops: u8,
         prefix_nops: u8,
+        section: Option<String>,
     ) -> Option<PatchableFunctionEntry> {
         if total_nops < prefix_nops {
             None
+        // Section name cannot contain null characters.
+        } else if section.as_ref().map(|x| x.contains('\0')).unwrap_or(false) {
+            None
         } else {
-            Some(Self { prefix: prefix_nops, entry: total_nops - prefix_nops })
+            Some(Self { prefix: prefix_nops, entry: total_nops - prefix_nops, section })
         }
     }
     pub fn prefix(&self) -> u8 {
@@ -3355,6 +3361,9 @@ impl PatchableFunctionEntry {
     pub fn entry(&self) -> u8 {
         self.entry
     }
+    pub fn section(&self) -> Option<&str> {
+        self.section.as_ref().map(|x| x.as_str())
+    }
 }
 
 /// `-Zpolonius` values, enabling the borrow checker polonius analysis, and which version: legacy,

compiler/rustc_session/src/options.rs

@@ -784,7 +784,7 @@ mod desc {
     pub(crate) const parse_passes: &str = "a space-separated list of passes, or `all`";
     pub(crate) const parse_panic_strategy: &str = "either `unwind`, `abort`, or `immediate-abort`";
     pub(crate) const parse_on_broken_pipe: &str = "either `kill`, `error`, or `inherit`";
-    pub(crate) const parse_patchable_function_entry: &str = "either two comma separated integers (total_nops,prefix_nops), with prefix_nops <= total_nops, or one integer (total_nops)";
+    pub(crate) const parse_patchable_function_entry: &str = "a comma separated list of (prefix_nops,total_nops,section_name), (prefix_nops,total_nops), or (total_nops). Where prefix_nops <= total_nops where 0 < total_nops <= 255 and prefix_nops <= total_nops";
     pub(crate) const parse_opt_panic_strategy: &str = parse_panic_strategy;
     pub(crate) const parse_relro_level: &str = "one of: `full`, `partial`, or `off`";
     pub(crate) const parse_sanitizers: &str = "comma separated list of sanitizers: `address`, `cfi`, `dataflow`, `hwaddress`, `kcfi`, `kernel-address`, `kernel-hwaddress`, `leak`, `memory`, `memtag`, `safestack`, `shadow-call-stack`, `thread`, or 'realtime'";
@@ -1204,20 +1204,24 @@ pub mod parse {
     ) -> bool {
         let mut total_nops = 0;
         let mut prefix_nops = 0;
+        let mut section = None;
 
         if !parse_number(&mut total_nops, v) {
-            let parts = v.and_then(|v| v.split_once(',')).unzip();
-            if !parse_number(&mut total_nops, parts.0) {
+            let parts: Vec<_> = v.unwrap_or("").split(',').collect();
+            if parts.len() < 2 || parts.len() > 3 {
                 return false;
             }
-            if !parse_number(&mut prefix_nops, parts.1) {
+
+            if !parse_number(&mut total_nops, Some(parts[0])) {
+                return false;
+            }
+            if !parse_number(&mut prefix_nops, Some(parts[1])) {
                 return false;
             }
+            section = parts.get(2).map(|x| x.to_string());
         }
 
-        if let Some(pfe) =
-            PatchableFunctionEntry::from_total_and_prefix_nops(total_nops, prefix_nops)
-        {
+        if let Some(pfe) = PatchableFunctionEntry::from_parts(total_nops, prefix_nops, section) {
             *slot = pfe;
             return true;
         }