1 // RUN: %clang_dfsan                         %s -o %t && %run %t | FileCheck %s
2 // RUN: %clang_dfsan  -mllvm -dfsan-args-abi %s -o %t && %run %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.
write_callback(int fd,const void * buf,size_t count)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 
write_string_to_stdout(char * string)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 
test_can_write_without_callback()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 
test_can_write_with_callback()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 
test_failing_write_runs_callback()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 
main(int argc,char * argv[])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 }
111