Fix handling of non-multiplexed (TTY) streams over upgraded socketsfix_upgraded_tty_streams

Signed-off-by: Joffrey F <joffrey@docker.com>

1 files changed, 20 insertions, 1 deletions

diff --git a/docker/utils/socket.py b/docker/utils/socket.py
index 4080f25..54392d2 100644
--- a/docker/utils/socket.py
+++ b/docker/utils/socket.py
@@ -75,5 +75,24 @@ def frames_iter(socket):
             break
         while n > 0:
             result = read(socket, n)
-            n -= len(result)
+            if result is None:
+                continue
+            data_length = len(result)
+            if data_length == 0:
+                # We have reached EOF
+                return
+            n -= data_length
             yield result
+
+
+def socket_raw_iter(socket):
+    """
+    Returns a generator of data read from the socket.
+    This is used for non-multiplexed streams.
+    """
+    while True:
+        result = read(socket)
+        if len(result) == 0:
+            # We have reached EOF
+            return
+        yield result