From 4e18d7969abc379de89cc52618408bb38a2c1857 Mon Sep 17 00:00:00 2001
From: Satoshi MITANI <s-mitani@mercari.com>
Date: Fri, 6 Sep 2024 09:20:00 +0900
Subject: [PATCH] Support get_server_public_key option  (#1377)

* caching_sha2_password requests secure connection if cache is not ready on server-side.
get_server_public_key option enables clients to create secure connection automatically even if connection is not SSL.

* return error if get_server_public_key option is not supported
---
 README.md             |  1 +
 ext/mysql2/client.c   | 16 ++++++++++++++++
 ext/mysql2/extconf.rb |  1 +
 lib/mysql2/client.rb  |  4 ++--
 4 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index cf35222d5..59ae2bc54 100644
--- a/README.md
+++ b/README.md
@@ -281,6 +281,7 @@ Mysql2::Client.new(
   :reconnect = true/false,
   :local_infile = true/false,
   :secure_auth = true/false,
+  :get_server_public_key = true/false,
   :default_file = '/path/to/my.cfg',
   :default_group = 'my.cfg section',
   :default_auth = 'authentication_windows_client'
diff --git a/ext/mysql2/client.c b/ext/mysql2/client.c
index 5a2fcc1bd..8fe33ea39 100644
--- a/ext/mysql2/client.c
+++ b/ext/mysql2/client.c
@@ -996,6 +996,13 @@ static VALUE _mysql_client_options(VALUE self, int opt, VALUE value) {
       retval  = charval;
       break;
 
+#ifdef HAVE_CONST_MYSQL_OPT_GET_SERVER_PUBLIC_KEY
+    case MYSQL_OPT_GET_SERVER_PUBLIC_KEY:
+      boolval = (value == Qfalse ? 0 : 1);
+      retval = &boolval;
+      break;
+#endif
+
 #ifdef HAVE_MYSQL_DEFAULT_AUTH
     case MYSQL_DEFAULT_AUTH:
       charval = (const char *)StringValueCStr(value);
@@ -1485,6 +1492,14 @@ static VALUE set_init_command(VALUE self, VALUE value) {
   return _mysql_client_options(self, MYSQL_INIT_COMMAND, value);
 }
 
+static VALUE set_get_server_public_key(VALUE self, VALUE value) {
+#ifdef HAVE_CONST_MYSQL_OPT_GET_SERVER_PUBLIC_KEY
+  return _mysql_client_options(self, MYSQL_OPT_GET_SERVER_PUBLIC_KEY, value);
+#else
+  rb_raise(cMysql2Error, "get-server-public-key is not available, you may need a newer MySQL client library");
+#endif
+}
+
 static VALUE set_default_auth(VALUE self, VALUE value) {
 #ifdef HAVE_MYSQL_DEFAULT_AUTH
   return _mysql_client_options(self, MYSQL_DEFAULT_AUTH, value);
@@ -1596,6 +1611,7 @@ void init_mysql2_client() {
   rb_define_private_method(cMysql2Client, "default_file=", set_read_default_file, 1);
   rb_define_private_method(cMysql2Client, "default_group=", set_read_default_group, 1);
   rb_define_private_method(cMysql2Client, "init_command=", set_init_command, 1);
+  rb_define_private_method(cMysql2Client, "get_server_public_key=", set_get_server_public_key, 1);
   rb_define_private_method(cMysql2Client, "default_auth=", set_default_auth, 1);
   rb_define_private_method(cMysql2Client, "ssl_set", set_ssl_options, 5);
   rb_define_private_method(cMysql2Client, "ssl_mode=", rb_set_ssl_mode_option, 1);
diff --git a/ext/mysql2/extconf.rb b/ext/mysql2/extconf.rb
index a1f6cb240..68e68bcb1 100644
--- a/ext/mysql2/extconf.rb
+++ b/ext/mysql2/extconf.rb
@@ -159,6 +159,7 @@ def add_ssl_defines(header)
 have_const('SERVER_QUERY_WAS_SLOW', mysql_h)
 have_const('MYSQL_OPTION_MULTI_STATEMENTS_ON', mysql_h)
 have_const('MYSQL_OPTION_MULTI_STATEMENTS_OFF', mysql_h)
+have_const('MYSQL_OPT_GET_SERVER_PUBLIC_KEY', mysql_h)
 
 # my_bool is replaced by C99 bool in MySQL 8.0, but we want
 # to retain compatibility with the typedef in earlier MySQLs.
diff --git a/lib/mysql2/client.rb b/lib/mysql2/client.rb
index 582b6e305..d952cec17 100644
--- a/lib/mysql2/client.rb
+++ b/lib/mysql2/client.rb
@@ -32,11 +32,11 @@ def initialize(opts = {})
       opts[:connect_timeout] = 120 unless opts.key?(:connect_timeout)
 
       # TODO: stricter validation rather than silent massaging
-      %i[reconnect connect_timeout local_infile read_timeout write_timeout default_file default_group secure_auth init_command automatic_close enable_cleartext_plugin default_auth].each do |key|
+      %i[reconnect connect_timeout local_infile read_timeout write_timeout default_file default_group secure_auth init_command automatic_close enable_cleartext_plugin default_auth get_server_public_key].each do |key|
         next unless opts.key?(key)
 
         case key
-        when :reconnect, :local_infile, :secure_auth, :automatic_close, :enable_cleartext_plugin
+        when :reconnect, :local_infile, :secure_auth, :automatic_close, :enable_cleartext_plugin, :get_server_public_key
           send(:"#{key}=", !!opts[key]) # rubocop:disable Style/DoubleNegation
         when :connect_timeout, :read_timeout, :write_timeout
           send(:"#{key}=", Integer(opts[key])) unless opts[key].nil?