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

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

6 files changed, 42 insertions, 11 deletions

diff --git a/docker/api/client.py b/docker/api/client.py
index 65b5baa..1de10c7 100644
--- a/docker/api/client.py
+++ b/docker/api/client.py
@@ -32,7 +32,7 @@ from ..errors import (
 from ..tls import TLSConfig
 from ..transport import SSLAdapter, UnixAdapter
 from ..utils import utils, check_resource, update_headers
-from ..utils.socket import frames_iter
+from ..utils.socket import frames_iter, socket_raw_iter
 from ..utils.json_stream import json_stream
 try:
     from ..transport import NpipeAdapter
@@ -362,13 +362,19 @@ class APIClient(
         for out in response.iter_content(chunk_size=1, decode_unicode=True):
             yield out
 
-    def _read_from_socket(self, response, stream):
+    def _read_from_socket(self, response, stream, tty=False):
         socket = self._get_raw_response_socket(response)
 
+        gen = None
+        if tty is False:
+            gen = frames_iter(socket)
+        else:
+            gen = socket_raw_iter(socket)
+
         if stream:
-            return frames_iter(socket)
+            return gen
         else:
-            return six.binary_type().join(frames_iter(socket))
+            return six.binary_type().join(gen)
 
     def _disable_socket_timeout(self, socket):
         """ Depending on the combination of python version and whether we're
@@ -398,9 +404,13 @@ class APIClient(
 
             s.settimeout(None)
 
-    def _get_result(self, container, stream, res):
+    @check_resource('container')
+    def _check_is_tty(self, container):
         cont = self.inspect_container(container)
-        return self._get_result_tty(stream, res, cont['Config']['Tty'])
+        return cont['Config']['Tty']
+
+    def _get_result(self, container, stream, res):
+        return self._get_result_tty(stream, res, self._check_is_tty(container))
 
     def _get_result_tty(self, stream, res, is_tty):
         # Stream multi-plexing was only introduced in API v1.6. Anything
diff --git a/docker/api/container.py b/docker/api/container.py
index 06c575d..dde1325 100644
--- a/docker/api/container.py
+++ b/docker/api/container.py
@@ -52,7 +52,9 @@ class ContainerApiMixin(object):
         u = self._url("/containers/{0}/attach", container)
         response = self._post(u, headers=headers, params=params, stream=stream)
 
-        return self._read_from_socket(response, stream)
+        return self._read_from_socket(
+            response, stream, self._check_is_tty(container)
+        )
 
     @utils.check_resource('container')
     def attach_socket(self, container, params=None, ws=False):
diff --git a/docker/api/exec_api.py b/docker/api/exec_api.py
index 2b407ce..6f42524 100644
--- a/docker/api/exec_api.py
+++ b/docker/api/exec_api.py
@@ -153,4 +153,4 @@ class ExecApiMixin(object):
             return self._result(res)
         if socket:
             return self._get_raw_response_socket(res)
-        return self._read_from_socket(res, stream)
+        return self._read_from_socket(res, stream, tty)
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
diff --git a/tests/integration/api_build_test.py b/tests/integration/api_build_test.py
index 209c1f2..d0aa5c2 100644
--- a/tests/integration/api_build_test.py
+++ b/tests/integration/api_build_test.py
@@ -244,8 +244,8 @@ class BuildTest(BaseAPIIntegrationTest):
         with pytest.raises(errors.NotFound):
             self.client.inspect_image('dockerpytest_nonebuild')
 
+    @requires_experimental(until=None)
     @requires_api_version('1.25')
-    @requires_experimental
     def test_build_squash(self):
         script = io.BytesIO('\n'.join([
             'FROM busybox',
diff --git a/tests/unit/api_test.py b/tests/unit/api_test.py
index 83848c5..6ac92c4 100644
--- a/tests/unit/api_test.py
+++ b/tests/unit/api_test.py
@@ -83,7 +83,7 @@ def fake_delete(self, url, *args, **kwargs):
     return fake_request('DELETE', url, *args, **kwargs)
 
 
-def fake_read_from_socket(self, response, stream):
+def fake_read_from_socket(self, response, stream, tty=False):
     return six.binary_type()