From 95ebbb3e8166787927e1ca14d94b12be8d138405 Mon Sep 17 00:00:00 2001
From: Inada Naoki <songofacandy@gmail.com>
Date: Wed, 29 Jan 2025 16:43:19 +0900
Subject: [PATCH 1/3] fix auth_switch_request handling

---
 pymysql/_auth.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/pymysql/_auth.py b/pymysql/_auth.py
index 8ce744fb..4790449b 100644
--- 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

From b1aa823ddd2e11ac733433a03717ad9f62a01847 Mon Sep 17 00:00:00 2001
From: Inada Naoki <songofacandy@gmail.com>
Date: Wed, 29 Jan 2025 17:05:43 +0900
Subject: [PATCH 2/3] update test

---
 pymysql/connections.py |  4 ++++
 tests/test_auth.py     | 28 +++++++++++++++++++++++++++-
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/pymysql/connections.py b/pymysql/connections.py
index 2ddcb3f7..99fcfcd0 100644
--- 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
index e5e2a64e..d7a0e82f 100644
--- 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

From 427ee2ef46e74745c1377822702c5790a7ddb736 Mon Sep 17 00:00:00 2001
From: Inada Naoki <songofacandy@gmail.com>
Date: Wed, 29 Jan 2025 17:56:21 +0900
Subject: [PATCH 3/3] test

---
 .coveragerc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.coveragerc b/.coveragerc
index a9ec9942..efa9a2ff 100644
--- a/.coveragerc
+++ b/.coveragerc
@@ -2,6 +2,7 @@
 branch = True
 source =
     pymysql
+    tests
 omit = pymysql/tests/*
        pymysql/tests/thirdparty/test_MySQLdb/*