Merge pull request #6478 from trailofbits/6477-add-stdout-stderr-as-optional-taint-sinks-behind-a-cli-flag

Adds `stdout` and `stderr` as taint sinks

Author: Marek Surovič

polytracker/include/taintdag/taint.hpp

@@ -46,7 +46,7 @@ const source_index_t source_index_mask = max_source_index;
 
 // Use the remaining bits for source file offset
 const size_t source_offset_bits = storage_bits - source_index_bits - 2;
-using source_offset_t = storage_t;
+using source_offset_t = int64_t;
 const source_offset_t max_source_offset =
     (static_cast<source_offset_t>(1) << source_offset_bits) - 1;
 

polytracker/src/polytracker/main.cpp

@@ -25,6 +25,8 @@
 DECLARE_EARLY_CONSTRUCT(std::unordered_set<std::string>, target_sources);
 DECLARE_EARLY_CONSTRUCT(taintdag::PolyTracker, polytracker_tdag);
 DECLARE_EARLY_CONSTRUCT(std::string, polytracker_db_name);
+DECLARE_EARLY_CONSTRUCT(std::string, polytracker_stderr_sink);
+DECLARE_EARLY_CONSTRUCT(std::string, polytracker_stdout_sink);
 
 uint64_t byte_start = 0;
 uint64_t byte_end = 0;
@@ -55,6 +57,9 @@ void polytracker_parse_env() {
   if (auto pdb = getenv("POLYDB")) {
     get_polytracker_db_name() = pdb;
   }
+
+  get_polytracker_stdout_sink() = getenv("POLYTRACKER_STDOUT_SINK");
+  get_polytracker_stderr_sink() = getenv("POLYTRACKER_STDERR_SINK");
 }
 
 /*
@@ -71,12 +76,22 @@ void polytracker_get_settings() {
 }
 
 void polytracker_end() {
+  if (int f = fileno(stdout); f >= 0) {
+    get_polytracker_tdag().close_file(f);
+  }
+  if (int f = fileno(stderr); f >= 0) {
+    get_polytracker_tdag().close_file(f);
+  }
   // Explicitly destroy the PolyTracker instance to flush mapping to disk
   get_polytracker_tdag().~PolyTracker();
 }
 
 void polytracker_print_settings() {
-  printf("POLYDB:        %s\n", get_polytracker_db_name().c_str());
+  printf("POLYDB: %s\n", get_polytracker_db_name().c_str());
+  printf("POLYTRACKER_STDOUT_SINK: %s\n",
+         get_polytracker_stdout_sink().c_str());
+  printf("POLYTRACKER_STDERR_SINK: %s\n",
+         get_polytracker_stderr_sink().c_str());
 }
 
 void polytracker_start(func_mapping const *globals, uint64_t globals_count,
@@ -96,6 +111,14 @@ void polytracker_start(func_mapping const *globals, uint64_t globals_count,
            "instrumentation.\n");
   }
 
+  if (int f = fileno(stdout); f >= 0 && get_polytracker_stdout_sink() == "1") {
+    get_polytracker_tdag().open_file(f, "/dev/stdout");
+  }
+
+  if (int f = fileno(stderr); f >= 0 && get_polytracker_stderr_sink() == "1") {
+    get_polytracker_tdag().open_file(f, "/dev/stderr");
+  }
+
   // Set up the atexit call
   atexit(polytracker_end);
 }

polytracker/src/taint_sources/write_taints.cpp

@@ -27,9 +27,9 @@ EXT_C_FUNC ssize_t __dfsw_write(int fd, void *buf, size_t count,
                                 dfsan_label *ret_label) {
   auto current_offset = lseek(fd, 0, SEEK_CUR);
   auto write_count = write(fd, buf, count);
-  if (write_count > 0)
+  if (write_count > 0) {
     get_polytracker_tdag().taint_sink(fd, current_offset, buf, write_count);
-
+  }
   *ret_label = 0;
   return write_count;
 }
@@ -42,9 +42,10 @@ EXT_C_FUNC size_t __dfsw_fwrite(void *buf, size_t size, size_t count,
   auto current_offset = ftell(stream);
   auto write_count = fwrite(buf, size, count, stream);
   auto fd = fileno(stream);
-  if (write_count > 0)
+  if (write_count > 0) {
     get_polytracker_tdag().taint_sink(fd, current_offset, buf,
                                       write_count * size);
+  }
   *ret_label = 0;
   return write_count;
 }

polytracker/taint_dag.py

@@ -12,7 +12,7 @@
 )
 from pathlib import Path
 from mmap import mmap, PROT_READ
-from ctypes import Structure, c_uint64, c_int32, c_uint32, c_uint8, sizeof
+from ctypes import Structure, c_int64, c_uint64, c_int32, c_uint32, c_uint8, sizeof
 
 from .plugins import Command
 from .repl import PolyTrackerREPL
@@ -102,7 +102,7 @@ def __repr__(self) -> str:
 
 class TDSink(Structure):
     _pack_ = 1
-    _fields_ = [("fdidx", c_uint8), ("offset", c_uint64), ("label", c_uint32)]
+    _fields_ = [("fdidx", c_uint8), ("offset", c_int64), ("label", c_uint32)]
 
     def __repr__(self) -> str:
         return f"TDSink fdidx: {self.fdidx} offset: {self.offset} label: {self.label}"

tests/test_stdio.cpp

@@ -0,0 +1,24 @@
+#include <stdint.h>
+#include <stdio.h>
+
+int main(int argc, char *argv[]) {
+  if (argc < 1) {
+    return 1;
+  }
+
+  FILE *f = fopen(argv[1], "r");
+  if (!f) {
+    return 2;
+  }
+
+  char arr[2];
+  fread(&arr, sizeof(arr), 1, f);
+  if (arr[0] == '{') {
+    uint16_t val = (uint16_t)*arr;
+    val += 1;
+    fwrite(&val, sizeof(char), 1, stdout);
+    fwrite(&val, sizeof(char), 1, stderr);
+  }
+
+  return 0;
+}

tests/test_stdio.py

@@ -0,0 +1,26 @@
+import pytest
+
+from pathlib import Path
+from polytracker import taint_dag, ProgramTrace
+
+
+@pytest.fixture
+def set_env_vars(monkeypatch):
+    monkeypatch.setenv("POLYTRACKER_STDOUT_SINK", "1")
+    monkeypatch.setenv("POLYTRACKER_STDERR_SINK", "1")
+
+
+@pytest.mark.program_trace("test_stdio.cpp")
+def test_stdio(set_env_vars, program_trace: ProgramTrace):
+    assert isinstance(program_trace, taint_dag.TDProgramTrace)
+    stdout = Path("/dev/stdout")
+    stderr = Path("/dev/stderr")
+    headers = list(program_trace.tdfile.fd_headers)
+    paths = list(map(lambda h: h[0], headers))
+    assert len(paths) == 3
+    assert stdout in paths
+    assert stderr in paths
+    sinks = list(program_trace.tdfile.sinks)
+    assert len(sinks) == 2
+    assert paths[sinks[0].fdidx] == stdout
+    assert paths[sinks[1].fdidx] == stderr

tests/tracing.py

@@ -1,3 +1,4 @@
+from os import getenv
 import pytest
 from shutil import copyfile
 from subprocess import CalledProcessError
@@ -102,6 +103,8 @@ def validate_execute_target(
         "POLYDB": to_native_path(db_path),
         "POLYTRACE": "1",
         "POLYFUNC": "1",
+        "POLYTRACKER_STDOUT_SINK": getenv("POLYTRACKER_STDOUT_SINK", "0"),
+        "POLYTRACKER_STDERR_SINK": getenv("POLYTRACKER_STDERR_SINK", "0"),
     }
     if taint_all:
         del env["POLYPATH"]