Merge pull request #1130 from aanand/support-tcp-upgrade

Support for TCP upgrade

docker/api/container.py

@@ -15,12 +15,18 @@ class ContainerApiMixin(object):
             'logs': logs and 1 or 0,
             'stdout': stdout and 1 or 0,
             'stderr': stderr and 1 or 0,
-            'stream': stream and 1 or 0,
+            'stream': stream and 1 or 0
         }
+
+        headers = {
+            'Connection': 'Upgrade',
+            'Upgrade': 'tcp'
+        }
+
         u = self._url("/containers/{0}/attach", container)
-        response = self._post(u, params=params, stream=stream)
+        response = self._post(u, headers=headers, params=params, stream=stream)
 
-        return self._get_result(container, stream, response)
+        return self._read_from_socket(response, stream)
 
     @utils.check_resource
     def attach_socket(self, container, params=None, ws=False):
@@ -34,9 +40,18 @@ class ContainerApiMixin(object):
         if ws:
             return self._attach_websocket(container, params)
 
+        headers = {
+            'Connection': 'Upgrade',
+            'Upgrade': 'tcp'
+        }
+
         u = self._url("/containers/{0}/attach", container)
-        return self._get_raw_response_socket(self.post(
-            u, None, params=self._attach_params(params), stream=True))
+        return self._get_raw_response_socket(
+            self.post(
+                u, None, params=self._attach_params(params), stream=True,
+                headers=headers
+            )
+        )
 
     @utils.check_resource
     def commit(self, container, repository=None, tag=None, message=None,

docker/api/exec_api.py

@@ -56,8 +56,6 @@ class ExecApiMixin(object):
     def exec_start(self, exec_id, detach=False, tty=False, stream=False,
                    socket=False):
         # we want opened socket if socket == True
-        if socket:
-            stream = True
         if isinstance(exec_id, dict):
             exec_id = exec_id.get('Id')
 
@@ -66,10 +64,18 @@ class ExecApiMixin(object):
             'Detach': detach
         }
 
+        headers = {} if detach else {
+            'Connection': 'Upgrade',
+            'Upgrade': 'tcp'
+        }
+
         res = self._post_json(
-            self._url('/exec/{0}/start', exec_id), data=data, stream=stream
+            self._url('/exec/{0}/start', exec_id),
+            headers=headers,
+            data=data,
+            stream=True
         )
 
         if socket:
             return self._get_raw_response_socket(res)
-        return self._get_result_tty(stream, res, tty)
+        return self._read_from_socket(res, stream)

docker/client.py

@@ -29,6 +29,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
 try:
     from .transport import NpipeAdapter
 except ImportError:
@@ -305,6 +306,14 @@ class Client(
         for out in response.iter_content(chunk_size=1, decode_unicode=True):
             yield out
 
+    def _read_from_socket(self, response, stream):
+        socket = self._get_raw_response_socket(response)
+
+        if stream:
+            return frames_iter(socket)
+        else:
+            return six.binary_type().join(frames_iter(socket))
+
     def _disable_socket_timeout(self, socket):
         """ Depending on the combination of python version and whether we're
         connecting over http or https, we might need to access _sock, which

docker/utils/socket.py

+import errno
+import os
+import select
+import struct
+
+import six
+
+
+class SocketError(Exception):
+    pass
+
+
+def read(socket, n=4096):
+    """
+    Reads at most n bytes from socket
+    """
+    recoverable_errors = (errno.EINTR, errno.EDEADLK, errno.EWOULDBLOCK)
+
+    # wait for data to become available
+    select.select([socket], [], [])
+
+    try:
+        if hasattr(socket, 'recv'):
+            return socket.recv(n)
+        return os.read(socket.fileno(), n)
+    except EnvironmentError as e:
+        if e.errno not in recoverable_errors:
+            raise
+
+
+def read_exactly(socket, n):
+    """
+    Reads exactly n bytes from socket
+    Raises SocketError if there isn't enough data
+    """
+    data = six.binary_type()
+    while len(data) < n:
+        next_data = read(socket, n - len(data))
+        if not next_data:
+            raise SocketError("Unexpected EOF")
+        data += next_data
+    return data
+
+
+def next_frame_size(socket):
+    """
+    Returns the size of the next frame of data waiting to be read from socket,
+    according to the protocol defined here:
+
+    https://docs.docker.com/engine/reference/api/docker_remote_api_v1.24/#/attach-to-a-container
+    """
+    try:
+        data = read_exactly(socket, 8)
+    except SocketError:
+        return 0
+
+    _, actual = struct.unpack('>BxxxL', data)
+    return actual
+
+
+def frames_iter(socket):
+    """
+    Returns a generator of frames read from socket
+    """
+    n = next_frame_size(socket)
+    while n > 0:
+        yield read(socket, n)
+        n = next_frame_size(socket)

tests/helpers.py

-import errno
 import os
 import os.path
-import select
 import shutil
-import struct
 import tarfile
 import tempfile
 import unittest
@@ -54,7 +51,7 @@ def exec_driver_is_native():
         c = docker_client()
         EXEC_DRIVER = c.info()['ExecutionDriver']
         c.close()
-    return EXEC_DRIVER.startswith('native')
+    return EXEC_DRIVER.startswith('native') or EXEC_DRIVER == ''
 
 
 def docker_client(**kwargs):
@@ -67,49 +64,6 @@ def docker_client_kwargs(**kwargs):
     return client_kwargs
 
 
-def read_socket(socket, n=4096):
-    """ Code stolen from dockerpty to read the socket """
-    recoverable_errors = (errno.EINTR, errno.EDEADLK, errno.EWOULDBLOCK)
-
-    # wait for data to become available
-    select.select([socket], [], [])
-
-    try:
-        if hasattr(socket, 'recv'):
-            return socket.recv(n)
-        return os.read(socket.fileno(), n)
-    except EnvironmentError as e:
-        if e.errno not in recoverable_errors:
-            raise
-
-
-def next_packet_size(socket):
-    """ Code stolen from dockerpty to get the next packet size """
-    data = six.binary_type()
-    while len(data) < 8:
-        next_data = read_socket(socket, 8 - len(data))
-        if not next_data:
-            return 0
-        data = data + next_data
-
-    if data is None:
-        return 0
-
-    if len(data) == 8:
-        _, actual = struct.unpack('>BxxxL', data)
-        return actual
-
-
-def read_data(socket, packet_size):
-    data = six.binary_type()
-    while len(data) < packet_size:
-        next_data = read_socket(socket, packet_size - len(data))
-        if not next_data:
-            assert False, "Failed trying to read in the dataz"
-        data += next_data
-    return data
-
-
 class BaseTestCase(unittest.TestCase):
     tmp_imgs = []
     tmp_containers = []

tests/integration/container_test.py

@@ -3,6 +3,8 @@ import signal
 import tempfile
 
 import docker
+from docker.utils.socket import next_frame_size
+from docker.utils.socket import read_exactly
 import pytest
 import six
 
@@ -1025,9 +1027,9 @@ class AttachContainerTest(helpers.BaseTestCase):
 
         self.client.start(ident)
 
-        next_size = helpers.next_packet_size(pty_stdout)
+        next_size = next_frame_size(pty_stdout)
         self.assertEqual(next_size, len(line))
-        data = helpers.read_data(pty_stdout, next_size)
+        data = read_exactly(pty_stdout, next_size)
         self.assertEqual(data.decode('utf-8'), line)
 
 

tests/integration/exec_test.py

 import pytest
 
+from docker.utils.socket import next_frame_size
+from docker.utils.socket import read_exactly
+
 from .. import helpers
 
 BUSYBOX = helpers.BUSYBOX
@@ -107,9 +110,9 @@ class ExecTest(helpers.BaseTestCase):
         socket = self.client.exec_start(exec_id, socket=True)
         self.addCleanup(socket.close)
 
-        next_size = helpers.next_packet_size(socket)
+        next_size = next_frame_size(socket)
         self.assertEqual(next_size, len(line))
-        data = helpers.read_data(socket, next_size)
+        data = read_exactly(socket, next_size)
         self.assertEqual(data.decode('utf-8'), line)
 
     def test_exec_inspect(self):

tests/unit/api_test.py

@@ -93,6 +93,10 @@ def fake_put(self, url, *args, **kwargs):
 def fake_delete(self, url, *args, **kwargs):
     return fake_request('DELETE', url, *args, **kwargs)
 
+
+def fake_read_from_socket(self, response, stream):
+    return six.binary_type()
+
 url_base = 'http+docker://localunixsocket/'
 url_prefix = '{0}v{1}/'.format(
     url_base,
@@ -103,7 +107,8 @@ class DockerClientTest(base.Cleanup, base.BaseTestCase):
     def setUp(self):
         self.patcher = mock.patch.multiple(
             'docker.Client', get=fake_get, post=fake_post, put=fake_put,
-            delete=fake_delete
+            delete=fake_delete,
+            _read_from_socket=fake_read_from_socket
         )
         self.patcher.start()
         self.client = docker.Client()

tests/unit/exec_test.py

@@ -51,8 +51,36 @@ class ExecTest(DockerClientTest):
             }
         )
 
-        self.assertEqual(args[1]['headers'],
-                         {'Content-Type': 'application/json'})
+        self.assertEqual(
+            args[1]['headers'], {
+                'Content-Type': 'application/json',
+                'Connection': 'Upgrade',
+                'Upgrade': 'tcp'
+            }
+        )
+
+    def test_exec_start_detached(self):
+        self.client.exec_start(fake_api.FAKE_EXEC_ID, detach=True)
+
+        args = fake_request.call_args
+        self.assertEqual(
+            args[0][1], url_prefix + 'exec/{0}/start'.format(
+                fake_api.FAKE_EXEC_ID
+            )
+        )
+
+        self.assertEqual(
+            json.loads(args[1]['data']), {
+                'Tty': False,
+                'Detach': True
+            }
+        )
+
+        self.assertEqual(
+            args[1]['headers'], {
+                'Content-Type': 'application/json'
+            }
+        )
 
     def test_exec_inspect(self):
         self.client.exec_inspect(fake_api.FAKE_EXEC_ID)