Fix kubernetes port forwarding from hanging in certain edge cases

iciclespider · iciclespider · commit 08b41ad806bc · 2022-04-29T21:40:32.000-10:00
It turns out SSL sockets can buffer data such that the select method is not aware of it. See: https://docs.python.org/3/library/ssl.html#notes-on-non-blocking-sockets https://docs.python.org/3/library/ssl.html#ssl.SSLSocket.pending
diff --git a/kubernetes/base/stream/ws_client.py b/kubernetes/base/stream/ws_client.py
@@ -370,34 +370,38 @@ def _proxy(self):
             r, w, _ = select.select(rlist, wlist, [])
             for sock in r:
                 if sock == self.websocket:
-                    opcode, frame = self.websocket.recv_data_frame(True)
-                    if opcode == ABNF.OPCODE_BINARY:
-                        if not frame.data:
-                            raise RuntimeError("Unexpected frame data size")
-                        channel = six.byte2int(frame.data)
-                        if channel >= len(channel_ports):
-                            raise RuntimeError("Unexpected channel number: %s" % channel)
-                        port = channel_ports[channel]
-                        if channel_initialized[channel]:
-                            if channel % 2:
-                                if port.error is None:
-                                    port.error = ''
-                                port.error += frame.data[1:].decode()
+                    pending = True
+                    while pending:
+                        opcode, frame = self.websocket.recv_data_frame(True)
+                        if opcode == ABNF.OPCODE_BINARY:
+                            if not frame.data:
+                                raise RuntimeError("Unexpected frame data size")
+                            channel = six.byte2int(frame.data)
+                            if channel >= len(channel_ports):
+                                raise RuntimeError("Unexpected channel number: %s" % channel)
+                            port = channel_ports[channel]
+                            if channel_initialized[channel]:
+                                if channel % 2:
+                                    if port.error is None:
+                                        port.error = ''
+                                    port.error += frame.data[1:].decode()
+                                else:
+                                    port.data += frame.data[1:]
                             else:
-                                port.data += frame.data[1:]
-                        else:
-                            if len(frame.data) != 3:
-                                raise RuntimeError(
-                                    "Unexpected initial channel frame data size"
-                                )
-                            port_number = six.byte2int(frame.data[1:2]) + (six.byte2int(frame.data[2:3]) * 256)
-                            if port_number != port.port_number:
-                                raise RuntimeError(
-                                    "Unexpected port number in initial channel frame: %s" % port_number
-                                )
-                            channel_initialized[channel] = True
-                    elif opcode not in (ABNF.OPCODE_PING, ABNF.OPCODE_PONG, ABNF.OPCODE_CLOSE):
-                        raise RuntimeError("Unexpected websocket opcode: %s" % opcode)
+                                if len(frame.data) != 3:
+                                    raise RuntimeError(
+                                        "Unexpected initial channel frame data size"
+                                    )
+                                port_number = six.byte2int(frame.data[1:2]) + (six.byte2int(frame.data[2:3]) * 256)
+                                if port_number != port.port_number:
+                                    raise RuntimeError(
+                                        "Unexpected port number in initial channel frame: %s" % port_number
+                                    )
+                                channel_initialized[channel] = True
+                        elif opcode not in (ABNF.OPCODE_PING, ABNF.OPCODE_PONG, ABNF.OPCODE_CLOSE):
+                            raise RuntimeError("Unexpected websocket opcode: %s" % opcode)
+                        if not isinstance(self.websocket.sock, ssl.SSLSocket) or not self.websocket.sock.pending():
+                            pending = False
                 else:
                     port = local_ports[sock]
                     data = port.python.recv(1024 * 1024)
diff --git a/kubernetes/e2e_test/test_client.py b/kubernetes/e2e_test/test_client.py
@@ -233,7 +233,7 @@ def test_exit_code(self):
     # Skipping this test as this flakes a lot
     # See: https://github.com/kubernetes-client/python/issues/1300
     # Re-enable the test once the flakiness is investigated
-    @unittest.skip("skipping due to extreme flakiness")
+    #@unittest.skip("skipping due to extreme flakiness")
     def test_portforward_raw(self):
         client = api_client.ApiClient(configuration=self.config)
         api = core_v1_api.CoreV1Api(client)
