Basic support for assert()

polytracker/custom_abi/dfsan_abilist.txt

@@ -154,7 +154,8 @@ fun:llrintl=functional
 
 # Functions that produce an output that does not depend on the input (shadow is
 # zeroed automatically).
-fun:__assert_fail=discard
+fun:__assert_fail=uninstrumented
+fun:__assert_fail=custom
 fun:__ctype_b_loc=discard
 fun:__cxa_atexit=discard
 fun:__errno_location=discard
@@ -547,7 +548,6 @@ fun:__asinl_finite=uninstrumented
 fun:__asprintf=uninstrumented
 fun:__asprintf_chk=uninstrumented
 fun:__assert=uninstrumented
-fun:__assert_fail=uninstrumented
 fun:__assert_perror_fail=uninstrumented
 fun:__atan2_finite=uninstrumented
 fun:__atan2f_finite=uninstrumented

polytracker/custom_abi/polytracker_abilist.txt

@@ -154,7 +154,8 @@ fun:llrintl=functional
 
 # Functions that produce an output that does not depend on the input (shadow is
 # zeroed automatically).
-fun:__assert_fail=discard
+fun:__assert_fail=uninstrumented
+fun:__assert_fail=custom
 fun:__ctype_b_loc=discard
 fun:__cxa_atexit=discard
 fun:__errno_location=discard
@@ -548,7 +549,6 @@ fun:__asinl_finite=uninstrumented
 fun:__asprintf=uninstrumented
 fun:__asprintf_chk=uninstrumented
 fun:__assert=uninstrumented
-fun:__assert_fail=uninstrumented
 fun:__assert_perror_fail=uninstrumented
 fun:__atan2_finite=uninstrumented
 fun:__atan2f_finite=uninstrumented

polytracker/include/taintdag/outputfile.h

@@ -88,7 +88,7 @@ template <Section... Sections> class OutputFile {
 
   // The FileHeader is constructed at offset zero of the mapped memory (file).
   // It contains the metadata required to parse a TDAG file. Initially each
-  // section is assumed to be allocation_size in size. On destruction of
+  // section is assumed to be 0 bytes in size. On destruction of
   // the OutputFile instance the size field of each section is updated to
   // reflect the actual, used size.
   struct FileHeader {
@@ -97,7 +97,7 @@ template <Section... Sections> class OutputFile {
         (SectionMeta{.tag = Sections::tag,
                      .align = Sections::align_of,
                      .offset = 0,
-                     .size = Sections::allocation_size})...};
+                     .size = 0})...};
   };
 
   OutputFile(std::filesystem::path const &filename)

polytracker/src/taint_sources/taint_sources.cpp

@@ -40,6 +40,9 @@ EARLY_CONSTRUCT_EXTERN_GETTER(taintdag::PolyTracker, polytracker_tdag);
 using util::Length;
 using util::Offset;
 
+// To allow abort to somewhat gracefully finalize polytracker logging
+extern void polytracker_end();
+
 // To create some label functions
 // Following the libc custom functions from custom.cc
 namespace {
@@ -391,6 +394,21 @@ EXT_C_FUNC void __dfsw_exit(int ret_code, dfsan_label ret_code_label) {
   exit(ret_code);
 }
 
+// Need this for some reason. Not sure why. Should already be in assert.h.
+extern "C" void __assert_fail(const char *, const char *, unsigned,
+                              const char *);
+// Capture calls to abort and explicitly invoke polytracker_end to do a best
+// effort at updating the tdag with correct sizes.
+EXT_C_FUNC void __dfsw___assert_fail(const char *msg, const char *file,
+                                     unsigned line, const char *pretty_func,
+                                     dfsan_label msg_label,
+                                     dfsan_label file_label,
+                                     dfsan_label line_label,
+                                     dfsan_label pretty_func_label) {
+  polytracker_end();
+  __assert_fail(msg, file, line, pretty_func);
+}
+
 // Socket functions
 namespace {
 static std::optional<std::string> connect_name(int socket) {

tests/test_assert.cpp

@@ -0,0 +1,12 @@
+#include <cassert>
+#include <unistd.h>
+
+int main(int argc, char *argv[]) {
+
+  char data[2];
+  read(0, data, sizeof(data));
+
+  // This will terminate the program unexpectedly (tdag sizes might not be updated).
+  assert(data[0] == data[1]);
+
+}

tests/test_assert.py

@@ -0,0 +1,32 @@
+from pathlib import Path
+import subprocess
+import pytest
+from polytracker import taint_dag, PolyTrackerTrace
+from typing import cast
+
+
+@pytest.mark.program_trace("test_assert.cpp")
+def test_assert(instrumented_binary: Path, trace_file: Path):
+    stdin_data = "ab"
+
+    subprocess.run(
+        [str(instrumented_binary)],
+        input=stdin_data.encode("utf-8"),
+        env={"POLYDB": str(trace_file), "POLYTRACKER_STDIN_SOURCE": str(1)},
+    )
+    program_trace = PolyTrackerTrace.load(trace_file)
+    assert isinstance(program_trace, taint_dag.TDProgramTrace)
+
+    tdfile = program_trace.tdfile
+    assert tdfile.label_count == 4
+
+    t1 = cast(taint_dag.TDSourceNode, tdfile.decode_node(1))
+    assert isinstance(t1, taint_dag.TDSourceNode)
+
+    t2 = cast(taint_dag.TDSourceNode, tdfile.decode_node(2))
+    assert isinstance(t2, taint_dag.TDSourceNode)
+
+    t3 = cast(taint_dag.TDSourceNode, tdfile.decode_node(3))
+    assert isinstance(t3, taint_dag.TDRangeNode)
+    assert t3.first == 1
+    assert t3.last == 2