tracing: annotate which instructions can branch

PiperOrigin-RevId: 560772931

tracing/BUILD

@@ -199,6 +199,18 @@ cc_library(
     ],
 )
 
+cc_test(
+    name = "disassembler_test",
+    srcs = ["disassembler_test.cc"],
+    deps = [
+        ":capstone_disassembler",
+        ":disassembler",
+        ":xed_disassembler",
+        "@silifuzz//util:arch",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "analysis",
     srcs = ["analysis.cc"],

tracing/capstone_disassembler.cc

@@ -17,16 +17,54 @@
 #include <stdint.h>
 
 #include <cstddef>
-#include <limits>
 #include <string>
 
 #include "absl/strings/str_cat.h"
+#include "third_party/capstone/arm64.h"
 #include "third_party/capstone/capstone.h"
+#include "third_party/capstone/x86.h"
 #include "./tracing/disassembler.h"
+#include "./util/arch.h"
 #include "./util/checks.h"
 
 namespace silifuzz {
 
+namespace {
+
+template <typename Arch>
+bool InstructionCanBranch(cs_insn* decoded_insn);
+
+template <>
+bool InstructionCanBranch<X86_64>(cs_insn* decoded_insn) {
+  cs_detail* detail = decoded_insn->detail;
+  for (size_t i = 0; i < detail->groups_count; i++) {
+    if (detail->groups[i] == X86_GRP_JUMP ||
+        detail->groups[i] == X86_GRP_CALL || detail->groups[i] == X86_GRP_RET ||
+        detail->groups[i] == X86_GRP_BRANCH_RELATIVE) {
+      return true;
+    }
+  }
+  return false;
+}
+
+template <>
+bool InstructionCanBranch<AArch64>(cs_insn* decoded_insn) {
+  cs_detail* detail = decoded_insn->detail;
+  for (size_t i = 0; i < detail->groups_count; i++) {
+    // In practice aarch64 labels all these instructions as "JUMP" but check all
+    // the groups in case this changes.
+    if (detail->groups[i] == ARM64_GRP_JUMP ||
+        detail->groups[i] == ARM64_GRP_CALL ||
+        detail->groups[i] == ARM64_GRP_RET ||
+        detail->groups[i] == ARM64_GRP_BRANCH_RELATIVE) {
+      return true;
+    }
+  }
+  return false;
+}
+
+}  // namespace
+
 template <typename Arch>
 CapstoneDisassembler<Arch>::CapstoneDisassembler() : valid_(false) {
   cs_arch arch;
@@ -43,8 +81,7 @@ CapstoneDisassembler<Arch>::CapstoneDisassembler() : valid_(false) {
     LOG_FATAL("Bad arch_id");
   }
   CHECK_EQ(cs_open(arch, mode, &capstone_handle_), CS_ERR_OK);
-  // TODO(ncbray): turn this on when it's needed.
-  // CHECK_EQ(cs_option(capstone_handle_, CS_OPT_DETAIL, CS_OPT_ON), CS_ERR_OK);
+  CHECK_EQ(cs_option(capstone_handle_, CS_OPT_DETAIL, CS_OPT_ON), CS_ERR_OK);
   decoded_insn_ = cs_malloc(capstone_handle_);
 }
 
@@ -71,6 +108,11 @@ size_t CapstoneDisassembler<Arch>::InstructionSize() const {
   return valid_ ? decoded_insn_->size : 0;
 }
 
+template <typename Arch>
+bool CapstoneDisassembler<Arch>::CanBranch() const {
+  return valid_ ? InstructionCanBranch<Arch>(decoded_insn_) : false;
+}
+
 template <typename Arch>
 std::string CapstoneDisassembler<Arch>::FullText() {
   return valid_

tracing/capstone_disassembler.h

@@ -53,6 +53,8 @@ class CapstoneDisassembler : public Disassembler {
   // How much data was consumed by the last call to Disassemble.
   [[nodiscard]] size_t InstructionSize() const override;
 
+  [[nodiscard]] virtual bool CanBranch() const override;
+
   // The textual representation of the last instruction that was disassembled.
   [[nodiscard]] std::string FullText() override;
 

tracing/disassembler.h

@@ -51,6 +51,10 @@ class Disassembler {
   // How much data was consumed by the last call to Disassemble.
   [[nodiscard]] virtual size_t InstructionSize() const = 0;
 
+  // Can this instruction modify the instruction pointer, in some way?
+  // Includes direct jumps, indirect jumps, calls, returns, etc.
+  [[nodiscard]] virtual bool CanBranch() const = 0;
+
   // The textual representation of the last instruction that was disassembled.
   [[nodiscard]] virtual std::string FullText() = 0;
 

tracing/disassembler_test.cc

@@ -0,0 +1,162 @@
+// Copyright 2023 The SiliFuzz Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "./tracing/disassembler.h"
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include "./tracing/capstone_disassembler.h"
+#include "./tracing/xed_disassembler.h"
+#include "./util/arch.h"
+
+namespace silifuzz {
+
+namespace {
+
+using testing::HasSubstr;
+
+struct DisassemblerTest {
+  std::string bytes;
+  std::string opcode;
+  bool partial_opcode;
+  bool can_branch;
+};
+
+std::string insn(uint32_t value) {
+  return std::string(reinterpret_cast<char*>(&value), sizeof(value));
+}
+
+std::vector<DisassemblerTest> DisassemblerTests_X86_64() {
+  return {
+      {
+          .bytes = {0x90},
+          .opcode = "nop",
+      },
+      {
+          .bytes = {0x51},
+          .opcode = "push",
+      },
+      {
+          .bytes = {0x59},
+          .opcode = "pop",
+      },
+      {
+          .bytes = {0xeb, 0x00},
+          .opcode = "jmp",
+          .can_branch = true,
+      },
+      {
+          .bytes = {0xe2, 0xfe},
+          .opcode = "loop",
+          .can_branch = true,
+      },
+      {
+          .bytes = {0xff, 0xd1},
+          .opcode = "call",
+          .partial_opcode = true,
+          .can_branch = true,
+      },
+      {
+          .bytes = {0xc3},
+          .opcode = "ret",
+          .partial_opcode = true,
+          .can_branch = true,
+      },
+  };
+}
+
+std::vector<DisassemblerTest> DisassemblerTests_AArch64() {
+  return {
+      {
+          .bytes = insn(0xd503201f),
+          .opcode = "nop",
+      },
+      {
+          .bytes = insn(0xa9bf07e0),
+          .opcode = "stp",
+      },
+      {
+          .bytes = insn(0xa8c107e0),
+          .opcode = "ldp",
+      },
+      {
+          .bytes = insn(0xb82850fe),
+          .opcode = "ldsmin",
+      },
+      {
+          .bytes = insn(0x14000000),
+          .opcode = "b",
+          .can_branch = true,
+      },
+      {
+          .bytes = insn(0xd63f0000),
+          .opcode = "blr",
+          .can_branch = true,
+      },
+      {
+          .bytes = insn(0xd65f03c0),
+          .opcode = "ret",
+          .can_branch = true,
+      },
+  };
+}
+
+void RunDisassemblerTest(Disassembler& disasm, const DisassemblerTest& test) {
+  SCOPED_TRACE(test.opcode);
+  constexpr uint64_t kArbitraryAddress = 0x10000;
+  ASSERT_TRUE(disasm.Disassemble(
+      kArbitraryAddress, reinterpret_cast<const uint8_t*>(test.bytes.data()),
+      test.bytes.size()));
+  SCOPED_TRACE(disasm.FullText());
+  EXPECT_EQ(disasm.InstructionSize(), test.bytes.size());
+  std::string opcode = disasm.InstructionIDName(disasm.InstructionID());
+  if (test.partial_opcode) {
+    EXPECT_THAT(opcode, HasSubstr(test.opcode));
+  } else {
+    EXPECT_EQ(opcode, test.opcode);
+  }
+  EXPECT_EQ(disasm.CanBranch(), test.can_branch);
+}
+
+TEST(DisassemblerTest, Xed) {
+  XedDisassembler disasm;
+  const std::vector<DisassemblerTest> tests = DisassemblerTests_X86_64();
+  for (const DisassemblerTest& test : tests) {
+    RunDisassemblerTest(disasm, test);
+  }
+}
+
+TEST(DisassemblerTest, Capstone_x86_64) {
+  CapstoneDisassembler<X86_64> disasm;
+  const std::vector<DisassemblerTest> tests = DisassemblerTests_X86_64();
+  for (const DisassemblerTest& test : tests) {
+    RunDisassemblerTest(disasm, test);
+  }
+}
+
+TEST(DisassemblerTest, Capstone_AArch64) {
+  CapstoneDisassembler<AArch64> disasm;
+  const std::vector<DisassemblerTest> tests = DisassemblerTests_AArch64();
+  for (const DisassemblerTest& test : tests) {
+    RunDisassemblerTest(disasm, test);
+  }
+}
+
+}  // namespace
+
+}  // namespace silifuzz

tracing/execution_trace.h

@@ -39,6 +39,11 @@ struct InstructionInfo {
   // The size of the instruction in bytes.
   uint8_t size;
 
+  // Is this the type of instruction that can branch?
+  // This is type information and it does not indicate if the instruction
+  // actually branches.
+  bool can_branch;
+
   // Did this instruction play a role in producing the final end state?
   // The exact definition of this field depends on the kind of fault injection
   // performed, but at the time of writing this means that the instruction
@@ -171,6 +176,7 @@ absl::Status CaptureTrace(Tracer& tracer, Disassembler& disasm,
     info.address = address;
     info.instruction_id = disasm.InstructionID();
     info.size = disasm.InstructionSize();
+    info.can_branch = disasm.CanBranch();
   });
   absl::Status result = tracer.Run(execution_trace.MaxInstructions());
   // Capture the final state.

tracing/trace_tool.cc

@@ -81,7 +81,8 @@ void LogTrace(Disassembler& disasm, ExecutionTrace<Arch>& execution_trace,
             " addr=", absl::Hex(info.address, absl::kZeroPad8), " offset=",
             absl::Dec(info.address - execution_trace.EntryAddress(),
                       absl::kZeroPad4),
-            " size=", absl::Dec(info.size, absl::kZeroPad2));
+            " size=", absl::Dec(info.size, absl::kZeroPad2), " ",
+            info.can_branch ? "B" : "-");
 
         // How many bits changed?
         UContext<Arch> diff;

tracing/xed_disassembler.cc

@@ -17,7 +17,6 @@
 #include <stdint.h>
 
 #include <cstddef>
-#include <limits>
 #include <string>
 
 #include "absl/base/call_once.h"
@@ -27,6 +26,8 @@
 #include "./util/checks.h"
 
 extern "C" {
+#include "third_party/libxed/xed-category-enum.h"
+#include "third_party/libxed/xed-decoded-inst-api.h"
 #include "third_party/libxed/xed-iclass-enum.h"
 }
 
@@ -61,6 +62,14 @@ size_t XedDisassembler::InstructionSize() const {
   return valid_ ? xed_decoded_inst_get_length(&xedd_) : 0;
 }
 
+bool XedDisassembler::CanBranch() const {
+  if (!valid_) return false;
+
+  xed_category_enum_t category = xed_decoded_inst_get_category(&xedd_);
+  return category == XED_CATEGORY_CALL || category == XED_CATEGORY_COND_BR ||
+         category == XED_CATEGORY_RET || category == XED_CATEGORY_UNCOND_BR;
+}
+
 std::string XedDisassembler::FullText() {
   if (valid_) {
     xed_print_info_t pi;

tracing/xed_disassembler.h

@@ -55,6 +55,8 @@ class XedDisassembler : public Disassembler {
   // How much data was consumed by the last call to Disassemble.
   [[nodiscard]] size_t InstructionSize() const override;
 
+  [[nodiscard]] virtual bool CanBranch() const override;
+
   // The textual representation of the last instruction that was disassembled.
   [[nodiscard]] std::string FullText() override;