diff --git a/contrib/ruby/ext/trilogy-ruby/cext.c b/contrib/ruby/ext/trilogy-ruby/cext.c index 83f095c7..f0ce7e1c 100644 --- a/contrib/ruby/ext/trilogy-ruby/cext.c +++ b/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; diff --git a/src/client.c b/src/client.c index 344574ca..9084a010 100644 --- a/src/client.c +++ b/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;