close connection._rfile in Connection._force_close (PyMySQL#1184)

cfbolz · Yevhenii Yatchenko · commit 9f461516804b · 2024-09-20T13:17:52.000+03:00
fix PyMySQL#1183.
diff --git a/pymysql/connections.py b/pymysql/connections.py
@@ -176,6 +176,7 @@ class Connection(object):
     """
 
     _sock = None
+    _rfile = None
     _auth_plugin_name = ''
     _closed = False
     _secure = False
@@ -372,6 +373,8 @@ def open(self):
 
     def _force_close(self):
         """Close connection without QUIT message"""
+        if self._rfile:
+            self._rfile.close()
         if self._sock:
             try:
                 self._sock.close()
@@ -600,12 +603,7 @@ def connect(self, sock=None):
             if self.autocommit_mode is not None:
                 self.autocommit(self.autocommit_mode)
         except BaseException as e:
-            self._rfile = None
-            if sock is not None:
-                try:
-                    sock.close()
-                except:  # noqa
-                    pass
+            self._force_close()
 
             if isinstance(e, (OSError, IOError, socket.error)):
                 exc = err.OperationalError(
diff --git a/pymysql/tests/test_connection.py b/pymysql/tests/test_connection.py
@@ -565,3 +565,40 @@ def test_commit_during_multi_result(self):
         con.commit()
         cur.execute("SELECT 3")
         self.assertEqual(cur.fetchone()[0], 3)
+
+    def test_force_close_closes_socketio(self):
+        con = self.connect()
+        sock = con._sock
+        fileno = sock.fileno()
+        rfile = con._rfile
+
+        con._force_close()
+        assert rfile.closed
+        assert sock._closed
+        assert sock.fileno() != fileno  # should be set to -1
+
+    def test_socket_closed_on_exception_in_connect(self):
+        con = self.connect(defer_connect=True)
+        sock = None
+        rfile = None
+        fileno = -1
+
+        def _request_authentication():
+            nonlocal sock, rfile, fileno
+            sock = con._sock
+            assert sock is not None
+            fileno = sock.fileno()
+            rfile = con._rfile
+            assert rfile is not None
+            raise TypeError
+
+        con._request_authentication = _request_authentication
+
+        with pytest.raises(TypeError):
+            con.connect()
+        assert not con.open
+        assert con._rfile is None
+        assert con._sock is None
+        assert rfile.closed
+        assert sock._closed
+        assert sock.fileno() != fileno  # should be set to -1
