validate max_allowed_packet on the client side

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

@@ -27,7 +27,7 @@ static ID id_socket, id_host, id_port, id_username, id_password, id_found_rows,
     id_ivar_affected_rows, id_ivar_fields, id_ivar_last_insert_id, id_ivar_rows, id_ivar_query_time, id_password,
     id_database, id_ssl_ca, id_ssl_capath, id_ssl_cert, id_ssl_cipher, id_ssl_crl, id_ssl_crlpath, id_ssl_key,
     id_ssl_mode, id_tls_ciphersuites, id_tls_min_version, id_tls_max_version, id_multi_statement, id_multi_result,
-    id_from_code, id_from_errno, id_connection_options;
+    id_from_code, id_from_errno, id_connection_options, id_max_allowed_packet;
 
 struct trilogy_ctx {
     trilogy_conn_t conn;
@@ -415,6 +415,11 @@ static VALUE rb_trilogy_initialize(VALUE self, VALUE encoding, VALUE charset, VA
         connopt.keepalive_interval = NUM2USHORT(val);
     }
 
+    if ((val = rb_hash_lookup(opts, ID2SYM(id_max_allowed_packet))) != Qnil) {
+        Check_Type(val, T_FIXNUM);
+        connopt.max_allowed_packet = NUM2SIZET(val);
+    }
+
     if ((val = rb_hash_lookup(opts, ID2SYM(id_host))) != Qnil) {
         Check_Type(val, T_STRING);
 
@@ -1114,6 +1119,7 @@ RUBY_FUNC_EXPORTED void Init_cext()
     id_connect_timeout = rb_intern("connect_timeout");
     id_read_timeout = rb_intern("read_timeout");
     id_write_timeout = rb_intern("write_timeout");
+    id_max_allowed_packet = rb_intern("max_allowed_packet");
     id_keepalive_enabled = rb_intern("keepalive_enabled");
     id_keepalive_idle = rb_intern("keepalive_idle");
     id_keepalive_count = rb_intern("keepalive_count");

contrib/ruby/test/client_test.rb

@@ -745,28 +745,116 @@ def test_connect_by_multiple_names
     Trilogy.new(host: "localhost")
   end
 
+  PADDED_QUERY_TEMPLATE = "SELECT LENGTH('%s')"
+  PROTOCOL_OVERHEAD = 2 # One byte for the 0x03 (COM_QUERY); one because the packet is actually required to be shorter than the "max"
+  PADDED_QUERY_OVERHEAD =
+    PADDED_QUERY_TEMPLATE.size - "%s".size + PROTOCOL_OVERHEAD
+
+  def query_for_target_packet_size(size)
+    PADDED_QUERY_TEMPLATE % ("x" * (size - PADDED_QUERY_OVERHEAD))
+  end
+
   def set_max_allowed_packet(size)
     client = new_tcp_client
     client.query "SET GLOBAL max_allowed_packet = #{size}"
   ensure
     ensure_closed client
   end
 
-  def test_packet_size
+  def test_packet_size_lower_than_trilogy_max_packet_len
+    set_max_allowed_packet(4 * 1024 * 1024) # TRILOGY_MAX_PACKET_LEN is 16MB
+
+    client = new_tcp_client(max_allowed_packet: 4 * 1024 * 1024)
+
+    assert client.query query_for_target_packet_size(1 * 1024 * 1024)
+
+    assert client.query query_for_target_packet_size(2 * 1024 * 1024)
+
+    assert client.query query_for_target_packet_size(4 * 1024 * 1024)
+
+    exception = assert_raises Trilogy::QueryError do
+      client.query query_for_target_packet_size(4 * 1024 * 1024 + 1)
+    end
+
+    assert_equal "trilogy_query_send: TRILOGY_MAX_PACKET_EXCEEDED", exception.message
+  ensure
+    ensure_closed client
+  end
+
+  def test_packet_size_greater_than_trilogy_max_packet_len
+    set_max_allowed_packet(32 * 1024 * 1024) # TRILOGY_MAX_PACKET_LEN is 16MB
+
+    client = new_tcp_client(max_allowed_packet: 32 * 1024 * 1024)
+
+    assert client.query query_for_target_packet_size(16 * 1024 * 1024)
+
+    assert client.query query_for_target_packet_size(32 * 1024 * 1024)
+
+    exception = assert_raises Trilogy::QueryError do
+      client.query query_for_target_packet_size(32 * 1024 * 1024 + 1)
+    end
+
+    assert_equal "trilogy_query_send: TRILOGY_MAX_PACKET_EXCEEDED", exception.message
+  ensure
+    ensure_closed client
+  end
+
+  def test_configured_max_packet_below_server
     set_max_allowed_packet(32 * 1024 * 1024)
 
-    client = new_tcp_client
+    client = new_tcp_client(max_allowed_packet: 24 * 1024 * 1024)
 
-    create_test_table(client)
-    client.query "TRUNCATE trilogy_test"
+    assert client.query query_for_target_packet_size(16 * 1024 * 1024)
 
-    result = client.query "INSERT INTO trilogy_test (blob_test) VALUES ('#{"x" * (15 * 1024 * 1024)}')"
-    assert result
-    assert_equal 1, client.last_insert_id
+    assert client.query query_for_target_packet_size(24 * 1024 * 1024)
 
-    result = client.query "INSERT INTO trilogy_test (blob_test) VALUES ('#{"x" * (31 * 1024 * 1024)}')"
-    assert result
-    assert_equal 2, client.last_insert_id
+    exception = assert_raises Trilogy::QueryError do
+      client.query query_for_target_packet_size(24 * 1024 * 1024 + 1)
+    end
+
+    assert_equal "trilogy_query_send: TRILOGY_MAX_PACKET_EXCEEDED", exception.message
+  ensure
+    ensure_closed client
+  end
+
+  def test_configured_max_packet_above_server
+    set_max_allowed_packet(24 * 1024 * 1024)
+
+    client = new_tcp_client(max_allowed_packet: 32 * 1024 * 1024)
+
+    assert client.query query_for_target_packet_size(16 * 1024 * 1024)
+
+    assert client.query query_for_target_packet_size(24 * 1024 * 1024)
+
+    exception = assert_raises Trilogy::QueryError do
+      client.query query_for_target_packet_size(32 * 1024 * 1024 + 1)
+    end
+
+    assert_equal "trilogy_query_send: TRILOGY_MAX_PACKET_EXCEEDED", exception.message
+
+    exception = assert_raises Trilogy::QueryError do
+      client.query query_for_target_packet_size(24 * 1024 * 1024 + 1)
+    end
+
+    refute_match(/TRILOGY_MAX_PACKET_EXCEEDED/, exception.message)
+  ensure
+    ensure_closed client
+  end
+
+  def test_absolute_maximum_packet_size
+    skip unless ENV["CI"]
+
+    set_max_allowed_packet(1024 * 1024 * 1024) # 1GB is the highest maximum allowed
+
+    client = new_tcp_client(max_allowed_packet: 1024 * 1024 * 1024)
+
+    assert client.query query_for_target_packet_size(1024 * 1024 * 1024)
+
+    exception = assert_raises Trilogy::QueryError do
+      client.query query_for_target_packet_size(1024 * 1024 * 1024 + 1)
+    end
+
+    assert_equal "trilogy_query_send: TRILOGY_MAX_PACKET_EXCEEDED", exception.message
   ensure
     ensure_closed client
   end

inc/trilogy/builder.h

@@ -16,6 +16,8 @@
 typedef struct {
     trilogy_buffer_t *buffer;
     size_t header_offset;
+    size_t packet_length;
+    size_t packet_max_length;
     uint32_t fragment_length;
     uint8_t seq;
 } trilogy_builder_t;
@@ -54,6 +56,8 @@ void trilogy_builder_finalize(trilogy_builder_t *builder);
  * Return values:
  *   TRILOGY_OK     - The value was appended to the packet buffer.
  *   TRILOGY_SYSERR - A system error occurred, check errno.
+ *   TRILOGY_MAX_PACKET_EXCEEDED - Appending this value would exceed the maximum
+ *                                 packet size.
  */
 int trilogy_builder_write_uint8(trilogy_builder_t *builder, uint8_t val);
 
@@ -66,6 +70,8 @@ int trilogy_builder_write_uint8(trilogy_builder_t *builder, uint8_t val);
  * Return values:
  *   TRILOGY_OK     - The value was appended to the packet buffer.
  *   TRILOGY_SYSERR - A system error occurred, check errno.
+ *   TRILOGY_MAX_PACKET_EXCEEDED - Appending this value would exceed the maximum
+ *                                 packet size.
  */
 int trilogy_builder_write_uint16(trilogy_builder_t *builder, uint16_t val);
 
@@ -78,6 +84,8 @@ int trilogy_builder_write_uint16(trilogy_builder_t *builder, uint16_t val);
  * Return values:
  *   TRILOGY_OK     - The value was appended to the packet buffer.
  *   TRILOGY_SYSERR - A system error occurred, check errno.
+ *   TRILOGY_MAX_PACKET_EXCEEDED - Appending this value would exceed the maximum
+ *                                 packet size.
  */
 int trilogy_builder_write_uint24(trilogy_builder_t *builder, uint32_t val);
 
@@ -90,6 +98,8 @@ int trilogy_builder_write_uint24(trilogy_builder_t *builder, uint32_t val);
  * Return values:
  *   TRILOGY_OK     - The value was appended to the packet buffer.
  *   TRILOGY_SYSERR - A system error occurred, check errno.
+ *   TRILOGY_MAX_PACKET_EXCEEDED - Appending this value would exceed the maximum
+ *                                 packet size.
  */
 int trilogy_builder_write_uint32(trilogy_builder_t *builder, uint32_t val);
 
@@ -102,6 +112,8 @@ int trilogy_builder_write_uint32(trilogy_builder_t *builder, uint32_t val);
  * Return values:
  *   TRILOGY_OK     - The value was appended to the packet buffer.
  *   TRILOGY_SYSERR - A system error occurred, check errno.
+ *   TRILOGY_MAX_PACKET_EXCEEDED - Appending this value would exceed the maximum
+ *                                 packet size.
  */
 int trilogy_builder_write_uint64(trilogy_builder_t *builder, uint64_t val);
 
@@ -117,6 +129,8 @@ int trilogy_builder_write_uint64(trilogy_builder_t *builder, uint64_t val);
  * Return values:
  *   TRILOGY_OK     - The value was appended to the packet buffer.
  *   TRILOGY_SYSERR - A system error occurred, check errno.
+ *   TRILOGY_MAX_PACKET_EXCEEDED - Appending this value would exceed the maximum
+ *                                 packet size.
  */
 int trilogy_builder_write_lenenc(trilogy_builder_t *builder, uint64_t val);
 
@@ -130,6 +144,8 @@ int trilogy_builder_write_lenenc(trilogy_builder_t *builder, uint64_t val);
  * Return values:
  *   TRILOGY_OK     - The value was appended to the packet buffer.
  *   TRILOGY_SYSERR - A system error occurred, check errno.
+ *   TRILOGY_MAX_PACKET_EXCEEDED - Appending this value would exceed the maximum
+ *                                 packet size.
  */
 int trilogy_builder_write_buffer(trilogy_builder_t *builder, const void *data, size_t len);
 
@@ -144,6 +160,8 @@ int trilogy_builder_write_buffer(trilogy_builder_t *builder, const void *data, s
  * Return values:
  *   TRILOGY_OK     - The value was appended to the packet buffer.
  *   TRILOGY_SYSERR - A system error occurred, check errno.
+ *   TRILOGY_MAX_PACKET_EXCEEDED - Appending this value would exceed the maximum
+ *                                 packet size.
  */
 int trilogy_builder_write_lenenc_buffer(trilogy_builder_t *builder, const void *data, size_t len);
 
@@ -155,7 +173,23 @@ int trilogy_builder_write_lenenc_buffer(trilogy_builder_t *builder, const void *
  * Return values:
  *   TRILOGY_OK     - The value was appended to the packet buffer.
  *   TRILOGY_SYSERR - A system error occurred, check errno.
+ *   TRILOGY_MAX_PACKET_EXCEEDED - Appending this value would exceed the maximum
+ *                                 packet size.
  */
 int trilogy_builder_write_string(trilogy_builder_t *builder, const char *data);
 
+/* trilogy_builder_set_max_packet_length - Set the maximum packet length for
+ * the builder. Writing data to the builder that would cause the packet length
+ * to exceed this value will cause the builder to error.
+ *
+ * builder    - A pre-initialized trilogy_builder_t pointer
+ * max_length - The new maximum packet length to set
+ *
+ * Return values:
+ *   TRILOGY_OK                  - The maximum packet length was set.
+ *   TRILOGY_MAX_PACKET_EXCEEDED - The current packet length is already
+ *                                 larger than the requested maximum.
+ */
+int trilogy_builder_set_max_packet_length(trilogy_builder_t *builder, size_t max_length);
+
 #endif

inc/trilogy/error.h

@@ -21,7 +21,8 @@
     XX(TRILOGY_OPENSSL_ERR, -16) /* check ERR_get_error() */                                                           \
     XX(TRILOGY_UNSUPPORTED, -17)                                                                                       \
     XX(TRILOGY_DNS_ERR, -18)                                                                                           \
-    XX(TRILOGY_AUTH_SWITCH, -19)
+    XX(TRILOGY_AUTH_SWITCH, -19)                                                                                       \
+    XX(TRILOGY_MAX_PACKET_EXCEEDED, -20)
 
 enum {
 #define XX(name, code) name = code,

inc/trilogy/socket.h

@@ -67,6 +67,8 @@ typedef struct {
     uint16_t keepalive_interval;
 
     TRILOGY_CAPABILITIES_t flags;
+
+    size_t max_allowed_packet;
 } trilogy_sockopt_t;
 
 typedef struct trilogy_sock_t {

src/builder.c

@@ -39,6 +39,8 @@ int trilogy_builder_init(trilogy_builder_t *builder, trilogy_buffer_t *buff, uin
     builder->buffer->len = 0;
 
     builder->seq = seq;
+    builder->packet_length = 0;
+    builder->packet_max_length = SIZE_MAX;
 
     return write_header(builder);
 }
@@ -60,10 +62,15 @@ void trilogy_builder_finalize(trilogy_builder_t *builder)
 
 int trilogy_builder_write_uint8(trilogy_builder_t *builder, uint8_t val)
 {
+    if (builder->packet_length >= builder->packet_max_length - 1) {
+        return TRILOGY_MAX_PACKET_EXCEEDED;
+    }
+
     CHECKED(trilogy_buffer_expand(builder->buffer, 1));
 
     builder->buffer->buff[builder->buffer->len++] = val;
     builder->fragment_length++;
+    builder->packet_length++;
 
     if (builder->fragment_length == TRILOGY_MAX_PACKET_LEN) {
         CHECKED(write_continuation_header(builder));
@@ -137,6 +144,10 @@ int trilogy_builder_write_buffer(trilogy_builder_t *builder, const void *data, s
 
     size_t fragment_remaining = TRILOGY_MAX_PACKET_LEN - builder->fragment_length;
 
+    if (builder->packet_length >= builder->packet_max_length - len) {
+        return TRILOGY_MAX_PACKET_EXCEEDED;
+    }
+
     // if this buffer write is not going to straddle a fragment boundary:
     if (len < fragment_remaining) {
         CHECKED(trilogy_buffer_expand(builder->buffer, len));
@@ -145,6 +156,7 @@ int trilogy_builder_write_buffer(trilogy_builder_t *builder, const void *data, s
 
         builder->buffer->len += len;
         builder->fragment_length += len;
+        builder->packet_length += len;
 
         return TRILOGY_OK;
     }
@@ -157,6 +169,7 @@ int trilogy_builder_write_buffer(trilogy_builder_t *builder, const void *data, s
 
         builder->buffer->len += fragment_remaining;
         builder->fragment_length += fragment_remaining;
+        builder->packet_length += fragment_remaining;
 
         ptr += fragment_remaining;
         len -= fragment_remaining;
@@ -172,6 +185,7 @@ int trilogy_builder_write_buffer(trilogy_builder_t *builder, const void *data, s
 
         builder->buffer->len += len;
         builder->fragment_length += len;
+        builder->packet_length += len;
     }
 
     return TRILOGY_OK;
@@ -195,4 +209,15 @@ int trilogy_builder_write_string(trilogy_builder_t *builder, const char *data)
     return TRILOGY_OK;
 }
 
+int trilogy_builder_set_max_packet_length(trilogy_builder_t *builder, size_t max_length)
+{
+    if (builder->packet_length > max_length) {
+        return TRILOGY_MAX_PACKET_EXCEEDED;
+    }
+
+    builder->packet_max_length = max_length;
+
+    return TRILOGY_OK;
+}
+
 #undef CHECKED

src/client.c

@@ -60,6 +60,10 @@ static int begin_command_phase(trilogy_builder_t *builder, trilogy_conn_t *conn,
         return rc;
     }
 
+    if (conn->socket->opts.max_allowed_packet > 0) {
+        trilogy_builder_set_max_packet_length(builder, conn->socket->opts.max_allowed_packet);
+    }
+
     conn->packet_parser.sequence_number = seq + 1;
 
     return 0;