Fix memory leaks, miscellaneous issues.

9 files changed, 118 insertions, 48 deletions

diff --git a/Makefile b/Makefile
index 796dec7..bde9466 100644
--- a/Makefile
+++ b/Makefile
@@ -1,4 +1,4 @@
-CFLAGS += -Wall -Wextra -Wwrite-strings -std=c99 -D_BSD_SOURCE -D_DEFAULT_SOURCE
+CFLAGS += -Wall -Wextra -Wwrite-strings -Wno-missing-field-initializers -std=c99 -D_BSD_SOURCE -D_DEFAULT_SOURCE
 
 LDLIBS := -lev
 
diff --git a/doc/TODO b/doc/TODO
index f55eca4..3b53c59 100644
--- a/doc/TODO
+++ b/doc/TODO
@@ -1,9 +1,15 @@
+- verify checksums
+
+- add LD_PRELOAD support
+
 - add 'TCP Fast Open' support
 
 - optionally return ACKs and retransmit if no ACK
 
 - change some of the tiny hash tables to linked lists
 
+- add pcap support
+
 - add fake TCP options
 
 - add Windows support (not likely)
diff --git a/src/checksum.c b/src/checksum.c
index 53d4e5a..48df2c7 100644
--- a/src/checksum.c
+++ b/src/checksum.c
@@ -35,6 +35,9 @@
 #include <endian.h>
 #include <netinet/in.h>
 #include <stdint.h>
+#include <stdlib.h>
+#include <sys/socket.h>
+
 #include "checksum.h"
 
 /* Revised by Kenneth Albanowski for m68knommu. Basic problem: unaligned access
diff --git a/src/checksum.h b/src/checksum.h
index 700ee1a..74fe69f 100644
--- a/src/checksum.h
+++ b/src/checksum.h
@@ -1,7 +1,6 @@
-#include <netinet/in.h>
-#include <stdarg.h>
 #include <stdint.h>
-#include <stdlib.h>
+
+struct sockaddr;
 
 /* calculates the checksum of len bytes at buff when combined with wsum.
  * return value is already in network order, but must be inverted before
diff --git a/src/client.c b/src/client.c
index 1c3fb4e..447bd0b 100644
--- a/src/client.c
+++ b/src/client.c
@@ -7,6 +7,7 @@
 #include <netinet/in.h>
 #include <netinet/ip.h>
 #include <netinet/tcp.h>
+#include <signal.h>
 #include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -33,7 +34,7 @@ struct c_data {
 };
 
 struct o_c_rsock {
-    struct sockaddr *r_addr;
+    struct sockaddr_storage r_addr;
     struct o_c_sock *o_socks_by_lport;
     struct c_data *c_data;
     unsigned int used_ports[32768 / PORTS_IN_INT];
@@ -45,7 +46,7 @@ struct o_c_rsock {
 };
 
 struct o_c_sock {
-    struct sockaddr *c_address;
+    struct sockaddr_storage c_address;
     struct o_c_rsock *rsock;
     char *pending_data;
     size_t pending_data_size;
@@ -124,13 +125,14 @@ static void free_port(unsigned int *used_ports, uint16_t port_num) {
 /* prepare server address in TCP header format */
 static void c_prep_s_addr(struct o_c_sock *sock, struct tcphdr *hdr) {
     hdr->th_sport = sock->l_port;
-    hdr->th_dport = IN_ADDR_PORT(sock->rsock->r_addr);
+    hdr->th_dport = IN_ADDR_PORT(&sock->rsock->r_addr);
     hdr->th_seq = htonl(sock->seq_num);
     hdr->th_off = 5;
 }
 
 /* clean up a socket, don't bother freeing anything if the program is stopping */
 static void c_sock_cleanup(EV_P_ struct o_c_sock *sock, int stopping) {
+    DBG("cleaning up sock @ %p", sock);
     if (sock->status != TCP_SYN_SENT) {
         struct tcphdr buf = {
             .th_flags = sock->status == TCP_ESTABLISHED ? TH_FIN : TH_RST
@@ -145,15 +147,15 @@ static void c_sock_cleanup(EV_P_ struct o_c_sock *sock, int stopping) {
         } else if ((size_t)sz != sizeof(buf)) {
             fprintf(stderr, "send %s our packet: tried %lu, sent %zd\n", (size_t)sz > sizeof(buf) ? "expanded" : "truncated", sizeof(buf), sz);
         }
-
-        return;
     }
 
-    if (!stopping) {
+    if (!stopping || free_mem_on_exit) {
+        DBG("freeing associated resources");
         free_port(sock->rsock->used_ports, sock->l_port);
         ev_timer_stop(EV_A_ &sock->tm_w);
 
         HASH_DELETE(hh_lp, sock->rsock->o_socks_by_lport, sock);
+        HASH_DELETE(hh_ca, sock->rsock->c_data->o_socks_by_caddr, sock);
 
         if (!sock->rsock->o_socks_by_lport) {
             close(sock->rsock->fd);
@@ -162,7 +164,6 @@ static void c_sock_cleanup(EV_P_ struct o_c_sock *sock, int stopping) {
 
             HASH_DEL(sock->rsock->c_data->o_rsocks, sock->rsock);
 
-            free(sock->rsock->r_addr);
             free(sock->rsock);
         }
 
@@ -236,7 +237,7 @@ static void cc_cb(struct ev_loop *loop, ev_io *w, int revents __attribute__((unu
 
         char *rptr = rbuf;
 
-        if (rsock->r_addr->sa_family == AF_INET) {
+        if (rsock->r_addr.ss_family == AF_INET) {
             if ((size_t)rsz < sizeof(struct iphdr)) {
                 DBG("packet is smaller than IP header, ignoring");
                 return;
@@ -324,7 +325,7 @@ static void cc_cb(struct ev_loop *loop, ev_io *w, int revents __attribute__((unu
             should_ssz = rsz - rhdr->th_off * 32 / CHAR_BIT;
             if (should_ssz > 0) {
                 DBG("sending %zd bytes to client", should_ssz);
-                ssz = sendto(rsock->c_data->s_sock, rptr + rhdr->th_off * 32 / CHAR_BIT, should_ssz, 0, sock->c_address, rsock->c_data->s_addrlen);
+                ssz = sendto(rsock->c_data->s_sock, rptr + rhdr->th_off * 32 / CHAR_BIT, should_ssz, 0, (struct sockaddr *)&sock->c_address, rsock->c_data->s_addrlen);
 
                 if (ssz < 0) {
                     perror("sendto");
@@ -348,20 +349,19 @@ static inline struct o_c_rsock * c_rsock_init(struct addrinfo *res) {
     struct o_c_rsock *rsock;
     rsock = malloc(sizeof(*rsock));
     memset(&rsock->used_ports, 0, sizeof(rsock->used_ports));
-    rsock->r_addr = malloc(res->ai_addrlen);
 
-    memcpy(rsock->r_addr, res->ai_addr, res->ai_addrlen);
+    memcpy(&rsock->r_addr, res->ai_addr, res->ai_addrlen);
     rsock->r_addrlen = res->ai_addrlen;
     freeaddrinfo(res);
     rsock->o_socks_by_lport = NULL;
 
-    rsock->fd = socket(rsock->r_addr->sa_family, SOCK_RAW, IPPROTO_TCP);
+    rsock->fd = socket(rsock->r_addr.ss_family, SOCK_RAW, IPPROTO_TCP);
     if (!rsock->fd) {
         perror("socket");
         return NULL;
     }
 
-    if (connect(rsock->fd, rsock->r_addr, rsock->r_addrlen) == -1) {
+    if (connect(rsock->fd, (struct sockaddr *)&rsock->r_addr, rsock->r_addrlen) == -1) {
         perror("connect");
         return NULL;
     }
@@ -381,12 +381,12 @@ static inline struct o_c_rsock * c_rsock_init(struct addrinfo *res) {
 
     char proto[] = { 0, IPPROTO_TCP };
 
-    if (rsock->r_addr->sa_family != our_addr.ss_family)
+    if (rsock->r_addr.ss_family != our_addr.ss_family)
         abort();
 
     rsock->csum_p = csum_partial(proto, sizeof(proto),
             csum_sockaddr_partial((struct sockaddr *)&our_addr, 0,
-            csum_sockaddr_partial(rsock->r_addr, 1, 0)));
+            csum_sockaddr_partial((struct sockaddr *)&rsock->r_addr, 1, 0)));
 
     return rsock;
 }
@@ -422,8 +422,7 @@ static void cs_cb(EV_P_ ev_io *w, int revents __attribute__((unused))) {
                 return;
             }
 
-            sock->c_address = malloc(addresslen);
-            memcpy(sock->c_address, &c_data->pkt_addr, addresslen);
+            memcpy(&sock->c_address, &c_data->pkt_addr, addresslen);
 
             HASH_FIND(hh, c_data->o_rsocks, res->ai_addr, res->ai_addrlen, sock->rsock);
 
@@ -440,7 +439,7 @@ static void cs_cb(EV_P_ ev_io *w, int revents __attribute__((unused))) {
                 sock->rsock->io_w.data = sock->rsock;
                 ev_io_start(EV_A_ &sock->rsock->io_w);
 
-                HASH_ADD_KEYPTR(hh, c_data->o_rsocks, sock->rsock->r_addr, sock->rsock->r_addrlen, sock->rsock);
+                HASH_ADD(hh, c_data->o_rsocks, r_addr, sock->rsock->r_addrlen, sock->rsock);
             }
 
             uint16_t l_port = reserve_port(sock->rsock->used_ports);
@@ -455,7 +454,7 @@ static void cs_cb(EV_P_ ev_io *w, int revents __attribute__((unused))) {
 
             sock->csum_p = csum_partial(&sock->l_port, sizeof(in_port_t), sock->rsock->csum_p);
 
-            HASH_ADD_KEYPTR(hh_ca, c_data->o_socks_by_caddr, sock->c_address, addresslen, sock);
+            HASH_ADD(hh_ca, c_data->o_socks_by_caddr, c_address, addresslen, sock);
             HASH_ADD(hh_lp, sock->rsock->o_socks_by_lport, l_port, sizeof(in_port_t), sock);
 
             sock->seq_num = random();
@@ -526,20 +525,19 @@ static void c_cleanup() {
         return;
 
     DBG("cleaning up");
-    struct o_c_sock *sock;
-    for (sock = global_c_data->o_socks_by_caddr; sock != NULL; sock = sock->hh_ca.next) {
-        switch (sock->status) {
-        case TCP_CLOSE:
-            break;
-        default:
-            c_sock_cleanup(EV_DEFAULT, sock, 1);
-        }
-        // don't bother freeing anything because we're about to exit anyways
+    struct o_c_sock *sock, *tmp;
+    HASH_ITER(hh_ca, global_c_data->o_socks_by_caddr, sock, tmp) {
+        c_sock_cleanup(EV_DEFAULT, sock, 1);
     }
 
     global_c_data = NULL;
 }
 
+static void c_finish(EV_P_ ev_signal *w __attribute__((unused)), int revents __attribute__((unused))) {
+    c_cleanup();
+    ev_break(EV_A_ EVBREAK_ALL);
+}
+
 int start_client(const char *s_host, const char *s_port, const char *r_host, const char *r_port) {
     struct addrinfo *res;
     int r = getaddrinfo(s_host, s_port, NULL, &res);
@@ -577,15 +575,24 @@ int start_client(const char *s_host, const char *s_port, const char *r_host, con
 
     struct ev_loop *loop = EV_DEFAULT;
     ev_io s_watcher;
+    ev_signal iwatcher, twatcher;
 
     s_watcher.data = &c_data;
 
     ev_io_init(&s_watcher, cs_cb, c_data.s_sock, EV_READ);
     ev_io_start(loop, &s_watcher);
+    ev_signal_init(&iwatcher, c_finish, SIGINT);
+    ev_signal_start(loop, &iwatcher);
+    ev_signal_init(&twatcher, c_finish, SIGTERM);
+    ev_signal_start(loop, &twatcher);
 
     DBG("initialization complete, starting event loop");
     r = ev_run(loop, 0);
 
     c_cleanup();
+
+    if (free_mem_on_exit)
+        ev_loop_destroy(loop);
+
     return r;
 }
diff --git a/src/common.h b/src/common.h
index 5800581..fea0db7 100644
--- a/src/common.h
+++ b/src/common.h
@@ -5,3 +5,5 @@
 #endif
 
 #define IN_ADDR_PORT(addr) (((struct sockaddr_in *)addr)->sin_port)
+
+extern int free_mem_on_exit;
diff --git a/src/server.c b/src/server.c
index b16ecf2..7980a92 100644
--- a/src/server.c
+++ b/src/server.c
@@ -6,6 +6,7 @@
 #include <netinet/in.h>
 #include <netinet/ip.h>
 #include <netinet/tcp.h>
+#include <signal.h>
 #include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -41,6 +42,8 @@ struct s_data {
     uint16_t csum_p;
 };
 
+struct s_data *global_s_data;
+
 static inline void s_prep_c_addr(struct o_s_sock *sock, struct tcphdr *hdr) {
     hdr->th_sport = ((struct sockaddr_in *)sock->s_data->s_addr)->sin_port;
     hdr->th_dport = ((struct sockaddr_in *)&sock->c_addr)->sin_port;
@@ -48,7 +51,7 @@ static inline void s_prep_c_addr(struct o_s_sock *sock, struct tcphdr *hdr) {
     hdr->th_off = 5;
 }
 
-static void s_sock_cleanup(EV_P_ struct o_s_sock *sock) {
+static void s_sock_cleanup(EV_P_ struct o_s_sock *sock, int stopping) {
     DBG("cleaning up socket %p", sock);
 
     if (sock->status == TCP_ESTABLISHED) {
@@ -58,7 +61,9 @@ static void s_sock_cleanup(EV_P_ struct o_s_sock *sock) {
         };
         s_prep_c_addr(sock, &buf);
         ssize_t sz;
-        if ((sz = sendto(sock->s_data->s_sock, &buf, sizeof(buf), 0, (struct sockaddr *)&sock->s_data->pkt_addr, sock->s_data->s_addrlen)) == -1) {
+        // don't need to save the real port because we're deleting the sock anyways
+        ((struct sockaddr_in *)&sock->c_addr)->sin_port = htons(0);
+        if ((sz = sendto(sock->s_data->s_sock, &buf, sizeof(buf), 0, (struct sockaddr *)&sock->c_addr, sock->s_data->s_addrlen)) == -1) {
             perror("sendto");
             ev_break(EV_A_ EVBREAK_ONE);
             return;
@@ -67,21 +72,26 @@ static void s_sock_cleanup(EV_P_ struct o_s_sock *sock) {
         }
     }
 
-    if (sock->c_sock != -1) {
-        close(sock->c_sock);
-    }
+    if (!stopping || free_mem_on_exit) {
+        DBG("freeing associated resources");
+
+        if (sock->c_sock != -1) {
+            close(sock->c_sock);
+        }
 
-    ev_timer_stop(EV_A_ &sock->tm_w);
-    ev_io_stop(EV_A_ &sock->io_w);
+        ev_timer_stop(EV_A_ &sock->tm_w);
+        if (sock->status == TCP_ESTABLISHED)
+            ev_io_stop(EV_A_ &sock->io_w);
 
-    HASH_DEL(sock->s_data->o_socks_by_caddr, sock);
+        HASH_DEL(sock->s_data->o_socks_by_caddr, sock);
 
-    free(sock);
+        free(sock);
+    }
 }
 
 static void s_tm_cb(EV_P_ ev_timer *w, int revents __attribute__((unused))) {
     DBG("timing out socket %p", w->data);
-    s_sock_cleanup(EV_A_ w->data);
+    s_sock_cleanup(EV_A_ w->data, 0);
 }
 
 static void sc_cb(EV_P_ ev_io *w, int revents __attribute__((unused))) {
@@ -346,6 +356,25 @@ static void ss_cb(EV_P_ ev_io *w, int revents __attribute__((unused))) {
     }
 }
 
+/* atexit cleanup */
+static void s_cleanup() {
+    if (!global_s_data)
+        return;
+
+    DBG("cleaning up");
+    struct o_s_sock *sock, *tmp;
+    HASH_ITER(hh, global_s_data->o_socks_by_caddr, sock, tmp) {
+        s_sock_cleanup(EV_DEFAULT, sock, 1);
+    }
+
+    global_s_data = NULL;
+}
+
+static void s_finish(EV_P_ ev_signal *w __attribute__((unused)), int revents __attribute__((unused))) {
+    s_cleanup();
+    ev_break(EV_A_ EVBREAK_ALL);
+}
+
 int start_server(const char *s_host, const char *s_port, const char *r_host, const char *r_port) {
     struct addrinfo *res;
     int r = getaddrinfo(s_host, s_port, NULL, &res);
@@ -383,17 +412,29 @@ int start_server(const char *s_host, const char *s_port, const char *r_host, con
         return 1;
     }
 
+    global_s_data = &s_data;
+
     struct ev_loop *loop = EV_DEFAULT;
     ev_io s_watcher;
+    ev_signal iwatcher, twatcher;
 
     ev_io_init(&s_watcher, ss_cb, s_data.s_sock, EV_READ);
     ev_io_start(EV_A_ &s_watcher);
+    ev_signal_init(&iwatcher, s_finish, SIGINT);
+    ev_signal_start(loop, &iwatcher);
+    ev_signal_init(&twatcher, s_finish, SIGTERM);
+    ev_signal_start(loop, &twatcher);
 
     s_watcher.data = &s_data;
 
     DBG("initialization complete, starting event loop");
     r = ev_run(loop, 0);
 
+    s_cleanup();
+
+    if (free_mem_on_exit)
+        ev_loop_destroy(loop);
+
     freeaddrinfo(res);
 
     return r;
diff --git a/src/udpastcp.c b/src/udpastcp.c
index 7230c93..fa8fea2 100644
--- a/src/udpastcp.c
+++ b/src/udpastcp.c
@@ -2,9 +2,12 @@
 #include <string.h>
 #include <stdio.h>
 #include <time.h>
+#include "common.h"
 #include "server.h"
 #include "client.h"
 
+int free_mem_on_exit = 0;
+
 int main(int argc, char *argv[]) {
     if (argc < 6) {
         puts("usage: udpintcp client|server LISTEN_HOST LISTEN_PORT REMOTE_HOST REMOTE_PORT");
@@ -13,9 +16,16 @@ int main(int argc, char *argv[]) {
 
     srandom((unsigned int)time(NULL));
 
+    if (getenv("UDPASTCP_RELEASE_MEMORY")) {
+        DBG("UDPASTCP_RELEASE_MEMORY is set, will free all memory on exit.");
+        free_mem_on_exit = 1;
+    }
+
     if (!strcmp(argv[1], "client")) {
+        DBG("starting client listening on [%s]:%s connecting to [%s]:%s", argv[2], argv[3], argv[4], argv[5]);
         return start_client(argv[2], argv[3], argv[4], argv[5]) == 0;
     } else if (!strcmp(argv[1], "server")) {
+        DBG("starting server listening on [%s]:%s connecting to [%s]:%s", argv[2], argv[3], argv[4], argv[5]);
         return start_server(argv[2], argv[3], argv[4], argv[5]) == 0;
     } else {
         fputs("invalid mode\n", stderr);
diff --git a/test.sh b/test.sh
index 842d530..a82020a 100755
--- a/test.sh
+++ b/test.sh
@@ -1,19 +1,21 @@
 #!/bin/sh
 # this script tests basic udpastcp functionality.
 
+: ${UDPASTCP:=./udpastcp}
+
 test_bidi() {
     (
         pids=
-        trap 'kill $pids' EXIT
-        ./udpastcp client "$1" 36563 "$1" 64109 &
+        trap 'kill $pids' INT TERM EXIT
+        $UDPASTCP client "$1" 36563 "$1" 64109 &
         pids="$!"
-        ./udpastcp server "$1" 64109 "$1" 41465 &
+        $UDPASTCP server "$1" 64109 "$1" 41465 &
         pids="$pids $!"
-        ( ( sleep 0.4; echo BBBBBBBB; ) | socat "udp-listen:41465,pf=${2}" - ) &
+        ( ( sleep 0.5; echo BBBBBBBB; ) | socat "udp-listen:41465,pf=${2}" - ) &
         pids="$pids $!"
-        ( ( sleep 0.2; echo AAAAAAAA; ) | socat - "udp-connect:localhost:36563,pf=${2}" ) &
+        ( ( sleep 0.4; echo AAAAAAAA; ) | socat - "udp-connect:localhost:36563,pf=${2}" ) &
         pids="$pids $!"
-        sleep 0.5
+        sleep 0.6
     )
 }