Implement a new kcfi_arity feature that encodes an indirect call target's arity (i.e., the number of live-in registers) in the function's __cfi header.

scottconstable · scottconstable · commit 4fb9709491f5 · 2025-01-27T22:32:07.000-08:00
diff --git a/clang/docs/ControlFlowIntegrity.rst b/clang/docs/ControlFlowIntegrity.rst
@@ -336,6 +336,15 @@ cross-DSO function address equality. These properties make KCFI easier to
 adopt in low-level software. KCFI is limited to checking only function
 pointers, and isn't compatible with executable-only memory.
 
+``-fsanitize-kcfi-arity``
+-----------------------------
+
+For supported targets, this feature extends kCFI by telling the compiler to
+record information about each indirect-callable function's arity (i.e., the
+number of arguments passed in registers) into the binary. Some kernel CFI
+techniques, such as FineIBT, may be able to use this information to provide
+enhanced security.
+
 Member Function Pointer Call Checking
 =====================================
 
diff --git a/clang/docs/UsersManual.rst b/clang/docs/UsersManual.rst
@@ -2220,6 +2220,12 @@ are listed below.
 
    This option is currently experimental.
 
+.. option:: -fsanitize-kcfi-arity
+
+   Extends kernel indirect call forward-edge control flow integrity with
+   additional function arity information (for supported targets). See
+   :doc:`ControlFlowIntegrity` for more details.
+
 .. option:: -fstrict-vtable-pointers
 
    Enable optimizations based on the strict rules for overwriting polymorphic
diff --git a/clang/include/clang/Basic/CodeGenOptions.def b/clang/include/clang/Basic/CodeGenOptions.def
@@ -277,6 +277,7 @@ CODEGENOPT(SanitizeCfiICallNormalizeIntegers, 1, 0) ///< Normalize integer types
                                                     ///< CFI icall function signatures
 CODEGENOPT(SanitizeCfiCanonicalJumpTables, 1, 0) ///< Make jump table symbols canonical
                                                  ///< instead of creating a local jump table.
+CODEGENOPT(SanitizeKcfiArity, 1, 0) ///< Embed arity in KCFI patchable function prefix
 CODEGENOPT(SanitizeCoverageType, 2, 0) ///< Type of sanitizer coverage
                                        ///< instrumentation.
 CODEGENOPT(SanitizeCoverageIndirectCalls, 1, 0) ///< Enable sanitizer coverage
diff --git a/clang/include/clang/Basic/Features.def b/clang/include/clang/Basic/Features.def
@@ -254,6 +254,7 @@ FEATURE(is_trivially_constructible, LangOpts.CPlusPlus)
 FEATURE(is_trivially_copyable, LangOpts.CPlusPlus)
 FEATURE(is_union, LangOpts.CPlusPlus)
 FEATURE(kcfi, LangOpts.Sanitize.has(SanitizerKind::KCFI))
+FEATURE(kcfi_arity, LangOpts.Sanitize.has(SanitizerKind::KCFI))
 FEATURE(modules, LangOpts.Modules)
 FEATURE(safe_stack, LangOpts.Sanitize.has(SanitizerKind::SafeStack))
 FEATURE(shadow_call_stack,
diff --git a/clang/include/clang/Driver/Options.td b/clang/include/clang/Driver/Options.td
@@ -2619,6 +2619,10 @@ defm sanitize_cfi_canonical_jump_tables : BoolOption<"f", "sanitize-cfi-canonica
           "Do not make">,
   BothFlags<[], [ClangOption], " the jump table addresses canonical in the symbol table">>,
   Group<f_clang_Group>;
+def fsanitize_kcfi_arity : Flag<["-"], "fsanitize-kcfi-arity">,
+                           Group<f_clang_Group>,
+                           HelpText<"Embed function arity information into the KCFI patchable function prefix">,
+                           MarshallingInfoFlag<CodeGenOpts<"SanitizeKcfiArity">>;
 defm sanitize_stats : BoolOption<"f", "sanitize-stats",
   CodeGenOpts<"SanitizeStats">, DefaultFalse,
   PosFlag<SetTrue, [], [ClangOption], "Enable">,
diff --git a/clang/include/clang/Driver/SanitizerArgs.h b/clang/include/clang/Driver/SanitizerArgs.h
@@ -43,6 +43,7 @@ class SanitizerArgs {
   bool CfiICallGeneralizePointers = false;
   bool CfiICallNormalizeIntegers = false;
   bool CfiCanonicalJumpTables = false;
+  bool KcfiArity = false;
   int AsanFieldPadding = 0;
   bool SharedRuntime = false;
   bool StableABI = false;
diff --git a/clang/lib/CodeGen/CodeGenModule.cpp b/clang/lib/CodeGen/CodeGenModule.cpp
@@ -1149,6 +1149,8 @@ void CodeGenModule::Release() {
     if (CodeGenOpts.PatchableFunctionEntryOffset)
       getModule().addModuleFlag(llvm::Module::Override, "kcfi-offset",
                                 CodeGenOpts.PatchableFunctionEntryOffset);
+    if (CodeGenOpts.SanitizeKcfiArity)
+      getModule().addModuleFlag(llvm::Module::Override, "kcfi-arity", 1);
   }
 
   if (CodeGenOpts.CFProtectionReturn &&
diff --git a/clang/lib/Driver/SanitizerArgs.cpp b/clang/lib/Driver/SanitizerArgs.cpp
@@ -829,6 +829,8 @@ SanitizerArgs::SanitizerArgs(const ToolChain &TC,
     CfiICallNormalizeIntegers =
         Args.hasArg(options::OPT_fsanitize_cfi_icall_normalize_integers);
 
+    KcfiArity = Args.hasArg(options::OPT_fsanitize_kcfi_arity);
+
     if (AllAddedKinds & SanitizerKind::CFI && DiagnoseErrors)
       D.Diag(diag::err_drv_argument_not_allowed_with)
           << "-fsanitize=kcfi"
@@ -1383,6 +1385,14 @@ void SanitizerArgs::addArgs(const ToolChain &TC, const llvm::opt::ArgList &Args,
   if (CfiICallNormalizeIntegers)
     CmdArgs.push_back("-fsanitize-cfi-icall-experimental-normalize-integers");
 
+  if (KcfiArity) {
+    if (!TC.getTriple().isOSLinux() || !TC.getTriple().isArch64Bit()) {
+      TC.getDriver().Diag(clang::diag::err_drv_kcfi_arity_unsupported_target)
+          << TC.getTriple().str();
+    }
+    CmdArgs.push_back("-fsanitize-kcfi-arity");
+  }
+
   if (CfiCanonicalJumpTables)
     CmdArgs.push_back("-fsanitize-cfi-canonical-jump-tables");
 
diff --git a/clang/test/CodeGen/kcfi-arity.c b/clang/test/CodeGen/kcfi-arity.c
@@ -0,0 +1,7 @@
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -emit-llvm -fsanitize=kcfi -fsanitize-kcfi-arity -o - %s | FileCheck %s
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -emit-llvm -fsanitize=kcfi -fsanitize-kcfi-arity -x c++ -o - %s | FileCheck %s
+#if !__has_feature(kcfi_arity)
+#error Missing kcfi_arity?
+#endif
+
+// CHECK: ![[#]] = !{i32 4, !"kcfi-arity", i32 1}
diff --git a/llvm/lib/Target/X86/X86AsmPrinter.cpp b/llvm/lib/Target/X86/X86AsmPrinter.cpp
@@ -181,8 +181,29 @@ void X86AsmPrinter::emitKCFITypeId(const MachineFunction &MF) {
   // Embed the type hash in the X86::MOV32ri instruction to avoid special
   // casing object file parsers.
   EmitKCFITypePadding(MF);
+  unsigned DestReg = X86::EAX;
+
+  if (F.getParent()->getModuleFlag("kcfi-arity")) {
+    // The ArityToRegMap assumes the 64-bit Linux kernel ABI
+    const auto &Triple = MF.getTarget().getTargetTriple();
+    assert(Triple.isArch64Bit() && Triple.isOSLinux());
+
+    // Determine the function's arity (i.e., the number of arguments) at the ABI
+    // level by counting the number of parameters that are passed
+    // as registers, such as pointers and 64-bit (or smaller) integers. The
+    // Linux x86-64 ABI allows up to 6 parameters to be passed in GPRs.
+    // Additional parameters or parameters larger than 64 bits may be passed on
+    // the stack, in which case the arity is denoted as 7.
+    const unsigned ArityToRegMap[8] = {X86::EAX, X86::ECX, X86::EDX, X86::EBX,
+                                       X86::ESP, X86::EBP, X86::ESI, X86::EDI};
+    int Arity = MF.getInfo<X86MachineFunctionInfo>()->getArgumentStackSize() > 0
+                    ? 7
+                    : MF.getRegInfo().liveins().size();
+    DestReg = ArityToRegMap[Arity];
+  }
+
   EmitAndCountInstruction(MCInstBuilder(X86::MOV32ri)
-                              .addReg(X86::EAX)
+                              .addReg(DestReg)
                               .addImm(MaskKCFIType(Type->getZExtValue())));
 
   if (MAI->hasDotTypeDotSizeDirective()) {
diff --git a/llvm/test/CodeGen/X86/kcfi-arity.ll b/llvm/test/CodeGen/X86/kcfi-arity.ll
@@ -0,0 +1,205 @@
+; RUN: llc -mtriple=x86_64-unknown-linux-gnu -verify-machineinstrs < %s | FileCheck %s --check-prefix=ASM
+; RUN: llc -mtriple=x86_64-unknown-linux-gnu -verify-machineinstrs -stop-after=finalize-isel < %s | FileCheck %s --check-prefixes=MIR,ISEL
+; RUN: llc -mtriple=x86_64-unknown-linux-gnu -verify-machineinstrs -stop-after=kcfi < %s | FileCheck %s --check-prefixes=MIR,KCFI
+
+; ASM:       .p2align 4
+; ASM:       .type __cfi_f1,@function
+; ASM-LABEL: __cfi_f1:
+; ASM-NEXT:    nop
+; ASM-NEXT:    nop
+; ASM-NEXT:    nop
+; ASM-NEXT:    nop
+; ASM-NEXT:    nop
+; ASM-NEXT:    nop
+; ASM-NEXT:    nop
+; ASM-NEXT:    nop
+; ASM-NEXT:    nop
+; ASM-NEXT:    nop
+; ASM-NEXT:    nop
+; ASM-NEXT:    movl $12345678, %ecx
+; ASM-LABEL: .Lcfi_func_end0:
+; ASM-NEXT:  .size   __cfi_f1, .Lcfi_func_end0-__cfi_f1
+define void @f1(ptr noundef %x) !kcfi_type !1 {
+; ASM-LABEL: f1:
+; ASM:       # %bb.0:
+; ASM:         movl $4282621618, %r10d # imm = 0xFF439EB2
+; ASM-NEXT:    addl -4(%rdi), %r10d
+; ASM-NEXT:    je .Ltmp0
+; ASM-NEXT:  .Ltmp1:
+; ASM-NEXT:    ud2
+; ASM-NEXT:    .section .kcfi_traps,"ao",@progbits,.text
+; ASM-NEXT:  .Ltmp2:
+; ASM-NEXT:    .long .Ltmp1-.Ltmp2
+; ASM-NEXT:    .text
+; ASM-NEXT:  .Ltmp0:
+; ASM-NEXT:    callq *%rdi
+
+; MIR-LABEL: name: f1
+; MIR: body:
+; ISEL: CALL64r %0, csr_64, implicit $rsp, implicit $ssp, implicit-def $rsp, implicit-def $ssp, cfi-type 12345678
+; KCFI:       BUNDLE{{.*}} {
+; KCFI-NEXT:    KCFI_CHECK $rdi, 12345678, implicit-def $r10, implicit-def $r11, implicit-def $eflags
+; KCFI-NEXT:    CALL64r killed $rdi, csr_64, implicit $rsp, implicit $ssp, implicit-def $rsp, implicit-def $ssp
+; KCFI-NEXT:  }
+  call void %x() [ "kcfi"(i32 12345678) ]
+  ret void
+}
+
+; ASM-NOT: __cfi_f2:
+define void @f2(ptr noundef %x) {
+; ASM-LABEL: f2:
+
+; MIR-LABEL: name: f2
+; MIR: body:
+; ISEL: TCRETURNri64 %0, 0, csr_64, implicit $rsp, implicit $ssp, cfi-type 12345678
+; KCFI:       BUNDLE{{.*}} {
+; KCFI-NEXT:    KCFI_CHECK $rdi, 12345678, implicit-def $r10, implicit-def $r11, implicit-def $eflags
+; KCFI-NEXT:    TAILJMPr64 killed $rdi, csr_64, implicit $rsp, implicit $ssp, implicit $rsp, implicit $ssp
+; KCFI-NEXT:  }
+  tail call void %x() [ "kcfi"(i32 12345678) ]
+  ret void
+}
+
+; ASM-NOT: __cfi_f3:
+define void @f3(ptr noundef %x) #0 {
+; ASM-LABEL: f3:
+; MIR-LABEL: name: f3
+; MIR: body:
+; ISEL: CALL64pcrel32 &__llvm_retpoline_r11, csr_64, implicit $rsp, implicit $ssp, implicit-def $rsp, implicit-def $ssp, implicit killed $r11, cfi-type 12345678
+; KCFI:       BUNDLE{{.*}} {
+; KCFI-NEXT:    KCFI_CHECK $r11, 12345678, implicit-def $r10, implicit-def $r11, implicit-def $eflags
+; KCFI-NEXT:    CALL64pcrel32 &__llvm_retpoline_r11, csr_64, implicit $rsp, implicit $ssp, implicit-def $rsp, implicit-def $ssp, implicit internal killed $r11
+; KCFI-NEXT:  }
+  call void %x() [ "kcfi"(i32 12345678) ]
+  ret void
+}
+
+; ASM-NOT: __cfi_f4:
+define void @f4(ptr noundef %x) #0 {
+; ASM-LABEL: f4:
+; MIR-LABEL: name: f4
+; MIR: body:
+; ISEL: TCRETURNdi64 &__llvm_retpoline_r11, 0, csr_64, implicit $rsp, implicit $ssp, implicit killed $r11, cfi-type 12345678
+; KCFI:       BUNDLE{{.*}} {
+; KCFI-NEXT:    KCFI_CHECK $r11, 12345678, implicit-def $r10, implicit-def $r11, implicit-def $eflags
+; KCFI-NEXT:    TAILJMPd64 &__llvm_retpoline_r11, csr_64, implicit $rsp, implicit $ssp, implicit $rsp, implicit $ssp, implicit internal killed $r11
+; KCFI-NEXT:  }
+  tail call void %x() [ "kcfi"(i32 12345678) ]
+  ret void
+}
+
+;; Ensure we emit Value + 1 for unwanted values (e.g. endbr64 == 4196274163).
+; ASM-LABEL: __cfi_f5:
+; ASM: movl $4196274164, %ecx # imm = 0xFA1E0FF4
+define void @f5(ptr noundef %x) !kcfi_type !2 {
+; ASM-LABEL: f5:
+; ASM: movl $98693132, %r10d # imm = 0x5E1F00C
+  tail call void %x() [ "kcfi"(i32 4196274163) ]
+  ret void
+}
+
+;; Ensure we emit Value + 1 for unwanted values (e.g. -endbr64 == 98693133).
+; ASM-LABEL: __cfi_f6:
+; ASM: movl $98693134, %ecx # imm = 0x5E1F00E
+define void @f6(ptr noundef %x) !kcfi_type !3 {
+; ASM-LABEL: f6:
+; ASM: movl $4196274162, %r10d # imm = 0xFA1E0FF2
+  tail call void %x() [ "kcfi"(i32 98693133) ]
+  ret void
+}
+
+@g = external local_unnamed_addr global ptr, align 8
+
+define void @f7() {
+; MIR-LABEL: name: f7
+; MIR: body:
+; ISEL: TCRETURNmi64 killed %0, 1, $noreg, 0, $noreg, 0, csr_64, implicit $rsp, implicit $ssp, cfi-type 12345678
+; KCFI: $r11 = MOV64rm killed renamable $rax, 1, $noreg, 0, $noreg
+; KCFI-NEXT:  BUNDLE{{.*}} {
+; KCFI-NEXT:    KCFI_CHECK $r11, 12345678, implicit-def $r10, implicit-def $r11, implicit-def $eflags
+; KCFI-NEXT:    TAILJMPr64 internal $r11, csr_64, implicit $rsp, implicit $ssp, implicit $rsp, implicit $ssp
+; KCFI-NEXT:  }
+  %1 = load ptr, ptr @g, align 8
+  tail call void %1() [ "kcfi"(i32 12345678) ]
+  ret void
+}
+
+define void @f8() {
+; MIR-LABEL: name: f8
+; MIR: body:
+; ISEL: CALL64m killed %0, 1, $noreg, 0, $noreg, csr_64, implicit $rsp, implicit $ssp, implicit-def $rsp, implicit-def $ssp, cfi-type 12345678
+; KCFI: $r11 = MOV64rm killed renamable $rax, 1, $noreg, 0, $noreg
+; KCFI-NEXT:  BUNDLE{{.*}} {
+; KCFI-NEXT:    KCFI_CHECK $r11, 12345678, implicit-def $r10, implicit-def $r11, implicit-def $eflags
+; KCFI-NEXT:    CALL64r internal $r11, csr_64, implicit $rsp, implicit $ssp, implicit-def $rsp, implicit-def $ssp
+; KCFI-NEXT:  }
+  %1 = load ptr, ptr @g, align 8
+  call void %1() [ "kcfi"(i32 12345678) ]
+  ret void
+}
+
+%struct.S9 = type { [10 x i64] }
+
+;; Ensure that functions with large (e.g., greater than 8 bytes) arguments passed on the stack are assigned arity=7
+; ASM-LABEL: __cfi_f9:
+; ASM: movl	$199571451, %edi                # imm = 0xBE537FB
+define dso_local void @f9(ptr noundef byval(%struct.S9) align 8 %s) !kcfi_type !4  {
+entry:
+  ret void
+}
+
+;; Ensure that functions with fewer than 7 register arguments and no stack arguments are assigned arity<7
+; ASM-LABEL: __cfi_f10:
+; ASM: movl	$1046421190, %esi               # imm = 0x3E5F1EC6
+define dso_local void @f10(i32 noundef %v1, i32 noundef %v2, i32 noundef %v3, i32 noundef %v4, i32 noundef %v5, i32 noundef %v6) #0 !kcfi_type !5 {
+entry:
+  %v1.addr = alloca i32, align 4
+  %v2.addr = alloca i32, align 4
+  %v3.addr = alloca i32, align 4
+  %v4.addr = alloca i32, align 4
+  %v5.addr = alloca i32, align 4
+  %v6.addr = alloca i32, align 4
+  store i32 %v1, ptr %v1.addr, align 4
+  store i32 %v2, ptr %v2.addr, align 4
+  store i32 %v3, ptr %v3.addr, align 4
+  store i32 %v4, ptr %v4.addr, align 4
+  store i32 %v5, ptr %v5.addr, align 4
+  store i32 %v6, ptr %v6.addr, align 4
+  ret void
+}
+
+;; Ensure that functions with greater than 7 register arguments and no stack arguments are assigned arity=7
+; ASM-LABEL: __cfi_f11:
+; ASM: movl	$1342488295, %edi               # imm = 0x5004BEE7
+define dso_local void @f11(i32 noundef %v1, i32 noundef %v2, i32 noundef %v3, i32 noundef %v4, i32 noundef %v5, i32 noundef %v6, i32 noundef %v7, i32 noundef %v8) #0 !kcfi_type !6 {
+entry:
+  %v1.addr = alloca i32, align 4
+  %v2.addr = alloca i32, align 4
+  %v3.addr = alloca i32, align 4
+  %v4.addr = alloca i32, align 4
+  %v5.addr = alloca i32, align 4
+  %v6.addr = alloca i32, align 4
+  %v7.addr = alloca i32, align 4
+  %v8.addr = alloca i32, align 4
+  store i32 %v1, ptr %v1.addr, align 4
+  store i32 %v2, ptr %v2.addr, align 4
+  store i32 %v3, ptr %v3.addr, align 4
+  store i32 %v4, ptr %v4.addr, align 4
+  store i32 %v5, ptr %v5.addr, align 4
+  store i32 %v6, ptr %v6.addr, align 4
+  store i32 %v7, ptr %v7.addr, align 4
+  store i32 %v8, ptr %v8.addr, align 4
+  ret void
+}
+
+attributes #0 = { "target-features"="+retpoline-indirect-branches,+retpoline-indirect-calls" }
+
+!llvm.module.flags = !{!0, !7}
+!0 = !{i32 4, !"kcfi", i32 1}
+!1 = !{i32 12345678}
+!2 = !{i32 4196274163}
+!3 = !{i32 98693133}
+!4 = !{i32 199571451}
+!5 = !{i32 1046421190}
+!6 = !{i32 1342488295}
+!7 = !{i32 4, !"kcfi-arity", i32 1}

Original file line number	Diff line number	Diff line change
`@@ -1149,6 +1149,8 @@ void CodeGenModule::Release() {`
`1149`	`1149`	`if (CodeGenOpts.PatchableFunctionEntryOffset)`
`1150`	`1150`	`getModule().addModuleFlag(llvm::Module::Override, "kcfi-offset",`
`1151`	`1151`	`CodeGenOpts.PatchableFunctionEntryOffset);`
	`1152`	`+ if (CodeGenOpts.SanitizeKcfiArity)`
	`1153`	`+ getModule().addModuleFlag(llvm::Module::Override, "kcfi-arity", 1);`
`1152`	`1154`	`}`
`1153`	`1155`
`1154`	`1156`	`if (CodeGenOpts.CFProtectionReturn &&`