write_callback.c revision 7c9150579ed0278492f51cc8434b1d63a44b9bd1 
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:
17static int last_callback_arg_fd;
18static const void *last_callback_arg_buf;
19static 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.
23static int count_unverified_callbacks = 0;
24
25// This callbact will be installed using dfsan_set_write_callback()
26// in tests below.
27static 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
36static 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
47static 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
58static 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
81static 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
106int 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