From 7519b186b353d7053ea93eb735c8914cf8bb79ba 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

---
 aiomysql/connection.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/aiomysql/connection.py b/aiomysql/connection.py
index f2b139b1..4b2befcc 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)
@@ -798,7 +802,7 @@ async def _request_authentication(self):
                 )
             # Else: empty password
         elif auth_plugin == 'sha256_password':
-            if self._ssl_context and self.server_capabilities & CLIENT.SSL:
+            if self._secure:
                 authresp = self._password.encode('latin1') + b'\0'
             elif self._password:
                 authresp = b'\1'  # request public key
@@ -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)