diff --git a/src/prompt_toolkit/contrib/telnet/server.py b/src/prompt_toolkit/contrib/telnet/server.py index 5347a3f98..2b75a96db 100644 --- a/src/prompt_toolkit/contrib/telnet/server.py +++ b/src/prompt_toolkit/contrib/telnet/server.py @@ -177,23 +177,33 @@ def ttype_received(ttype: str) -> None: self.parser = TelnetProtocolParser(data_received, size_received, ttype_received) self.context: contextvars.Context | None = None + def _handle_incoming_data(self) -> None: + """Read and process data that is ready on the client socket.""" + try: + data = self.conn.recv(1024) + except OSError as e: + # A reset or another socket error means that the client disconnected. + logger.info( + "Connection closed by client. {!r} {!r}: {}".format(*self.addr, e) + ) + self.close() + return + + if data: + self.feed(data) + else: + # Connection closed by client. + logger.info("Connection closed by client. {!r} {!r}".format(*self.addr)) + self.close() + async def run_application(self) -> None: """ Run application. """ - def handle_incoming_data() -> None: - data = self.conn.recv(1024) - if data: - self.feed(data) - else: - # Connection closed by client. - logger.info("Connection closed by client. {!r} {!r}".format(*self.addr)) - self.close() - # Add reader. loop = get_running_loop() - loop.add_reader(self.conn, handle_incoming_data) + loop.add_reader(self.conn, self._handle_incoming_data) try: # Wait for v100_output to be properly instantiated diff --git a/tests/test_telnet.py b/tests/test_telnet.py new file mode 100644 index 000000000..ff3c644cb --- /dev/null +++ b/tests/test_telnet.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +from unittest.mock import Mock + +import prompt_toolkit.contrib.telnet.server as telnet_server +from prompt_toolkit.contrib.telnet.server import TelnetConnection + + +def _create_connection() -> TelnetConnection: + connection = object.__new__(TelnetConnection) + connection.conn = Mock() + connection.addr = ("127.0.0.1", 1234) + connection.parser = Mock() + connection.vt100_input = Mock() + connection.stdout = Mock() + connection._closed = False + return connection + + +def test_handle_incoming_data() -> None: + connection = _create_connection() + connection.conn.recv.return_value = b"hello" + + connection._handle_incoming_data() + + connection.parser.feed.assert_called_once_with(b"hello") + assert not connection._closed + + +def test_handle_incoming_data_eof(monkeypatch) -> None: + connection = _create_connection() + connection.conn.recv.return_value = b"" + loop = Mock() + monkeypatch.setattr(telnet_server, "get_running_loop", lambda: loop) + + connection._handle_incoming_data() + + assert connection._closed + loop.remove_reader.assert_called_once_with(connection.conn) + connection.conn.close.assert_called_once_with() + + +def test_handle_incoming_data_connection_reset(monkeypatch) -> None: + connection = _create_connection() + connection.conn.recv.side_effect = ConnectionResetError(104, "Connection reset") + loop = Mock() + monkeypatch.setattr(telnet_server, "get_running_loop", lambda: loop) + + connection._handle_incoming_data() + + assert connection._closed + loop.remove_reader.assert_called_once_with(connection.conn) + connection.conn.close.assert_called_once_with()