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