Merge pull request #2651 from jojinkb/fix-ws-stdout-truncation

Fix long exec stdout truncated at websocket frame boundaries

Author: kubernetes-prow[bot]

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:
@@ -231,7 +246,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
@@ -266,6 +281,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."""

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 .ws_client import STDOUT_CHANNEL
 from kubernetes.client.configuration import Configuration
 import os
+import select
 import socket
 import threading
 import pytest
@@ -128,6 +129,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 <long file>`.
+            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"""
 
@@ -213,6 +347,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()
@@ -227,7 +362,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)