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