support server_public_key_path option. (#744)

fix #682
PyMySQL · Nov 12, 2024 · feea60c · feea60c
1 parent 89511ee
commit feea60c
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 17 deletions.
diff --git a/doc/user_guide.rst b/doc/user_guide.rst
@@ -393,6 +393,10 @@ connect(parameters...)
             an exception is raised. *This must be a keyword
             parameter.*
 
+        server_public_key_path
+            specifies path to a RSA public key used by caching sha2 password authentication.
+            See https://dev.mysql.com/doc/refman/9.0/en/caching-sha2-pluggable-authentication.html
+
 .. _mysql_ssl_set: http://dev.mysql.com/doc/refman/en/mysql-ssl-set.html
 
 

diff --git a/src/MySQLdb/_mysql.c b/src/MySQLdb/_mysql.c
@@ -36,13 +36,18 @@ PERFORMANCE OF THIS SOFTWARE.
 #endif
 
 #if ((MYSQL_VERSION_ID >= 50555 && MYSQL_VERSION_ID <= 50599) || \
-(MYSQL_VERSION_ID >= 50636 && MYSQL_VERSION_ID <= 50699) || \
-(MYSQL_VERSION_ID >= 50711 && MYSQL_VERSION_ID <= 50799) || \
-(MYSQL_VERSION_ID >= 80000)) && \
-!defined(MARIADB_BASE_VERSION) && !defined(MARIADB_VERSION_ID)
+     (MYSQL_VERSION_ID >= 50636 && MYSQL_VERSION_ID <= 50699) || \
+     (MYSQL_VERSION_ID >= 50711 && MYSQL_VERSION_ID <= 50799) || \
+     (MYSQL_VERSION_ID >= 80000)) && \
+     !defined(MARIADB_BASE_VERSION) && !defined(MARIADB_VERSION_ID)
 #define HAVE_ENUM_MYSQL_OPT_SSL_MODE
 #endif
 
+#if defined(MARIADB_VERSION_ID) && MARIADB_VERSION_ID >= 100403 || \
+    !defined(MARIADB_VERSION_ID) && MYSQL_VERSION_ID >= 50723
+#define HAVE_MYSQL_SERVER_PUBLIC_KEY
+#endif
+
 #define PY_SSIZE_T_CLEAN 1
 #include "Python.h"
 
@@ -431,7 +436,7 @@ _mysql_ConnectionObject_Initialize(
                   "client_flag", "ssl", "ssl_mode",
                   "local_infile",
                   "read_timeout", "write_timeout", "charset",
-                  "auth_plugin",
+                  "auth_plugin", "server_public_key_path",
                   NULL } ;
     int connect_timeout = 0;
     int read_timeout = 0;
@@ -442,14 +447,15 @@ _mysql_ConnectionObject_Initialize(
          *read_default_file=NULL,
          *read_default_group=NULL,
          *charset=NULL,
-         *auth_plugin=NULL;
+         *auth_plugin=NULL,
+         *server_public_key_path=NULL;
 
     self->converter = NULL;
     self->open = false;
     self->reconnect = false;
 
     if (!PyArg_ParseTupleAndKeywords(args, kwargs,
-                "|ssssisOiiisssiOsiiiss:connect",
+                "|ssssisOiiisssiOsiiisss:connect",
                 kwlist,
                 &host, &user, &passwd, &db,
                 &port, &unix_socket, &conv,
@@ -462,10 +468,19 @@ _mysql_ConnectionObject_Initialize(
                 &read_timeout,
                 &write_timeout,
                 &charset,
-                &auth_plugin
+                &auth_plugin,
+                &server_public_key_path
     ))
         return -1;
 
+#ifndef HAVE_MYSQL_SERVER_PUBLIC_KEY
+    if (server_public_key_path) {
+        PyErr_SetString(_mysql_NotSupportedError, "server_public_key_path is not supported");
+        return -1;
+    }
+#endif
+    // For compatibility with PyPy, we need to keep strong reference
+    // to unicode objects until we use UTF8.
 #define _stringsuck(d,t,s) {t=PyMapping_GetItemString(s,#d);\
         if(t){d=PyUnicode_AsUTF8(t);ssl_keepref[n_ssl_keepref++]=t;}\
         PyErr_Clear();}
@@ -542,6 +557,10 @@ _mysql_ConnectionObject_Initialize(
         mysql_options(&(self->connection), MYSQL_OPT_SSL_CAPATH, capath);
         mysql_options(&(self->connection), MYSQL_OPT_SSL_CIPHER, cipher);
     }
+    for (int i=0 ; i<n_ssl_keepref; i++) {
+        Py_DECREF(ssl_keepref[i]);
+        ssl_keepref[i] = NULL;
+    }
 
 #ifdef HAVE_ENUM_MYSQL_OPT_SSL_MODE
     if (ssl_mode_set) {
@@ -557,14 +576,14 @@ _mysql_ConnectionObject_Initialize(
     my_bool my_true = 1;
     my_bool my_false = 0;
     if (ssl_mode_num >= SSLMODE_REQUIRED) {
-        mysql_optionsv(&(self->connection), MYSQL_OPT_SSL_ENFORCE, (void *)&my_true);
+        mysql_options(&(self->connection), MYSQL_OPT_SSL_ENFORCE, (void *)&my_true);
     } else {
-        mysql_optionsv(&(self->connection), MYSQL_OPT_SSL_ENFORCE, (void *)&my_false);
+        mysql_options(&(self->connection), MYSQL_OPT_SSL_ENFORCE, (void *)&my_false);
     }
     if (ssl_mode_num >= SSLMODE_VERIFY_CA) {
-        mysql_optionsv(&(self->connection), MYSQL_OPT_SSL_VERIFY_SERVER_CERT, (void *)&my_true);
+        mysql_options(&(self->connection), MYSQL_OPT_SSL_VERIFY_SERVER_CERT, (void *)&my_true);
     } else {
-        mysql_optionsv(&(self->connection), MYSQL_OPT_SSL_VERIFY_SERVER_CERT, (void *)&my_false);
+        mysql_options(&(self->connection), MYSQL_OPT_SSL_VERIFY_SERVER_CERT, (void *)&my_false);
     }
 #endif
 
@@ -574,17 +593,17 @@ _mysql_ConnectionObject_Initialize(
     if (auth_plugin) {
         mysql_options(&(self->connection), MYSQL_DEFAULT_AUTH, auth_plugin);
     }
+#ifdef HAVE_MYSQL_SERVER_PUBLIC_KEY
+    if (server_public_key_path) {
+        mysql_options(&(self->connection), MYSQL_SERVER_PUBLIC_KEY, server_public_key_path);
+    }
+#endif
 
     Py_BEGIN_ALLOW_THREADS
     conn = mysql_real_connect(&(self->connection), host, user, passwd, db,
                   port, unix_socket, client_flag);
     Py_END_ALLOW_THREADS
 
-    for (int i=0; i<n_ssl_keepref; i++) {
-        Py_DECREF(ssl_keepref[i]);
-        ssl_keepref[i] = NULL;
-    }
-
     if (!conn) {
         _mysql_Exception(self);
         return -1;

diff --git a/src/MySQLdb/connections.py b/src/MySQLdb/connections.py
@@ -137,6 +137,10 @@ class object, used to create cursors (keyword only)
             Since mysqlclient 2.2.4, ssl=True is alias of ssl_mode=REQUIRED
             for better compatibility with PyMySQL and MariaDB.
 
+        :param str server_public_key_path:
+            specify the path to a file RSA public key file for caching_sha2_password.
+            See https://dev.mysql.com/doc/refman/9.0/en/caching-sha2-pluggable-authentication.html
+
         :param bool local_infile:
             enables LOAD LOCAL INFILE; zero disables