fix(connection): unpack_from error caused by network issues (#185)

* fix(connection): Raise InterfaceError for network issue related to closed server side socket (#164)

* add unit test and integration test
* rectify unit test
* rectify integration test
* raise InterfaceError with Broken Pipe for timeout and blocking modes

* fix(connection): add empty buffer check upon connect and unit tests

* fix(connection): unpack_from caused by network issues
* rectify unit tests when network disabled

---------

Co-authored-by: Soksamnang Lim <soklim@amazon.com>

Author: soksamnanglim

redshift_connector/core.py

@@ -739,7 +739,21 @@ def get_calling_module() -> str:
             # Each time will read 5 bytes, the first byte, the code, inform the type of message
             # following 4 bytes inform the message's length
             # then can use this length to minus 4 to get the real data.
-            code, data_len = ci_unpack(self._read(5))
+            buffer = self._read(5)
+
+            if len(buffer) == 0:
+                if self._usock.timeout is not None:
+                    raise InterfaceError(
+                        "BrokenPipe: server socket closed. We noticed a timeout is set for this connection. Consider "
+                        "raising the timeout or defaulting timeout to none."
+                    )
+                else:
+                    raise InterfaceError(
+                        "BrokenPipe: server socket closed. Please check that client side networking configurations such "
+                        "as Proxies, firewalls, VPN, etc. are not affecting your network connection."
+                    )
+
+            code, data_len = ci_unpack(buffer)
             self.message_types[code](self._read(data_len - 4), None)
         if self.error is not None:
             raise self.error
@@ -2007,7 +2021,22 @@ def handle_messages(self: "Connection", cursor: Cursor) -> None:
         code = self.error = None
 
         while code != READY_FOR_QUERY:
-            code, data_len = ci_unpack(self._read(5))
+            buffer = self._read(5)
+
+            if len(buffer) == 0:
+                if self._usock.timeout is not None:
+                    raise InterfaceError(
+                        "BrokenPipe: server socket closed. We noticed a timeout is set for this connection. Consider "
+                        "raising the timeout or defaulting timeout to none."
+                    )
+                else:
+                    raise InterfaceError(
+                        "BrokenPipe: server socket closed. Please check that client side networking configurations such "
+                        "as Proxies, firewalls, VPN, etc. are not affecting your network connection."
+                    )
+
+            code, data_len = ci_unpack(buffer)
+
             self.message_types[code](self._read(data_len - 4), cursor)
 
         if self.error is not None:
@@ -2028,7 +2057,20 @@ def handle_messages_merge_socket_read(self: "Connection", cursor: Cursor):
         """
         code = self.error = None
         # read 5 bytes of message firstly
-        code, data_len = ci_unpack(self._read(5))
+        buffer = self._read(5)
+        if len(buffer) == 0:
+            if self._usock.timeout is not None:
+                raise InterfaceError(
+                    "BrokenPipe: server socket closed. We noticed a timeout is set for this connection. Consider "
+                    "raising the timeout or defaulting timeout to none."
+                )
+            else:
+                raise InterfaceError(
+                    "BrokenPipe: server socket closed. Please check that client side networking configurations such "
+                    "as Proxies, firewalls, VPN, etc. are not affecting your network connection."
+                )
+
+        code, data_len = ci_unpack(buffer)
 
         while True:
             if code == READY_FOR_QUERY:

test/integration/test_connection.py

@@ -2,9 +2,7 @@
 import logging
 import os
 import random
-import socket
 import string
-import struct
 import sys
 
 import pytest  # type: ignore
@@ -178,10 +176,31 @@ def test_broken_pipe(con, db_kwargs):
             pid1 = cur1.fetchone()[0]
 
             cur2.execute("select pg_terminate_backend(%s)", (pid1,))
-            try:
+            with pytest.raises(
+                redshift_connector.InterfaceError,
+                match="BrokenPipe: server socket closed. Please check that client side networking configurations such "
+                "as Proxies, firewalls, VPN, etc. are not affecting your network connection.",
+            ):
+                cur1.execute("select 1")
+
+
+# case 2: same connector configuration, but should throw an error since the timeout is set,
+def test_broken_pipe_timeout(con, db_kwargs):
+    db_kwargs["timeout"] = 60
+    with redshift_connector.connect(**db_kwargs) as db1:
+        with db1.cursor() as cur1, con.cursor() as cur2:
+            print(db1._usock.timeout)
+            cur1.execute("select pg_backend_pid()")
+            pid1 = cur1.fetchone()[0]
+
+            cur2.execute("select pg_terminate_backend(%s)", (pid1,))
+            with pytest.raises(
+                redshift_connector.InterfaceError,
+                match="BrokenPipe: server socket closed. We noticed a timeout is "
+                "set for this connection. Consider raising the timeout or "
+                "defaulting timeout to none.",
+            ):
                 cur1.execute("select 1")
-            except Exception as e:
-                assert isinstance(e, (socket.error, struct.error))
 
 
 def test_application_name_integer(db_kwargs):

test/unit/test_connection.py

@@ -1,9 +1,9 @@
+import socket
 import typing
 from collections import deque
 from decimal import Decimal
-from unittest.mock import patch
-import socket
 from unittest import mock
+from unittest.mock import patch
 
 import pytest  # type: ignore
 
@@ -13,8 +13,8 @@
     Error,
     IntegrityError,
     InterfaceError,
+    OperationalError,
     ProgrammingError,
-    OperationalError
 )
 from redshift_connector.config import (
     ClientProtocolVersion,
@@ -28,7 +28,9 @@
 from redshift_connector.utils.type_utils import py_types as PY_TYPES
 from redshift_connector.utils.type_utils import redshift_types as REDSHIFT_TYPES
 
-test_error_responses_data: typing.List[typing.Tuple[bytes, typing.Dict, typing.Type[Error]]] = [
+test_error_responses_data: typing.List[
+    typing.Tuple[bytes, typing.Dict, typing.Type[Error]]
+] = [
     (
         (
             b"SERROR\x00"
@@ -332,8 +334,141 @@ def test_client_os_version_is_not_present():
     with patch("platform.platform", side_effect=Exception("not for you")):
         assert mock_connection.client_os_version == "unknown"
 
+
 def test_socket_timeout_error():
     with mock.patch('socket.socket.connect') as mock_socket:
         mock_socket.side_effect = (socket.timeout)
         with pytest.raises(OperationalError):
             Connection(user='mock_user', password='mock_password', host='localhost', port=8080, database='mocked')
+
+
+def mock_read(*args, **kwargs):
+    return b""
+
+
+def test_handle_messages_broken_pipe_blocking():
+    # mock the connection and mock the read attribute
+    mock_connection: Connection = Connection.__new__(Connection)
+    mock_connection._read = mock_read
+
+    # we only need to mock the raw socket
+    mock_usock = mock.Mock()
+    mock_usock.timeout = None
+    mock_connection._usock = mock_usock
+
+    mock_cursor: Cursor = Cursor.__new__(Cursor)
+    mock_cursor.ps = None
+
+    with pytest.raises(
+            InterfaceError,
+            match="BrokenPipe: server socket closed. Please check that client side networking configurations such "
+                  "as Proxies, firewalls, VPN, etc. are not affecting your network connection.",
+    ):
+        mock_connection.handle_messages(mock_cursor)
+
+
+def test_handle_messages_broken_pipe_timeout():
+    # mock the connection and mock the read attribute
+    mock_connection: Connection = Connection.__new__(Connection)
+    mock_connection._read = mock_read
+
+    # we only need to mock the raw socket
+    mock_usock = mock.Mock()
+    mock_usock.timeout = 47
+    mock_connection._usock = mock_usock
+
+    mock_cursor: Cursor = Cursor.__new__(Cursor)
+    mock_cursor.ps = None
+
+    with pytest.raises(
+            InterfaceError,
+            match="BrokenPipe: server socket closed. We noticed a timeout is set for this connection. Consider "
+                  "raising the timeout or defaulting timeout to none.",
+    ):
+        mock_connection.handle_messages(mock_cursor)
+
+
+def test_handle_messages_merge_socket_read_broken_pipe_blocking():
+    # mock the connection and mock the read attribute
+    mock_connection: Connection = Connection.__new__(Connection)
+    mock_connection._read = mock_read
+
+    # we only need to mock the raw socket
+    mock_usock = mock.Mock()
+    mock_usock.timeout = None
+    mock_connection._usock = mock_usock
+
+    mock_cursor: Cursor = Cursor.__new__(Cursor)
+    mock_cursor.ps = None
+
+    with pytest.raises(
+            InterfaceError,
+            match="BrokenPipe: server socket closed. Please check that client side networking configurations such "
+                  "as Proxies, firewalls, VPN, etc. are not affecting your network connection.",
+    ):
+        mock_connection.handle_messages_merge_socket_read(mock_cursor)
+
+
+def test_handle_messages_merge_socket_read_broken_pipe_timeout():
+    # mock the connection and mock the read attribute
+    mock_connection: Connection = Connection.__new__(Connection)
+    mock_connection._read = mock_read
+
+    # we only need to mock the raw socket
+    mock_usock = mock.Mock()
+    mock_usock.timeout = 47
+    mock_connection._usock = mock_usock
+
+    mock_cursor: Cursor = Cursor.__new__(Cursor)
+    mock_cursor.ps = None
+
+    with pytest.raises(
+        InterfaceError,
+        match="BrokenPipe: server socket closed. We noticed a timeout is set for this connection. Consider "
+        "raising the timeout or defaulting timeout to none.",
+    ):
+        mock_connection.handle_messages_merge_socket_read(mock_cursor)
+
+
+def test_broken_pipe_on_connect(db_kwargs):
+    db_kwargs["ssl"] = False
+
+    with mock.patch("socket.getaddrinfo") as mock_getaddrinfo:
+        addr_tuple = [(0, 1, 2, "", ('3.226.18.73', 5439)), (2, 1, 6, '', ('3.226.18.73', 5439))]
+        mock_getaddrinfo.return_value = addr_tuple
+        with mock.patch('socket.socket.connect') as mock_usock:
+            mock_usock.side_effect = lambda *args, **kwargs: None
+            with mock.patch("socket.socket.makefile") as mock_sock:
+                mock_file = mock_sock.return_value
+                mock_file._read.return_value = b""
+                with pytest.raises(
+                    InterfaceError,
+                    match="BrokenPipe: server socket closed. Please check that client side networking configurations such "
+                          "as Proxies, firewalls, VPN, etc. are not affecting your network connection.",
+                ):
+                    db_kwargs.pop("region")
+                    db_kwargs.pop("cluster_identifier")
+                    Connection(**db_kwargs)
+
+def test_broken_pipe_timeout_on_connect(db_kwargs):
+    db_kwargs["ssl"] = False
+    db_kwargs["timeout"] = 60
+
+
+    with mock.patch("socket.getaddrinfo") as mock_getaddrinfo:
+        addr_tuple =[(0, 1, 2, "", ('3.226.18.73', 5439)), (2, 1, 6, '', ('3.226.18.73', 5439))]
+        mock_getaddrinfo.return_value= addr_tuple
+        with mock.patch('socket.socket.connect') as mock_usock:
+            mock_usock.side_effect = lambda *args, **kwargs: None
+
+            with mock.patch("socket.socket.makefile") as mock_sock:
+                mock_file = mock_sock.return_value
+                mock_file._read.return_value = b""
+                with pytest.raises(
+                    InterfaceError,
+                    match="BrokenPipe: server socket closed. We noticed a timeout is set for this connection. Consider "
+                    "raising the timeout or defaulting timeout to none.",
+                ):
+                    db_kwargs.pop("region")
+                    db_kwargs.pop("cluster_identifier")
+                    Connection(**db_kwargs)