Add user-defined callback on write() calls.

Add dfsan_set_write_callback(), which sets a callback to be invoked when
a write() call is invoked within DFSan instrumented code.

Patch by Sam Kerner!

Differential Revision: http://reviews.llvm.org/D3268

llvm-svn: 207131

compiler-rt/include/sanitizer/dfsan_interface.h

@@ -39,6 +39,9 @@ struct dfsan_label_info {
   void *userdata;
 };
 
+/// Signature of the callback argument to dfsan_set_write_callback().
+typedef void (*dfsan_write_callback_t)(int fd, const void *buf, size_t count);
+
 /// Computes the union of \c l1 and \c l2, possibly creating a union label in
 /// the process.
 dfsan_label dfsan_union(dfsan_label l1, dfsan_label l2);
@@ -77,6 +80,11 @@ dfsan_label dfsan_has_label_with_desc(dfsan_label label, const char *desc);
 /// 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);
+
 #ifdef __cplusplus
 }  // extern "C"
 

compiler-rt/lib/dfsan/dfsan_custom.cc

@@ -800,4 +800,43 @@ __dfsw_socketpair(int domain, int type, int protocol, int sv[2],
   }
   return ret;
 }
+
+// Type of the trampoline function passed to the custom version of
+// dfsan_set_write_callback.
+typedef void (*write_trampoline_t)(
+    void *callback,
+    int fd, const void *buf, ssize_t count,
+    dfsan_label fd_label, dfsan_label buf_label, dfsan_label count_label);
+
+// Calls to dfsan_set_write_callback() set the values in this struct.
+// Calls to the custom version of write() read (and invoke) them.
+static struct {
+  write_trampoline_t write_callback_trampoline = NULL;
+  void *write_callback = NULL;
+} write_callback_info;
+
+SANITIZER_INTERFACE_ATTRIBUTE void
+__dfsw_dfsan_set_write_callback(
+    write_trampoline_t write_callback_trampoline,
+    void *write_callback,
+    dfsan_label write_callback_label,
+    dfsan_label *ret_label) {
+  write_callback_info.write_callback_trampoline = write_callback_trampoline;
+  write_callback_info.write_callback = write_callback;
+}
+
+SANITIZER_INTERFACE_ATTRIBUTE int
+__dfsw_write(int fd, const void *buf, size_t count,
+             dfsan_label fd_label, dfsan_label buf_label,
+             dfsan_label count_label, dfsan_label *ret_label) {
+  if (write_callback_info.write_callback != NULL) {
+    write_callback_info.write_callback_trampoline(
+        write_callback_info.write_callback,
+        fd, buf, count,
+        fd_label, buf_label, count_label);
+  }
+
+  *ret_label = 0;
+  return write(fd, buf, count);
+}
 }

compiler-rt/lib/dfsan/done_abilist.txt

@@ -24,6 +24,8 @@ fun:dfsan_has_label=uninstrumented
 fun:dfsan_has_label=discard
 fun:dfsan_has_label_with_desc=uninstrumented
 fun:dfsan_has_label_with_desc=discard
+fun:dfsan_set_write_callback=uninstrumented
+fun:dfsan_set_write_callback=custom
 
 ###############################################################################
 # glibc
@@ -142,7 +144,6 @@ fun:symlink=discard
 fun:syscall=discard
 fun:unlink=discard
 fun:uselocale=discard
-fun:write=discard
 
 # Functions that produce output does not depend on the input (need to zero the
 # shadow manually).
@@ -192,6 +193,10 @@ fun:strncmp=custom
 fun:strrchr=custom
 fun:strstr=custom
 
+# Functions which take action based on global state, such as running a callback
+# set by a sepperate function.
+fun:write=custom
+
 # Functions that take a callback (wrap the callback manually).
 fun:dl_iterate_phdr=custom
 

compiler-rt/test/dfsan/custom.c

@@ -30,6 +30,7 @@
 
 dfsan_label i_label = 0;
 dfsan_label j_label = 0;
+dfsan_label k_label = 0;
 dfsan_label i_j_label = 0;
 
 #define ASSERT_ZERO_LABEL(data) \
@@ -339,6 +340,54 @@ void test_ctime_r() {
   ASSERT_READ_ZERO_LABEL(buf, strlen(buf) + 1);
 }
 
+static int write_callback_count = 0;
+static int last_fd;
+static const void *last_buf;
+static size_t last_count;
+
+void write_callback(int fd, const void *buf, size_t count) {
+  write_callback_count++;
+
+  last_fd = fd;
+  last_buf = buf;
+  last_count = count;
+}
+
+void test_dfsan_set_write_callback() {
+  char buf[] = "Sample chars";
+  int buf_len = strlen(buf);
+
+  int fd = open("/dev/null", O_WRONLY);
+
+  dfsan_set_write_callback(write_callback);
+
+  write_callback_count = 0;
+
+  // Callback should be invoked on every call to write().
+  int res = write(fd, buf, buf_len);
+  assert(write_callback_count == 1);
+  ASSERT_READ_ZERO_LABEL(&res, sizeof(res));
+  ASSERT_READ_ZERO_LABEL(&last_fd, sizeof(last_fd));
+  ASSERT_READ_ZERO_LABEL(last_buf, sizeof(last_buf));
+  ASSERT_READ_ZERO_LABEL(&last_count, sizeof(last_count));
+
+  // Add a label to write() arguments.  Check that the labels are readable from
+  // the values passed to the callback.
+  dfsan_set_label(i_label, &fd, sizeof(fd));
+  dfsan_set_label(j_label, &(buf[3]), 1);
+  dfsan_set_label(k_label, &buf_len, sizeof(buf_len));
+  
+  res = write(fd, buf, buf_len);
+  assert(write_callback_count == 2);
+  ASSERT_READ_ZERO_LABEL(&res, sizeof(res));
+  ASSERT_READ_LABEL(&last_fd, sizeof(last_fd), i_label);
+  ASSERT_READ_LABEL(&last_buf[3], sizeof(last_buf[3]), j_label);
+  ASSERT_READ_LABEL(last_buf, sizeof(last_buf), j_label);
+  ASSERT_READ_LABEL(&last_count, sizeof(last_count), k_label);
+
+  dfsan_set_write_callback(NULL);
+}
+
 void test_fgets() {
   char *buf = (char*) malloc(128);
   FILE *f = fopen("/etc/passwd", "r");
@@ -702,14 +751,39 @@ void test_socketpair() {
   ASSERT_READ_ZERO_LABEL(fd, sizeof(fd));
 }
 
+void test_write() {
+  int fd = open("/dev/null", O_WRONLY);
+
+  char buf[] = "a string";
+  int len = strlen(buf);
+
+  // The result of a write always unlabeled.
+  int res = write(fd, buf, len);
+  assert(res > 0);
+  ASSERT_ZERO_LABEL(res);
+
+  // Label all arguments to write().
+  dfsan_set_label(i_label, &(buf[3]), 1);
+  dfsan_set_label(j_label, &fd, sizeof(fd));
+  dfsan_set_label(i_label, &len, sizeof(len));
+
+  // The value returned by write() should have no label.
+  res = write(fd, buf, len);
+  ASSERT_ZERO_LABEL(res);
+
+  close(fd);
+}
+
 int main(void) {
   i_label = dfsan_create_label("i", 0);
   j_label = dfsan_create_label("j", 0);
+  k_label = dfsan_create_label("k", 0);
   i_j_label = dfsan_union(i_label, j_label);
 
   test_calloc();
   test_clock_gettime();
   test_ctime_r();
+  test_dfsan_set_write_callback();
   test_dl_iterate_phdr();
   test_dlopen();
   test_fgets();
@@ -755,4 +829,5 @@ int main(void) {
   test_strtoul();
   test_strtoull();
   test_time();
+  test_write();
 }

compiler-rt/test/dfsan/write_callback.c

@@ -0,0 +1,110 @@
+// RUN: %clang_dfsan                         -m64 %s -o %t && %t | FileCheck %s
+// RUN: %clang_dfsan  -mllvm -dfsan-args-abi -m64 %s -o %t && %t | FileCheck %s
+
+// Tests that the custom implementation of write() does writes with or without
+// a callback set using dfsan_set_write_callback().
+
+#include <sanitizer/dfsan_interface.h>
+
+#include <assert.h>
+#include <fcntl.h>
+#include <stdio.h>
+#include <string.h>
+#include <unistd.h>
+
+// Check write callback arguments by having the callback store them in
+// the following variables:
+static int last_callback_arg_fd;
+static const void *last_callback_arg_buf;
+static size_t last_callback_arg_count;
+
+// Allow tests to check the number of callbacks made by incrementing
+// this count.  When callbacks are verified, the count is reset.
+static int count_unverified_callbacks = 0;
+
+// This callbact will be installed using dfsan_set_write_callback()
+// in tests below.
+static void write_callback(int fd, const void *buf, size_t count) {
+  // Do not do anything in this function that might call write().
+  count_unverified_callbacks++;
+
+  last_callback_arg_fd = fd;
+  last_callback_arg_buf = buf;
+  last_callback_arg_count = count;
+}
+
+static void write_string_to_stdout(char *string) {
+  char *cur = string;
+  int bytes_left = strlen(string);
+  while (bytes_left > 0) {
+    int res = write(fileno(stdout), cur, bytes_left);
+    assert (res >= 0);
+    cur += res;
+    bytes_left -= res;
+  }
+}
+
+static void test_can_write_without_callback() {
+  dfsan_set_write_callback(NULL);
+  count_unverified_callbacks = 0;
+
+  char aString[] = "Test that writes work without callback.\n";
+  // CHECK: Test that writes work without callback.
+  write_string_to_stdout(aString);
+
+  assert(count_unverified_callbacks == 0);
+}
+
+static void test_can_write_with_callback() {
+  dfsan_set_write_callback(write_callback);
+
+  count_unverified_callbacks = 0;
+
+  char stringWithCallback[] = "Test that writes work with callback.\n";
+  // CHECK: Test that writes work with callback.
+  write_string_to_stdout(stringWithCallback);
+
+  // Data was written, so at least one call to write() was made.
+  // Because a write may not process all the bytes it is passed, there
+  // may have been several calls to write().
+  assert(count_unverified_callbacks > 0);
+  count_unverified_callbacks = 0;
+
+  dfsan_set_write_callback(NULL);
+
+  char stringWithoutCallback[] = "Writes work after the callback is removed.\n";
+  // CHECK: Writes work after the callback is removed.
+  write_string_to_stdout(stringWithoutCallback);
+  assert(count_unverified_callbacks == 0);
+}
+
+static void test_failing_write_runs_callback() {
+  // Open /dev/null in read-only mode.  Calling write() on fd will fail.
+  int fd = open("/dev/null", O_RDONLY);
+  assert(fd != -1);
+
+  // Install a callback.
+  dfsan_set_write_callback(write_callback);
+
+  // Write to the read-only file handle.  The write will fail, but the callback
+  // should still be invoked.
+  char aString[] = "This text will fail to be written.\n";
+  int len = strlen(aString);
+  int write_result = write(fd, aString, len);
+  assert(write_result == -1);
+
+  assert(count_unverified_callbacks == 1);
+  count_unverified_callbacks = 0;
+
+  assert(fd == last_callback_arg_fd);
+  assert(aString == last_callback_arg_buf);
+  assert(len == last_callback_arg_count);
+
+  close(fd);
+}
+
+int main(int argc, char* argv[]) {
+  test_can_write_without_callback();
+  test_can_write_with_callback();
+  test_failing_write_runs_callback();
+}