Fix closure of sockets after failed connection (#794)

Author: robsdedude

neo4j/io/_socket.py

@@ -117,8 +117,7 @@ def sendall(self, data):
         return self._wait_for_io(self._socket.sendall, data)
 
     def close(self):
-        self._socket.shutdown(SHUT_RDWR)
-        self._socket.close()
+        self.close_socket(self._socket)
 
     @classmethod
     def _connect(cls, resolved_address, timeout, keep_alive):
@@ -160,7 +159,7 @@ def _connect(cls, resolved_address, timeout, keep_alive):
             log.debug("[#0000]  C: <ERROR> %s %s", type(error).__name__,
                       " ".join(map(repr, error.args)))
             log.debug("[#0000]  C: <CLOSE> %s", resolved_address)
-            s.close()
+            cls.close_socket(s)
             raise ServiceUnavailable(
                 "Failed to establish connection to {!r} (reason {})".format(
                     resolved_address, error))
@@ -175,6 +174,7 @@ def _secure(cls, s, host, ssl_context):
                 sni_host = host if HAS_SNI and host else None
                 s = ssl_context.wrap_socket(s, server_hostname=sni_host)
             except (OSError, SSLError, CertificateError) as cause:
+                cls.close_socket(s)
                 raise BoltSecurityError(
                     message="Failed to establish encrypted connection.",
                     address=(host, local_port)
@@ -233,20 +233,20 @@ def _handshake(cls, s, resolved_address):
             # If no data is returned after a successful select
             # response, the server has closed the connection
             log.debug("[#%04X]  S: <CLOSE>", local_port)
-            BoltSocket.close_socket(s)
+            cls.close_socket(s)
             raise ServiceUnavailable(
                 "Connection to {address} closed without handshake response".format(
                     address=resolved_address))
         if data_size != 4:
             # Some garbled data has been received
             log.debug("[#%04X]  S: @*#!", local_port)
-            s.close()
+            cls.close_socket(s)
             raise BoltProtocolError(
                 "Expected four byte Bolt handshake response from %r, received %r instead; check for incorrect port number" % (
                 resolved_address, data), address=resolved_address)
         elif data == b"HTTP":
             log.debug("[#%04X]  S: <CLOSE>", local_port)
-            BoltSocket.close_socket(s)
+            cls.close_socket(s)
             raise ServiceUnavailable(
                 "Cannot to connect to Bolt service on {!r} "
                 "(looks like HTTP)".format(resolved_address))
@@ -257,12 +257,14 @@ def _handshake(cls, s, resolved_address):
 
     @classmethod
     def close_socket(cls, socket_):
+        if isinstance(socket_, BoltSocket):
+            socket_ = socket_._socket
         try:
-            if isinstance(socket_, BoltSocket):
-                socket_.close()
-            else:
-                socket_.shutdown(SHUT_RDWR)
-                socket_.close()
+            socket_.shutdown(SHUT_RDWR)
+        except OSError:
+            pass
+        try:
+            socket_.close()
         except OSError:
             pass
 
@@ -296,11 +298,11 @@ def connect(cls, address, *, timeout, custom_resolver, ssl_context,
                 log.debug("[#%04X]  C: <CONNECTION FAILED> %s", local_port,
                           err_str)
                 if s:
-                    BoltSocket.close_socket(s)
+                    cls.close_socket(s)
                 errors.append(error)
             except Exception:
                 if s:
-                    BoltSocket.close_socket(s)
+                    cls.close_socket(s)
                 raise
         if not errors:
             raise ServiceUnavailable(