test/dfsan/write_callback.c

   1 // RUN: %clang_dfsan                         -m64 %s -o %t && %t | FileCheck %s
   2 // RUN: %clang_dfsan  -mllvm -dfsan-args-abi -m64 %s -o %t && %t | FileCheck %s
   3
   4 // Tests that the custom implementation of write() does writes with or without
   5 // a callback set using dfsan_set_write_callback().
   6
   7 #include <sanitizer/dfsan_interface.h>
   8
   9 #include <assert.h>
  10 #include <fcntl.h>
  11 #include <stdio.h>
  12 #include <string.h>
  13 #include <unistd.h>
  14
  15 // Check write callback arguments by having the callback store them in
  16 // the following variables:
  17 static int last_callback_arg_fd;
  18 static const void *last_callback_arg_buf;
  19 static size_t last_callback_arg_count;
  20
  21 // Allow tests to check the number of callbacks made by incrementing
  22 // this count.  When callbacks are verified, the count is reset.
  23 static int count_unverified_callbacks = 0;
  24
  25 // This callbact will be installed using dfsan_set_write_callback()
  26 // in tests below.
  27 static void write_callback(int fd, const void *buf, size_t count) {
  28   // Do not do anything in this function that might call write().
  29   count_unverified_callbacks++;
  30
  31   last_callback_arg_fd = fd;
  32   last_callback_arg_buf = buf;
  33   last_callback_arg_count = count;
  34 }
  35
  36 static void write_string_to_stdout(char *string) {
  37   char *cur = string;
  38   int bytes_left = strlen(string);
  39   while (bytes_left > 0) {
  40     int res = write(fileno(stdout), cur, bytes_left);
  41     assert (res >= 0);
  42     cur += res;
  43     bytes_left -= res;
  44   }
  45 }
  46
  47 static void test_can_write_without_callback() {
  48   dfsan_set_write_callback(NULL);
  49   count_unverified_callbacks = 0;
  50
  51   char aString[] = "Test that writes work without callback.\n";
  52   // CHECK: Test that writes work without callback.
  53   write_string_to_stdout(aString);
  54
  55   assert(count_unverified_callbacks == 0);
  56 }
  57
  58 static void test_can_write_with_callback() {
  59   dfsan_set_write_callback(write_callback);
  60
  61   count_unverified_callbacks = 0;
  62
  63   char stringWithCallback[] = "Test that writes work with callback.\n";
  64   // CHECK: Test that writes work with callback.
  65   write_string_to_stdout(stringWithCallback);
  66
  67   // Data was written, so at least one call to write() was made.
  68   // Because a write may not process all the bytes it is passed, there
  69   // may have been several calls to write().
  70   assert(count_unverified_callbacks > 0);
  71   count_unverified_callbacks = 0;
  72
  73   dfsan_set_write_callback(NULL);
  74
  75   char stringWithoutCallback[] = "Writes work after the callback is removed.\n";
  76   // CHECK: Writes work after the callback is removed.
  77   write_string_to_stdout(stringWithoutCallback);
  78   assert(count_unverified_callbacks == 0);
  79 }
  80
  81 static void test_failing_write_runs_callback() {
  82   // Open /dev/null in read-only mode.  Calling write() on fd will fail.
  83   int fd = open("/dev/null", O_RDONLY);
  84   assert(fd != -1);
  85
  86   // Install a callback.
  87   dfsan_set_write_callback(write_callback);
  88
  89   // Write to the read-only file handle.  The write will fail, but the callback
  90   // should still be invoked.
  91   char aString[] = "This text will fail to be written.\n";
  92   int len = strlen(aString);
  93   int write_result = write(fd, aString, len);
  94   assert(write_result == -1);
  95
  96   assert(count_unverified_callbacks == 1);
  97   count_unverified_callbacks = 0;
  98
  99   assert(fd == last_callback_arg_fd);
 100   assert(aString == last_callback_arg_buf);
 101   assert(len == last_callback_arg_count);
 102
 103   close(fd);
 104 }
 105
 106 int main(int argc, char* argv[]) {
 107   test_can_write_without_callback();
 108   test_can_write_with_callback();
 109   test_failing_write_runs_callback();
 110 }