Disable buffering based on presence of Connection Upgrade headers

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

2 files changed, 35 insertions, 37 deletions

diff --git a/Makefile b/Makefile
index 991b93a..efa4232 100644
--- a/Makefile
+++ b/Makefile
@@ -27,19 +27,19 @@ test: flake8 unit-test unit-test-py3 integration-dind integration-dind-ssl
 
 .PHONY: unit-test
 unit-test: build
-	docker run --rm docker-sdk-python py.test tests/unit
+	docker run -t --rm docker-sdk-python py.test tests/unit
 
 .PHONY: unit-test-py3
 unit-test-py3: build-py3
-	docker run --rm docker-sdk-python3 py.test tests/unit
+	docker run -t --rm docker-sdk-python3 py.test tests/unit
 
 .PHONY: integration-test
 integration-test: build
-	docker run --rm -v /var/run/docker.sock:/var/run/docker.sock docker-sdk-python py.test tests/integration/${file}
+	docker run -t --rm -v /var/run/docker.sock:/var/run/docker.sock docker-sdk-python py.test -v tests/integration/${file}
 
 .PHONY: integration-test-py3
 integration-test-py3: build-py3
-	docker run --rm -v /var/run/docker.sock:/var/run/docker.sock docker-sdk-python3 py.test tests/integration/${file}
+	docker run -t --rm -v /var/run/docker.sock:/var/run/docker.sock docker-sdk-python3 py.test tests/integration/${file}
 
 TEST_API_VERSION ?= 1.30
 TEST_ENGINE_VERSION ?= 17.06.0-ce
@@ -49,9 +49,9 @@ integration-dind: build build-py3
 	docker rm -vf dpy-dind || :
 	docker run -d --name dpy-dind --privileged dockerswarm/dind:${TEST_ENGINE_VERSION} dockerd\
 		-H tcp://0.0.0.0:2375 --experimental
-	docker run --rm --env="DOCKER_HOST=tcp://docker:2375" --env="DOCKER_TEST_API_VERSION=${TEST_API_VERSION}"\
+	docker run -t --rm --env="DOCKER_HOST=tcp://docker:2375" --env="DOCKER_TEST_API_VERSION=${TEST_API_VERSION}"\
 		--link=dpy-dind:docker docker-sdk-python py.test tests/integration
-	docker run --rm --env="DOCKER_HOST=tcp://docker:2375" --env="DOCKER_TEST_API_VERSION=${TEST_API_VERSION}"\
+	docker run -t --rm --env="DOCKER_HOST=tcp://docker:2375" --env="DOCKER_TEST_API_VERSION=${TEST_API_VERSION}"\
 		--link=dpy-dind:docker docker-sdk-python3 py.test tests/integration
 	docker rm -vf dpy-dind
 
@@ -63,21 +63,21 @@ integration-dind-ssl: build-dind-certs build build-py3
 		-v /tmp --privileged dockerswarm/dind:${TEST_ENGINE_VERSION} dockerd --tlsverify\
 		--tlscacert=/certs/ca.pem --tlscert=/certs/server-cert.pem\
 		--tlskey=/certs/server-key.pem -H tcp://0.0.0.0:2375 --experimental
-	docker run --rm --volumes-from dpy-dind-ssl --env="DOCKER_HOST=tcp://docker:2375"\
+	docker run -t --rm --volumes-from dpy-dind-ssl --env="DOCKER_HOST=tcp://docker:2375"\
 		--env="DOCKER_TLS_VERIFY=1" --env="DOCKER_CERT_PATH=/certs" --env="DOCKER_TEST_API_VERSION=${TEST_API_VERSION}"\
 		--link=dpy-dind-ssl:docker docker-sdk-python py.test tests/integration
-	docker run --rm --volumes-from dpy-dind-ssl --env="DOCKER_HOST=tcp://docker:2375"\
+	docker run -t --rm --volumes-from dpy-dind-ssl --env="DOCKER_HOST=tcp://docker:2375"\
 		--env="DOCKER_TLS_VERIFY=1" --env="DOCKER_CERT_PATH=/certs" --env="DOCKER_TEST_API_VERSION=${TEST_API_VERSION}"\
 		--link=dpy-dind-ssl:docker docker-sdk-python3 py.test tests/integration
 	docker rm -vf dpy-dind-ssl dpy-dind-certs
 
 .PHONY: flake8
 flake8: build
-	docker run --rm docker-sdk-python flake8 docker tests
+	docker run -t --rm docker-sdk-python flake8 docker tests
 
 .PHONY: docs
 docs: build-docs
-	docker run --rm -it -v `pwd`:/src docker-sdk-python-docs sphinx-build docs docs/_build
+	docker run --rm -t -v `pwd`:/src docker-sdk-python-docs sphinx-build docs docs/_build
 
 .PHONY: shell
 shell: build
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):