MachineFunction: -fsanitize={function,kcfi}: ensure 4-byte alignment

[MaskRay]

Summary

Fix https://github.com/llvm/llvm-project/issues/63579

% cat a.c
void foo() {}
% clang --target=arm-none-eabi -mthumb -mno-unaligned-access -fsanitize=kcfi a.c -S -o - | grep p2align
        .p2align        1
% clang --target=armv6m-none-eabi -fsanitize=function a.c -S -o - | grep p2align
        .p2align        1

Ensure that -fsanitize={function,kcfi} instrumented functions are aligned by at
least 4, so that loading the type hash before the function label will not cause
a misaligned access. This is especially important for -mno-unaligned-access
configurations that doesn't set setMinFunctionAlignment to 4 or greater.

With this patch, the generated assembly for the examples above will contain .p2align 2.

If __attribute__((aligned(N))) or -falign-functions=N is specified, the
larger alignment will be used.

[samitolvanen]

Thanks for fixing this! This approach looks good to me.

llvm/test/CodeGen/RISCV/kcfi-patchable-function-prefix.ll

3–7

Nit: These checks are the same now, so we can replace both lines with one CHECK: line. Also the other instance below.

[MaskRay]

improve llvm/test/CodeGen/RISCV/kcfi-patchable-function-prefix.ll

[simon_tatham]

Thanks – this looks much more like what I was imagining!

(I didn't just fix it myself this way because I was nervous about it affecting non-Arm architectures whose requirements I didn't know anything about. But @efriedma's comment on D154043 seems like a good argument that it's the right thing everywhere.)

llvm/lib/CodeGen/MachineFunction.cpp

void MachineFunction::init() {
   Alignment = STI->getTargetLowering()->getMinFunctionAlignment();
 
   // FIXME: Shouldn't use pref alignment if explicit alignment is set on F.
   // FIXME: Use Function::hasOptSize().
   if (!F.hasFnAttribute(Attribute::OptimizeForSize))
     Alignment = std::max(Alignment,
                          STI->getTargetLowering()->getPrefFunctionAlignment());
 
+  // -fsanitize=function and -fsanitize=kcfi instrument indirect function calls
+  // to load a type hash before the function label. Ensure functions are aligned
+  // by a least 4 to avoid unaligned access, which is especially important for
+  // -mno-unaligned-access.
+  if (F.hasMetadata(LLVMContext::MD_func_sanitize) ||
+      F.getMetadata(LLVMContext::MD_kcfi_type))
+    Alignment = std::max(Alignment, Align(4));
+
   if (AlignAllFunctions)
     Alignment = Align(1ULL << AlignAllFunctions);
 
   JumpTableInfo = nullptr;
 
   if (isFuncletEHPersonality(classifyEHPersonality(
           F.hasPersonalityFn() ? F.getPersonalityFn() : nullptr))) {
     WinEHInfo = new (Allocator) WinEHFuncInfo();

llvm/test/CodeGen/ARM/func-sanitizer.ll

+; RUN: llc -mtriple=thumbv6m-none-eabi < %s | FileCheck %s
+
+; CHECK-LABEL: .globl nosan
+; CHECK-NEXT:  .p2align 1
+; CHECK-NEXT:  .type nosan,%function
+; CHECK-NEXT:  .code 16
+; CHECK-NEXT:  .thumb_func
+; CHECK-NEXT:  nosan:
+define dso_local void @nosan() nounwind {
+  ret void
+}
+
+;; The alignment is at least 4 to avoid unaligned type hash loads when this
+;; instrumented function is indirectly called.
+; CHECK-LABEL: .globl foo
+; CHECK-NEXT:  .p2align 2
+; CHECK-NEXT:  .type foo,%function
+; CHECK-NEXT:  .long 3238382334
+; CHECK-NEXT:  .long 3170468932
+; CHECK-NEXT:  .code 16
+; CHECK-NEXT:  .thumb_func
+; CHECK-NEXT:  foo:
+define dso_local void @foo() nounwind !func_sanitize !0 {
+  ret void
+}
+
+;; If "align" is smaller than 4 (required alignment from !func_sanitize), use 4.
+; CHECK-LABEL: .globl align2
+; CHECK-NEXT:  .p2align 2
+define dso_local void @align2() nounwind align 2 !func_sanitize !0 {
+  ret void
+}
+
+;; If "align" is larger than 4, use its value.
+; CHECK-LABEL: .globl align8
+; CHECK-NEXT:  .p2align 3
+define dso_local void @align8() nounwind align 8 !func_sanitize !0 {
+  ret void
+}
+
+!0 = !{i32 -1056584962, i32 -1124498364}

llvm/test/CodeGen/ARM/kcfi.ll

+; RUN: llc -mtriple=thumbv6m-none-eabi < %s | FileCheck %s
+
+; CHECK-LABEL: .globl nosan
+; CHECK-NEXT:  .p2align 1
+; CHECK-NEXT:  .type nosan,%function
+; CHECK-NEXT:  .code 16
+; CHECK-NEXT:  .thumb_func
+; CHECK-NEXT:  nosan:
+define dso_local void @nosan() nounwind {
+  ret void
+}
+
+;; The alignment is at least 4 to avoid unaligned type hash loads when this
+;; instrumented function is indirectly called.
+; CHECK-LABEL: .globl f1
+; CHECK-NEXT:  .p2align 2
+; CHECK-NEXT:  .type f1,%function
+; CHECK-NEXT:  .long 3170468932
+; CHECK-NEXT:  .code 16
+; CHECK-NEXT:  .thumb_func
+; CHECK-NEXT:  f1:
+define void @f1(ptr noundef %x) !kcfi_type !1 {
+  ret void
+}
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 4, !"kcfi", i32 1}
+!1 = !{i32 -1124498364}

llvm/test/CodeGen/RISCV/kcfi-patchable-function-prefix.ll

 ; RUN: llc -mtriple=riscv64 -verify-machineinstrs < %s | FileCheck %s --check-prefixes=CHECK,NOC
 ; RUN: llc -mtriple=riscv64 -mattr=+c -verify-machineinstrs < %s | FileCheck %s --check-prefixes=CHECK,C
 
-; NOC:            .p2align 2
-; C:              .p2align 1
+;; The alignment is at least 4 to avoid unaligned type hash loads when this
+;; instrumented function is indirectly called.
+; CHECK-LABEL:    .globl f1
+; CHECK:          .p2align 2
 ; CHECK-NOT:        nop
 ; CHECK:          .word   12345678
 ; CHECK-LABEL:    f1:
 define void @f1(ptr noundef %x) !kcfi_type !1 {
 ; CHECK:            lw      t1, -4(a0)
   call void %x() [ "kcfi"(i32 12345678) ]
   ret void
 }
 
+; CHECK-LABEL:    .globl f2
 ; NOC:            .p2align 2
 ; C:              .p2align 1
 ; CHECK-NOT:       .word
 ; CHECK-NOT:        nop
 ; CHECK-LABEL:    f2:
 define void @f2(ptr noundef %x) {
 ; CHECK:            lw      t1, -4(a0)
   call void %x() [ "kcfi"(i32 12345678) ]
   ret void
 }
 
-; NOC:            .p2align 2
-; C:              .p2align 1
+; CHECK-LABEL:    .globl f3
+; CHECK:          .p2align 2
 ; CHECK:          .word   12345678
 ; CHECK-COUNT-11:   nop
 ; CHECK-LABEL:    f3:
 define void @f3(ptr noundef %x) #0 !kcfi_type !1 {
 ; NOC:              lw      t1, -48(a0)
 ; C:                lw      t1, -26(a0)
   call void %x() [ "kcfi"(i32 12345678) ]
   ret void
 }
 
+; CHECK-LABEL:    .globl f4
 ; NOC:            .p2align 2
 ; C:              .p2align 1
 ; CHECK-NOT:      .word
 ; CHECK-COUNT-11:   nop
 ; CHECK-LABEL:    f4:
 define void @f4(ptr noundef %x) #0 {
 ; NOC:            lw      t1, -48(a0)
 ; C:              lw      t1, -26(a0)