[BOLT] Add validation for direct call/branch targets (#165406)

In some edge cases, a binary may contain direct `branch` or `call`
instructions whose target do not point to a valid executable
instruction. This can occur due to compiler bugs, hand-written assembly,
obfuscation technique, **or when control flow targets a data by
mistake.**

We also encountered the problems as described in this
[issue](#149382), where "data
in code" within OpenSSL's hand-written assembly was misidentified as
instructions(island identification seems fail due to the absence of a
corresponding data symbol). The problem occurred because a data sequence
was incorrectly disassembled as a "jb" instruction.

The point here is that the data should not be pointed to by any edge, so
this patch tries to address this by validating the destination address
for **direct branches and calls**. If the target instruction is
invalid(implies a corrupted control flow), this function will be set
ignored.

Although this approach appears helpful for addressing the 'data in code'
problem, its validation might be compromised if the data can be
disassembled as normal instruction.

Author: Jinjie-Huang

bolt/include/bolt/Core/BinaryContext.h

@@ -943,10 +943,11 @@ class BinaryContext {
   /// that should be used by the branch. For example, main or secondary entry
   /// point.
   ///
-  /// If \p Address is an invalid destination, such as a constant island, return
-  /// nullptr and mark \p BF as ignored, since we cannot properly handle a
-  /// branch to a constant island.
-  MCSymbol *handleExternalBranchTarget(uint64_t Address, BinaryFunction &BF);
+  /// This function also performs validations: If \p Address points to an
+  /// invalid instruction or lies within a constant island, return nullptr and
+  /// mark both \p Source and \p Target as ignored.
+  MCSymbol *handleExternalBranchTarget(uint64_t Address, BinaryFunction &Source,
+                                       BinaryFunction &Target);
 
   /// Analyze memory contents at the given \p Address and return the type of
   /// memory contents (such as a possible jump table).

bolt/include/bolt/Core/BinaryFunction.h

@@ -385,6 +385,10 @@ class BinaryFunction {
   /// True if the function should not have an associated symbol table entry.
   bool IsAnonymous{false};
 
+  /// Indicates whether branch validation has already been performed,
+  /// to avoid redundant processing.
+  bool NeedBranchValidation{true};
+
   /// Name for the section this function code should reside in.
   std::string CodeSectionName;
 
@@ -2320,6 +2324,11 @@ class BinaryFunction {
   /// zero-value bytes.
   bool isZeroPaddingAt(uint64_t Offset) const;
 
+  /// Validate if the target of any internal direct branch/call is a valid
+  /// executable instruction.
+  /// Return true if all the targets are valid, false otherwise.
+  bool validateInternalBranches();
+
   /// Check that entry points have an associated instruction at their
   /// offsets after disassembly.
   void postProcessEntryPoints();

bolt/lib/Core/BinaryContext.cpp

@@ -531,20 +531,40 @@ BinaryContext::handleAddressRef(uint64_t Address, BinaryFunction &BF,
 }
 
 MCSymbol *BinaryContext::handleExternalBranchTarget(uint64_t Address,
-                                                    BinaryFunction &BF) {
-  if (BF.isInConstantIsland(Address)) {
-    BF.setIgnored();
-    this->outs() << "BOLT-WARNING: ignoring entry point at address 0x"
-                 << Twine::utohexstr(Address)
-                 << " in constant island of function " << BF << '\n';
-    return nullptr;
+                                                    BinaryFunction &Source,
+                                                    BinaryFunction &Target) {
+  const uint64_t Offset = Address - Target.getAddress();
+  assert(Offset < Target.getSize() &&
+         "Address should be inside the referenced function");
+
+  bool IsValid = true;
+  if (Source.NeedBranchValidation) {
+    if (Target.CurrentState == BinaryFunction::State::Disassembled &&
+        !Target.getInstructionAtOffset(Offset)) {
+      this->errs()
+          << "BOLT-WARNING: corrupted control flow detected in function "
+          << Source
+          << ": an external branch/call targets an invalid instruction "
+          << "in function " << Target << " at address 0x"
+          << Twine::utohexstr(Address) << "; ignoring both functions\n";
+      IsValid = false;
+    }
+    if (Target.isInConstantIsland(Address)) {
+      this->errs() << "BOLT-WARNING: ignoring entry point at address 0x"
+                   << Twine::utohexstr(Address)
+                   << " in constant island of function " << Target << '\n';
+      IsValid = false;
+    }
   }
 
-  const uint64_t Offset = Address - BF.getAddress();
-  assert(Offset < BF.getSize() &&
-         "Address should be inside the referenced function");
+  if (!IsValid) {
+    Source.NeedBranchValidation = false;
+    Source.setIgnored();
+    Target.setIgnored();
+    return nullptr;
+  }
 
-  return Offset ? BF.addEntryPointAtOffset(Offset) : BF.getSymbol();
+  return Offset ? Target.addEntryPointAtOffset(Offset) : Target.getSymbol();
 }
 
 MemoryContentsType BinaryContext::analyzeMemoryAt(uint64_t Address,
@@ -1433,7 +1453,7 @@ void BinaryContext::processInterproceduralReferences() {
 
       // Create an extra entry point if needed. Can also render the target
       // function ignored if the reference is invalid.
-      handleExternalBranchTarget(Address, *TargetFunction);
+      handleExternalBranchTarget(Address, Function, *TargetFunction);
 
       continue;
     }

bolt/lib/Core/BinaryFunction.cpp

@@ -1038,8 +1038,10 @@ MCSymbol *BinaryFunction::getOrCreateLocalLabel(uint64_t Address) {
 
   // For AArch64, check if this address is part of a constant island.
   if (BC.isAArch64()) {
-    if (MCSymbol *IslandSym = getOrCreateIslandAccess(Address))
+    if (MCSymbol *IslandSym = getOrCreateIslandAccess(Address)) {
+      Labels[Offset] = IslandSym;
       return IslandSym;
+    }
   }
 
   if (Offset == getSize())
@@ -1692,7 +1694,7 @@ bool BinaryFunction::scanExternalRefs() {
       // Get a reference symbol for the function when address is a valid code
       // reference.
       BranchTargetSymbol =
-          BC.handleExternalBranchTarget(TargetAddress, *TargetFunction);
+          BC.handleExternalBranchTarget(TargetAddress, *this, *TargetFunction);
       if (!BranchTargetSymbol)
         continue;
     }
@@ -1900,6 +1902,36 @@ bool BinaryFunction::scanExternalRefs() {
   return Success;
 }
 
+bool BinaryFunction::validateInternalBranches() {
+  if (!isSimple() || TrapsOnEntry)
+    return true;
+
+  for (const auto &KV : Labels) {
+    MCSymbol *Label = KV.second;
+    if (getSecondaryEntryPointSymbol(Label))
+      continue;
+
+    const uint32_t Offset = KV.first;
+    // Skip empty functions and out-of-bounds offsets,
+    // as they may not be disassembled.
+    if (!Offset || (Offset > getSize()))
+      continue;
+
+    if (!getInstructionAtOffset(Offset) ||
+        isInConstantIsland(getAddress() + Offset)) {
+      BC.errs() << "BOLT-WARNING: corrupted control flow detected in function "
+                << *this << ": an internal branch/call targets an invalid "
+                << "instruction at address 0x"
+                << Twine::utohexstr(getAddress() + Offset)
+                << "; ignoring this function\n";
+      setIgnored();
+      return false;
+    }
+  }
+
+  return true;
+}
+
 void BinaryFunction::postProcessEntryPoints() {
   if (!isSimple())
     return;

bolt/lib/Rewrite/RewriteInstance.cpp

@@ -3689,6 +3689,7 @@ void RewriteInstance::disassembleFunctions() {
     if (!shouldDisassemble(Function))
       continue;
 
+    Function.validateInternalBranches();
     Function.postProcessEntryPoints();
     Function.postProcessJumpTables();
   }

bolt/test/AArch64/constant-island-alignment.s

@@ -53,6 +53,7 @@ _start:
   blr x0
   mov x0, #1
   ret
+.size _start,.-_start
   nop
 # CHECK: {{0|8}} <$d>:
 .Lci:

bolt/test/AArch64/validate-branch-target.s

@@ -0,0 +1,36 @@
+## Test that BOLT errs when detecting the target 
+## of a direct call/branch is a invalid instruction
+
+# REQUIRES: system-linux
+# RUN: rm -rf %t && mkdir -p %t && cd %t
+# RUN: llvm-mc -filetype=obj -triple aarch64-unknown-linux %s -o main.o
+# RUN: %clang %cflags %t/main.o -o main.exe -Wl,-q
+# RUN: llvm-bolt %t/main.exe -o %t/main.exe.bolt -lite=0 2>&1 | FileCheck %s --check-prefix=CHECK-TARGETS
+
+# CHECK-TARGETS: BOLT-WARNING: corrupted control flow detected in function external_corrupt: an external branch/call targets an invalid instruction in function external_func at address 0x{{[0-9a-f]+}}; ignoring both functions
+# CHECK-TARGETS: BOLT-WARNING: corrupted control flow detected in function internal_corrupt: an internal branch/call targets an invalid instruction at address 0x{{[0-9a-f]+}}; ignoring this function
+
+
+.globl  internal_corrupt
+.type   internal_corrupt,@function
+internal_corrupt:
+    b constant_island_0  // targeting the data in code
+constant_island_0:
+    .word 0xffffffff
+.size   internal_corrupt,.-internal_corrupt
+
+
+.globl  external_corrupt
+.type   external_corrupt,@function
+external_corrupt:
+    b   constant_island_1  // targeting the data in code externally
+.size   external_corrupt,.-external_corrupt
+
+.globl  external_func
+.type   external_func,@function
+external_func:
+    add x0, x0, x1
+constant_island_1:
+    .word 0xffffffff // data in code
+    ret
+.size   external_func,.-external_func

bolt/test/X86/validate-branch-target.s

@@ -0,0 +1,33 @@
+## Test that BOLT errs when detecting the target 
+## of a direct call/branch is a invalid instruction
+
+# REQUIRES: system-linux
+# RUN: rm -rf %t && mkdir -p %t && cd %t
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-linux %s -o main.o
+# RUN: %clang %cflags -pie -Wl,-q %t/main.o -o main.exe
+# RUN: llvm-bolt %t/main.exe -o %t/main.exe.bolt -lite=0 2>&1 | FileCheck %s --check-prefix=CHECK-TARGETS
+
+# CHECK-TARGETS: BOLT-WARNING: corrupted control flow detected in function external_corrupt: an external branch/call targets an invalid instruction in function external_func at address 0x{{[0-9a-f]+}}; ignoring both functions
+# CHECK-TARGETS: BOLT-WARNING: corrupted control flow detected in function internal_corrupt: an internal branch/call targets an invalid instruction at address 0x{{[0-9a-f]+}}; ignoring this function
+
+
+.globl	internal_corrupt
+.type	internal_corrupt,@function
+internal_corrupt:
+	jb  data_in_code + 1  # targeting the data in code, and jump into the middle of 'xorb' instruction
+data_in_code:
+	.byte 0x34, 0x01 # data in code, will be disassembled as 'xorb 0x1, %al'
+.size	internal_corrupt,.-internal_corrupt
+
+
+.globl	external_corrupt
+.type	external_corrupt,@function
+external_corrupt:
+	jb  external_func + 1  # targeting the middle of normal instruction externally
+.size	external_corrupt,.-external_corrupt
+
+.globl	external_func
+.type	external_func,@function
+external_func:
+	addq  $1, %rax  # normal instruction
+.size	external_func,.-external_func