[KCFI] Fix hash offset calculation in Thumb mode

[samitolvanen]

Summary

ARM stores the Thumb state in the least significant bit of the
function pointers. When compiling for ARM or Thumb, as all
instructions are at least 16-bit aligned, ignore the LSB when
computing the prefix hash location, so we can support both
pure Thumb and mixed ARM/Thumb binaries.

[MaskRay]

@simon_tatham

[simon_tatham]

Thanks! This should fix https://github.com/llvm/llvm-project/issues/62936.

I agree that it's better to do the fix here, instead of adding dedicated support in the backend the way AArch64 does it, because as far as I can see backend KCFI checks are inserted after register allocation and have to use registers that are already dead. And in at least one Arm situation (namely Armv6-M) no such registers are conveniently available, so the check would become more expensive.

llvm/lib/Transforms/Instrumentation/KCFI.cpp

PreservedAnalyses KCFIPass::run(Function &F, FunctionAnalysisManager &AM) {
   if (F.hasFnAttribute("patchable-function-prefix"))
     Ctx.diagnose(
         DiagnosticInfoKCFI("-fpatchable-function-entry=N,M, where M>0 is not "
                            "compatible with -fsanitize=kcfi on this target"));
 
   IntegerType *Int32Ty = Type::getInt32Ty(Ctx);
   MDNode *VeryUnlikelyWeights =
       MDBuilder(Ctx).createBranchWeights(1, (1U << 20) - 1);
+  Triple T(M.getTargetTriple());
 
   for (CallInst *CI : KCFICalls) {
     // Get the expected hash value.
     const uint32_t ExpectedHash =
         cast<ConstantInt>(CI->getOperandBundle(LLVMContext::OB_kcfi)->Inputs[0])
             ->getZExtValue();
 
     // Drop the KCFI operand bundle.
     CallBase *Call =
         CallBase::removeOperandBundle(CI, LLVMContext::OB_kcfi, CI);
     assert(Call != CI);
     Call->copyMetadata(*CI);
     CI->replaceAllUsesWith(Call);
     CI->eraseFromParent();
 
     if (!Call->isIndirectCall())
       continue;
 
     // Emit a check and trap if the target hash doesn't match.
     IRBuilder<> Builder(Call);
-    Value *HashPtr = Builder.CreateConstInBoundsGEP1_32(
-        Int32Ty, Call->getCalledOperand(), -1);
+    Value *FuncPtr = Call->getCalledOperand();
+    // ARM uses the least significant bit of the function pointer to select
+    // between ARM and Thumb modes for the callee. Instructions are always
+    // at least 16-bit aligned, so clear the LSB before we compute the hash
+    // location.
+    if (T.isARM() || T.isThumb()) {
+      FuncPtr = Builder.CreateIntToPtr(
+          Builder.CreateAnd(Builder.CreatePtrToInt(FuncPtr, Int32Ty),
+                            ConstantInt::get(Int32Ty, -2)),
+          FuncPtr->getType());
+    }
+    Value *HashPtr = Builder.CreateConstInBoundsGEP1_32(Int32Ty, FuncPtr, -1);
     Value *Test = Builder.CreateICmpNE(Builder.CreateLoad(Int32Ty, HashPtr),
                                        ConstantInt::get(Int32Ty, ExpectedHash));
     Instruction *ThenTerm =
         SplitBlockAndInsertIfThen(Test, Call, false, VeryUnlikelyWeights);
     Builder.SetInsertPoint(ThenTerm);
     Builder.CreateCall(Intrinsic::getDeclaration(&M, Intrinsic::debugtrap));
     ++NumKCFIChecks;
   }
 
   return PreservedAnalyses::none();
 }

llvm/test/Transforms/KCFI/kcfi.ll

-; RUN: opt -S -passes=kcfi %s | FileCheck %s
+; RUN: opt -S -passes=kcfi %s | FileCheck --check-prefixes=CHECK,NOARM %s
+; RUN: %if arm-registered-target %{ opt -S -passes=kcfi -mtriple=thumbv7m-unknown-linux-gnu %s | FileCheck --check-prefixes=CHECK,ARM %s %}
 
 ; CHECK-LABEL: define void @f1(
 define void @f1(ptr noundef %x) {
-  ; CHECK:      %[[#GEPI:]] = getelementptr inbounds i32, ptr %x, i32 -1
+  ; ARM:        %[[#FINT:]] = ptrtoint ptr %x to i32
+  ; ARM-NEXT:   %[[#FAND:]] = and i32 %[[#FINT]], -2
+  ; ARM-NEXT:   %[[#FPTR:]] = inttoptr i32 %[[#FAND]] to ptr
+  ; ARM-NEXT:   %[[#GEPI:]] = getelementptr inbounds i32, ptr %[[#FPTR]], i32 -1
+  ; NOARM:      %[[#GEPI:]] = getelementptr inbounds i32, ptr %x, i32 -1
   ; CHECK-NEXT: %[[#LOAD:]] = load i32, ptr %[[#GEPI]], align 4
   ; CHECK-NEXT: %[[#ICMP:]] = icmp ne i32 %[[#LOAD]], 12345678
   ; CHECK-NEXT: br i1 %[[#ICMP]], label %[[#TRAP:]], label %[[#CALL:]], !prof ![[#WEIGHTS:]]
   ; CHECK:      [[#TRAP]]:
   ; CHECK-NEXT: call void @llvm.debugtrap()
   ; CHECK-NEXT: br label %[[#CALL]]
   ; CHECK:      [[#CALL]]:
   ; CHECK-NEXT: call void %x()
   ; CHECK-NOT:  [ "kcfi"(i32 12345678) ]
   call void %x() [ "kcfi"(i32 12345678) ]
   ret void
 }
 
 !llvm.module.flags = !{!0}
 !0 = !{i32 4, !"kcfi", i32 1}
 ; CHECK: ![[#WEIGHTS]] = !{!"branch_weights", i32 1, i32 1048575}