Cross-DSO control flow integrity (LLVM part)

[eugenis]

Summary

An LTO pass that inserts a __cfi_check() function that validates a call based on a hash of the call-site-known type and the target pointer.

[pcc]

Please also add design documentation to docs/ControlFlowIntegrityDesign.rst.

include/llvm/Transforms/IPO/CrossDSOCFI.h

21 ↗

(On Diff #42257)

Why is this in a header? Were you planning to add unit tests for this function?

lib/Transforms/IPO/CrossDSOCFI.cpp

2

module's

11

Please do.

49

This will yield a different result depending on the endianness of the host machine.

64

This function doesn't look like it is defined anywhere.

90

This is already happening, no?

101

This won't be necessary if we teach the clang frontend to not emit bitset entries for declarations when doing cross-DSO.

105

This isn't exactly right; MDNode::isDistinct() corresponds directly to whether the bitset entry has global scope. I would check isDistinct here and assert that the identifier is an MDString (since we don't support hashing anything else).

lib/Transforms/IPO/PassManagerBuilder.cpp

603

I'd prefer if this stated what the function does (i.e. "create function that ...").

test/Transforms/CrossDSOCFI/basic.ll

6

Probably should check for hash values here since we know what they're going to be.

37–44

The body of the test could be made a lot simpler. E.g. these don't need to be vtables.

[pcc]

include/llvm/Transforms/IPO/CrossDSOCFI.h

22 ↗

(On Diff #42258)

Oh, I see that you need this in the frontend. I wonder if there's a better way to do this that doesn't involve the frontend and the pass conspiring to use the same hash values? Maybe the bitset identifiers should be hash values in cross-DSO mode?

[eugenis]

include/llvm/Transforms/IPO/CrossDSOCFI.h

22 ↗

(On Diff #42258)

Switched to hash values in bitset identifiers. The frontend now emits 2 bitset entries for each function or vtable entry: one with a string identifier for local checks, and one with a hash value for cross-DSO checks. We don't use hash values for local checks because of possible hash function collisions that could weaken the checks.

lib/Transforms/IPO/CrossDSOCFI.cpp

50

Moved to clang (see D15367) and fixed.

91

You mean insertion at the end of the module? I think so, but I wonder if anything else can be done to ensure that nothing gets reordered later.
We will catch any such problem at runtime.

102

Moved to a separate function and replaced a "continue" with an assert, just in case.

106

This is gone.

[pcc]

lib/Transforms/IPO/CrossDSOCFI.cpp

85

Early return here.

92

Sure, so let's word the FIXME like that.

112

I think this will break if the ConstantInts are not of type i64. Can you handle this gracefully?

141

Isn't this non-deterministic? I think it will depend on the ConstantInt addresses.

test/Transforms/CrossDSOCFI/basic.ll

18

I think you can make these tests more specific once the non-determinism has been fixed in the code.

[eugenis]

lib/Transforms/IPO/CrossDSOCFI.cpp

112

Do we really need to handle this case? Frontend never inserts constantints in the bitset metadata that are not i64, and we can make this a requirement.

141

Good catch. Fixed.

[pcc]

lib/Transforms/IPO/CrossDSOCFI.cpp

113

It can be a requirement but we shouldn't produce invalid IR if it happens. I think the easiest thing to do is to return nullptr from CrossDSOCFI::extractBitSetTypeId if the ConstantInt is not i64.

test/Transforms/CrossDSOCFI/basic.ll

6

You can test for the label names as well (here and below).

[pcc]

LGTM with nit.

lib/Transforms/IPO/CrossDSOCFI.cpp

98

Can return void

[pcc]

Sorry, a few other nits I noticed.

lib/Transforms/IPO/CrossDSOCFI.cpp

122

You don't need to cast this.

125

Same here.

145

Don't shadow TypeId.

[eugenis]

r255693, thanks!

include/llvm/InitializePasses.h

 void initializeCFGViewerPass(PassRegistry&);
 void initializeConstantHoistingPass(PassRegistry&);
 void initializeCodeGenPreparePass(PassRegistry&);
 void initializeConstantMergePass(PassRegistry&);
 void initializeConstantPropagationPass(PassRegistry&);
 void initializeMachineCopyPropagationPass(PassRegistry&);
 void initializeCostModelAnalysisPass(PassRegistry&);
 void initializeCorrelatedValuePropagationPass(PassRegistry&);
+void initializeCrossDSOCFIPass(PassRegistry&);
 void initializeDAEPass(PassRegistry&);
 void initializeDAHPass(PassRegistry&);
 void initializeDCEPass(PassRegistry&);
 void initializeDSEPass(PassRegistry&);
 void initializeDeadInstEliminationPass(PassRegistry&);
 void initializeDeadMachineInstructionElimPass(PassRegistry&);
 void initializeDelinearizationPass(PassRegistry &);
 void initializeDependenceAnalysisPass(PassRegistry&);

include/llvm/Transforms/IPO.h

 /// createBarrierNoopPass - This pass is purely a module pass barrier in a pass
 /// manager.
 ModulePass *createBarrierNoopPass();
 
 /// \brief This pass lowers bitset metadata and the llvm.bitset.test intrinsic
 /// to bitsets.
 ModulePass *createLowerBitSetsPass();
 
+/// \brief This pass export CFI checks for use by external modules.
+ModulePass *createCrossDSOCFIPass();
+
 //===----------------------------------------------------------------------===//
 // SampleProfilePass - Loads sample profile data from disk and generates
 // IR metadata to reflect the profile.
 ModulePass *createSampleProfileLoaderPass();
 ModulePass *createSampleProfileLoaderPass(StringRef Name);
 
 } // End llvm namespace
 
 #endif

lib/Transforms/IPO/CMakeLists.txt

 add_llvm_library(LLVMipo
   ArgumentPromotion.cpp
   BarrierNoopPass.cpp
   ConstantMerge.cpp
+  CrossDSOCFI.cpp
   DeadArgumentElimination.cpp
   ElimAvailExtern.cpp
   ExtractGV.cpp
   FunctionAttrs.cpp
   FunctionImport.cpp
   GlobalDCE.cpp
   GlobalOpt.cpp
   IPConstantPropagation.cpp

lib/Transforms/IPO/CrossDSOCFI.cpp

+//===-- CrossDSOCFI.cpp - Externalize this module's CFI checks ------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass exports all llvm.bitset's found in the module in the form of a
+// __cfi_check function, which can be used to verify cross-DSO call targets.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/IPO.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/EquivalenceClasses.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/IR/Constant.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/GlobalObject.h"
+#include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/MDBuilder.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Operator.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "cross-dso-cfi"
+
+STATISTIC(TypeIds, "Number of unique type identifiers");
+
+namespace {
+
+struct CrossDSOCFI : public ModulePass {
+  static char ID;
+  CrossDSOCFI() : ModulePass(ID) {
+    initializeCrossDSOCFIPass(*PassRegistry::getPassRegistry());
+  }
+
+  Module *M;
+  MDNode *VeryLikelyWeights;
+
+  ConstantInt *extractBitSetTypeId(MDNode *MD);
+  void buildCFICheck();
+
+  bool doInitialization(Module &M) override;
+  bool runOnModule(Module &M) override;
+};
+
+} // anonymous namespace
+
+INITIALIZE_PASS_BEGIN(CrossDSOCFI, "cross-dso-cfi", "Cross-DSO CFI", false,
+                      false)
+INITIALIZE_PASS_END(CrossDSOCFI, "cross-dso-cfi", "Cross-DSO CFI", false, false)
+char CrossDSOCFI::ID = 0;
+
+ModulePass *llvm::createCrossDSOCFIPass() { return new CrossDSOCFI; }
+
+bool CrossDSOCFI::doInitialization(Module &Mod) {
+  M = &Mod;
+  VeryLikelyWeights =
+      MDBuilder(M->getContext()).createBranchWeights((1U << 20) - 1, 1);
+
+  return false;
+}
+
+/// extractBitSetTypeId - Extracts TypeId from a hash-based bitset MDNode.
+ConstantInt *CrossDSOCFI::extractBitSetTypeId(MDNode *MD) {
+  // This check excludes vtables for classes inside anonymous namespaces.
+  auto TM = dyn_cast<ValueAsMetadata>(MD->getOperand(0));
+  if (!TM)
+    return nullptr;
+  auto C = dyn_cast_or_null<ConstantInt>(TM->getValue());
+  if (!C) return nullptr;
+  // We are looking for i64 constants.
+  if (C->getBitWidth() != 64) return nullptr;
+
+  // Sanity check.
+  auto FM = dyn_cast_or_null<ValueAsMetadata>(MD->getOperand(1));
+  // Can be null if a function was removed by an optimization.
+  if (FM) {
+    auto F = dyn_cast<Function>(FM->getValue());
+    // But can never be a function declaration.
+    assert(!F || !F->isDeclaration());
+  }
+  return C;
+}
+
+/// buildCFICheck - emits __cfi_check for the current module.
+void CrossDSOCFI::buildCFICheck() {
+  // FIXME: verify that __cfi_check ends up near the end of the code section,
+  // but before the jump slots created in LowerBitSets.
+  llvm::DenseSet<uint64_t> BitSetIds;
+  NamedMDNode *BitSetNM = M->getNamedMetadata("llvm.bitsets");
+
+  if (BitSetNM)
+    for (unsigned I = 0, E = BitSetNM->getNumOperands(); I != E; ++I)
+      if (ConstantInt *TypeId = extractBitSetTypeId(BitSetNM->getOperand(I)))
+        BitSetIds.insert(TypeId->getZExtValue());
+
+  LLVMContext &Ctx = M->getContext();
+  Constant *C = M->getOrInsertFunction(
+      "__cfi_check",
+      FunctionType::get(
+          Type::getVoidTy(Ctx),
+          {Type::getInt64Ty(Ctx), PointerType::getUnqual(Type::getInt8Ty(Ctx))},
+          false));
+  Function *F = dyn_cast<Function>(C);
+  F->setAlignment(4096);
+  auto args = F->arg_begin();
+  Argument &CallSiteTypeId = *(args++);
+  CallSiteTypeId.setName("CallSiteTypeId");
+  Argument &Addr = *(args++);
+  Addr.setName("Addr");
+  assert(args == F->arg_end());
+
+  BasicBlock *BB = BasicBlock::Create(Ctx, "entry", F);
+
+  BasicBlock *TrapBB = BasicBlock::Create(Ctx, "trap", F);
+  IRBuilder<> IRBTrap(TrapBB);
+  Function *TrapFn = Intrinsic::getDeclaration(M, Intrinsic::trap);
+  llvm::CallInst *TrapCall = IRBTrap.CreateCall(TrapFn);
+  TrapCall->setDoesNotReturn();
+  TrapCall->setDoesNotThrow();
+  IRBTrap.CreateUnreachable();
+
+  BasicBlock *ExitBB = BasicBlock::Create(Ctx, "exit", F);
+  IRBuilder<> IRBExit(ExitBB);
+  IRBExit.CreateRetVoid();
+
+  IRBuilder<> IRB(BB);
+  SwitchInst *SI = IRB.CreateSwitch(&CallSiteTypeId, TrapBB, BitSetIds.size());
+  for (uint64_t TypeId : BitSetIds) {
+    ConstantInt *CaseTypeId = ConstantInt::get(Type::getInt64Ty(Ctx), TypeId);
+    BasicBlock *TestBB = BasicBlock::Create(Ctx, "test", F);
+    IRBuilder<> IRBTest(TestBB);
+    Function *BitsetTestFn =
+        Intrinsic::getDeclaration(M, Intrinsic::bitset_test);
+
+    Value *Test = IRBTest.CreateCall(
+        BitsetTestFn, {&Addr, MetadataAsValue::get(
+                                  Ctx, ConstantAsMetadata::get(CaseTypeId))});
+    BranchInst *BI = IRBTest.CreateCondBr(Test, ExitBB, TrapBB);
+    BI->setMetadata(LLVMContext::MD_prof, VeryLikelyWeights);
+
+    SI->addCase(CaseTypeId, TestBB);
+    ++TypeIds;
+  }
+}
+
+bool CrossDSOCFI::runOnModule(Module &M) {
+  if (M.getModuleFlag("Cross-DSO CFI") == nullptr)
+    return false;
+  buildCFICheck();
+  return true;
+}

lib/Transforms/IPO/IPO.cpp

 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/Transforms/IPO.h"
 
 using namespace llvm;
 
 void llvm::initializeIPO(PassRegistry &Registry) {
   initializeArgPromotionPass(Registry);
   initializeConstantMergePass(Registry);
+  initializeCrossDSOCFIPass(Registry);
   initializeDAEPass(Registry);
   initializeDAHPass(Registry);
   initializeFunctionAttrsPass(Registry);
   initializeGlobalDCEPass(Registry);
   initializeGlobalOptPass(Registry);
   initializeIPCPPass(Registry);
   initializeAlwaysInlinerPass(Registry);
   initializeSimpleInlinerPass(Registry);

lib/Transforms/IPO/PassManagerBuilder.cpp

if (LibraryInfo)
     PM.add(new TargetLibraryInfoWrapperPass(*LibraryInfo));
 
   if (VerifyInput)
     PM.add(createVerifierPass());
 
   if (OptLevel > 1)
     addLTOOptimizationPasses(PM);
 
+  // Create a function that performs CFI checks for cross-DSO calls with targets
+  // in the current module.
+  PM.add(createCrossDSOCFIPass());
+
   // Lower bit sets to globals. This pass supports Clang's control flow
   // integrity mechanisms (-fsanitize=cfi*) and needs to run at link time if CFI
   // is enabled. The pass does nothing if CFI is disabled.
   PM.add(createLowerBitSetsPass());
 
   if (OptLevel != 0)
     addLateLTOOptimizationPasses(PM);
 

test/Transforms/CrossDSOCFI/basic.ll

+; RUN: opt -S -cross-dso-cfi < %s | FileCheck %s
+
+; CHECK:     define void @__cfi_check(i64 %[[TYPE:.*]], i8* %[[ADDR:.*]]) align 4096
+; CHECK:     switch i64 %[[TYPE]], label %[[TRAP:.*]] [
+; CHECK-NEXT:   i64 111, label %[[L1:.*]]
+; CHECK-NEXT:   i64 222, label %[[L2:.*]]
+; CHECK-NEXT:   i64 333, label %[[L3:.*]]
+; CHECK-NEXT:   i64 444, label %[[L4:.*]]
+; CHECK-NEXT: {{]$}}
+
+; CHECK:     [[TRAP]]:
+; CHECK-NEXT:   call void @llvm.trap()
+; CHECK-MEXT:   unreachable
+
+; CHECK:     [[EXIT:.*]]:
+; CHECK-NEXT:   ret void
+
+; CHECK:     [[L1]]:
+; CHECK-NEXT:   call i1 @llvm.bitset.test(i8* %[[ADDR]], metadata i64 111)
+; CHECK-NEXT:   br {{.*}} label %[[EXIT]], label %[[TRAP]]
+
+; CHECK:     [[L2]]:
+; CHECK-NEXT:   call i1 @llvm.bitset.test(i8* %[[ADDR]], metadata i64 222)
+; CHECK-NEXT:   br {{.*}} label %[[EXIT]], label %[[TRAP]]
+
+; CHECK:     [[L3]]:
+; CHECK-NEXT:   call i1 @llvm.bitset.test(i8* %[[ADDR]], metadata i64 333)
+; CHECK-NEXT:   br {{.*}} label %[[EXIT]], label %[[TRAP]]
+
+; CHECK:     [[L4]]:
+; CHECK-NEXT:   call i1 @llvm.bitset.test(i8* %[[ADDR]], metadata i64 444)
+; CHECK-NEXT:   br {{.*}} label %[[EXIT]], label %[[TRAP]]
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+@_ZTV1A = constant i8 0
+@_ZTI1A = constant i8 0
+@_ZTS1A = constant i8 0
+@_ZTV1B = constant i8 0
+@_ZTI1B = constant i8 0
+@_ZTS1B = constant i8 0
+
+define signext i8 @f11() {
+entry:
+  ret i8 1
+}
+
+define signext i8 @f12() {
+entry:
+  ret i8 2
+}
+
+define signext i8 @f13() {
+entry:
+  ret i8 3
+}
+
+define i32 @f21() {
+entry:
+  ret i32 4
+}
+
+define i32 @f22() {
+entry:
+  ret i32 5
+}
+
+!llvm.bitsets = !{!0, !1, !2, !3, !4, !7, !8, !9, !10, !11, !12, !13, !14, !15}
+!llvm.module.flags = !{!17}
+
+!0 = !{!"_ZTSFcvE", i8 ()* @f11, i64 0}
+!1 = !{i64 111, i8 ()* @f11, i64 0}
+!2 = !{!"_ZTSFcvE", i8 ()* @f12, i64 0}
+!3 = !{i64 111, i8 ()* @f12, i64 0}
+!4 = !{!"_ZTSFcvE", i8 ()* @f13, i64 0}
+!5 = !{i64 111, i8 ()* @f13, i64 0}
+!6 = !{!"_ZTSFivE", i32 ()* @f21, i64 0}
+!7 = !{i64 222, i32 ()* @f21, i64 0}
+!8 = !{!"_ZTSFivE", i32 ()* @f22, i64 0}
+!9 = !{i64 222, i32 ()* @f22, i64 0}
+!10 = !{!"_ZTS1A", i8* @_ZTV1A, i64 16}
+!11 = !{i64 333, i8* @_ZTV1A, i64 16}
+!12 = !{!"_ZTS1A", i8* @_ZTV1B, i64 16}
+!13 = !{i64 333, i8* @_ZTV1B, i64 16}
+!14 = !{!"_ZTS1B", i8* @_ZTV1B, i64 16}
+!15 = !{i64 444, i8* @_ZTV1B, i64 16}
+!17= !{i32 4, !"Cross-DSO CFI", i32 1}