Repository: trailofbits/polytracker (language: python)

Issue:
Applications that exit with `assert` kills the `polytracker info` command
If an application exits due to an assert (or otherwise unexpected exit). The size fields for each section remains at it's [allocated size](https://github.com/trailofbits/polytracker/blob/master/polytracker/include/taintdag/outputfile.h#L100). Under normal circumstances, [that size is updated when the application exits](https://github.com/trailofbits/polytracker/blob/master/polytracker/include/taintdag/outputfile.h#L119-L131).

A first mitigation to prevent crash and just return an empty output file could be to initialize each section size to zero. Other options to explore is to have each section update the size as soon as anything is written (requiring an additional write for each write to a section).


Discussion / hints:

Fix patch:
diff --git a/polytracker/custom_abi/dfsan_abilist.txt b/polytracker/custom_abi/dfsan_abilist.txt
index 1fcdcef9..6ddf8a97 100644
--- a/polytracker/custom_abi/dfsan_abilist.txt
+++ b/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
diff --git a/polytracker/custom_abi/polytracker_abilist.txt b/polytracker/custom_abi/polytracker_abilist.txt
index 51b71f07..3d67e62d 100644
--- a/polytracker/custom_abi/polytracker_abilist.txt
+++ b/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
diff --git a/polytracker/include/taintdag/outputfile.h b/polytracker/include/taintdag/outputfile.h
index 2782036b..752bba1f 100644
--- a/polytracker/include/taintdag/outputfile.h
+++ b/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)
diff --git a/polytracker/src/taint_sources/taint_sources.cpp b/polytracker/src/taint_sources/taint_sources.cpp
index b1a25a4d..4d4ea846 100644
--- a/polytracker/src/taint_sources/taint_sources.cpp
+++ b/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) {


Test patch:
diff --git a/tests/test_assert.cpp b/tests/test_assert.cpp
new file mode 100644
index 00000000..9a0d9ef2
--- /dev/null
+++ b/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]);
+
+}
\ No newline at end of file
diff --git a/tests/test_assert.py b/tests/test_assert.py
new file mode 100644
index 00000000..05ea3bab
--- /dev/null
+++ b/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