git.ozlabs.org Git - ccan/blob - 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 "65015"
  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_idle());
  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 }