Handle telnet client socket resets gracefully

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

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()