diff --git a/kubernetes/base/stream/ws_client.py b/kubernetes/base/stream/ws_client.py index bd80480a6d..50d0b31a3a 100644 --- a/kubernetes/base/stream/ws_client.py +++ b/kubernetes/base/stream/ws_client.py @@ -197,8 +197,23 @@ def write_stdin(self, data): """The same as write_channel with channel=0.""" self.write_channel(STDIN_CHANNEL, data) + def _frames_immediately_available(self): + """Return True if at least one more frame is already waiting on + the socket and can be read without blocking.""" + # Prefer poll over select for the same reasons as in update(). + if hasattr(select, "poll"): + poll = select.poll() + poll.register(self.sock.sock, select.POLLIN) + r = poll.poll(0) + poll.unregister(self.sock.sock) + else: + r, _, _ = select.select((self.sock.sock, ), (), (), 0) + return bool(r) + def update(self, timeout=0): - """Update channel buffers with at most one complete frame of input.""" + """Update channel buffers with all complete frames of input that + are available, waiting at most `timeout` seconds for the first + one.""" if not self.is_open(): return if not self.sock.connected: @@ -229,7 +244,7 @@ def update(self, timeout=0): r, _, _ = select.select( (self.sock.sock, ), (), (), timeout) - if r: + while r: op_code, frame = self.sock.recv_data_frame(True) if op_code == ABNF.OPCODE_CLOSE: self._connected = False @@ -264,6 +279,13 @@ def update(self, timeout=0): else: self._channels[channel] += data + # Output larger than the websocket frame size (e.g. long exec + # stdout) arrives as a sequence of frames. Consume every frame + # that is already available before returning, otherwise callers + # reading a channel after a single update() would see the + # output truncated at a frame boundary. + r = self._frames_immediately_available() + def run_forever(self, timeout=None): """Wait till connection is closed or timeout reached. Buffer any input received during this time.""" diff --git a/kubernetes/base/stream/ws_client_test.py b/kubernetes/base/stream/ws_client_test.py index 2099624672..3144fa9945 100644 --- a/kubernetes/base/stream/ws_client_test.py +++ b/kubernetes/base/stream/ws_client_test.py @@ -16,10 +16,11 @@ 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 select import socket import threading import pytest @@ -127,6 +128,139 @@ def test_websocket_proxycare(self): assert dictval(connect_opts, 'http_no_proxy') == expect_noproxy +class WSClientMultiFrameReadTest(unittest.TestCase): + """Tests that reads spanning multiple websocket frames are not + truncated at a frame boundary (issue #2226)""" + + 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) + + def test_read_stdout_returns_all_available_frames(self): + """Verify a single peek/read returns output spanning several frames. + + The server sends long exec stdout as a sequence of websocket + messages (32768 bytes each). Reading only one frame per update() + truncated the output at a frame boundary for callers using the + peek_stdout()/read_stdout() pattern.""" + frame_payload = 32768 + total = 70000 + payload = b'x' * total + + mock_ws = MagicMock() + mock_ws.subprotocol = V5_CHANNEL_PROTOCOL + mock_ws.connected = True + client = self._make_client(mock_ws) + + server_sock, client_sock = socket.socketpair() + try: + # Make sure the whole simulated server output fits in the + # socket buffers so sendall() below cannot block. + client_sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, + 4 * frame_payload) + server_sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, + 4 * frame_payload) + + ws = websocket.WebSocket() + ws.sock = client_sock + ws.connected = True + client.sock = ws + + # The stdout messages (server frames are unmasked), followed by + # a close frame, just like an exec of `cat `. + frames = b'' + for offset in range(0, total, frame_payload): + chunk = (bytes([STDOUT_CHANNEL]) + + payload[offset:offset + frame_payload]) + frames += websocket.ABNF( + 1, 0, 0, 0, websocket.ABNF.OPCODE_BINARY, 0, + chunk).format() + frames += websocket.ABNF( + 1, 0, 0, 0, websocket.ABNF.OPCODE_CLOSE, 0, + b'\x03\xe8').format() + server_sock.sendall(frames) + + out = '' + if client.peek_stdout(timeout=5): + out = client.read_stdout() + + self.assertEqual(len(out), total) + # The close frame was consumed as well + self.assertFalse(client.is_open()) + finally: + server_sock.close() + client_sock.close() + + def test_update_consumes_all_immediately_available_frames(self): + """Verify update() drains every frame that is already readable""" + with patch('select.poll') as mock_poll, \ + patch('select.select') as mock_select: + # The socket reports readable twice, then has no more data. + mock_poll.return_value.poll.side_effect = [ + [(10, select.POLLIN)], [(10, select.POLLIN)], []] + mock_select.side_effect = [ + ([10], [], []), ([10], [], []), ([], [], [])] + + frame1 = MagicMock() + frame1.data = bytes([STDOUT_CHANNEL]) + b'first' + frame2 = MagicMock() + frame2.data = bytes([STDOUT_CHANNEL]) + b'second' + + mock_ws = MagicMock() + mock_ws.subprotocol = V5_CHANNEL_PROTOCOL + mock_ws.connected = True + mock_ws.recv_data_frame.side_effect = [ + (websocket.ABNF.OPCODE_BINARY, frame1), + (websocket.ABNF.OPCODE_BINARY, frame2)] + + client = self._make_client(mock_ws) + client.update(timeout=0) + + self.assertEqual(mock_ws.recv_data_frame.call_count, 2) + self.assertEqual(client._channels.get(STDOUT_CHANNEL), + 'firstsecond') + + def test_update_stops_draining_on_close_frame(self): + """Verify the drain loop terminates when a close frame arrives""" + with patch('select.poll') as mock_poll, \ + patch('select.select') as mock_select: + # The socket always reports readable. + mock_poll.return_value.poll.return_value = [(10, select.POLLIN)] + mock_select.return_value = ([10], [], []) + + frame1 = MagicMock() + frame1.data = bytes([STDOUT_CHANNEL]) + b'output' + close_frame = MagicMock() + close_frame.data = b'\x03\xe8' + + mock_ws = MagicMock() + mock_ws.subprotocol = V5_CHANNEL_PROTOCOL + mock_ws.connected = True + mock_ws.recv_data_frame.side_effect = [ + (websocket.ABNF.OPCODE_BINARY, frame1), + (websocket.ABNF.OPCODE_CLOSE, close_frame)] + + client = self._make_client(mock_ws) + client.update(timeout=0) + + self.assertEqual(mock_ws.recv_data_frame.call_count, 2) + self.assertEqual(client._channels.get(STDOUT_CHANNEL), 'output') + self.assertFalse(client.is_open()) + + class WSClientProtocolTest(unittest.TestCase): """Tests for WSClient V5 protocol handling""" @@ -211,6 +345,7 @@ def test_update_receives_close_v5(self): def test_update_ignores_close_signal_v4(self): """Verify update treats 0xFF as regular data (or ignores signal interpretation) when v4""" with patch.object(ws_client_module, 'create_websocket') as mock_create, \ + patch('select.poll') as mock_poll, \ patch('select.select') as mock_select: mock_ws = MagicMock() @@ -224,7 +359,11 @@ def test_update_ignores_close_signal_v4(self): mock_ws.recv_data_frame.return_value = (websocket.ABNF.OPCODE_BINARY, frame) mock_create.return_value = mock_ws - mock_select.return_value = ([mock_ws.sock], [], []) + # The frame is readable once, then there is no more data. + mock_poll.return_value.poll.side_effect = [ + [(10, select.POLLIN)], []] + mock_select.side_effect = [ + ([mock_ws.sock], [], []), ([], [], [])] client = WSClient(self.config_mock, "ws://test", headers=None, capture_all=True, binary=True) # binary=True to avoid decode errors client.update(timeout=0)