Add flag to dump the labels when the program terminates.

[martignlo]

I previously used the wrong diff. Sorry!

[martignlo]

Added test.

[pcc]

lib/dfsan/dfsan.cc

274

You can skip index 0.

347

The other sanitizers use atexit for this type of thing.

test/dfsan/dump_labels.c

4

You should also test the case where the program terminates due to running out of labels.

[martignlo]

lib/dfsan/dfsan.cc

274

Done.

347

Moved the registration to dfsan_init.

test/dfsan/dump_labels.c

4

Done.

[pcc]

LGTM

test/dfsan/dump_labels.c

59

You could add a CHECK-OOL-NOT line to make sure this line is not printed.

[martignlo]

test/dfsan/dump_labels.c

59

Dropped since the exitstatus already captures this.

[martignlo]

Closed by commit rL222425 (authored by @martignlo).

include/sanitizer/dfsan_interface.h

 /// Returns the number of labels allocated.
 size_t dfsan_get_label_count(void);
 
 /// Sets a callback to be invoked on calls to write().  The callback is invoked
 /// before the write is done.  The write is not guaranteed to succeed when the
 /// callback executes.  Pass in NULL to remove any callback.
 void dfsan_set_write_callback(dfsan_write_callback_t labeled_write_callback);
 
+/// Writes the labels currently used by the program to the given file
+/// descriptor. The lines of the output have the following format:
+///
+/// <label> <parent label 1> <parent label 2> <label description if any>
+void dfsan_dump_labels(int fd);
+
 #ifdef __cplusplus
 }  // extern "C"
 
 template <typename T>
 void dfsan_set_label(dfsan_label label, T &data) {  // NOLINT
   dfsan_set_label(label, (void *)&data, sizeof(T));
 }
 
 #endif
 
 #endif  // DFSAN_INTERFACE_H

lib/dfsan/dfsan.h

struct Flags {
   bool warn_unimplemented;
   // Whether to warn on non-zero labels.
   bool warn_nonzero_labels;
   // Whether to propagate labels only when there is an obvious data dependency
   // (e.g., when comparing strings, ignore the fact that the output of the
   // comparison might be data-dependent on the content of the strings). This
   // applies only to the custom functions defined in 'custom.c'.
   bool strict_data_dependencies;
+  // The path of the file where to dump the labels when the program terminates.
+  const char* dump_labels_at_exit;
 };
 
 extern Flags flags_data;
 inline Flags &flags() {
   return flags_data;
 }
 
 }  // namespace __dfsan
 
 #endif  // DFSAN_H

lib/dfsan/dfsan.cc

 static const uptr kUnionTableAddr = 0x200000000000;
 static const uptr kUnusedAddr = kUnionTableAddr + sizeof(dfsan_union_table_t);
 static const uptr kAppAddr = 0x700000008000;
 
 static atomic_dfsan_label *union_table(dfsan_label l1, dfsan_label l2) {
   return &(*(dfsan_union_table_t *) kUnionTableAddr)[l1][l2];
 }
 
+// Checks we do not run out of labels.
+static void dfsan_check_label(dfsan_label label) {
+  if (label == kInitializingLabel) {
+    Report("FATAL: DataFlowSanitizer: out of labels\n");
+    Die();
+  }
+}
+
 // Resolves the union of two unequal labels.  Nonequality is a precondition for
 // this function (the instrumentation pass inlines the equality test).
 extern "C" SANITIZER_INTERFACE_ATTRIBUTE
 dfsan_label __dfsan_union(dfsan_label l1, dfsan_label l2) {
   DCHECK_NE(l1, l2);
 
   if (l1 == 0)
     return l2;
if (atomic_compare_exchange_strong(table_ent, &label, kInitializingLabel,
     // in the cases we are interested in) a label may only subsume labels
     // created earlier (i.e. with a lower numerical value).
     if (__dfsan_label_info[l2].l1 == l1 ||
         __dfsan_label_info[l2].l2 == l1) {
       label = l2;
     } else {
       label =
         atomic_fetch_add(&__dfsan_last_label, 1, memory_order_relaxed) + 1;
-      CHECK_NE(label, kInitializingLabel);
+      dfsan_check_label(label);
       __dfsan_label_info[label].l1 = l1;
       __dfsan_label_info[label].l2 = l2;
     }
     atomic_store(table_ent, label, memory_order_release);
   } else if (label == kInitializingLabel) {
     // Another thread is initializing the entry.  Wait until it is finished.
     do {
       internal_sched_yield();
if (l1 == l2)
     return l1;
   return __dfsan_union(l1, l2);
 }
 
 extern "C" SANITIZER_INTERFACE_ATTRIBUTE
 dfsan_label dfsan_create_label(const char *desc, void *userdata) {
   dfsan_label label =
     atomic_fetch_add(&__dfsan_last_label, 1, memory_order_relaxed) + 1;
-  CHECK_NE(label, kInitializingLabel);
+  dfsan_check_label(label);
   __dfsan_label_info[label].l1 = __dfsan_label_info[label].l2 = 0;
   __dfsan_label_info[label].desc = desc;
   __dfsan_label_info[label].userdata = userdata;
   return label;
 }
 
 extern "C" SANITIZER_INTERFACE_ATTRIBUTE
 void __dfsan_set_label(dfsan_label label, void *addr, uptr size) {

 extern "C" SANITIZER_INTERFACE_ATTRIBUTE uptr
 dfsan_get_label_count(void) {
   dfsan_label max_label_allocated =
       atomic_load(&__dfsan_last_label, memory_order_relaxed);
 
   return static_cast<uptr>(max_label_allocated);
 }
 
+extern "C" SANITIZER_INTERFACE_ATTRIBUTE void
+dfsan_dump_labels(int fd) {
+  dfsan_label last_label =
+      atomic_load(&__dfsan_last_label, memory_order_relaxed);
+
+  for (uptr l = 1; l <= last_label; ++l) {
+    char buf[64];
+    internal_snprintf(buf, sizeof(buf), "%u %u %u ", l,
+                      __dfsan_label_info[l].l1, __dfsan_label_info[l].l2);
+    internal_write(fd, buf, internal_strlen(buf));
+    if (__dfsan_label_info[l].l1 == 0 && __dfsan_label_info[l].desc) {
+      internal_write(fd, __dfsan_label_info[l].desc,
+                     internal_strlen(__dfsan_label_info[l].desc));
+    }
+    internal_write(fd, "\n", 1);
+  }
+}
+
 static void InitializeFlags(Flags &f, const char *env) {
   f.warn_unimplemented = true;
   f.warn_nonzero_labels = false;
   f.strict_data_dependencies = true;
+  f.dump_labels_at_exit = "";
 
   ParseFlag(env, &f.warn_unimplemented, "warn_unimplemented", "");
   ParseFlag(env, &f.warn_nonzero_labels, "warn_nonzero_labels", "");
   ParseFlag(env, &f.strict_data_dependencies, "strict_data_dependencies", "");
+  ParseFlag(env, &f.dump_labels_at_exit, "dump_labels_at_exit", "");
+}
+
+static void dfsan_fini() {
+  if (internal_strcmp(flags().dump_labels_at_exit, "") != 0) {
+    fd_t fd = OpenFile(flags().dump_labels_at_exit, true /* write */);
+    if (fd == kInvalidFd) {
+      Report("WARNING: DataFlowSanitizer: unable to open output file %s\n",
+             flags().dump_labels_at_exit);
+      return;
+    }
+
+    Report("INFO: DataFlowSanitizer: dumping labels to %s\n",
+           flags().dump_labels_at_exit);
+    dfsan_dump_labels(fd);
+    internal_close(fd);
+  }
 }
 
 #ifdef DFSAN_NOLIBC
 extern "C" void dfsan_init() {
 #else
 static void dfsan_init(int argc, char **argv, char **envp) {
 #endif
   MmapFixedNoReserve(kShadowAddr, kUnusedAddr - kShadowAddr);
 
   // Protect the region of memory we don't use, to preserve the one-to-one
   // mapping from application to shadow memory. But if ASLR is disabled, Linux
   // will load our executable in the middle of our unused region. This mostly
   // works so long as the program doesn't use too much memory. We support this
   // case by disabling memory protection when ASLR is disabled.
   uptr init_addr = (uptr)&dfsan_init;
   if (!(init_addr >= kUnusedAddr && init_addr < kAppAddr))
     Mprotect(kUnusedAddr, kAppAddr - kUnusedAddr);
 
   InitializeFlags(flags(), GetEnv("DFSAN_OPTIONS"));
 
   InitializeInterceptors();
+
+  // Register the fini callback to run when the program terminates successfully
+  // or it is killed by the runtime.
+  Atexit(dfsan_fini);
+  SetDieCallback(dfsan_fini);
+
+  __dfsan_label_info[kInitializingLabel].desc = "<init label>";
 }
 
 #if !defined(DFSAN_NOLIBC) && SANITIZER_CAN_USE_PREINIT_ARRAY
 __attribute__((section(".preinit_array"), used))
 static void (*dfsan_init_ptr)(int, char **, char **) = dfsan_init;
 #endif

test/dfsan/dump_labels.c

+// RUN: %clang_dfsan -m64 %s -o %t
+// RUN: DFSAN_OPTIONS=dump_labels_at_exit=/dev/stdout %run %t 2>&1 | FileCheck %s
+// RUN: DFSAN_OPTIONS=dump_labels_at_exit=/dev/stdout not %run %t c 2>&1 | FileCheck %s --check-prefix=CHECK-OOL
+// RUN: DFSAN_OPTIONS=dump_labels_at_exit=/dev/stdout not %run %t u 2>&1 | FileCheck %s --check-prefix=CHECK-OOL
+
+// Tests that labels are properly dumped at program termination.
+
+#include <sanitizer/dfsan_interface.h>
+#include <assert.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+int main(int argc, char** argv) {
+  int i = 1;
+  dfsan_label i_label = dfsan_create_label("i", 0);
+  dfsan_set_label(i_label, &i, sizeof(i));
+
+  int j = 2;
+  dfsan_label j_label = dfsan_create_label("j", 0);
+  dfsan_set_label(j_label, &j, sizeof(j));
+
+  int k = 3;
+  dfsan_label k_label = dfsan_create_label("k", 0);
+  dfsan_set_label(k_label, &k, sizeof(k));
+
+  dfsan_label ij_label = dfsan_get_label(i + j);
+  dfsan_label ijk_label = dfsan_get_label(i + j + k);
+
+  fprintf(stderr, "i %d j %d k %d ij %d ijk %d\n", i_label, j_label, k_label,
+          ij_label, ijk_label);
+
+  // CHECK: 1 0 0 i
+  // CHECK: 2 0 0 j
+  // CHECK: 3 0 0 k
+  // CHECK: 4 1 2
+  // CHECK: 5 3 4
+
+  if (argc > 1) {
+    // Exhaust the labels.
+    unsigned long num_labels = 1 << (sizeof(dfsan_label) * 8);
+    for (unsigned long i =  ijk_label + 1; i < num_labels - 2; ++i) {
+      dfsan_label l = dfsan_create_label("l", 0);
+      assert(l == i);
+    }
+
+    // Consume the last available label.
+    dfsan_label l = dfsan_union(5, 6);
+    assert(l == num_labels - 2);
+
+    // Try to allocate another label (either explicitly or by unioning two
+    // existing labels), but expect a crash.
+    if (argv[1][0] == 'c') {
+      l = dfsan_create_label("l", 0);
+    } else {
+      l = dfsan_union(6, 7);
+    }
+
+    // This should not happen if the program correctly runs out of labels.
+    fprintf(stderr, "l %d\n", l);
+
+    // CHECK-OOL: FATAL: DataFlowSanitizer: out of labels
+    // CHECK-OOL: 1 0 0 i
+    // CHECK-OOL: 2 0 0 j
+    // CHECK-OOL: 3 0 0 k
+    // CHECK-OOL: 4 1 2
+    // CHECK-OOL: 5 3 4
+    // CHECK-OOL: 6 0 0
+    // CHECK-OOL: 65534 5 6
+    // CHECK-OOL: 65535 0 0 <init label>
+  }
+
+  return 0;
+}