Refine memory management (#579)

- Added one macro to avoid dangling pointers
- Added two functions to perform NULL pointer check
  since the allocation is not guaranteed by C library, although it
  is a rare case, just for sanity
- Add NULL pointer check to brealloc() and bfree() and for sanity as well

Signed-off-by: Syrone Wong <wong.syrone@gmail.com>

src/cache.c

@@ -28,6 +28,7 @@
 #include <errno.h>
 #include <stdlib.h>
 #include "cache.h"
+#include "mm-wrapper.h"
 
 #ifdef __MINGW32__
 #include "win32.h"
@@ -91,13 +92,12 @@ int cache_delete(struct cache *cache, int keep_data)
                     cache->free_cb(entry->data);
                 }
             }
-            free(entry->key);
-            free(entry);
+            SS_SAFEFREE(entry->key);
+            SS_SAFEFREE(entry);
         }
     }
 
-    free(cache);
-    cache = NULL;
+    SS_SAFEFREE(cache);
     return 0;
 }
 
@@ -130,11 +130,11 @@ int cache_remove(struct cache *cache, char *key, size_t key_len)
             if (cache->free_cb) {
                 cache->free_cb(tmp->data);
             } else {
-                free(tmp->data);
+                SS_SAFEFREE(tmp->data);
             }
         }
-        free(tmp->key);
-        free(tmp);
+        SS_SAFEFREE(tmp->key);
+        SS_SAFEFREE(tmp);
     }
 
     return 0;
@@ -230,7 +230,7 @@ int cache_insert(struct cache *cache, char *key, size_t key_len, void *data)
         return ENOMEM;
     }
 
-    entry->key = malloc(key_len);
+    entry->key = SS_SAFEMALLOC(key_len);
     memcpy(entry->key, key, key_len);
     entry->data = data;
     HASH_ADD_KEYPTR(hh, cache->entries, entry->key, key_len, entry);
@@ -242,11 +242,11 @@ int cache_insert(struct cache *cache, char *key, size_t key_len, void *data)
                 if (cache->free_cb) {
                     cache->free_cb(entry->data);
                 } else {
-                    free(entry->data);
+                    SS_SAFEFREE(entry->data);
                 }
             }
-            free(entry->key);
-            free(entry);
+            SS_SAFEFREE(entry->key);
+            SS_SAFEFREE(entry);
             break;
         }
     }

src/encrypt.c

@@ -76,6 +76,7 @@
 #include "cache.h"
 #include "encrypt.h"
 #include "utils.h"
+#include "mm-wrapper.h"
 
 #define OFFSET_ROL(p, o) ((uint64_t)(*(p + o)) << (8 * o))
 
@@ -213,16 +214,17 @@ static int safe_memcmp(const void *s1, const void *s2, size_t n)
 int balloc(buffer_t *ptr, size_t capacity)
 {
     memset(ptr, 0, sizeof(buffer_t));
-    ptr->array    = malloc(capacity);
+    ptr->array    = SS_SAFEMALLOC(capacity);
     ptr->capacity = capacity;
     return capacity;
 }
 
 int brealloc(buffer_t *ptr, size_t len, size_t capacity)
 {
+    if (ptr == NULL) return -1;
     size_t real_capacity = max(len, capacity);
     if (ptr->capacity < real_capacity) {
-        ptr->array    = realloc(ptr->array, real_capacity);
+        ptr->array    = SS_SAFEREALLOC(ptr->array, real_capacity);
         ptr->capacity = real_capacity;
     }
     return real_capacity;
@@ -230,12 +232,12 @@ int brealloc(buffer_t *ptr, size_t len, size_t capacity)
 
 void bfree(buffer_t *ptr)
 {
+    if (ptr == NULL) return;
     ptr->idx      = 0;
     ptr->len      = 0;
     ptr->capacity = 0;
     if (ptr->array != NULL) {
-        free(ptr->array);
-        ptr->array = NULL;
+        SS_SAFEFREE(ptr->array);
     }
 }
 
@@ -306,8 +308,8 @@ static void merge(uint8_t *left, int llength, uint8_t *right,
         }
     }
 
-    free(ltmp);
-    free(rtmp);
+    SS_SAFEFREE(ltmp);
+    SS_SAFEFREE(rtmp);
 }
 
 static void merge_sort(uint8_t array[], int length,
@@ -365,8 +367,8 @@ void enc_table_init(const char *pass)
     uint64_t key = 0;
     uint8_t *digest;
 
-    enc_table = malloc(256);
-    dec_table = malloc(256);
+    enc_table = SS_SAFEMALLOC(256);
+    dec_table = SS_SAFEMALLOC(256);
 
     digest = enc_md5((const uint8_t *)pass, strlen(pass), NULL);
 

src/jconf.c

@@ -29,6 +29,7 @@
 #include "jconf.h"
 #include "json.h"
 #include "string.h"
+#include "mm-wrapper.h"
 
 #include <libcork/core.h>
 
@@ -49,10 +50,8 @@ static char *to_string(const json_value *value)
 
 void free_addr(ss_addr_t *addr)
 {
-    free(addr->host);
-    free(addr->port);
-    addr->host = NULL;
-    addr->port = NULL;
+    SS_SAFEFREE(addr->host);
+    SS_SAFEFREE(addr->port);
 }
 
 void parse_addr(const char *str, ss_addr_t *addr)
@@ -117,7 +116,7 @@ jconf_t *read_jconf(const char *file)
         FATAL("Too large config file.");
     }
 
-    buf = malloc(pos + 1);
+    buf = SS_SAFEMALLOC(pos + 1);
     if (buf == NULL) {
         FATAL("No enough memory.");
     }
@@ -199,7 +198,7 @@ jconf_t *read_jconf(const char *file)
         FATAL("Invalid config file");
     }
 
-    free(buf);
+    SS_SAFEFREE(buf);
     json_value_free(obj);
     return &conf;
 }

src/json.c

@@ -28,6 +28,7 @@
  */
 
 #include "json.h"
+#include "mm-wrapper.h"
 
 #ifdef _MSC_VER
 #ifndef _CRT_SECURE_NO_WARNINGS
@@ -90,12 +91,12 @@ typedef struct {
 
 static void *default_alloc(size_t size, int zero, void *user_data)
 {
-    return zero ? calloc(1, size) : malloc(size);
+    return zero ? calloc(1, size) : SS_SAFEMALLOC(size);
 }
 
 static void default_free(void *ptr, void *user_data)
 {
-    free(ptr);
+    SS_SAFEFREE(ptr);
 }
 
 static void *json_alloc(json_state *state, unsigned long size, int zero)

src/local.c

@@ -66,6 +66,7 @@
 #include "socks5.h"
 #include "acl.h"
 #include "local.h"
+#include "mm-wrapper.h"
 
 #ifndef EAGAIN
 #define EAGAIN EWOULDBLOCK
@@ -774,13 +775,13 @@ static void remote_send_cb(EV_P_ ev_io *w, int revents)
 static remote_t *new_remote(int fd, int timeout)
 {
     remote_t *remote;
-    remote = malloc(sizeof(remote_t));
+    remote = SS_SAFEMALLOC(sizeof(remote_t));
 
     memset(remote, 0, sizeof(remote_t));
 
-    remote->buf                 = malloc(sizeof(buffer_t));
-    remote->recv_ctx            = malloc(sizeof(remote_ctx_t));
-    remote->send_ctx            = malloc(sizeof(remote_ctx_t));
+    remote->buf                 = SS_SAFEMALLOC(sizeof(buffer_t));
+    remote->recv_ctx            = SS_SAFEMALLOC(sizeof(remote_ctx_t));
+    remote->send_ctx            = SS_SAFEMALLOC(sizeof(remote_ctx_t));
     remote->recv_ctx->connected = 0;
     remote->send_ctx->connected = 0;
     remote->fd                  = fd;
@@ -806,11 +807,11 @@ static void free_remote(remote_t *remote)
     }
     if (remote->buf != NULL) {
         bfree(remote->buf);
-        free(remote->buf);
+        SS_SAFEFREE(remote->buf);
     }
-    free(remote->recv_ctx);
-    free(remote->send_ctx);
-    free(remote);
+    SS_SAFEFREE(remote->recv_ctx);
+    SS_SAFEFREE(remote->send_ctx);
+    SS_SAFEFREE(remote);
 }
 
 static void close_and_free_remote(EV_P_ remote_t *remote)
@@ -828,13 +829,13 @@ static void close_and_free_remote(EV_P_ remote_t *remote)
 static server_t *new_server(int fd, int method)
 {
     server_t *server;
-    server = malloc(sizeof(server_t));
+    server = SS_SAFEMALLOC(sizeof(server_t));
 
     memset(server, 0, sizeof(server_t));
 
-    server->recv_ctx            = malloc(sizeof(server_ctx_t));
-    server->send_ctx            = malloc(sizeof(server_ctx_t));
-    server->buf                 = malloc(sizeof(buffer_t));
+    server->recv_ctx            = SS_SAFEMALLOC(sizeof(server_ctx_t));
+    server->send_ctx            = SS_SAFEMALLOC(sizeof(server_ctx_t));
+    server->buf                 = SS_SAFEMALLOC(sizeof(buffer_t));
     server->recv_ctx->connected = 0;
     server->send_ctx->connected = 0;
     server->fd                  = fd;
@@ -842,8 +843,8 @@ static server_t *new_server(int fd, int method)
     server->send_ctx->server    = server;
 
     if (method) {
-        server->e_ctx = malloc(sizeof(struct enc_ctx));
-        server->d_ctx = malloc(sizeof(struct enc_ctx));
+        server->e_ctx = SS_SAFEMALLOC(sizeof(struct enc_ctx));
+        server->d_ctx = SS_SAFEMALLOC(sizeof(struct enc_ctx));
         enc_ctx_init(method, server->e_ctx, 1);
         enc_ctx_init(method, server->d_ctx, 0);
     } else {
@@ -870,19 +871,19 @@ static void free_server(server_t *server)
     }
     if (server->e_ctx != NULL) {
         cipher_context_release(&server->e_ctx->evp);
-        free(server->e_ctx);
+        SS_SAFEFREE(server->e_ctx);
     }
     if (server->d_ctx != NULL) {
         cipher_context_release(&server->d_ctx->evp);
-        free(server->d_ctx);
+        SS_SAFEFREE(server->d_ctx);
     }
     if (server->buf != NULL) {
         bfree(server->buf);
-        free(server->buf);
+        SS_SAFEFREE(server->buf);
     }
-    free(server->recv_ctx);
-    free(server->send_ctx);
-    free(server);
+    SS_SAFEFREE(server->recv_ctx);
+    SS_SAFEFREE(server->send_ctx);
+    SS_SAFEFREE(server);
 }
 
 static void close_and_free_server(EV_P_ server_t *server)
@@ -1188,12 +1189,12 @@ int main(int argc, char **argv)
     // Setup proxy context
     listen_ctx_t listen_ctx;
     listen_ctx.remote_num  = remote_num;
-    listen_ctx.remote_addr = malloc(sizeof(struct sockaddr *) * remote_num);
+    listen_ctx.remote_addr = SS_SAFEMALLOC(sizeof(struct sockaddr *) * remote_num);
     for (i = 0; i < remote_num; i++) {
         char *host = remote_addr[i].host;
         char *port = remote_addr[i].port == NULL ? remote_port :
                      remote_addr[i].port;
-        struct sockaddr_storage *storage = malloc(sizeof(struct sockaddr_storage));
+        struct sockaddr_storage *storage = SS_SAFEMALLOC(sizeof(struct sockaddr_storage));
         memset(storage, 0, sizeof(struct sockaddr_storage));
         if (get_sockaddr(host, port, storage, 1) == -1) {
             FATAL("failed to resolve the provided hostname");
@@ -1256,8 +1257,8 @@ int main(int argc, char **argv)
     }
 
     for (i = 0; i < remote_num; i++)
-        free(listen_ctx.remote_addr[i]);
-    free(listen_ctx.remote_addr);
+        SS_SAFEFREE(listen_ctx.remote_addr[i]);
+    SS_SAFEFREE(listen_ctx.remote_addr);
 
 #ifdef __MINGW32__
     winsock_cleanup();
@@ -1323,7 +1324,7 @@ int start_ss_local_server(profile_t profile)
     LOGI("initialize ciphers... %s", method);
     int m = enc_init(password, method);
 
-    struct sockaddr_storage *storage = malloc(sizeof(struct sockaddr_storage));
+    struct sockaddr_storage *storage = SS_SAFEMALLOC(sizeof(struct sockaddr_storage));
     memset(storage, 0, sizeof(struct sockaddr_storage));
     if (get_sockaddr(remote_host, remote_port_str, storage, 1) == -1) {
         return -1;
@@ -1334,7 +1335,7 @@ int start_ss_local_server(profile_t profile)
     listen_ctx_t listen_ctx;
 
     listen_ctx.remote_num     = 1;
-    listen_ctx.remote_addr    = malloc(sizeof(struct sockaddr *));
+    listen_ctx.remote_addr    = SS_SAFEMALLOC(sizeof(struct sockaddr *));
     listen_ctx.remote_addr[0] = (struct sockaddr *)storage;
     listen_ctx.timeout        = timeout;
     listen_ctx.method         = m;
@@ -1387,7 +1388,7 @@ int start_ss_local_server(profile_t profile)
     free_connections(loop);
     close(listen_ctx.fd);
 
-    free(listen_ctx.remote_addr);
+    SS_SAFEFREE(listen_ctx.remote_addr);
 
 #ifdef __MINGW32__
     winsock_cleanup();

src/manager.c

@@ -66,6 +66,7 @@
 #include "json.h"
 #include "utils.h"
 #include "manager.h"
+#include "mm-wrapper.h"
 
 #ifndef BUF_SIZE
 #define BUF_SIZE 65535
@@ -323,7 +324,7 @@ static void remove_server(char *prefix, char *port)
     cork_hash_table_delete(server_table, (void *)port, (void **)&old_port, (void **)&old_server);
 
     if (old_server != NULL) {
-        free(old_server);
+        SS_SAFEFREE(old_server);
     }
 
     stop_server(prefix, port);
@@ -371,7 +372,7 @@ static void manager_recv_cb(EV_P_ ev_io *w, int revents)
         if (server == NULL || server->port[0] == 0 || server->password[0] == 0) {
             LOGE("invalid command: %s:%s", buf, get_data(buf, r));
             if (server != NULL) {
-                free(server);
+                SS_SAFEFREE(server);
             }
             goto ERROR_MSG;
         }
@@ -389,13 +390,13 @@ static void manager_recv_cb(EV_P_ ev_io *w, int revents)
         if (server == NULL || server->port[0] == 0) {
             LOGE("invalid command: %s:%s", buf, get_data(buf, r));
             if (server != NULL) {
-                free(server);
+                SS_SAFEFREE(server);
             }
             goto ERROR_MSG;
         }
 
         remove_server(working_dir, server->port);
-        free(server);
+        SS_SAFEFREE(server);
 
         char msg[3] = "ok";
         if (sendto(manager->fd, msg, 3, 0, (struct sockaddr *)&claddr, len) != 3) {

src/mm-wrapper.h

@@ -0,0 +1,52 @@
+/*
+ * mm-wrapper.h - Define safe memory management wrapper
+ *
+ * Copyright (C) 2013 - 2016, Max Lv <max.c.lv@gmail.com>
+ *
+ * This file is part of the shadowsocks-libev.
+ * shadowsocks-libev is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * shadowsocks-libev is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with shadowsocks-libev; see the file COPYING. If not, see
+ * <http://www.gnu.org/licenses/>.
+ */
+
+#ifndef _MM_WRAPPER_H
+#define _MM_WRAPPER_H
+
+#include <stdlib.h>
+
+#define SS_SAFEFREE(ptr) \
+    do {                 \
+        free(ptr);       \
+        ptr = NULL;      \
+    } while(0)
+
+static inline void *SS_SAFEMALLOC(size_t size);
+static inline void *SS_SAFEREALLOC(void *ptr, size_t new_size);
+
+static inline void *SS_SAFEMALLOC(size_t size) {
+    void *tmp = malloc(size);
+    if (tmp == NULL)
+        exit(EXIT_FAILURE);
+    return tmp;
+}
+
+static inline void *SS_SAFEREALLOC(void *ptr, size_t new_size) {
+    void *new = realloc(ptr, new_size);
+    if (new == NULL) {
+        free(ptr); ptr = NULL;
+        exit(EXIT_FAILURE);
+    }
+    return new;
+}
+
+#endif // _MM_WRAPPER_H

src/redir.c

@@ -49,6 +49,7 @@
 #include "utils.h"
 #include "common.h"
 #include "redir.h"
+#include "mm-wrapper.h"
 
 #ifndef EAGAIN
 #define EAGAIN EWOULDBLOCK
@@ -458,13 +459,13 @@ static void remote_send_cb(EV_P_ ev_io *w, int revents)
 static remote_t *new_remote(int fd, int timeout)
 {
     remote_t *remote;
-    remote = malloc(sizeof(remote_t));
+    remote = SS_SAFEMALLOC(sizeof(remote_t));
 
     memset(remote, 0, sizeof(remote_t));
 
-    remote->recv_ctx            = malloc(sizeof(remote_ctx_t));
-    remote->send_ctx            = malloc(sizeof(remote_ctx_t));
-    remote->buf                 = malloc(sizeof(buffer_t));
+    remote->recv_ctx            = SS_SAFEMALLOC(sizeof(remote_ctx_t));
+    remote->send_ctx            = SS_SAFEMALLOC(sizeof(remote_ctx_t));
+    remote->buf                 = SS_SAFEMALLOC(sizeof(buffer_t));
     remote->fd                  = fd;
     remote->recv_ctx->remote    = remote;
     remote->recv_ctx->connected = 0;
@@ -489,11 +490,11 @@ static void free_remote(remote_t *remote)
         }
         if (remote->buf != NULL) {
             bfree(remote->buf);
-            free(remote->buf);
+            SS_SAFEFREE(remote->buf);
         }
-        free(remote->recv_ctx);
-        free(remote->send_ctx);
-        free(remote);
+        SS_SAFEFREE(remote->recv_ctx);
+        SS_SAFEFREE(remote->send_ctx);
+        SS_SAFEFREE(remote);
     }
 }
 
@@ -511,11 +512,11 @@ static void close_and_free_remote(EV_P_ remote_t *remote)
 static server_t *new_server(int fd, int method)
 {
     server_t *server;
-    server = malloc(sizeof(server_t));
+    server = SS_SAFEMALLOC(sizeof(server_t));
 
-    server->recv_ctx            = malloc(sizeof(server_ctx_t));
-    server->send_ctx            = malloc(sizeof(server_ctx_t));
-    server->buf                 = malloc(sizeof(buffer_t));
+    server->recv_ctx            = SS_SAFEMALLOC(sizeof(server_ctx_t));
+    server->send_ctx            = SS_SAFEMALLOC(sizeof(server_ctx_t));
+    server->buf                 = SS_SAFEMALLOC(sizeof(buffer_t));
     server->fd                  = fd;
     server->recv_ctx->server    = server;
     server->recv_ctx->connected = 0;
@@ -523,8 +524,8 @@ static server_t *new_server(int fd, int method)
     server->send_ctx->connected = 0;
 
     if (method) {
-        server->e_ctx = malloc(sizeof(enc_ctx_t));
-        server->d_ctx = malloc(sizeof(enc_ctx_t));
+        server->e_ctx = SS_SAFEMALLOC(sizeof(enc_ctx_t));
+        server->d_ctx = SS_SAFEMALLOC(sizeof(enc_ctx_t));
         enc_ctx_init(method, server->e_ctx, 1);
         enc_ctx_init(method, server->d_ctx, 0);
     } else {
@@ -548,19 +549,19 @@ static void free_server(server_t *server)
         }
         if (server->e_ctx != NULL) {
             cipher_context_release(&server->e_ctx->evp);
-            free(server->e_ctx);
+            SS_SAFEFREE(server->e_ctx);
         }
         if (server->d_ctx != NULL) {
             cipher_context_release(&server->d_ctx->evp);
-            free(server->d_ctx);
+            SS_SAFEFREE(server->d_ctx);
         }
         if (server->buf != NULL) {
             bfree(server->buf);
-            free(server->buf);
+            SS_SAFEFREE(server->buf);
         }
-        free(server->recv_ctx);
-        free(server->send_ctx);
-        free(server);
+        SS_SAFEFREE(server->recv_ctx);
+        SS_SAFEFREE(server->send_ctx);
+        SS_SAFEFREE(server);
     }
 }
 
@@ -803,12 +804,12 @@ int main(int argc, char **argv)
     // Setup proxy context
     listen_ctx_t listen_ctx;
     listen_ctx.remote_num  = remote_num;
-    listen_ctx.remote_addr = malloc(sizeof(struct sockaddr *) * remote_num);
+    listen_ctx.remote_addr = SS_SAFEMALLOC(sizeof(struct sockaddr *) * remote_num);
     for (int i = 0; i < remote_num; i++) {
         char *host = remote_addr[i].host;
         char *port = remote_addr[i].port == NULL ? remote_port :
                      remote_addr[i].port;
-        struct sockaddr_storage *storage = malloc(sizeof(struct sockaddr_storage));
+        struct sockaddr_storage *storage = SS_SAFEMALLOC(sizeof(struct sockaddr_storage));
         memset(storage, 0, sizeof(struct sockaddr_storage));
         if (get_sockaddr(host, port, storage, 1) == -1) {
             FATAL("failed to resolve the provided hostname");

src/resolv.c

@@ -46,6 +46,7 @@
 
 #include "resolv.h"
 #include "utils.h"
+#include "mm-wrapper.h"
 
 /*
  * Implement DNS resolution interface using libudns
@@ -146,7 +147,7 @@ struct ResolvQuery *resolv_query(const char *hostname, void (*client_cb)(struct
     /*
      * Wrap udns's call back in our own
      */
-    struct ResolvQuery *cb_data = malloc(sizeof(struct ResolvQuery));
+    struct ResolvQuery *cb_data = SS_SAFEMALLOC(sizeof(struct ResolvQuery));
     if (cb_data == NULL) {
         LOGE("Failed to allocate memory for DNS query callback data.");
         return NULL;
@@ -184,8 +185,7 @@ struct ResolvQuery *resolv_query(const char *hostname, void (*client_cb)(struct
         if (cb_data->client_free_cb != NULL) {
             cb_data->client_free_cb(cb_data->client_cb_data);
         }
-        free(cb_data);
-        cb_data = NULL;
+        SS_SAFEFREE(cb_data);
     }
 
     return cb_data;
@@ -200,15 +200,14 @@ void resolv_cancel(struct ResolvQuery *query_handle)
          i++)
         if (cb_data->queries[i] != NULL) {
             dns_cancel(ctx, cb_data->queries[i]);
-            free(cb_data->queries[i]);
-            cb_data->queries[i] = NULL;
+            SS_SAFEFREE(cb_data->queries[i]);
         }
 
     if (cb_data->client_free_cb != NULL) {
         cb_data->client_free_cb(cb_data->client_cb_data);
     }
 
-    free(cb_data);
+    SS_SAFEFREE(cb_data);
 }
 
 /*
@@ -235,7 +234,7 @@ static void dns_query_v4_cb(struct dns_ctx *ctx, struct dns_rr_a4 *result, void
             LOGI("IPv4 resolv: %s", dns_strerror(dns_status(ctx)));
         }
     } else if (result->dnsa4_nrr > 0) {
-        struct sockaddr **new_responses = realloc(cb_data->responses,
+        struct sockaddr **new_responses = SS_SAFEREALLOC(cb_data->responses,
                                                   (cb_data->response_count +
                                                    result->dnsa4_nrr) *
                                                   sizeof(struct sockaddr *));
@@ -263,7 +262,7 @@ static void dns_query_v4_cb(struct dns_ctx *ctx, struct dns_rr_a4 *result, void
         }
     }
 
-    free(result);
+    SS_SAFEFREE(result);
     cb_data->queries[0] = NULL; /* mark A query as being completed */
 
     /* Once all queries have completed, call client callback */
@@ -281,7 +280,7 @@ static void dns_query_v6_cb(struct dns_ctx *ctx, struct dns_rr_a6 *result, void
             LOGI("IPv6 resolv: %s", dns_strerror(dns_status(ctx)));
         }
     } else if (result->dnsa6_nrr > 0) {
-        struct sockaddr **new_responses = realloc(cb_data->responses,
+        struct sockaddr **new_responses = SS_SAFEREALLOC(cb_data->responses,
                                                   (cb_data->response_count +
                                                    result->dnsa6_nrr) *
                                                   sizeof(struct sockaddr *));
@@ -309,7 +308,7 @@ static void dns_query_v6_cb(struct dns_ctx *ctx, struct dns_rr_a6 *result, void
         }
     }
 
-    free(result);
+    SS_SAFEFREE(result);
     cb_data->queries[1] = NULL; /* mark AAAA query as being completed */
 
     /* Once all queries have completed, call client callback */
@@ -336,13 +335,13 @@ static void process_client_callback(struct ResolvQuery *cb_data)
     cb_data->client_cb(best_address, cb_data->client_cb_data);
 
     for (int i = 0; i < cb_data->response_count; i++)
-        free(cb_data->responses[i]);
+        SS_SAFEFREE(cb_data->responses[i]);
 
-    free(cb_data->responses);
+    SS_SAFEFREE(cb_data->responses);
     if (cb_data->client_free_cb != NULL) {
         cb_data->client_free_cb(cb_data->client_cb_data);
     }
-    free(cb_data);
+    SS_SAFEFREE(cb_data);
 }
 
 static struct sockaddr *choose_ipv4_first(struct ResolvQuery *cb_data)

src/server.c

@@ -64,6 +64,7 @@
 #include "utils.h"
 #include "acl.h"
 #include "server.h"
+#include "mm-wrapper.h"
 
 #ifndef EAGAIN
 #define EAGAIN EWOULDBLOCK
@@ -1107,10 +1108,10 @@ static remote_t *new_remote(int fd)
 
     remote_t *remote;
 
-    remote                      = malloc(sizeof(remote_t));
-    remote->recv_ctx            = malloc(sizeof(remote_ctx_t));
-    remote->send_ctx            = malloc(sizeof(remote_ctx_t));
-    remote->buf                 = malloc(sizeof(buffer_t));
+    remote                      = SS_SAFEMALLOC(sizeof(remote_t));
+    remote->recv_ctx            = SS_SAFEMALLOC(sizeof(remote_ctx_t));
+    remote->send_ctx            = SS_SAFEMALLOC(sizeof(remote_ctx_t));
+    remote->buf                 = SS_SAFEMALLOC(sizeof(buffer_t));
     remote->fd                  = fd;
     remote->recv_ctx->remote    = remote;
     remote->recv_ctx->connected = 0;
@@ -1133,11 +1134,11 @@ static void free_remote(remote_t *remote)
     }
     if (remote->buf != NULL) {
         bfree(remote->buf);
-        free(remote->buf);
+        SS_SAFEFREE(remote->buf);
     }
-    free(remote->recv_ctx);
-    free(remote->send_ctx);
-    free(remote);
+    SS_SAFEFREE(remote->recv_ctx);
+    SS_SAFEFREE(remote->send_ctx);
+    SS_SAFEFREE(remote);
 }
 
 static void close_and_free_remote(EV_P_ remote_t *remote)
@@ -1161,13 +1162,13 @@ static server_t *new_server(int fd, listen_ctx_t *listener)
     }
 
     server_t *server;
-    server = malloc(sizeof(server_t));
+    server = SS_SAFEMALLOC(sizeof(server_t));
 
     memset(server, 0, sizeof(server_t));
 
-    server->recv_ctx            = malloc(sizeof(server_ctx_t));
-    server->send_ctx            = malloc(sizeof(server_ctx_t));
-    server->buf                 = malloc(sizeof(buffer_t));
+    server->recv_ctx            = SS_SAFEMALLOC(sizeof(server_ctx_t));
+    server->send_ctx            = SS_SAFEMALLOC(sizeof(server_ctx_t));
+    server->buf                 = SS_SAFEMALLOC(sizeof(buffer_t));
     server->fd                  = fd;
     server->recv_ctx->server    = server;
     server->recv_ctx->connected = 0;
@@ -1179,8 +1180,8 @@ static server_t *new_server(int fd, listen_ctx_t *listener)
     server->remote              = NULL;
 
     if (listener->method) {
-        server->e_ctx = malloc(sizeof(enc_ctx_t));
-        server->d_ctx = malloc(sizeof(enc_ctx_t));
+        server->e_ctx = SS_SAFEMALLOC(sizeof(enc_ctx_t));
+        server->d_ctx = SS_SAFEMALLOC(sizeof(enc_ctx_t));
         enc_ctx_init(listener->method, server->e_ctx, 1);
         enc_ctx_init(listener->method, server->d_ctx, 0);
     } else {
@@ -1197,7 +1198,7 @@ static server_t *new_server(int fd, listen_ctx_t *listener)
 
     server->chunk = (chunk_t *)malloc(sizeof(chunk_t));
     memset(server->chunk, 0, sizeof(chunk_t));
-    server->chunk->buf = malloc(sizeof(buffer_t));
+    server->chunk->buf = SS_SAFEMALLOC(sizeof(buffer_t));
     memset(server->chunk->buf, 0, sizeof(buffer_t));
 
     cork_dllist_add(&connections, &server->entries);
@@ -1212,31 +1213,29 @@ static void free_server(server_t *server)
     if (server->chunk != NULL) {
         if (server->chunk->buf != NULL) {
             bfree(server->chunk->buf);
-            free(server->chunk->buf);
-            server->chunk->buf = NULL;
+            SS_SAFEFREE(server->chunk->buf);
         }
-        free(server->chunk);
-        server->chunk = NULL;
+        SS_SAFEFREE(server->chunk);
     }
     if (server->remote != NULL) {
         server->remote->server = NULL;
     }
     if (server->e_ctx != NULL) {
         cipher_context_release(&server->e_ctx->evp);
-        free(server->e_ctx);
+        SS_SAFEFREE(server->e_ctx);
     }
     if (server->d_ctx != NULL) {
         cipher_context_release(&server->d_ctx->evp);
-        free(server->d_ctx);
+        SS_SAFEFREE(server->d_ctx);
     }
     if (server->buf != NULL) {
         bfree(server->buf);
-        free(server->buf);
+        SS_SAFEFREE(server->buf);
     }
 
-    free(server->recv_ctx);
-    free(server->send_ctx);
-    free(server);
+    SS_SAFEFREE(server->recv_ctx);
+    SS_SAFEFREE(server->send_ctx);
+    SS_SAFEFREE(server);
 }
 
 static void close_and_free_server(EV_P_ server_t *server)

src/tunnel.c

@@ -58,6 +58,7 @@
 #include "netutils.h"
 #include "utils.h"
 #include "tunnel.h"
+#include "mm-wrapper.h"
 
 #ifndef EAGAIN
 #define EAGAIN EWOULDBLOCK
@@ -506,13 +507,13 @@ static void remote_send_cb(EV_P_ ev_io *w, int revents)
 static remote_t *new_remote(int fd, int timeout)
 {
     remote_t *remote;
-    remote = malloc(sizeof(remote_t));
+    remote = SS_SAFEMALLOC(sizeof(remote_t));
 
     memset(remote, 0, sizeof(remote_t));
 
-    remote->buf                 = malloc(sizeof(buffer_t));
-    remote->recv_ctx            = malloc(sizeof(remote_ctx_t));
-    remote->send_ctx            = malloc(sizeof(remote_ctx_t));
+    remote->buf                 = SS_SAFEMALLOC(sizeof(buffer_t));
+    remote->recv_ctx            = SS_SAFEMALLOC(sizeof(remote_ctx_t));
+    remote->send_ctx            = SS_SAFEMALLOC(sizeof(remote_ctx_t));
     remote->fd                  = fd;
     remote->recv_ctx->remote    = remote;
     remote->recv_ctx->connected = 0;
@@ -537,11 +538,11 @@ static void free_remote(remote_t *remote)
         }
         if (remote->buf) {
             bfree(remote->buf);
-            free(remote->buf);
+            SS_SAFEFREE(remote->buf);
         }
-        free(remote->recv_ctx);
-        free(remote->send_ctx);
-        free(remote);
+        SS_SAFEFREE(remote->recv_ctx);
+        SS_SAFEFREE(remote->send_ctx);
+        SS_SAFEFREE(remote);
     }
 }
 
@@ -560,10 +561,10 @@ static server_t *new_server(int fd, int method)
 {
     server_t *server;
 
-    server                      = malloc(sizeof(server_t));
-    server->buf                 = malloc(sizeof(buffer_t));
-    server->recv_ctx            = malloc(sizeof(server_ctx_t));
-    server->send_ctx            = malloc(sizeof(server_ctx_t));
+    server                      = SS_SAFEMALLOC(sizeof(server_t));
+    server->buf                 = SS_SAFEMALLOC(sizeof(buffer_t));
+    server->recv_ctx            = SS_SAFEMALLOC(sizeof(server_ctx_t));
+    server->send_ctx            = SS_SAFEMALLOC(sizeof(server_ctx_t));
     server->fd                  = fd;
     server->recv_ctx->server    = server;
     server->recv_ctx->connected = 0;
@@ -571,8 +572,8 @@ static server_t *new_server(int fd, int method)
     server->send_ctx->connected = 0;
 
     if (method) {
-        server->e_ctx = malloc(sizeof(struct enc_ctx));
-        server->d_ctx = malloc(sizeof(struct enc_ctx));
+        server->e_ctx = SS_SAFEMALLOC(sizeof(struct enc_ctx));
+        server->d_ctx = SS_SAFEMALLOC(sizeof(struct enc_ctx));
         enc_ctx_init(method, server->e_ctx, 1);
         enc_ctx_init(method, server->d_ctx, 0);
     } else {
@@ -596,19 +597,19 @@ static void free_server(server_t *server)
         }
         if (server->e_ctx != NULL) {
             cipher_context_release(&server->e_ctx->evp);
-            free(server->e_ctx);
+            SS_SAFEFREE(server->e_ctx);
         }
         if (server->d_ctx != NULL) {
             cipher_context_release(&server->d_ctx->evp);
-            free(server->d_ctx);
+            SS_SAFEFREE(server->d_ctx);
         }
         if (server->buf) {
             bfree(server->buf);
-            free(server->buf);
+            SS_SAFEFREE(server->buf);
         }
-        free(server->recv_ctx);
-        free(server->send_ctx);
-        free(server);
+        SS_SAFEFREE(server->recv_ctx);
+        SS_SAFEFREE(server->send_ctx);
+        SS_SAFEFREE(server);
     }
 }
 
@@ -882,12 +883,12 @@ int main(int argc, char **argv)
     struct listen_ctx listen_ctx;
     listen_ctx.tunnel_addr = tunnel_addr;
     listen_ctx.remote_num  = remote_num;
-    listen_ctx.remote_addr = malloc(sizeof(struct sockaddr *) * remote_num);
+    listen_ctx.remote_addr = SS_SAFEMALLOC(sizeof(struct sockaddr *) * remote_num);
     for (i = 0; i < remote_num; i++) {
         char *host = remote_addr[i].host;
         char *port = remote_addr[i].port == NULL ? remote_port :
                      remote_addr[i].port;
-        struct sockaddr_storage *storage = malloc(sizeof(struct sockaddr_storage));
+        struct sockaddr_storage *storage = SS_SAFEMALLOC(sizeof(struct sockaddr_storage));
         memset(storage, 0, sizeof(struct sockaddr_storage));
         if (get_sockaddr(host, port, storage, 1) == -1) {
             FATAL("failed to resolve the provided hostname");

src/udprelay.c

@@ -60,6 +60,7 @@
 #include "netutils.h"
 #include "cache.h"
 #include "udprelay.h"
+#include "mm-wrapper.h"
 
 #ifdef MODULE_REMOTE
 #define MAX_UDP_CONN_NUM 512
@@ -475,7 +476,7 @@ int create_server_socket(const char *host, const char *port)
 
 remote_ctx_t *new_remote(int fd, server_ctx_t *server_ctx)
 {
-    remote_ctx_t *ctx = malloc(sizeof(remote_ctx_t));
+    remote_ctx_t *ctx = SS_SAFEMALLOC(sizeof(remote_ctx_t));
     memset(ctx, 0, sizeof(remote_ctx_t));
 
     ctx->fd         = fd;
@@ -490,7 +491,7 @@ remote_ctx_t *new_remote(int fd, server_ctx_t *server_ctx)
 
 server_ctx_t *new_server_ctx(int fd)
 {
-    server_ctx_t *ctx = malloc(sizeof(server_ctx_t));
+    server_ctx_t *ctx = SS_SAFEMALLOC(sizeof(server_ctx_t));
     memset(ctx, 0, sizeof(server_ctx_t));
 
     ctx->fd = fd;
@@ -503,9 +504,9 @@ server_ctx_t *new_server_ctx(int fd)
 #ifdef MODULE_REMOTE
 struct query_ctx *new_query_ctx(char *buf, size_t len)
 {
-    struct query_ctx *ctx = malloc(sizeof(struct query_ctx));
+    struct query_ctx *ctx = SS_SAFEMALLOC(sizeof(struct query_ctx));
     memset(ctx, 0, sizeof(struct query_ctx));
-    ctx->buf = malloc(sizeof(buffer_t));
+    ctx->buf = SS_SAFEMALLOC(sizeof(buffer_t));
     balloc(ctx->buf, len);
     memcpy(ctx->buf->array, buf, len);
     ctx->buf->len = len;
@@ -521,9 +522,9 @@ void close_and_free_query(EV_P_ struct query_ctx *ctx)
         }
         if (ctx->buf != NULL) {
             bfree(ctx->buf);
-            free(ctx->buf);
+            SS_SAFEFREE(ctx->buf);
         }
-        free(ctx);
+        SS_SAFEFREE(ctx);
     }
 }
 
@@ -535,7 +536,7 @@ void close_and_free_remote(EV_P_ remote_ctx_t *ctx)
         ev_timer_stop(EV_A_ & ctx->watcher);
         ev_io_stop(EV_A_ & ctx->io);
         close(ctx->fd);
-        free(ctx);
+        SS_SAFEFREE(ctx);
     }
 }
 
@@ -655,7 +656,7 @@ static void remote_recv_cb(EV_P_ ev_io *w, int revents)
     socklen_t src_addr_len = sizeof(src_addr);
     memset(&src_addr, 0, src_addr_len);
 
-    buffer_t *buf = malloc(sizeof(buffer_t));
+    buffer_t *buf = SS_SAFEMALLOC(sizeof(buffer_t));
     balloc(buf, BUF_SIZE);
 
     // recv
@@ -795,7 +796,7 @@ static void remote_recv_cb(EV_P_ ev_io *w, int revents)
 CLEAN_UP:
 
     bfree(buf);
-    free(buf);
+    SS_SAFEFREE(buf);
 }
 
 static void server_recv_cb(EV_P_ ev_io *w, int revents)
@@ -804,7 +805,7 @@ static void server_recv_cb(EV_P_ ev_io *w, int revents)
     struct sockaddr_storage src_addr;
     memset(&src_addr, 0, sizeof(struct sockaddr_storage));
 
-    buffer_t *buf = malloc(sizeof(buffer_t));
+    buffer_t *buf = SS_SAFEMALLOC(sizeof(buffer_t));
     balloc(buf, BUF_SIZE);
 
     socklen_t src_addr_len = sizeof(struct sockaddr_storage);
@@ -1237,7 +1238,7 @@ static void server_recv_cb(EV_P_ ev_io *w, int revents)
 
 CLEAN_UP:
     bfree(buf);
-    free(buf);
+    SS_SAFEFREE(buf);
 }
 
 void free_cb(void *element)
@@ -1309,7 +1310,7 @@ void free_udprelay()
         ev_io_stop(loop, &server_ctx->io);
         close(server_ctx->fd);
         cache_delete(server_ctx->conn_cache, 0);
-        free(server_ctx);
+        SS_SAFEFREE(server_ctx);
         server_ctx_list[server_num] = NULL;
     }
 }