fix auth_switch_request handling (#1200)

methane · web-flow · commit 01af30fea088 · 2025-01-29T18:06:45.000+09:00
diff --git a/.coveragerc b/.coveragerc
@@ -2,6 +2,7 @@
 branch = True
 source =
     pymysql
+    tests
 omit = pymysql/tests/*
        pymysql/tests/thirdparty/test_MySQLdb/*
 
diff --git a/pymysql/_auth.py b/pymysql/_auth.py
@@ -166,6 +166,8 @@ def sha256_password_auth(conn, pkt):
 
     if pkt.is_auth_switch_request():
         conn.salt = pkt.read_all()
+        if conn.salt.endswith(b"\0"):
+            conn.salt = conn.salt[:-1]
         if not conn.server_public_key and conn.password:
             # Request server public key
             if DEBUG:
@@ -215,9 +217,11 @@ def caching_sha2_password_auth(conn, pkt):
 
     if pkt.is_auth_switch_request():
         # Try from fast auth
-        if DEBUG:
-            print("caching sha2: Trying fast path")
         conn.salt = pkt.read_all()
+        if conn.salt.endswith(b"\0"):  # str.removesuffix is available in 3.9
+            conn.salt = conn.salt[:-1]
+        if DEBUG:
+            print(f"caching sha2: Trying fast path. salt={conn.salt.hex()!r}")
         scrambled = scramble_caching_sha2(conn.password, conn.salt)
         pkt = _roundtrip(conn, scrambled)
     # else: fast auth is tried in initial handshake
diff --git a/pymysql/connections.py b/pymysql/connections.py
@@ -47,6 +47,7 @@
     DEFAULT_USER = None
 
 DEBUG = False
+_DEFAULT_AUTH_PLUGIN = None  # if this is not None, use it instead of server's default.
 
 TEXT_TYPES = {
     FIELD_TYPE.BIT,
@@ -1158,6 +1159,9 @@ def _get_server_information(self):
             else:
                 self._auth_plugin_name = data[i:server_end].decode("utf-8")
 
+        if _DEFAULT_AUTH_PLUGIN is not None:  # for tests
+            self._auth_plugin_name = _DEFAULT_AUTH_PLUGIN
+
     def get_server_info(self):
         return self.server_version
 
diff --git a/tests/test_auth.py b/tests/test_auth.py
@@ -71,6 +71,19 @@ def test_caching_sha2_password():
     con.query("FLUSH PRIVILEGES")
     con.close()
 
+    # Fast path after auth_switch_request
+    pymysql.connections._DEFAULT_AUTH_PLUGIN = "mysql_native_password"
+    con = pymysql.connect(
+        user="user_caching_sha2",
+        password=pass_caching_sha2,
+        host=host,
+        port=port,
+        ssl=ssl,
+    )
+    con.query("FLUSH PRIVILEGES")
+    con.close()
+    pymysql.connections._DEFAULT_AUTH_PLUGIN = None
+
 
 def test_caching_sha2_password_ssl():
     con = pymysql.connect(
@@ -88,7 +101,20 @@ def test_caching_sha2_password_ssl():
         password=pass_caching_sha2,
         host=host,
         port=port,
-        ssl=None,
+        ssl=ssl,
+    )
+    con.query("FLUSH PRIVILEGES")
+    con.close()
+
+    # Fast path after auth_switch_request
+    pymysql.connections._DEFAULT_AUTH_PLUGIN = "mysql_native_password"
+    con = pymysql.connect(
+        user="user_caching_sha2",
+        password=pass_caching_sha2,
+        host=host,
+        port=port,
+        ssl=ssl,
     )
     con.query("FLUSH PRIVILEGES")
     con.close()
+    pymysql.connections._DEFAULT_AUTH_PLUGIN = None
