diff --git a/src/epplibwrapper/client.py b/src/epplibwrapper/client.py index a7856298b..9d203b246 100644 --- a/src/epplibwrapper/client.py +++ b/src/epplibwrapper/client.py @@ -1,6 +1,7 @@ """Provide a wrapper around epplib to handle authentication and errors.""" import logging +from gevent.lock import BoundedSemaphore try: from epplib.client import Client @@ -52,10 +53,16 @@ class EPPLibWrapper: "urn:ietf:params:xml:ns:contact-1.0", ], ) + # We should only ever have one active connection at a time + self.connection_lock = BoundedSemaphore(1) + + self.connection_lock.acquire() try: self._initialize_client() except Exception: - logger.warning("Unable to configure epplib. Registrar cannot contact registry.") + logger.warning("Unable to configure the connection to the registry.") + finally: + self.connection_lock.release() def _initialize_client(self) -> None: """Initialize a client, assuming _login defined. Sets _client to initialized @@ -74,11 +81,7 @@ class EPPLibWrapper: ) try: # use the _client object to connect - self._client.connect() # type: ignore - response = self._client.send(self._login) # type: ignore - if response.code >= 2000: # type: ignore - self._client.close() # type: ignore - raise LoginError(response.msg) # type: ignore + self._connect() except TransportError as err: message = "_initialize_client failed to execute due to a connection error." logger.error(f"{message} Error: {err}") @@ -90,13 +93,33 @@ class EPPLibWrapper: logger.error(f"{message} Error: {err}") raise RegistryError(message) from err + def _connect(self) -> None: + """Connects to EPP. Sends a login command. If an invalid response is returned, + the client will be closed and a LoginError raised.""" + self._client.connect() # type: ignore + response = self._client.send(self._login) # type: ignore + if response.code >= 2000: # type: ignore + self._client.close() # type: ignore + raise LoginError(response.msg) # type: ignore + def _disconnect(self) -> None: - """Close the connection.""" + """Close the connection. Sends a logout command and closes the connection.""" + self._send_logout_command() + self._close_client() + + def _send_logout_command(self): + """Sends a logout command to epp""" try: self._client.send(commands.Logout()) # type: ignore - self._client.close() # type: ignore - except Exception: - logger.warning("Connection to registry was not cleanly closed.") + except Exception as err: + logger.warning(f"Logout command not sent successfully: {err}") + + def _close_client(self): + """Closes an active client connection""" + try: + self._client.close() + except Exception as err: + logger.warning(f"Connection to registry was not cleanly closed: {err}") def _send(self, command): """Helper function used by `send`.""" @@ -146,6 +169,8 @@ class EPPLibWrapper: cmd_type = command.__class__.__name__ if not cleaned: raise ValueError("Please sanitize user input before sending it.") + + self.connection_lock.acquire() try: return self._send(command) except RegistryError as err: @@ -161,6 +186,8 @@ class EPPLibWrapper: return self._retry(command) else: raise err + finally: + self.connection_lock.release() try: diff --git a/src/epplibwrapper/tests/test_client.py b/src/epplibwrapper/tests/test_client.py index f95b37dcd..d30ec4865 100644 --- a/src/epplibwrapper/tests/test_client.py +++ b/src/epplibwrapper/tests/test_client.py @@ -1,5 +1,9 @@ +import datetime +from dateutil.tz import tzlocal # type: ignore from unittest.mock import MagicMock, patch +from pathlib import Path from django.test import TestCase +from gevent.exceptions import ConcurrentObjectUseError from epplibwrapper.client import EPPLibWrapper from epplibwrapper.errors import RegistryError, LoginError from .common import less_console_noise @@ -8,6 +12,9 @@ import logging try: from epplib.exceptions import TransportError from epplib.responses import Result + from epplib.transport import SocketTransport + from epplib import commands + from epplib.models import common, info except ImportError: pass @@ -255,3 +262,116 @@ class TestClient(TestCase): mock_close.assert_called_once() # send() is called 5 times: send(login), send(command) fail, send(logout), send(login), send(command) self.assertEquals(mock_send.call_count, 5) + + def fake_failure_send_concurrent_threads(self, command=None, cleaned=None): + """ + Raises a ConcurrentObjectUseError, which gevent throws when accessing + the same thread from two different locations. + """ + # This error is thrown when two threads are being used concurrently + raise ConcurrentObjectUseError("This socket is already used by another greenlet") + + def do_nothing(self, command=None): + """ + A placeholder method that performs no action. + """ + pass # noqa + + def fake_success_send(self, command=None, cleaned=None): + """ + Simulates receiving a success response from EPP. + """ + mock = MagicMock( + code=1000, + msg="Command completed successfully", + res_data=None, + cl_tr_id="xkw1uo#2023-10-17T15:29:09.559376", + sv_tr_id="5CcH4gxISuGkq8eqvr1UyQ==-35a", + extensions=[], + msg_q=None, + ) + return mock + + def fake_info_domain_received(self, command=None, cleaned=None): + """ + Simulates receiving a response by reading from a predefined XML file. + """ + location = Path(__file__).parent / "utility" / "infoDomain.xml" + xml = (location).read_bytes() + return xml + + def get_fake_epp_result(self): + """Mimics a return from EPP by returning a dictionary in the same format""" + result = { + "cl_tr_id": None, + "code": 1000, + "extensions": [], + "msg": "Command completed successfully", + "msg_q": None, + "res_data": [ + info.InfoDomainResultData( + roid="DF1340360-GOV", + statuses=[ + common.Status( + state="serverTransferProhibited", + description=None, + lang="en", + ), + common.Status(state="inactive", description=None, lang="en"), + ], + cl_id="gov2023-ote", + cr_id="gov2023-ote", + cr_date=datetime.datetime(2023, 8, 15, 23, 56, 36, tzinfo=tzlocal()), + up_id="gov2023-ote", + up_date=datetime.datetime(2023, 8, 17, 2, 3, 19, tzinfo=tzlocal()), + tr_date=None, + name="test3.gov", + registrant="TuaWnx9hnm84GCSU", + admins=[], + nsset=None, + keyset=None, + ex_date=datetime.date(2024, 8, 15), + auth_info=info.DomainAuthInfo(pw="2fooBAR123fooBaz"), + ) + ], + "sv_tr_id": "wRRNVhKhQW2m6wsUHbo/lA==-29a", + } + return result + + def test_send_command_close_failure_recovers(self): + """ + Validates the resilience of the connection handling mechanism + during command execution on retry. + + Scenario: + - Initialization of the connection is successful. + - An attempt to send a command fails with a specific error code (ConcurrentObjectUseError) + - The client attempts to retry. + - Subsequently, the client re-initializes the connection. + - A retry of the command execution post-reinitialization succeeds. + """ + + expected_result = self.get_fake_epp_result() + wrapper = None + # Trigger a retry + # Do nothing on connect, as we aren't testing it and want to connect while + # mimicking the rest of the client as closely as possible (which is not entirely possible with MagicMock) + with patch.object(EPPLibWrapper, "_connect", self.do_nothing): + with patch.object(SocketTransport, "send", self.fake_failure_send_concurrent_threads): + wrapper = EPPLibWrapper() + tested_command = commands.InfoDomain(name="test.gov") + try: + wrapper.send(tested_command, cleaned=True) + except RegistryError as err: + expected_error = "InfoDomain failed to execute due to an unknown error." + self.assertEqual(err.args[0], expected_error) + else: + self.fail("Registry error was not thrown") + + # After a retry, try sending again to see if the connection recovers + with patch.object(EPPLibWrapper, "_connect", self.do_nothing): + with patch.object(SocketTransport, "send", self.fake_success_send), patch.object( + SocketTransport, "receive", self.fake_info_domain_received + ): + result = wrapper.send(tested_command, cleaned=True) + self.assertEqual(expected_result, result.__dict__)