UBet specific adaptations

polytracker/build.py

@@ -147,6 +147,26 @@ def _optimize_bitcode(input_bitcode: Path, output_bitcode: Path) -> None:
     subprocess.check_call(cmd)
 
 
+def _preopt_instrument_bitcode(input_bitcode: Path, output_bitcode: Path) -> None:
+    POLY_PASS_PATH: Path = _ensure_path_exists(
+        _compiler_dir_path() / "pass" / "libPolytrackerPass.so"
+    )
+
+    cmd = [
+        "opt",
+        "-load",
+        str(POLY_PASS_PATH),
+        "-load-pass-plugin",
+        str(POLY_PASS_PATH),
+        "-passes=pt-tcf",
+        str(input_bitcode),
+        "-o",
+        str(output_bitcode),
+    ]
+    # execute `cmd`
+    subprocess.check_call(cmd)
+
+
 def _instrument_bitcode(
     input_bitcode: Path,
     output_bitcode: Path,
@@ -398,16 +418,30 @@ def __init_arguments__(self, parser: argparse.ArgumentParser):
             help="specify additional ignore lists to polytracker",
         )
 
+        parser.add_argument(
+            "--cflog",
+            action="store_true",
+            help="instrument with control affecting dataflow logging",
+        )
+
     def run(self, args: argparse.Namespace):
         for target in args.targets:
             blight_cmds = _read_blight_journal(args.journal_path)
             target_cmd, target_path = _find_target(target, blight_cmds)
             bc_path = target_path.with_suffix(".bc")
+            opt_bc = bc_path.with_suffix(".opt.bc")
             _extract_bitcode(target_path, bc_path)
-            _optimize_bitcode(bc_path, bc_path)
+            if args.cflog:
+                # Control affecting data flow logging happens before optimization
+                pre_opt = bc_path.with_suffix(".preopt.bc")
+                _preopt_instrument_bitcode(bc_path, pre_opt)
+
+                _optimize_bitcode(pre_opt, opt_bc)
+            else:
+                _optimize_bitcode(bc_path, opt_bc)
             inst_bc_path = Path(f"{bc_path.stem}.instrumented.bc")
             _instrument_bitcode(
-                bc_path,
+                opt_bc,
                 inst_bc_path,
                 args.ignore_lists,
                 args.taint,

polytracker/custom_abi/dfsan_abilist.txt

@@ -30,6 +30,10 @@ fun:open64=custom
 ##########################################
 # Polytracker functions
 #########################################
+fun:__polytracker_leave_function=uninstrumented
+fun:__polytracker_leave_function=discard
+fun:__polytracker_enter_function=uninstrumented
+fun:__polytracker_enter_function=discard
 fun:__polytracker_log_func_entry=uninstrumented
 fun:__polytracker_log_func_entry=discard
 fun:__polytracker_log_func_exit=uninstrumented
@@ -43,6 +47,8 @@ fun:__polytracker_log_taint_op=uninstrumented
 fun:__polytracker_log_taint_op=custom
 fun:__polytracker_log_conditional_branch=uninstrumented
 fun:__polytracker_log_conditional_branch=custom
+fun:__polytracker_log_tainted_control_flow=uninstrumented
+fun:__polytracker_log_tainted_control_flow=custom
 # -- end
 fun:__polytracker_dump=uninstrumented
 fun:__polytracker_dump=discard

polytracker/include/polytracker/passes/tainted_control_flow.h

@@ -0,0 +1,66 @@
+/*
+ * Copyright (c) 2023-present, Trail of Bits, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed in accordance with the terms specified in
+ * the LICENSE file found in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <llvm/IR/InstVisitor.h>
+#include <llvm/IR/PassManager.h>
+#include <unordered_map>
+
+namespace polytracker {
+namespace detail {
+struct FunctionMappingJSONWriter;
+}
+
+class TaintedControlFlowPass
+    : public llvm::PassInfoMixin<TaintedControlFlowPass>,
+      public llvm::InstVisitor<TaintedControlFlowPass> {
+  //
+  llvm::IntegerType *label_ty{nullptr};
+  // Taint tracking startup
+  llvm::FunctionCallee taint_start_fn;
+  // Log taint label affecting control flow
+  llvm::FunctionCallee cond_br_log_fn;
+  // Log enter/leave functions
+  llvm::FunctionCallee fn_enter_log_fn;
+  llvm::FunctionCallee fn_leave_log_fn;
+
+  // Helpers
+  void insertCondBrLogCall(llvm::Instruction &inst, llvm::Value *val);
+  void insertTaintStartupCall(llvm::Module &mod);
+  void declareLoggingFunctions(llvm::Module &mod);
+
+  llvm::ConstantInt *get_function_id_const(llvm::Function &f);
+  llvm::ConstantInt *get_function_id_const(llvm::Instruction &i);
+
+public:
+  using function_id = uint32_t;
+
+  TaintedControlFlowPass();
+  TaintedControlFlowPass(TaintedControlFlowPass &&);
+  ~TaintedControlFlowPass();
+
+  llvm::PreservedAnalyses run(llvm::Module &mod,
+                              llvm::ModuleAnalysisManager &mam);
+  void visitGetElementPtrInst(llvm::GetElementPtrInst &gep);
+  void visitBranchInst(llvm::BranchInst &bi);
+  void visitSwitchInst(llvm::SwitchInst &si);
+  void visitSelectInst(llvm::SelectInst &si);
+
+  void instrumentFunctionEnter(llvm::Function &func);
+  void visitReturnInst(llvm::ReturnInst &ri);
+
+  function_id function_mapping(llvm::Function &func);
+
+  std::unordered_map<uintptr_t, function_id> function_ids_;
+  function_id function_counter_{0};
+
+  std::unique_ptr<detail::FunctionMappingJSONWriter> function_mapping_writer_;
+};
+
+} // namespace polytracker

polytracker/include/taintdag/control_flow_log.h

@@ -0,0 +1,88 @@
+
+/*
+ * Copyright (c) 2022-present, Trail of Bits, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed in accordance with the terms specified in
+ * the LICENSE file found in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include "taintdag/outputfile.h"
+#include "taintdag/section.h"
+#include "taintdag/taint.h"
+#include "taintdag/util.h"
+
+namespace taintdag {
+
+namespace details {
+// A uint32_t varint encoded by setting highest bit for all but the final byte.
+// Requires up to 5 bytes of storage as each output byte uses 7 input bits.
+// Total maximum need is floor(32/7) = 5. Returns number of bytes required.
+size_t varint_encode(uint32_t val, uint8_t *buffer) {
+  auto orig_buffer = buffer;
+  while (val >= 0x80) {
+    *buffer++ = 0x80 | (val & 0x7f);
+    val >>= 7;
+  }
+  *buffer++ = val & 0x7f;
+  return buffer - orig_buffer;
+}
+// TODO (hbrodin): Should probably used std::span
+} // namespace details
+
+struct ControlFlowLog : public SectionBase {
+  enum EventType {
+    EnterFunction = 0,
+    LeaveFunction = 1,
+    TaintedControlFlow = 2,
+  };
+
+  static constexpr uint8_t tag{8};
+  static constexpr size_t align_of{1};
+  static constexpr size_t allocation_size{1024 * 1024 * 1024};
+
+  template <typename OF>
+  ControlFlowLog(SectionArg<OF> of) : SectionBase(of.range) {}
+
+  void function_event(EventType evt, uint32_t function_id) {
+    uint8_t buffer[6];
+    buffer[0] = static_cast<uint8_t>(evt);
+    auto used = details::varint_encode(function_id, &buffer[1]);
+    auto total = used + 1;
+
+    if (auto wctx = write(total)) {
+      std::copy(&buffer[0], &buffer[total], wctx->mem.begin());
+    } else {
+      error_exit("Failed to write ", total,
+                 " bytes of output to the ControlFlowLog Section.");
+    }
+  }
+  void enter_function(uint32_t function_id) {
+    function_event(EnterFunction, function_id);
+  }
+
+  void leave_function(uint32_t function_id) {
+    function_event(LeaveFunction, function_id);
+  }
+
+  void tainted_control_flow(label_t label, uint32_t function_id) {
+    // 1 byte event, <= 5 bytes function id, <= 5 bytes label
+    uint8_t buffer[11];
+    buffer[0] = static_cast<uint8_t>(TaintedControlFlow);
+    auto used = details::varint_encode(function_id, &buffer[1]);
+    auto total = used + 1;
+    used = details::varint_encode(label, &buffer[total]);
+    total += used;
+
+    if (auto wctx = write(total)) {
+      std::copy(&buffer[0], &buffer[total], wctx->mem.begin());
+    } else {
+      error_exit("Failed to write ", total,
+                 " bytes of output to the ControlFlowLog Section.");
+    }
+  }
+};
+
+} // namespace taintdag

polytracker/include/taintdag/polytracker.h

@@ -12,6 +12,7 @@
 #include <span>
 
 #include "taintdag/bitmap_section.h"
+#include "taintdag/control_flow_log.h"
 #include "taintdag/fnmapping.h"
 #include "taintdag/fntrace.h"
 #include "taintdag/labels.h"
@@ -54,6 +55,21 @@ class PolyTracker {
   // Update the label, it affects control flow
   void affects_control_flow(label_t taint_label);
 
+  // Instrumentation callback for when control flow is influenced by a
+  // a tainted value
+  void log_tainted_control_flow(label_t taint_label, uint32_t function_id);
+
+  // Instrumentation callback for when execution enters a function
+  // NOTE: There is a overlap in functionality between this and `function_entry`
+  // they will co-exist for now as they operate slightly different. The
+  // underlying reason is that this was developed separately to support the
+  // Tainted Control Flow logging mechanism.
+  void enter_function(uint32_t function_id);
+
+  // Instrumentation callback for when execution leaves a function
+  // NOTE: Se `enter_function` comment about overlap.
+  void leave_function(uint32_t function_id);
+
   // Log tainted data flowed into the sink
   void taint_sink(int fd, util::Offset offset, void const *mem, size_t length);
   // Same as before, but use same label for all data
@@ -79,7 +95,7 @@ class PolyTracker {
   // sections and in which order they appear.
   using ConcreteOutputFile =
       OutputFile<Sources, Labels, StringTable, TaintSink,
-                 SourceLabelIndexSection, Functions, Events>;
+                 SourceLabelIndexSection, Functions, Events, ControlFlowLog>;
   ConcreteOutputFile output_file_;
 
   // Tracking source offsets for streams (where offsets can be determined by

polytracker/src/passes/CMakeLists.txt

@@ -6,7 +6,7 @@ endif(APPLE)
 
 add_library(
   PolytrackerPass SHARED
-  taint_tracking.cpp remove_fn_attr.cpp function_tracing.cpp
+  taint_tracking.cpp remove_fn_attr.cpp function_tracing.cpp tainted_control_flow.cpp
   DataFlowSanitizer.cpp utils.cpp pass_plugin.cpp)
 
 target_link_libraries(

polytracker/src/passes/pass_plugin.cpp

@@ -13,6 +13,7 @@
 #include "polytracker/passes/function_tracing.h"
 #include "polytracker/passes/remove_fn_attr.h"
 #include "polytracker/passes/taint_tracking.h"
+#include "polytracker/passes/tainted_control_flow.h"
 
 llvm::PassPluginLibraryInfo getPolyTrackerPluginInfo() {
   return {LLVM_PLUGIN_API_VERSION, "PolyTracker", "",
@@ -36,6 +37,10 @@ llvm::PassPluginLibraryInfo getPolyTrackerPluginInfo() {
                     mpm.addPass(polytracker::FunctionTracingPass());
                     return true;
                   }
+                  if (name == "pt-tcf") {
+                    mpm.addPass(polytracker::TaintedControlFlowPass());
+                    return true;
+                  }
                   return false;
                 });
           }};