Fix WSClient.update dropping frames buffered in the SSL socket

WSClient.update() waits for the underlying socket to become readable
via poll()/select() before reading a frame. SSL sockets decrypt an
entire TLS record at a time, so when several websocket frames arrive
in one TLS record, the first recv_data_frame() call consumes the whole
record from the socket and the remaining frames sit decrypted inside
the SSLSocket's internal buffer, where poll()/select() cannot see
them. Those frames are only delivered once new data arrives on the
connection and are lost if it never does, which randomly truncates
large outputs read through the stream API.

Check SSLSocket.pending() before polling, mirroring what
PortForward._proxy() already does, so buffered frames are consumed
without waiting for socket readability.

The OPCODE_CONT handling suggested in the issue is not needed:
websocket-client reassembles continuation frames inside
recv_data_frame() and returns the opcode of the initial frame, so
update() never observes OPCODE_CONT. A regression test documents that
fragmented messages are delivered in full.

Signed-off-by: Jojin <[email protected]>

Author: jojinkb

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"):

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