Disable buffering based on presence of Connection Upgrade headerscjh1-attach-buf-fix

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

1 files changed, 25 insertions, 27 deletions

diff --git a/docker/transport/unixconn.py b/docker/transport/unixconn.py
index 16e22a8..7cb8771 100644
--- a/docker/transport/unixconn.py
+++ b/docker/transport/unixconn.py
@@ -18,7 +18,20 @@ except ImportError:
 RecentlyUsedContainer = urllib3._collections.RecentlyUsedContainer
 
 
+class UnixHTTPResponse(httplib.HTTPResponse, object):
+    def __init__(self, sock, *args, **kwargs):
+        disable_buffering = kwargs.pop('disable_buffering', False)
+        super(UnixHTTPResponse, self).__init__(sock, *args, **kwargs)
+        if disable_buffering is True:
+            # We must first create a new pointer then close the old one
+            # to avoid closing the underlying socket.
+            new_fp = sock.makefile('rb', 0)
+            self.fp.close()
+            self.fp = new_fp
+
+
 class UnixHTTPConnection(httplib.HTTPConnection, object):
+
     def __init__(self, base_url, unix_socket, timeout=60):
         super(UnixHTTPConnection, self).__init__(
             'localhost', timeout=timeout
@@ -26,6 +39,7 @@ class UnixHTTPConnection(httplib.HTTPConnection, object):
         self.base_url = base_url
         self.unix_socket = unix_socket
         self.timeout = timeout
+        self.disable_buffering = False
 
     def connect(self):
         sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
@@ -33,24 +47,16 @@ class UnixHTTPConnection(httplib.HTTPConnection, object):
         sock.connect(self.unix_socket)
         self.sock = sock
 
+    def putheader(self, header, *values):
+        super(UnixHTTPConnection, self).putheader(header, *values)
+        if header == 'Connection' and 'Upgrade' in values:
+            self.disable_buffering = True
 
-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)
-
+    def response_class(self, sock, *args, **kwargs):
+        if self.disable_buffering:
+            kwargs['disable_buffering'] = True
 
-class AttachUnixHTTPConnection(UnixHTTPConnection):
-    '''
-    A HTTPConnection that returns responses that don't used buffering.
-    '''
-    response_class = AttachHTTPResponse
+        return UnixHTTPResponse(sock, *args, **kwargs)
 
 
 class UnixHTTPConnectionPool(urllib3.connectionpool.HTTPConnectionPool):
@@ -63,17 +69,9 @@ class UnixHTTPConnectionPool(urllib3.connectionpool.HTTPConnectionPool):
         self.timeout = timeout
 
     def _new_conn(self):
-        # 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
-            )
+        return UnixHTTPConnection(
+            self.base_url, self.socket_path, self.timeout
+        )
 
 
 class UnixAdapter(requests.adapters.HTTPAdapter):