Prevent data loss when attaching to container

The use of buffering within httplib.HTTPResponse can cause data
to be lost. socket.makefile() is called without a bufsize, which
causes a buffer to be used when recieving data. The attach
methods do a HTTP upgrade to tcp before the raw socket is using
to stream data from the container. The problem is that if the
container starts stream data while httplib/http.client is reading
the response to the attach request part of the data ends will end
up in the buffer of fileobject created within the HTTPResponse
object. This data is lost as after the attach request data is
read directly from the raw socket.

Signed-off-by: Chris Harris <chris.harris@kitware.com>

1 files changed, 30 insertions, 3 deletions

diff --git a/docker/transport/unixconn.py b/docker/transport/unixconn.py
index 3565cfb..16e22a8 100644
--- a/docker/transport/unixconn.py
+++ b/docker/transport/unixconn.py
@@ -34,6 +34,25 @@ class UnixHTTPConnection(httplib.HTTPConnection, object):
         self.sock = sock
 
 
+class AttachHTTPResponse(httplib.HTTPResponse):
+    '''
+    A HTTPResponse object that doesn't use a buffered fileobject.
+    '''
+    def __init__(self, sock, *args, **kwargs):
+        # Delegate to super class
+        httplib.HTTPResponse.__init__(self, sock, *args, **kwargs)
+
+        # Override fp with a fileobject that doesn't buffer
+        self.fp = sock.makefile('rb', 0)
+
+
+class AttachUnixHTTPConnection(UnixHTTPConnection):
+    '''
+    A HTTPConnection that returns responses that don't used buffering.
+    '''
+    response_class = AttachHTTPResponse
+
+
 class UnixHTTPConnectionPool(urllib3.connectionpool.HTTPConnectionPool):
     def __init__(self, base_url, socket_path, timeout=60, maxsize=10):
         super(UnixHTTPConnectionPool, self).__init__(
@@ -44,9 +63,17 @@ class UnixHTTPConnectionPool(urllib3.connectionpool.HTTPConnectionPool):
         self.timeout = timeout
 
     def _new_conn(self):
-        return UnixHTTPConnection(
-            self.base_url, self.socket_path, self.timeout
-        )
+        # Special case for attach url, as we do a http upgrade to tcp and
+        # a buffered connection can cause data loss.
+        path = urllib3.util.parse_url(self.base_url).path
+        if path.endswith('attach'):
+            return AttachUnixHTTPConnection(
+                self.base_url, self.socket_path, self.timeout
+            )
+        else:
+            return UnixHTTPConnection(
+                self.base_url, self.socket_path, self.timeout
+            )
 
 
 class UnixAdapter(requests.adapters.HTTPAdapter):