Fix use-after-close race causing SEGV

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

@@ -119,7 +119,8 @@ static inline buffer_pool *get_buffer_pool(void)
 static void buffer_checkout(trilogy_buffer_t *buffer, size_t initial_capacity)
 {
     buffer_pool * pool = get_buffer_pool();
-    if (pool->len) {
+
+    if (pool && pool->len) {
         pool->len--;
         buffer->buff = pool->entries[pool->len].buff;
         buffer->cap = pool->entries[pool->len].cap;
@@ -133,6 +134,13 @@ static bool buffer_checkin(trilogy_buffer_t *buffer)
 {
     buffer_pool * pool = get_buffer_pool();
 
+    if (pool == NULL) {
+        xfree(buffer->buff);
+        buffer->buff = NULL;
+        buffer->cap = 0;
+        return false;
+    }
+
     if (pool->len >= BUFFER_POOL_MAX_SIZE) {
         xfree(buffer->buff);
         buffer->buff = NULL;
@@ -348,6 +356,10 @@ static VALUE allocate_trilogy(VALUE klass)
 
 static int flush_writes(struct trilogy_ctx *ctx)
 {
+    if (ctx->conn.socket == NULL) {
+        return TRILOGY_CLOSED_CONNECTION;
+    }
+
     while (1) {
         int rc = trilogy_flush_writes(&ctx->conn);
 
@@ -931,6 +943,10 @@ static VALUE read_query_response(VALUE vargs)
         if (rc != TRILOGY_OK) {
             handle_trilogy_error(ctx, rc, "trilogy_query_recv");
         }
+
+        if (ctx->conn.socket == NULL) {
+            rb_raise(Trilogy_ConnectionClosedError, "Connection closed while waiting for response");
+        }
     }
 
     struct timespec finish;

src/client.c

@@ -11,6 +11,10 @@
 
 static inline TRILOGY_PACKET_TYPE_t current_packet_type(trilogy_conn_t *conn)
 {
+    if (conn->packet_buffer.buff == NULL || conn->packet_buffer.len == 0) {
+        return TRILOGY_PACKET_UNKNOWN;
+    }
+
     return (TRILOGY_PACKET_TYPE_t)conn->packet_buffer.buff[0];
 }
 
@@ -71,6 +75,10 @@ static int begin_command_phase(trilogy_builder_t *builder, trilogy_conn_t *conn,
 
 static int read_packet(trilogy_conn_t *conn)
 {
+    if (conn->socket == NULL || conn->packet_buffer.buff == NULL) {
+        return TRILOGY_CLOSED_CONNECTION;
+    }
+
     if (conn->recv_buff_pos == conn->recv_buff_len) {
         ssize_t nread = trilogy_sock_read(conn->socket, conn->recv_buff, sizeof(conn->recv_buff));
 
@@ -154,6 +162,10 @@ int trilogy_init(trilogy_conn_t *conn)
 
 int trilogy_flush_writes(trilogy_conn_t *conn)
 {
+    if (conn->socket == NULL) {
+        return TRILOGY_CLOSED_CONNECTION;
+    }
+
     void *ptr = conn->packet_buffer.buff + conn->packet_buffer_written;
     size_t len = conn->packet_buffer.len - conn->packet_buffer_written;