From 7761a8cf3568f25243e7881b315b9fc922237673 Mon Sep 17 00:00:00 2001
From: Tomas Tomecek <ttomecek@redhat.com>
Date: Thu, 10 May 2018 16:10:36 +0200
Subject: [PATCH] asd

---
 docker/api/container.py |  4 +++-
 docker/api/exec_api.py  |  2 +-
 docker/client.py        | 22 ++++++++++++++++------
 docker/utils/socket.py  | 31 +++++++++++++++++++++++++++----
 tests/unit/api_test.py  |  2 +-
 5 files changed, 48 insertions(+), 13 deletions(-)

diff --git a/docker/api/container.py b/docker/api/container.py
index b8507d8..a164e7d 100644
--- a/docker/api/container.py
+++ b/docker/api/container.py
@@ -26,7 +26,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
     def attach_socket(self, container, params=None, ws=False):
diff --git a/docker/api/exec_api.py b/docker/api/exec_api.py
index 6e49996..9468f46 100644
--- a/docker/api/exec_api.py
+++ b/docker/api/exec_api.py
@@ -78,4 +78,4 @@ class ExecApiMixin(object):
 
         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/client.py b/docker/client.py
index 3fa19e0..eb928c6 100644
--- a/docker/client.py
+++ b/docker/client.py
@@ -16,7 +16,7 @@ from .ssladapter import ssladapter
 from .tls import TLSConfig
 from .transport import UnixAdapter
 from .utils import utils, check_resource, update_headers, kwargs_from_env
-from .utils.socket import frames_iter
+from .utils.socket import frames_iter, socket_raw_iter
 try:
     from .transport import NpipeAdapter
 except ImportError:
@@ -317,13 +317,19 @@ class Client(
         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
@@ -353,9 +359,13 @@ class Client(
 
             s.settimeout(None)
 
-    def _get_result(self, container, stream, res):
+    @check_resource
+    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/utils/socket.py b/docker/utils/socket.py
index 164b845..c3a5f90 100644
--- a/docker/utils/socket.py
+++ b/docker/utils/socket.py
@@ -59,7 +59,7 @@ def next_frame_size(socket):
     try:
         data = read_exactly(socket, 8)
     except SocketError:
-        return 0
+        return -1
 
     _, actual = struct.unpack('>BxxxL', data)
     return actual
@@ -69,7 +69,30 @@ def frames_iter(socket):
     """
     Returns a generator of frames read from socket
     """
-    n = next_frame_size(socket)
-    while n > 0:
-        yield read(socket, n)
+    while True:
         n = next_frame_size(socket)
+        if n < 0:
+            break
+        while n > 0:
+            result = read(socket, n)
+            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/unit/api_test.py b/tests/unit/api_test.py
index 8faca6b..3aeff85 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()
 
 url_base = '{0}/'.format(fake_api.prefix)
-- 
2.17.0