io: change io_idle() to io_wait()
[ccan] / ccan / io / test / run-set_alloc.c
1 #include <ccan/tap/tap.h>
2 #include <stdlib.h>
3 #include <stdio.h>
4 #include <signal.h>
5 #include <sys/types.h>
6 #include <sys/wait.h>
7
8 /* Make sure we override these! */
9 static void *no_malloc(size_t size)
10 {
11         abort();
12 }
13 static void *no_realloc(void *p, size_t size)
14 {
15         abort();
16 }
17 static void no_free(void *p)
18 {
19         abort();
20 }
21 #define malloc no_malloc
22 #define realloc no_realloc
23 #define free no_free
24
25 #include <ccan/io/poll.c>
26 #include <ccan/io/io.c>
27
28 #undef malloc
29 #undef realloc
30 #undef free
31
32 static unsigned int alloc_count, realloc_count, free_count;
33 static void *ptrs[100];
34
35 static void **find_ptr(void *p)
36 {
37         unsigned int i;
38
39         for (i = 0; i < 100; i++)
40                 if (ptrs[i] == p)
41                         return ptrs + i;
42         return NULL;
43 }
44
45 static void *allocfn(size_t size)
46 {
47         alloc_count++;
48         return *find_ptr(NULL) = malloc(size);
49 }
50
51 static void *reallocfn(void *ptr, size_t size)
52 {
53         realloc_count++;
54         if (!ptr)
55                 alloc_count++;
56
57         return *find_ptr(ptr) = realloc(ptr, size);
58 }
59
60 static void freefn(void *ptr)
61 {
62         free_count++;
63         free(ptr);
64         *find_ptr(ptr) = NULL;
65 }
66
67 #ifndef PORT
68 #define PORT "65115"
69 #endif
70
71 struct data {
72         int state;
73         int timeout_usec;
74         bool timed_out;
75         char buf[4];
76 };
77
78
79 static struct io_plan no_timeout(struct io_conn *conn, struct data *d)
80 {
81         ok1(d->state == 1);
82         d->state++;
83         return io_close();
84 }
85
86 static struct io_plan timeout(struct io_conn *conn, struct data *d)
87 {
88         ok1(d->state == 1);
89         d->state++;
90         d->timed_out = true;
91         return io_close();
92 }
93
94 static void finish_ok(struct io_conn *conn, struct data *d)
95 {
96         ok1(d->state == 2);
97         d->state++;
98         io_break(d, io_never());
99 }
100
101 static void init_conn(int fd, struct data *d)
102 {
103         struct io_conn *conn;
104
105         ok1(d->state == 0);
106         d->state++;
107
108         conn = io_new_conn(fd, io_read(d->buf, sizeof(d->buf), no_timeout, d));
109         io_set_finish(conn, finish_ok, d);
110         io_timeout(conn, time_from_usec(d->timeout_usec), timeout, d);
111 }
112
113 static int make_listen_fd(const char *port, struct addrinfo **info)
114 {
115         int fd, on = 1;
116         struct addrinfo *addrinfo, hints;
117
118         memset(&hints, 0, sizeof(hints));
119         hints.ai_family = AF_UNSPEC;
120         hints.ai_socktype = SOCK_STREAM;
121         hints.ai_flags = AI_PASSIVE;
122         hints.ai_protocol = 0;
123
124         if (getaddrinfo(NULL, port, &hints, &addrinfo) != 0)
125                 return -1;
126
127         fd = socket(addrinfo->ai_family, addrinfo->ai_socktype,
128                     addrinfo->ai_protocol);
129         if (fd < 0)
130                 return -1;
131
132         setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on));
133         if (bind(fd, addrinfo->ai_addr, addrinfo->ai_addrlen) != 0) {
134                 close(fd);
135                 return -1;
136         }
137         if (listen(fd, 1) != 0) {
138                 close(fd);
139                 return -1;
140         }
141         *info = addrinfo;
142         return fd;
143 }
144
145 int main(void)
146 {
147         struct data *d = allocfn(sizeof(*d));
148         struct addrinfo *addrinfo;
149         struct io_listener *l;
150         int fd, status;
151
152         io_set_alloc(allocfn, reallocfn, freefn);
153
154         /* This is how many tests you plan to run */
155         plan_tests(25);
156         d->state = 0;
157         d->timed_out = false;
158         d->timeout_usec = 100000;
159         fd = make_listen_fd(PORT, &addrinfo);
160         ok1(fd >= 0);
161         l = io_new_listener(fd, init_conn, d);
162         ok1(l);
163         fflush(stdout);
164
165         if (!fork()) {
166                 int i;
167
168                 io_close_listener(l);
169                 fd = socket(addrinfo->ai_family, addrinfo->ai_socktype,
170                             addrinfo->ai_protocol);
171                 if (fd < 0)
172                         exit(1);
173                 if (connect(fd, addrinfo->ai_addr, addrinfo->ai_addrlen) != 0)
174                         exit(2);
175                 signal(SIGPIPE, SIG_IGN);
176                 usleep(500000);
177                 for (i = 0; i < strlen("hellothere"); i++) {
178                         if (write(fd, "hellothere" + i, 1) != 1)
179                                 break;
180                 }
181                 close(fd);
182                 freeaddrinfo(addrinfo);
183                 free(d);
184                 exit(i);
185         }
186         ok1(io_loop() == d);
187         ok1(d->state == 3);
188         ok1(d->timed_out == true);
189         ok1(wait(&status));
190         ok1(WIFEXITED(status));
191         ok1(WEXITSTATUS(status) < sizeof(d->buf));
192
193         /* This one shouldn't time out. */
194         d->state = 0;
195         d->timed_out = false;
196         d->timeout_usec = 500000;
197         fflush(stdout);
198
199         if (!fork()) {
200                 int i;
201
202                 io_close_listener(l);
203                 fd = socket(addrinfo->ai_family, addrinfo->ai_socktype,
204                             addrinfo->ai_protocol);
205                 if (fd < 0)
206                         exit(1);
207                 if (connect(fd, addrinfo->ai_addr, addrinfo->ai_addrlen) != 0)
208                         exit(2);
209                 signal(SIGPIPE, SIG_IGN);
210                 usleep(100000);
211                 for (i = 0; i < strlen("hellothere"); i++) {
212                         if (write(fd, "hellothere" + i, 1) != 1)
213                                 break;
214                 }
215                 close(fd);
216                 freeaddrinfo(addrinfo);
217                 free(d);
218                 exit(i);
219         }
220         ok1(io_loop() == d);
221         ok1(d->state == 3);
222         ok1(d->timed_out == false);
223         ok1(wait(&status));
224         ok1(WIFEXITED(status));
225         ok1(WEXITSTATUS(status) >= sizeof(d->buf));
226
227         io_close_listener(l);
228         freeaddrinfo(addrinfo);
229
230         /* We should have tested each one at least once! */
231         ok1(realloc_count);
232         ok1(alloc_count);
233         ok1(free_count);
234
235         ok1(free_count < alloc_count);
236         freefn(d);
237         ok1(free_count == alloc_count);
238
239         return exit_status();
240 }