Support custom allocators

And hook Ruby's allocator in the Ruby extension.

This is largely copied from how prism do it.

We had several bugs in the past where we were freeing memory
allocated by Ruby with the system allocator and vice-versa.

It is much simpler if we only use one allocator, and it also
has the benefit of feeding Ruby GC with malloc statistics.

Author: byroot

contrib/ruby/ext/trilogy-ruby/cext.c

@@ -41,9 +41,7 @@ static void buffer_pool_free(void *data)
     buffer_pool *pool = (buffer_pool *)data;
     if (pool->capa) {
         for (size_t index = 0; index < pool->len; index++) {
-            // NB: buff was allocated by trilogy/buffer.h using raw `malloc`
-            // hence we must use raw `free`.
-            free(pool->entries[index].buff);
+            xfree(pool->entries[index].buff);
         }
         xfree(pool->entries);
     }
@@ -124,7 +122,7 @@ static void buffer_checkout(trilogy_buffer_t *buffer, size_t initial_capacity)
         buffer->buff = pool->entries[pool->len].buff;
         buffer->cap = pool->entries[pool->len].cap;
     } else {
-        buffer->buff = malloc(initial_capacity);
+        buffer->buff = xmalloc(initial_capacity);
         buffer->cap = initial_capacity;
     }
 }
@@ -134,7 +132,7 @@ static bool buffer_checkin(trilogy_buffer_t *buffer)
     buffer_pool * pool = get_buffer_pool();
 
     if (pool->len >= BUFFER_POOL_MAX_SIZE) {
-        free(buffer->buff);
+        xfree(buffer->buff);
         buffer->buff = NULL;
         buffer->cap = 0;
         return false;

example/trilogy_query.c

@@ -38,7 +38,7 @@ static int execute_query(trilogy_conn_t *conn, const char *sql)
         rc = trilogy_read_full_column(conn, &column);
 
         if (rc < 0) {
-            free(binary_columns);
+            xfree(binary_columns);
             return rc;
         }
 
@@ -58,7 +58,7 @@ static int execute_query(trilogy_conn_t *conn, const char *sql)
 
     // shut scan-build up
     if (column_count == 0) {
-        free(binary_columns);
+        xfree(binary_columns);
         return TRILOGY_OK;
     }
 
@@ -84,8 +84,8 @@ static int execute_query(trilogy_conn_t *conn, const char *sql)
         printf("\n");
     }
 
-    free(binary_columns);
-    free(values);
+    xfree(binary_columns);
+    xfree(values);
 
     if (rc == TRILOGY_EOF) {
         rc = TRILOGY_OK;
@@ -145,7 +145,7 @@ int main(int argc, char *argv[])
         switch (opt) {
         case 'h':
             if (optarg) {
-                free(connopt.hostname);
+                xfree(connopt.hostname);
                 connopt.hostname = strdup(optarg);
             }
             break;
@@ -156,25 +156,25 @@ int main(int argc, char *argv[])
             break;
         case 's':
             if (optarg) {
-                free(sql);
+                xfree(sql);
                 sql = strdup(optarg);
             }
             break;
         case 'd':
             if (optarg) {
-                free(connopt.database);
+                xfree(connopt.database);
                 connopt.database = strdup(optarg);
             }
             break;
         case 'u':
             if (optarg) {
-                free(connopt.username);
+                xfree(connopt.username);
                 connopt.username = strdup(optarg);
             }
             break;
         case 'p':
             if (optarg) {
-                free(connopt.password);
+                xfree(connopt.password);
                 connopt.password = strdup(optarg);
                 connopt.password_len = strlen(optarg);
             }
@@ -224,11 +224,11 @@ int main(int argc, char *argv[])
     fail_on_error(&conn, err, "closing connection");
     fprintf(stderr, "connection closed\n");
 
-    free(connopt.hostname);
-    free(connopt.username);
-    free(sql);
-    free(connopt.database);
-    free(connopt.password);
+    xfree(connopt.hostname);
+    xfree(connopt.username);
+    xfree(sql);
+    xfree(connopt.database);
+    xfree(connopt.password);
 
     trilogy_free(&conn);
     exit(EXIT_SUCCESS);

inc/trilogy.h

@@ -1,6 +1,7 @@
 #ifndef TRILOGY_H
 #define TRILOGY_H
 
+#include "trilogy/allocator.h"
 #include "trilogy/blocking.h"
 #include "trilogy/client.h"
 #include "trilogy/error.h"

inc/trilogy/allocator.h

@@ -0,0 +1,48 @@
+#ifndef TRILOGY_INTERNAL_ALLOCATOR_H
+#define TRILOGY_INTERNAL_ALLOCATOR_H
+
+/* If you build Trilogy with a custom allocator, configure it with
+ * "-D TRILOGY_XALLOCATOR" to use your own allocator that defines xmalloc,
+ * xrealloc, xcalloc, and xfree.
+ *
+ * For example, your `trilogy_xallocator.h` file could look like this:
+ *
+ * ```
+ * #ifndef TRILOGY_XALLOCATOR_H
+ * #define TRILOGY_XALLOCATOR_H
+ * #define xmalloc          my_malloc
+ * #define xrealloc         my_realloc
+ * #define xcalloc          my_calloc
+ * #define xfree            my_free
+ * #endif
+ * ```
+ */
+#ifdef TRILOGY_XALLOCATOR
+    #include "trilogy_xallocator.h"
+#else
+    #ifndef xmalloc
+        /* The malloc function that should be used. This can be overridden with
+         * the TRILOGY_XALLOCATOR define. */
+        #define xmalloc malloc
+    #endif
+
+    #ifndef xrealloc
+        /* The realloc function that should be used. This can be overridden with
+         * the TRILOGY_XALLOCATOR define. */
+        #define xrealloc realloc
+    #endif
+
+    #ifndef xcalloc
+        /* The calloc function that should be used. This can be overridden with
+         * the TRILOGY_XALLOCATOR define. */
+        #define xcalloc calloc
+    #endif
+
+    #ifndef xfree
+        /* The free function that should be used. This can be overridden with
+         * the TRILOGY_XALLOCATOR define. */
+        #define xfree free
+    #endif
+#endif
+
+#endif

src/buffer.c

@@ -2,14 +2,15 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include "trilogy/allocator.h"
 #include "trilogy/buffer.h"
 #include "trilogy/error.h"
 
 int trilogy_buffer_init(trilogy_buffer_t *buffer, size_t initial_capacity)
 {
     buffer->len = 0;
     buffer->cap = initial_capacity;
-    buffer->buff = malloc(initial_capacity);
+    buffer->buff = xmalloc(initial_capacity);
 
     if (buffer->buff == NULL) {
         return TRILOGY_SYSERR;
@@ -37,7 +38,7 @@ int trilogy_buffer_expand(trilogy_buffer_t *buffer, size_t needed)
             new_cap *= EXPAND_MULTIPLIER;
         }
 
-        uint8_t *new_buff = realloc(buffer->buff, new_cap);
+        uint8_t *new_buff = xrealloc(buffer->buff, new_cap);
         if (new_buff == NULL)
             return TRILOGY_SYSERR;
 
@@ -77,7 +78,7 @@ int trilogy_buffer_write(trilogy_buffer_t *buffer, const uint8_t *ptr, size_t le
 void trilogy_buffer_free(trilogy_buffer_t *buffer)
 {
     if (buffer->buff) {
-        free(buffer->buff);
+        xfree(buffer->buff);
         buffer->buff = NULL;
         buffer->len = buffer->cap = 0;
     }

src/client.c

@@ -7,6 +7,7 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include "trilogy/allocator.h"
 #include "trilogy/client.h"
 #include "trilogy/error.h"
 
@@ -566,7 +567,7 @@ static int encrypt_password_with_public_key(const uint8_t *scramble, size_t scra
         return TRILOGY_MEM_ERROR;
     }
     size_t plaintext_len = password_len + 1;
-    uint8_t *plaintext = malloc(plaintext_len);
+    uint8_t *plaintext = xmalloc(plaintext_len);
 
     if (plaintext == NULL) {
         return TRILOGY_MEM_ERROR;
@@ -585,7 +586,7 @@ static int encrypt_password_with_public_key(const uint8_t *scramble, size_t scra
 
     BIO *bio = BIO_new_mem_buf((void *)key_data, (int)key_data_len);
     if (bio == NULL) {
-        free(plaintext);
+        xfree(plaintext);
         return TRILOGY_OPENSSL_ERR;
     }
 
@@ -600,7 +601,7 @@ static int encrypt_password_with_public_key(const uint8_t *scramble, size_t scra
     if (public_key == NULL) {
         ERR_clear_error();
         memset(plaintext, 0, plaintext_len);
-        free(plaintext);
+        xfree(plaintext);
         return TRILOGY_AUTH_PLUGIN_ERROR;
     }
 
@@ -609,7 +610,7 @@ static int encrypt_password_with_public_key(const uint8_t *scramble, size_t scra
     if (key_size <= 0) {
         EVP_PKEY_free(public_key);
         memset(plaintext, 0, plaintext_len);
-        free(plaintext);
+        xfree(plaintext);
         return TRILOGY_AUTH_PLUGIN_ERROR;
     }
     ciphertext_len = (size_t)key_size;
@@ -628,11 +629,11 @@ static int encrypt_password_with_public_key(const uint8_t *scramble, size_t scra
         RSA_free(public_key);
 #endif
         memset(plaintext, 0, plaintext_len);
-        free(plaintext);
+        xfree(plaintext);
         return TRILOGY_AUTH_PLUGIN_ERROR;
     }
 
-    ciphertext = malloc(ciphertext_len);
+    ciphertext = xmalloc(ciphertext_len);
 
     if (ciphertext == NULL) {
 #if OPENSSL_VERSION_NUMBER >= 0x30000000L
@@ -641,7 +642,7 @@ static int encrypt_password_with_public_key(const uint8_t *scramble, size_t scra
         RSA_free(public_key);
 #endif
         memset(plaintext, 0, plaintext_len);
-        free(plaintext);
+        xfree(plaintext);
         return TRILOGY_MEM_ERROR;
     }
 
@@ -676,13 +677,13 @@ static int encrypt_password_with_public_key(const uint8_t *scramble, size_t scra
 #endif
 
     memset(plaintext, 0, plaintext_len);
-    free(plaintext);
+    xfree(plaintext);
 
     if (rc == TRILOGY_OK) {
         *encrypted_out = ciphertext;
     } else {
         memset(ciphertext, 0, ciphertext_len);
-        free(ciphertext);
+        xfree(ciphertext);
     }
 
     return rc;
@@ -755,7 +756,7 @@ static int handle_fast_auth_fail(trilogy_conn_t *conn, trilogy_handshake_t *hand
 
     rc = send_auth_buffer(conn, encrypted, encrypted_len);
     memset(encrypted, 0, encrypted_len);
-    free(encrypted);
+    xfree(encrypted);
 
     if (rc < 0) {
         return rc;

src/socket.c

@@ -11,6 +11,7 @@
 #include <stdlib.h>
 #include <unistd.h>
 
+#include "trilogy/allocator.h"
 #include "trilogy/error.h"
 #include "trilogy/socket.h"
 
@@ -111,28 +112,28 @@ static int _cb_raw_close(trilogy_sock_t *_sock)
     if (sock->addr) {
         if (sock->base.opts.hostname == NULL && sock->base.opts.path != NULL) {
             /* We created these with calloc so must free them instead of calling freeaddrinfo */
-            free(sock->addr->ai_addr);
-            free(sock->addr);
+            xfree(sock->addr->ai_addr);
+            xfree(sock->addr);
         } else {
             freeaddrinfo(sock->addr);
         }
     }
 
-    free(sock->base.opts.hostname);
-    free(sock->base.opts.path);
-    free(sock->base.opts.database);
-    free(sock->base.opts.username);
-    free(sock->base.opts.password);
-    free(sock->base.opts.ssl_ca);
-    free(sock->base.opts.ssl_capath);
-    free(sock->base.opts.ssl_cert);
-    free(sock->base.opts.ssl_cipher);
-    free(sock->base.opts.ssl_crl);
-    free(sock->base.opts.ssl_crlpath);
-    free(sock->base.opts.ssl_key);
-    free(sock->base.opts.tls_ciphersuites);
+    xfree(sock->base.opts.hostname);
+    xfree(sock->base.opts.path);
+    xfree(sock->base.opts.database);
+    xfree(sock->base.opts.username);
+    xfree(sock->base.opts.password);
+    xfree(sock->base.opts.ssl_ca);
+    xfree(sock->base.opts.ssl_capath);
+    xfree(sock->base.opts.ssl_cert);
+    xfree(sock->base.opts.ssl_cipher);
+    xfree(sock->base.opts.ssl_crl);
+    xfree(sock->base.opts.ssl_crlpath);
+    xfree(sock->base.opts.ssl_key);
+    xfree(sock->base.opts.tls_ciphersuites);
 
-    free(sock);
+    xfree(sock);
     return rc;
 }
 
@@ -311,7 +312,7 @@ static char *strdupnullok(const char *str)
 
 trilogy_sock_t *trilogy_sock_new(const trilogy_sockopt_t *opts)
 {
-    struct trilogy_sock *sock = malloc(sizeof(struct trilogy_sock));
+    struct trilogy_sock *sock = xmalloc(sizeof(struct trilogy_sock));
 
     sock->base.connect_cb = _cb_raw_connect;
     sock->base.read_cb = _cb_raw_read;
@@ -328,7 +329,7 @@ trilogy_sock_t *trilogy_sock_new(const trilogy_sockopt_t *opts)
     sock->base.opts.username = strdupnullok(opts->username);
 
     if (sock->base.opts.password) {
-        sock->base.opts.password = malloc(opts->password_len);
+        sock->base.opts.password = xmalloc(opts->password_len);
         memcpy(sock->base.opts.password, opts->password, opts->password_len);
     }
 

test/builder_test.c

@@ -60,7 +60,7 @@ TEST test_builder_write_uint8_split_packet()
     const uint8_t expected_header2[] = {0x01, 0x00, 0x00, 0x01};
     ASSERT_MEM_EQ(builder.buffer->buff + TRILOGY_MAX_PACKET_LEN + 4, expected_header2, sizeof(expected_header2));
 
-    free(bytes);
+    xfree(bytes);
 
     trilogy_buffer_free(&buff);
     PASS();
@@ -123,7 +123,7 @@ TEST test_builder_write_uint8_exceeds_large_max()
     err = trilogy_builder_write_uint8(&builder, 0x03);
     ASSERT_EQ(TRILOGY_MAX_PACKET_EXCEEDED, err);
 
-    free(bytes);
+    xfree(bytes);
 
     trilogy_buffer_free(&buff);
     PASS();
@@ -331,7 +331,7 @@ TEST test_builder_write_large_buffer()
     const uint8_t expected_header2[] = {0x0A, 0x00, 0x00, 0x01};
     ASSERT_MEM_EQ(builder.buffer->buff + TRILOGY_MAX_PACKET_LEN + 4, expected_header2, sizeof(expected_header2));
 
-    free(bytes);
+    xfree(bytes);
 
     trilogy_buffer_free(&buff);
     PASS();