trilogy 2.10.0 → 2.11.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 6e84050730e47a4bd826b9c15b226d0f825a46df8849c3c4a89c01011e27fb6c
-  data.tar.gz: eddfedb0481b09220bb95f1de423b473aa78f2ad8058de692b85c52bba1c0480
+  metadata.gz: 0625a7ff415d68c8f5e421188e66fe7c41ebadd2afc39a7d6b89c6318ea9ef9d
+  data.tar.gz: c6b10922d195af6ec3b5140963749b59d4ef4a967728aeeb826377bc66992bc0
 SHA512:
-  metadata.gz: a1fa8ecb46aa1b1595c6d6f1a20b5dd4267452d5c59b0571fe31c4cc854907aeb8f8ad1496099e53ba83b29c0594f88dc6f4f45586ede8b1ad53d70b0ff2a774
-  data.tar.gz: a18b393f8b035c8ee3315d42b23d8a46e1597f4ccbac0347acab0161a8c041572eeb8a13d4d323c97a63040f9f18baa04e0742c73e65df33d5763fdd27b6c581
+  metadata.gz: 5da895e095a69094c7f1844d341ea2bf0c4461f7744f892a5a6f3aff1552761ebb8f099a5df151450c6c32e894e88cfebf72052c8190aef749be0d8d56ad0611
+  data.tar.gz: 60830ae194994fa09a975d3cea8146632300bfef62e98726c7cd6edfdb18a91bd3bce55a9fbf0d5b5d2fa0932d2a264333d87ccb57ae222e59a46bd3c5740fd7

data/Rakefile

@@ -19,10 +19,25 @@ Rake::TestTask.new do |t|
   t.test_files = FileList['test/*_test.rb']
   t.verbose = true
 end
-task :test => :compile
+
+task :test => [:compile, "db:clean"]
 
 task :default => :test
 
 task :console => :compile do
   sh "ruby -I lib -r trilogy -S irb"
 end
+
+namespace :db do
+  task :create do
+    mysql_command = "mysql -uroot -e"
+    %x( #{mysql_command} "create DATABASE IF NOT EXISTS test DEFAULT CHARACTER SET utf8mb4" )
+  end
+
+  task :drop do
+    mysql_command = "mysql -uroot -e"
+    %x( #{mysql_command} "drop DATABASE IF EXISTS test" )
+  end
+
+  task :clean => ["db:drop", "db:create"]
+end 

data/ext/trilogy-ruby/cext.c

@@ -124,7 +124,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 = RB_ALLOC_N(uint8_t, initial_capacity);
+        buffer->buff = malloc(initial_capacity);
         buffer->cap = initial_capacity;
     }
 }
@@ -134,7 +134,7 @@ static bool buffer_checkin(trilogy_buffer_t *buffer)
     buffer_pool * pool = get_buffer_pool();
 
     if (pool->len >= BUFFER_POOL_MAX_SIZE) {
-        xfree(buffer->buff);
+        free(buffer->buff);
         buffer->buff = NULL;
         buffer->cap = 0;
         return false;
@@ -172,9 +172,9 @@ static ID id_socket, id_host, id_port, id_username, id_password, id_found_rows,
 
 struct trilogy_ctx {
     trilogy_conn_t conn;
-    char server_version[TRILOGY_SERVER_VERSION_SIZE + 1];
+    rb_encoding *encoding;
     unsigned int query_flags;
-    VALUE encoding;
+    char server_version[TRILOGY_SERVER_VERSION_SIZE + 1];
 };
 
 static void rb_trilogy_acquire_buffer(struct trilogy_ctx *ctx)
@@ -191,12 +191,6 @@ static void rb_trilogy_release_buffer(struct trilogy_ctx *ctx)
     }
 }
 
-static void mark_trilogy(void *ptr)
-{
-    struct trilogy_ctx *ctx = ptr;
-    rb_gc_mark(ctx->encoding);
-}
-
 static void free_trilogy(void *ptr)
 {
     struct trilogy_ctx *ctx = ptr;
@@ -218,7 +212,7 @@ static size_t trilogy_memsize(const void *ptr) {
 static const rb_data_type_t trilogy_data_type = {
     .wrap_struct_name = "trilogy",
     .function = {
-        .dmark = mark_trilogy,
+        .dmark = NULL,
         .dfree = free_trilogy,
         .dsize = trilogy_memsize,
     },
@@ -455,42 +449,29 @@ static int _cb_ruby_wait(trilogy_sock_t *sock, trilogy_wait_t wait)
     return TRILOGY_OK;
 }
 
-struct nogvl_sock_args {
-    int rc;
-    trilogy_sock_t *sock;
-};
-
-static void *no_gvl_resolve(void *data)
+static int try_connect(struct trilogy_ctx *ctx, trilogy_handshake_t *handshake, const trilogy_sockopt_t *opts, int fd)
 {
-    struct nogvl_sock_args *args = data;
-    args->rc = trilogy_sock_resolve(args->sock);
-    return NULL;
-}
+    if (fd < 0) {
+        return TRILOGY_ERR;
+    }
 
-static int try_connect(struct trilogy_ctx *ctx, trilogy_handshake_t *handshake, const trilogy_sockopt_t *opts)
-{
     trilogy_sock_t *sock = trilogy_sock_new(opts);
     if (sock == NULL) {
         return TRILOGY_ERR;
     }
 
-    struct nogvl_sock_args args = {.rc = 0, .sock = sock};
-
-    // Do the DNS resolving with the GVL unlocked. At this point all
-    // configuration data is copied and available to the trilogy socket.
-    rb_thread_call_without_gvl(no_gvl_resolve, (void *)&args, RUBY_UBF_IO, NULL);
-
-    int rc = args.rc;
-
-    if (rc != TRILOGY_OK) {
-        trilogy_sock_close(sock);
-        return rc;
-    }
+    int rc;
 
     /* replace the default wait callback with our GVL-aware callback so we can
 escape the GVL on each wait operation without going through call_without_gvl */
     sock->wait_cb = _cb_ruby_wait;
-    rc = trilogy_connect_send_socket(&ctx->conn, sock);
+
+    int newfd = dup(fd);
+    if (newfd < 0) {
+        return TRILOGY_ERR;
+    }
+
+    rc = trilogy_connect_set_fd(&ctx->conn, sock, newfd);
     if (rc < 0) {
         trilogy_sock_close(sock);
         return rc;
@@ -534,7 +515,12 @@ static void auth_switch(struct trilogy_ctx *ctx, trilogy_handshake_t *handshake)
         }
 
         if (rc != TRILOGY_AGAIN) {
-            handle_trilogy_error(ctx, rc, "trilogy_auth_recv");
+            if (rc == TRILOGY_UNSUPPORTED) {
+                handle_trilogy_error(ctx, rc, "trilogy_auth_recv: caching_sha2_password requires either TCP with TLS or a unix socket");
+            }
+            else {
+                handle_trilogy_error(ctx, rc, "trilogy_auth_recv");
+            }
         }
 
         rc = trilogy_sock_wait_read(ctx->conn.socket);
@@ -588,12 +574,7 @@ static void authenticate(struct trilogy_ctx *ctx, trilogy_handshake_t *handshake
         }
 
         if (rc != TRILOGY_AGAIN) {
-            if (rc == TRILOGY_UNSUPPORTED) {
-                handle_trilogy_error(ctx, rc, "trilogy_auth_recv: caching_sha2_password requires either TCP with TLS or a unix socket");
-            }
-            else {
-                handle_trilogy_error(ctx, rc, "trilogy_auth_recv");
-            }
+            handle_trilogy_error(ctx, rc, "trilogy_auth_recv");
         }
 
         rc = trilogy_sock_wait_read(ctx->conn.socket);
@@ -607,14 +588,24 @@ static void authenticate(struct trilogy_ctx *ctx, trilogy_handshake_t *handshake
     }
 }
 
-static VALUE rb_trilogy_connect(VALUE self, VALUE encoding, VALUE charset, VALUE opts)
+#ifndef HAVE_RB_IO_DESCRIPTOR /* Ruby < 3.1 */
+static int rb_io_descriptor(VALUE io)
+{
+    rb_io_t *fptr;
+    GetOpenFile(io, fptr);
+    rb_io_check_closed(fptr);
+    return fptr->fd;
+}
+#endif
+
+static VALUE rb_trilogy_connect(VALUE self, VALUE raw_socket, VALUE encoding, VALUE charset, VALUE opts)
 {
     struct trilogy_ctx *ctx = get_ctx(self);
     trilogy_sockopt_t connopt = {0};
     trilogy_handshake_t handshake;
     VALUE val;
 
-    RB_OBJ_WRITE(self, &ctx->encoding, encoding);
+    ctx->encoding = rb_to_encoding(encoding);
     connopt.encoding = NUM2INT(charset);
 
     Check_Type(opts, T_HASH);
@@ -762,9 +753,17 @@ static VALUE rb_trilogy_connect(VALUE self, VALUE encoding, VALUE charset, VALUE
         connopt.tls_max_version = NUM2INT(val);
     }
 
-    rb_trilogy_acquire_buffer(ctx);
+    VALUE io = rb_io_get_io(raw_socket);
+
+    rb_io_t *fptr;
+    GetOpenFile(io, fptr);
+    rb_io_check_readable(fptr);
+    rb_io_check_writable(fptr);
 
-    int rc = try_connect(ctx, &handshake, &connopt);
+    int fd = rb_io_descriptor(io);
+
+    rb_trilogy_acquire_buffer(ctx);
+    int rc = try_connect(ctx, &handshake, &connopt, fd);
     if (rc != TRILOGY_OK) {
         if (connopt.path) {
             handle_trilogy_error(ctx, rc, "trilogy_connect - unable to connect to %s", connopt.path);
@@ -983,10 +982,10 @@ static VALUE read_query_response(VALUE vargs)
             }
         }
 
-#ifdef HAVE_RB_INTERNED_STR
-        VALUE column_name = rb_interned_str(column.name, column.name_len);
+#ifdef HAVE_RB_ENC_INTERNED_STR
+        VALUE column_name = rb_enc_interned_str(column.name, column.name_len, ctx->encoding);
 #else
-        VALUE column_name = rb_str_new(column.name, column.name_len);
+        VALUE column_name = rb_enc_str_new(column.name, column.name_len, ctx->encoding);
         OBJ_FREEZE(column_name);
 #endif
 
@@ -1093,12 +1092,37 @@ static VALUE rb_trilogy_more_results_exist(VALUE self)
     }
 }
 
+static VALUE rb_trilogy_abandon_results(VALUE self)
+{
+    struct trilogy_ctx *ctx = get_open_ctx(self);
+
+    long count = 0;
+    while (ctx->conn.server_status & TRILOGY_SERVER_STATUS_MORE_RESULTS_EXISTS) {
+        count++;
+        int rc = trilogy_drain_results(&ctx->conn);
+        while (rc == TRILOGY_AGAIN) {
+            rc = trilogy_sock_wait_read(ctx->conn.socket);
+            if (rc != TRILOGY_OK) {
+                handle_trilogy_error(ctx, rc, "trilogy_sock_wait_read");
+            }
+
+            rc = trilogy_drain_results(&ctx->conn);
+        }
+
+        if (rc != TRILOGY_OK) {
+            handle_trilogy_error(ctx, rc, "trilogy_drain_results");
+        }
+    }
+
+    return LONG2NUM(count);
+}
+
 static VALUE rb_trilogy_query(VALUE self, VALUE query)
 {
     struct trilogy_ctx *ctx = get_open_ctx(self);
 
     StringValue(query);
-    query = rb_str_export_to_enc(query, rb_to_encoding(ctx->encoding));
+    query = rb_str_export_to_enc(query, ctx->encoding);
 
     rb_trilogy_acquire_buffer(ctx);
 
@@ -1333,7 +1357,7 @@ RUBY_FUNC_EXPORTED void Init_cext(void)
     VALUE Trilogy = rb_const_get(rb_cObject, rb_intern("Trilogy"));
     rb_define_alloc_func(Trilogy, allocate_trilogy);
 
-    rb_define_private_method(Trilogy, "_connect", rb_trilogy_connect, 3);
+    rb_define_private_method(Trilogy, "_connect", rb_trilogy_connect, 4);
     rb_define_method(Trilogy, "change_db", rb_trilogy_change_db, 1);
     rb_define_alias(Trilogy, "select_db", "change_db");
     rb_define_method(Trilogy, "query", rb_trilogy_query, 1);
@@ -1357,6 +1381,7 @@ RUBY_FUNC_EXPORTED void Init_cext(void)
     rb_define_method(Trilogy, "server_version", rb_trilogy_server_version, 0);
     rb_define_method(Trilogy, "more_results_exist?", rb_trilogy_more_results_exist, 0);
     rb_define_method(Trilogy, "next_result", rb_trilogy_next_result, 0);
+    rb_define_method(Trilogy, "abandon_results!", rb_trilogy_abandon_results, 0);
     rb_define_method(Trilogy, "set_server_option", rb_trilogy_set_server_option, 1);
     rb_define_const(Trilogy, "TLS_VERSION_10", INT2NUM(TRILOGY_TLS_VERSION_10));
     rb_define_const(Trilogy, "TLS_VERSION_11", INT2NUM(TRILOGY_TLS_VERSION_11));

data/ext/trilogy-ruby/extconf.rb

@@ -12,10 +12,12 @@ File.binwrite("trilogy.c",
 $objs = %w[trilogy.o cast.o cext.o]
 append_cflags(["-I #{__dir__}/inc", "-std=gnu99", "-fvisibility=hidden"])
 
-dir_config("openssl")
+dir_config("openssl").any? || pkg_config("openssl")
 
 have_library("crypto", "CRYPTO_malloc")
 have_library("ssl", "SSL_new")
-have_func("rb_interned_str", "ruby.h")
 have_func("rb_ractor_local_storage_value_newkey", "ruby.h")
+have_func("rb_enc_interned_str", "ruby.h")
+have_func("rb_io_descriptor", "ruby.h") # Ruby 3.1+
+
 create_makefile "trilogy/cext"

data/ext/trilogy-ruby/inc/trilogy/client.h

@@ -159,6 +159,8 @@ int trilogy_connect_send(trilogy_conn_t *conn, const trilogy_sockopt_t *opts);
  */
 int trilogy_connect_send_socket(trilogy_conn_t *conn, trilogy_sock_t *sock);
 
+int trilogy_connect_set_fd(trilogy_conn_t *conn, trilogy_sock_t *sock, int fd);
+
 /* trilogy_connect_recv - Read the initial handshake from the server.
  *
  * This should be called after trilogy_connect_send returns TRILOGY_OK. Calling

data/ext/trilogy-ruby/inc/trilogy/socket.h

@@ -88,6 +88,8 @@ typedef struct trilogy_sock_t {
 
 static inline int trilogy_sock_connect(trilogy_sock_t *sock) { return sock->connect_cb(sock); }
 
+void trilogy_sock_set_fd(trilogy_sock_t *sock, int fd);
+
 static inline ssize_t trilogy_sock_read(trilogy_sock_t *sock, void *buf, size_t n)
 {
     return sock->read_cb(sock, buf, n);

data/ext/trilogy-ruby/src/client.c

@@ -1,4 +1,10 @@
 #include <fcntl.h>
+#include <limits.h>
+#include <openssl/err.h>
+#include <openssl/evp.h>
+#include <openssl/pem.h>
+#include <openssl/rsa.h>
+#include <stdlib.h>
 #include <string.h>
 
 #include "trilogy/client.h"
@@ -118,6 +124,38 @@ static int begin_write(trilogy_conn_t *conn)
     return trilogy_flush_writes(conn);
 }
 
+static int flush_current_packet(trilogy_conn_t *conn)
+{
+    int rc = begin_write(conn);
+
+    while (rc == TRILOGY_AGAIN) {
+        rc = trilogy_sock_wait_write(conn->socket);
+
+        if (rc != TRILOGY_OK) {
+            return rc;
+        }
+
+        rc = trilogy_flush_writes(conn);
+    }
+
+    return rc;
+}
+
+static int read_packet_blocking(trilogy_conn_t *conn)
+{
+    int rc;
+
+    while ((rc = read_packet(conn)) == TRILOGY_AGAIN) {
+        rc = trilogy_sock_wait_read(conn->socket);
+
+        if (rc != TRILOGY_OK) {
+            return rc;
+        }
+    }
+
+    return rc;
+}
+
 int trilogy_init_no_buffer(trilogy_conn_t *conn)
 {
     conn->affected_rows = 0;
@@ -279,7 +317,8 @@ static int read_auth_switch_packet(trilogy_conn_t *conn, trilogy_handshake_t *ha
     return TRILOGY_AUTH_SWITCH;
 }
 
-static int handle_generic_response(trilogy_conn_t *conn) {
+static int handle_generic_response(trilogy_conn_t *conn)
+{
     switch (current_packet_type(conn)) {
     case TRILOGY_PACKET_OK:
         return read_ok_packet(conn);
@@ -330,6 +369,16 @@ int trilogy_connect_send_socket(trilogy_conn_t *conn, trilogy_sock_t *sock)
     return TRILOGY_OK;
 }
 
+int trilogy_connect_set_fd(trilogy_conn_t *conn, trilogy_sock_t *sock, int fd)
+{
+    trilogy_sock_set_fd(sock, fd);
+
+    conn->socket = sock;
+    conn->packet_parser.sequence_number = 0;
+
+    return TRILOGY_OK;
+}
+
 int trilogy_connect_recv(trilogy_conn_t *conn, trilogy_handshake_t *handshake_out)
 {
     int rc = read_packet(conn);
@@ -427,9 +476,294 @@ void trilogy_auth_clear_password(trilogy_conn_t *conn)
     }
 }
 
+#define CACHING_SHA2_REQUEST_PUBLIC_KEY 2
+#define CACHING_SHA2_SCRAMBLE_LEN 20
 #define FAST_AUTH_OK 3
 #define FAST_AUTH_FAIL 4
 
+static int read_auth_result(trilogy_conn_t *conn)
+{
+    int rc = read_packet_blocking(conn);
+
+    if (rc < 0) {
+        return rc;
+    }
+
+    trilogy_auth_clear_password(conn);
+    return handle_generic_response(conn);
+}
+
+static int send_cleartext_password(trilogy_conn_t *conn)
+{
+    trilogy_builder_t builder;
+
+    int rc = begin_command_phase(&builder, conn, conn->packet_parser.sequence_number);
+
+    if (rc < 0) {
+        return rc;
+    }
+
+    if (conn->socket->opts.password_len == 0) {
+        rc = trilogy_builder_write_uint8(&builder, 0);
+
+        if (rc < 0) {
+            return rc;
+        }
+
+        trilogy_builder_finalize(&builder);
+        return flush_current_packet(conn);
+    }
+
+    rc = trilogy_build_auth_clear_password(&builder, conn->socket->opts.password, conn->socket->opts.password_len);
+
+    if (rc < 0) {
+        return rc;
+    }
+
+    return flush_current_packet(conn);
+}
+
+static int send_auth_buffer(trilogy_conn_t *conn, const void *buff, size_t buff_len)
+{
+    trilogy_builder_t builder;
+    int rc = begin_command_phase(&builder, conn, conn->packet_parser.sequence_number);
+
+    if (rc < 0) {
+        return rc;
+    }
+
+    rc = trilogy_builder_write_buffer(&builder, buff, buff_len);
+    if (rc < 0) {
+        return rc;
+    }
+
+    trilogy_builder_finalize(&builder);
+
+    return flush_current_packet(conn);
+}
+
+static int send_public_key_request(trilogy_conn_t *conn)
+{
+    uint8_t request = CACHING_SHA2_REQUEST_PUBLIC_KEY;
+
+    return send_auth_buffer(conn, &request, sizeof(request));
+}
+
+static int encrypt_password_with_public_key(const uint8_t *scramble, size_t scramble_len, trilogy_conn_t *conn,
+                                            const uint8_t *key_data, size_t key_data_len, uint8_t **encrypted_out,
+                                            size_t *encrypted_len)
+{
+    int rc = TRILOGY_OK;
+    uint8_t *ciphertext = NULL;
+    size_t ciphertext_len = 0;
+
+    if (key_data_len == 0 || key_data_len > INT_MAX) {
+        return TRILOGY_AUTH_PLUGIN_ERROR;
+    }
+
+    size_t password_len = conn->socket->opts.password_len;
+    if (password_len == SIZE_MAX) {
+        return TRILOGY_MEM_ERROR;
+    }
+    size_t plaintext_len = password_len + 1;
+    uint8_t *plaintext = malloc(plaintext_len);
+
+    if (plaintext == NULL) {
+        return TRILOGY_MEM_ERROR;
+    }
+
+    if (password_len > 0) {
+        memcpy(plaintext, conn->socket->opts.password, password_len);
+    }
+    plaintext[plaintext_len - 1] = '\0';
+
+    if (scramble_len > 0) {
+        for (size_t i = 0; i < plaintext_len; i++) {
+            plaintext[i] ^= scramble[i % scramble_len];
+        }
+    }
+
+    BIO *bio = BIO_new_mem_buf((void *)key_data, (int)key_data_len);
+    if (bio == NULL) {
+        free(plaintext);
+        return TRILOGY_OPENSSL_ERR;
+    }
+
+#if OPENSSL_VERSION_NUMBER >= 0x30000000L
+    EVP_PKEY *public_key = PEM_read_bio_PUBKEY(bio, NULL, NULL, NULL);
+#else
+    RSA *public_key = PEM_read_bio_RSA_PUBKEY(bio, NULL, NULL, NULL);
+#endif
+
+    BIO_free(bio);
+
+    if (public_key == NULL) {
+        ERR_clear_error();
+        memset(plaintext, 0, plaintext_len);
+        free(plaintext);
+        return TRILOGY_AUTH_PLUGIN_ERROR;
+    }
+
+#if OPENSSL_VERSION_NUMBER >= 0x30000000L
+    int key_size = EVP_PKEY_get_size(public_key);
+    if (key_size <= 0) {
+        EVP_PKEY_free(public_key);
+        memset(plaintext, 0, plaintext_len);
+        free(plaintext);
+        return TRILOGY_AUTH_PLUGIN_ERROR;
+    }
+    ciphertext_len = (size_t)key_size;
+#else
+    ciphertext_len = (size_t)RSA_size(public_key);
+#endif
+
+    /*
+       When using RSA_PKCS1_OAEP_PADDING the password length must be less
+       than RSA_size(rsa) - 41.
+     */
+    if (ciphertext_len == 0 || plaintext_len + 41 >= ciphertext_len) {
+#if OPENSSL_VERSION_NUMBER >= 0x30000000L
+        EVP_PKEY_free(public_key);
+#else
+        RSA_free(public_key);
+#endif
+        memset(plaintext, 0, plaintext_len);
+        free(plaintext);
+        return TRILOGY_AUTH_PLUGIN_ERROR;
+    }
+
+    ciphertext = malloc(ciphertext_len);
+
+    if (ciphertext == NULL) {
+#if OPENSSL_VERSION_NUMBER >= 0x30000000L
+        EVP_PKEY_free(public_key);
+#else
+        RSA_free(public_key);
+#endif
+        memset(plaintext, 0, plaintext_len);
+        free(plaintext);
+        return TRILOGY_MEM_ERROR;
+    }
+
+#if OPENSSL_VERSION_NUMBER >= 0x30000000L
+    EVP_PKEY_CTX *ctx = EVP_PKEY_CTX_new(public_key, NULL);
+    if (ctx == NULL || EVP_PKEY_encrypt_init(ctx) <= 0 ||
+        EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_PKCS1_OAEP_PADDING) <= 0) {
+        rc = TRILOGY_OPENSSL_ERR;
+    } else {
+        size_t out_len = ciphertext_len;
+
+        if (EVP_PKEY_encrypt(ctx, ciphertext, &out_len, plaintext, plaintext_len) <= 0) {
+            rc = TRILOGY_OPENSSL_ERR;
+        } else {
+            *encrypted_len = out_len;
+        }
+    }
+
+    if (ctx) {
+        EVP_PKEY_CTX_free(ctx);
+    }
+    EVP_PKEY_free(public_key);
+#else
+    int out_len = RSA_public_encrypt((int)plaintext_len, plaintext, ciphertext, public_key, RSA_PKCS1_OAEP_PADDING);
+    RSA_free(public_key);
+
+    if (out_len < 0) {
+        rc = TRILOGY_OPENSSL_ERR;
+    } else {
+        *encrypted_len = (size_t)out_len;
+    }
+#endif
+
+    memset(plaintext, 0, plaintext_len);
+    free(plaintext);
+
+    if (rc == TRILOGY_OK) {
+        *encrypted_out = ciphertext;
+    } else {
+        memset(ciphertext, 0, ciphertext_len);
+        free(ciphertext);
+    }
+
+    return rc;
+}
+
+static int handle_fast_auth_fail(trilogy_conn_t *conn, trilogy_handshake_t *handshake, const uint8_t *auth_data,
+                                 size_t auth_data_len)
+{
+    int rc;
+    bool use_ssl = (conn->socket->opts.flags & TRILOGY_CAPABILITIES_SSL) != 0;
+    bool has_unix_socket = (conn->socket->opts.path != NULL);
+
+    // No password to send, so we can safely respond even without TLS.
+    if (conn->socket->opts.password_len == 0) {
+        rc = send_cleartext_password(conn);
+        if (rc < 0) {
+            return rc;
+        }
+
+        return read_auth_result(conn);
+    }
+
+    if (use_ssl || has_unix_socket) {
+        rc = send_cleartext_password(conn);
+        if (rc < 0) {
+            return rc;
+        }
+
+        return read_auth_result(conn);
+    }
+
+    const uint8_t *public_key_data = NULL;
+    size_t public_key_len = 0;
+
+    if (auth_data_len > 1) {
+        public_key_data = auth_data + 1;
+        public_key_len = auth_data_len - 1;
+    } else {
+        rc = send_public_key_request(conn);
+        if (rc < 0) {
+            return rc;
+        }
+
+        rc = read_packet_blocking(conn);
+        if (rc < 0) {
+            return rc;
+        }
+
+        if (current_packet_type(conn) == TRILOGY_PACKET_ERR) {
+            return read_err_packet(conn);
+        }
+
+        if (current_packet_type(conn) != TRILOGY_PACKET_AUTH_MORE_DATA || conn->packet_buffer.len < 2) {
+            return TRILOGY_PROTOCOL_VIOLATION;
+        }
+
+        public_key_data = conn->packet_buffer.buff + 1;
+        public_key_len = conn->packet_buffer.len - 1;
+    }
+
+    uint8_t *encrypted = NULL;
+    size_t encrypted_len = 0;
+
+    rc = encrypt_password_with_public_key((const uint8_t *)handshake->scramble, CACHING_SHA2_SCRAMBLE_LEN, conn,
+                                          public_key_data, public_key_len, &encrypted, &encrypted_len);
+
+    if (rc < 0) {
+        return rc;
+    }
+
+    rc = send_auth_buffer(conn, encrypted, encrypted_len);
+    memset(encrypted, 0, encrypted_len);
+    free(encrypted);
+
+    if (rc < 0) {
+        return rc;
+    }
+
+    return read_auth_result(conn);
+}
+
 int trilogy_auth_recv(trilogy_conn_t *conn, trilogy_handshake_t *handshake)
 {
     int rc = read_packet(conn);
@@ -440,67 +774,24 @@ int trilogy_auth_recv(trilogy_conn_t *conn, trilogy_handshake_t *handshake)
 
     switch (current_packet_type(conn)) {
     case TRILOGY_PACKET_AUTH_MORE_DATA: {
-        bool use_ssl = (conn->socket->opts.flags & TRILOGY_CAPABILITIES_SSL) != 0;
-        bool has_unix_socket = (conn->socket->opts.path != NULL);
+        const uint8_t *auth_data = conn->packet_buffer.buff + 1;
+        size_t auth_data_len = conn->packet_buffer.len - 1;
 
-        if (!use_ssl && !has_unix_socket) {
-            return TRILOGY_UNSUPPORTED;
+        if (auth_data_len < 1) {
+            return TRILOGY_PROTOCOL_VIOLATION;
         }
 
-        uint8_t byte = conn->packet_buffer.buff[1];
+        uint8_t byte = auth_data[0];
         switch (byte) {
-            case FAST_AUTH_OK:
-                break;
-            case FAST_AUTH_FAIL:
-                {
-                    trilogy_builder_t builder;
-                    int err = begin_command_phase(&builder, conn, conn->packet_parser.sequence_number);
-
-                    if (err < 0) {
-                        return err;
-                    }
-
-                    err = trilogy_build_auth_clear_password(&builder, conn->socket->opts.password, conn->socket->opts.password_len);
-
-                    if (err < 0) {
-                        return err;
-                    }
-
-                    int rc = begin_write(conn);
-
-                    while (rc == TRILOGY_AGAIN) {
-                        rc = trilogy_sock_wait_write(conn->socket);
-                        if (rc != TRILOGY_OK) {
-                            return rc;
-                        }
-
-                        rc = trilogy_flush_writes(conn);
-                    }
-                    if (rc != TRILOGY_OK) {
-                        return rc;
-                    }
-
-                    break;
-                }
-            default:
-                return TRILOGY_UNEXPECTED_PACKET;
-        }
-        while (1) {
-            rc = read_packet(conn);
+        case FAST_AUTH_OK:
+            return read_auth_result(conn);
 
-            if (rc == TRILOGY_OK) {
-                break;
-            }
-            else if (rc == TRILOGY_AGAIN) {
-                rc = trilogy_sock_wait_read(conn->socket);
-            }
+        case FAST_AUTH_FAIL:
+            return handle_fast_auth_fail(conn, handshake, auth_data, auth_data_len);
 
-            if (rc != TRILOGY_OK) {
-                return rc;
-            }
+        default:
+            return TRILOGY_UNEXPECTED_PACKET;
         }
-        trilogy_auth_clear_password(conn);
-        return handle_generic_response(conn);
     }
 
     case TRILOGY_PACKET_EOF:
@@ -555,7 +846,8 @@ int trilogy_set_option_send(trilogy_conn_t *conn, const uint16_t option)
     return begin_write(conn);
 }
 
-int trilogy_set_option_recv(trilogy_conn_t *conn) {
+int trilogy_set_option_recv(trilogy_conn_t *conn)
+{
     int rc = read_packet(conn);
 
     if (rc < 0) {
@@ -575,7 +867,6 @@ int trilogy_set_option_recv(trilogy_conn_t *conn) {
     }
 }
 
-
 int trilogy_ping_send(trilogy_conn_t *conn)
 {
     trilogy_builder_t builder;
@@ -727,6 +1018,7 @@ int trilogy_drain_results(trilogy_conn_t *conn)
         }
 
         if (current_packet_type(conn) == TRILOGY_PACKET_EOF && conn->packet_buffer.len < 9) {
+            read_eof_packet(conn);
             return TRILOGY_OK;
         }
     }
@@ -998,9 +1290,7 @@ int trilogy_stmt_reset_send(trilogy_conn_t *conn, trilogy_stmt_t *stmt)
     return begin_write(conn);
 }
 
-int trilogy_stmt_reset_recv(trilogy_conn_t *conn) {
-    return read_generic_response(conn);
-}
+int trilogy_stmt_reset_recv(trilogy_conn_t *conn) { return read_generic_response(conn); }
 
 int trilogy_stmt_close_send(trilogy_conn_t *conn, trilogy_stmt_t *stmt)
 {

data/ext/trilogy-ruby/src/packet_parser.c

@@ -74,7 +74,7 @@ size_t trilogy_packet_parser_execute(trilogy_packet_parser_t *parser, const uint
             break;
         }
         case S_SEQ: {
-            if (cur_byte != parser->sequence_number) {
+            if (cur_byte != parser->sequence_number && cur_byte > 0) {
                 *error = TRILOGY_INVALID_SEQUENCE_ID;
                 return i;
             }

data/ext/trilogy-ruby/src/socket.c

@@ -25,6 +25,12 @@ struct trilogy_sock {
     SSL *ssl;
 };
 
+void trilogy_sock_set_fd(trilogy_sock_t *_sock, int fd)
+{
+    struct trilogy_sock *sock = (struct trilogy_sock *)_sock;
+    sock->fd = fd;
+}
+
 static int _cb_raw_fd(trilogy_sock_t *_sock)
 {
     struct trilogy_sock *sock = (struct trilogy_sock *)_sock;

data/lib/trilogy/error.rb

@@ -54,6 +54,12 @@ class Trilogy
     include ConnectionError
   end
 
+  class SynchronizationError < BaseError
+    def initialize(message = "This connection is already in use by another thread or fiber")
+      super
+    end
+  end
+
   # Trilogy::ClientError is the base error type for invalid queries or parameters
   # that shouldn't be retried.
   class ClientError < BaseError
@@ -94,6 +100,7 @@ class Trilogy
       1160 => BaseConnectionError, # ER_NET_ERROR_ON_WRITE
       1161 => BaseConnectionError, # ER_NET_WRITE_INTERRUPTED
       1927 => BaseConnectionError, # ER_CONNECTION_KILLED
+      4031 => BaseConnectionError, # Disconnected by server
     }
     class << self
       def from_code(message, code)

data/lib/trilogy/version.rb

@@ -1,3 +1,3 @@
 class Trilogy
-  VERSION = "2.10.0"
+  VERSION = "2.11.0"
 end

data/lib/trilogy.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require "socket"
 require "trilogy/version"
 require "trilogy/error"
 require "trilogy/result"
@@ -7,7 +8,34 @@ require "trilogy/cext"
 require "trilogy/encoding"
 
 class Trilogy
+  IO_TIMEOUT_ERROR =
+    if defined?(IO::TimeoutError)
+      IO::TimeoutError
+    else
+      Class.new(StandardError)
+    end
+  private_constant :IO_TIMEOUT_ERROR
+
+  Synchronization = Module.new
+
+  source = public_instance_methods(false).flat_map do |method|
+    [
+      "def #{method}(...)",
+        "raise SynchronizationError unless @mutex.try_lock",
+        "begin",
+          "super",
+        "ensure",
+          "@mutex.unlock",
+        "end",
+      "end",
+    ]
+  end
+  Synchronization.class_eval(source.join(";"), __FILE__, __LINE__)
+
+  prepend(Synchronization)
+
   def initialize(options = {})
+    @mutex = Mutex.new
     options[:port] = options[:port].to_i if options[:port]
     mysql_encoding = options[:encoding] || "utf8mb4"
     encoding = Trilogy::Encoding.find(mysql_encoding)
@@ -15,7 +43,54 @@ class Trilogy
     @connection_options = options
     @connected_host = nil
 
-    _connect(encoding, charset, options)
+    socket = nil
+    begin
+      if host = options[:host]
+        port = options[:port] || 3306
+        connect_timeout = options[:connect_timeout] || options[:write_timeout]
+
+        socket = TCPSocket.new(host, port, connect_timeout: connect_timeout)
+
+        socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+
+        if keepalive_enabled = options[:keepalive_enabled]
+          keepalive_idle = options[:keepalive_idle]
+          keepalive_interval = options[:keepalive_interval]
+          keepalive_count = options[:keepalive_count]
+
+          socket.setsockopt(Socket::SOL_SOCKET, Socket::SO_KEEPALIVE, true)
+
+          if keepalive_idle > 0 && defined?(Socket::TCP_KEEPIDLE)
+            socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_KEEPIDLE, keepalive_idle)
+          end
+          if keepalive_interval > 0 && defined?(Socket::TCP_KEEPINTVL)
+            socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_KEEPINTVL, keepalive_interval)
+          end
+          if keepalive_count > 0 && defined?(Socket::TCP_KEEPCNT)
+            socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_KEEPCNT, keepalive_count)
+          end
+        end
+      else
+        path = options[:socket] ||= "/tmp/mysql.sock"
+        socket = UNIXSocket.new(path)
+      end
+    rescue Errno::ETIMEDOUT, IO_TIMEOUT_ERROR => e
+      raise Trilogy::TimeoutError, e.message
+    rescue SocketError => e
+      connection_str = host ? "#{host}:#{port}" : path
+      raise Trilogy::BaseConnectionError, "unable to connect to \"#{connection_str}\": #{e.message}"
+    rescue => e
+      if e.respond_to?(:errno)
+        raise Trilogy::SyscallError.from_errno(e.errno, e.message)
+      else
+        raise
+      end
+    end
+
+    _connect(socket, encoding, charset, options)
+  ensure
+    # Socket's fd will be dup'd in C
+    socket&.close
   end
 
   def connection_options

data/trilogy.gemspec

@@ -25,7 +25,4 @@ Gem::Specification.new do |s|
   s.required_ruby_version = ">= 3.0"
 
   s.add_dependency "bigdecimal"
-
-  s.add_development_dependency "rake-compiler", "~> 1.0"
-  s.add_development_dependency "minitest", "~> 5.5"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: trilogy
 version: !ruby/object:Gem::Version
-  version: 2.10.0
+  version: 2.11.0
 platform: ruby
 authors:
 - GitHub Engineering
@@ -23,34 +23,6 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '0'
-- !ruby/object:Gem::Dependency
-  name: rake-compiler
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - "~>"
-      - !ruby/object:Gem::Version
-        version: '1.0'
-  type: :development
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - "~>"
-      - !ruby/object:Gem::Version
-        version: '1.0'
-- !ruby/object:Gem::Dependency
-  name: minitest
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - "~>"
-      - !ruby/object:Gem::Version
-        version: '5.5'
-  type: :development
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - "~>"
-      - !ruby/object:Gem::Version
-        version: '5.5'
 email: opensource+trilogy@github.com
 executables: []
 extensions: