diff --git a/kubernetes/base/stream/ws_client.py b/kubernetes/base/stream/ws_client.py index bd80480a6d..68b0f892f7 100644 --- a/kubernetes/base/stream/ws_client.py +++ b/kubernetes/base/stream/ws_client.py @@ -205,6 +205,16 @@ def update(self, timeout=0): self._connected = False return + # SSL sockets decrypt an entire TLS record at a time, so a previous + # recv_data_frame() call may have pulled the bytes of the next frames + # off the socket already: those frames sit decrypted in the + # SSLSocket's internal buffer, invisible to select()/poll() on the + # underlying socket. Without this check the buffered frames would + # only be delivered once new data arrives on the socket, and would be + # lost if it never does. + if (isinstance(self.sock.sock, ssl.SSLSocket) + and self.sock.sock.pending()): + r = True # The options here are: # select.select() - this will work on most OS, however, it has a # limitation of only able to read fd numbers up to 1024. @@ -214,7 +224,7 @@ def update(self, timeout=0): # efficient as epoll. Will work for fd numbers above 1024. # select.epoll() - newest and most efficient way of polling. # However, only works on linux. - if hasattr(select, "poll"): + elif hasattr(select, "poll"): poll = select.poll() poll.register(self.sock.sock, select.POLLIN) if timeout is not None and timeout != float("inf"): diff --git a/kubernetes/base/stream/ws_client_test.py b/kubernetes/base/stream/ws_client_test.py index 2099624672..740625d73b 100644 --- a/kubernetes/base/stream/ws_client_test.py +++ b/kubernetes/base/stream/ws_client_test.py @@ -16,11 +16,12 @@ from unittest.mock import MagicMock, patch from . import ws_client as ws_client_module -from .ws_client import get_websocket_url, WSClient, V5_CHANNEL_PROTOCOL, V4_CHANNEL_PROTOCOL, CLOSE_CHANNEL, STDIN_CHANNEL +from .ws_client import get_websocket_url, WSClient, V5_CHANNEL_PROTOCOL, V4_CHANNEL_PROTOCOL, CLOSE_CHANNEL, STDIN_CHANNEL, STDOUT_CHANNEL from .ws_client import websocket_proxycare from kubernetes.client.configuration import Configuration import os import socket +import ssl import threading import pytest from kubernetes import stream, client, config @@ -385,6 +386,108 @@ def test_readline_channel_returns_empty_bytes_on_expired_timeout(self): self.assertEqual(line, b"") +class WSClientUpdateTest(unittest.TestCase): + """Tests for WSClient.update() frame consumption (issue #2375)""" + + def setUp(self): + # Mock configuration to avoid real connections in WSClient.__init__ + self.config_mock = MagicMock() + self.config_mock.assert_hostname = False + self.config_mock.api_key = {} + self.config_mock.proxy = None + self.config_mock.ssl_ca_cert = None + self.config_mock.cert_file = None + self.config_mock.key_file = None + self.config_mock.verify_ssl = True + + def _make_client(self, mock_ws): + with patch.object(ws_client_module, 'create_websocket') as mock_create: + mock_create.return_value = mock_ws + return WSClient(self.config_mock, "wss://test", headers=None, + capture_all=True, binary=True) + + def test_update_reads_frames_pending_in_ssl_buffer(self): + """Verify update reads a frame buffered inside the SSL socket even + when the underlying socket does not report as readable. + + SSL sockets decrypt a whole TLS record at a time, so frames that + share a TLS record with a previously read frame sit decrypted in + the SSLSocket's buffer where select()/poll() cannot see them.""" + with patch('select.poll') as mock_poll, \ + patch('select.select') as mock_select: + # Nothing is readable on the underlying socket. + mock_poll.return_value.poll.return_value = [] + mock_select.return_value = ([], [], []) + + mock_ws = MagicMock() + mock_ws.subprotocol = V4_CHANNEL_PROTOCOL + mock_ws.connected = True + # A decrypted frame is waiting inside the SSL socket. + mock_ws.sock = MagicMock(spec=ssl.SSLSocket) + mock_ws.sock.pending.return_value = 6 + frame = MagicMock() + frame.data = bytes([STDOUT_CHANNEL]) + b'hello' + mock_ws.recv_data_frame.return_value = (websocket.ABNF.OPCODE_BINARY, frame) + + client = self._make_client(mock_ws) + client.update(timeout=0) + + self.assertEqual(client._channels.get(STDOUT_CHANNEL), b'hello') + + def test_update_polls_when_no_ssl_data_pending(self): + """Verify update falls back to poll/select when the SSL socket has no + buffered data""" + with patch('select.poll') as mock_poll, \ + patch('select.select') as mock_select: + mock_poll.return_value.poll.return_value = [] + mock_select.return_value = ([], [], []) + + mock_ws = MagicMock() + mock_ws.subprotocol = V4_CHANNEL_PROTOCOL + mock_ws.connected = True + mock_ws.sock = MagicMock(spec=ssl.SSLSocket) + mock_ws.sock.pending.return_value = 0 + + client = self._make_client(mock_ws) + client.update(timeout=0) + + mock_ws.recv_data_frame.assert_not_called() + + def test_update_receives_fragmented_message(self): + """Verify a message fragmented into continuation frames is delivered + in full. + + websocket-client reassembles continuation (OPCODE_CONT) frames inside + recv_data_frame() and returns the opcode of the initial frame, so + update() must buffer the complete message.""" + mock_ws = MagicMock() + mock_ws.subprotocol = V4_CHANNEL_PROTOCOL + mock_ws.connected = True + client = self._make_client(mock_ws) + + server_sock, client_sock = socket.socketpair() + try: + ws = websocket.WebSocket() + ws.sock = client_sock + ws.connected = True + client.sock = ws + + # A stdout message fragmented into an initial data frame (fin=0) + # and a continuation frame (fin=1); server frames are unmasked. + initial = websocket.ABNF(0, 0, 0, 0, websocket.ABNF.OPCODE_BINARY, + 0, bytes([STDOUT_CHANNEL]) + b'A' * 10) + cont = websocket.ABNF(1, 0, 0, 0, websocket.ABNF.OPCODE_CONT, + 0, b'B' * 10) + server_sock.sendall(initial.format() + cont.format()) + + client.update(timeout=5) + + self.assertEqual(client._channels.get(STDOUT_CHANNEL), + b'A' * 10 + b'B' * 10) + finally: + server_sock.close() + client_sock.close() + @pytest.fixture(scope="module") def dummy_proxy():