From 8bb1e6f117bb21440ea238160e4ba641389e71bf Mon Sep 17 00:00:00 2001
From: Richard Schwab <gitrichardschwab-7a2qxq42kj@central-intelligence.agency>
Date: Fri, 28 Jan 2022 17:50:43 +0100
Subject: [PATCH] implement a shortcut for determining secure connections, now
 supporting unix sockets ports https://github.com/PyMySQL/PyMySQL/pull/696
 ports check for server tls support for secure connections from
 https://github.com/PyMySQL/PyMySQL/pull/353

---
 CHANGES.txt            |  1 +
 aiomysql/connection.py | 12 ++++++++----
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/CHANGES.txt b/CHANGES.txt
index b5b276b1..b537f279 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -10,6 +10,7 @@ To be included in 1.0.0 (unreleased)
 * Fix timed out MySQL 8.0 connections being returned from Pool #660
 * Ensure connections are properly closed before raising an OperationalError when the server connection is lost #660
 * Ensure connections are properly closed before raising an InternalError when packet sequence numbers are out of sync #660
+* Unix sockets are now internally considered secure, allowing sha256_password and caching_sha2_password auth methods to be used #695
 
 
 0.0.22 (2021-11-14)
diff --git a/aiomysql/connection.py b/aiomysql/connection.py
index f2b139b1..1bd279d3 100644
--- a/aiomysql/connection.py
+++ b/aiomysql/connection.py
@@ -229,6 +229,7 @@ def __init__(self, host="localhost", user=None, password="",
         self._client_auth_plugin = auth_plugin
         self._server_auth_plugin = ""
         self._auth_plugin_used = ""
+        self._secure = False
         self.server_public_key = server_public_key
         self.salt = None
 
@@ -526,7 +527,7 @@ async def _connect(self):
         # raise OperationalError(CR.CR_SERVER_GONE_ERROR,
         # "MySQL server has gone away (%r)" % (e,))
         try:
-            if self._unix_socket and self._host in ('localhost', '127.0.0.1'):
+            if self._unix_socket:
                 self._reader, self._writer = await \
                     asyncio.wait_for(
                         _open_unix_connection(
@@ -534,6 +535,7 @@ async def _connect(self):
                         timeout=self.connect_timeout)
                 self.host_info = "Localhost via UNIX socket: " + \
                                  self._unix_socket
+                self._secure = True
             else:
                 self._reader, self._writer = await \
                     asyncio.wait_for(
@@ -743,7 +745,7 @@ async def _request_authentication(self):
         if self.user is None:
             raise ValueError("Did not specify a username")
 
-        if self._ssl_context:
+        if self._ssl_context and self.server_capabilities & CLIENT.SSL:
             # capablities, max packet, charset
             data = struct.pack('<IIB', self.client_flag, 16777216, 33)
             data += b'\x00' * (32 - len(data))
@@ -770,6 +772,8 @@ async def _request_authentication(self):
                 server_hostname=self._host
             )
 
+            self._secure = True
+
         charset_id = charset_by_name(self.charset).id
         if isinstance(self.user, str):
             _user = self.user.encode(self.encoding)
@@ -960,7 +964,7 @@ async def caching_sha2_password_auth(self, pkt):
 
         logger.debug("caching sha2: Trying full auth...")
 
-        if self._ssl_context:
+        if self._secure:
             logger.debug("caching sha2: Sending plain "
                          "password via secure connection")
             self.write_packet(self._password.encode('latin1') + b'\0')
@@ -991,7 +995,7 @@ async def caching_sha2_password_auth(self, pkt):
         pkt.check_error()
 
     async def sha256_password_auth(self, pkt):
-        if self._ssl_context:
+        if self._secure:
             logger.debug("sha256: Sending plain password")
             data = self._password.encode('latin1') + b'\0'
             self.write_packet(data)