Read upgraded streams through the buffered reader

docker/api/client.py

@@ -1,3 +1,4 @@
+import io
 import json
 import struct
 import urllib
@@ -428,6 +429,22 @@ def _read_from_socket(self, response, stream, tty=True, demux=False):
         """
         socket = self._get_raw_response_socket(response)
 
+        # The daemon answers an upgraded request with the response headers and
+        # then writes the stream on the same connection. http.client parses
+        # those headers through a buffered reader, which reads up to a whole
+        # buffer at a time, so the first frames of the stream can already be
+        # sitting in that buffer by the time we get here. Reading from the
+        # socket would skip them, and they would be lost.
+        # See https://github.com/docker/docker-py/issues/3332 and
+        # https://github.com/docker/docker-py/issues/2042.
+        reader = getattr(response.raw._fp, 'fp', None)
+        if isinstance(reader, io.BufferedReader):
+            # docker.utils.socket.read() cannot wait on a buffered reader the
+            # way it waits on a socket, so a quiet stream would now end on the
+            # socket timeout. Disable it, as the other streaming helpers do.
+            self._disable_socket_timeout(socket)
+            socket = reader
+
         gen = frames_iter(socket, tty)
 
         if demux:

docker/utils/socket.py

@@ -1,4 +1,5 @@
 import errno
+import io
 import os
 import select
 import socket as pysocket
@@ -23,14 +24,28 @@ class SocketError(Exception):
 NPIPE_ENDED = 109
 
 
+def _is_pipe_ended(socket, exception):
+    """
+    Whether exception is the npipe equivalent of a closed connection.
+    """
+    if isinstance(socket, io.BufferedReader):
+        # NpipeSocket.makefile() wraps the socket in a raw stream, which the
+        # buffered reader then wraps in turn.
+        socket = getattr(getattr(socket, 'raw', None), 'sock', None)
+
+    return (isinstance(socket, NpipeSocket) and
+            len(exception.args) > 0 and
+            exception.args[0] == NPIPE_ENDED)
+
+
 def read(socket, n=4096):
     """
     Reads at most n bytes from socket
     """
 
     recoverable_errors = (errno.EINTR, errno.EDEADLK, errno.EWOULDBLOCK)
 
-    if not isinstance(socket, NpipeSocket):
+    if not isinstance(socket, (NpipeSocket, io.BufferedReader)):
         if not hasattr(select, "poll"):
             # Limited to 1024
             select.select([socket], [], [])
@@ -40,6 +55,15 @@ def read(socket, n=4096):
             poll.poll()
 
     try:
+        if isinstance(socket, io.BufferedReader):
+            # A buffered reader is not waited on above: data that it has
+            # already buffered would not show up in a poll of the file
+            # descriptor, and the wait would block until more data arrived.
+            # read1() returns what is buffered and only reads the descriptor
+            # once the buffer is empty, which is the contract read() expects
+            # here. Plain read() would instead block until it had n bytes,
+            # holding back a frame that is complete but shorter than that.
+            return socket.read1(n)
         if hasattr(socket, 'recv'):
             return socket.recv(n)
         if isinstance(socket, pysocket.SocketIO):
@@ -49,10 +73,7 @@ def read(socket, n=4096):
         if e.errno not in recoverable_errors:
             raise
     except Exception as e:
-        is_pipe_ended = (isinstance(socket, NpipeSocket) and
-                         len(e.args) > 0 and
-                         e.args[0] == NPIPE_ENDED)
-        if is_pipe_ended:
+        if _is_pipe_ended(socket, e):
             # npipes don't support duplex sockets, so we interpret
             # a PIPE_ENDED error as a close operation (0-length read).
             return ''

tests/unit/api_test.py

@@ -588,6 +588,182 @@ def test_read_from_socket_no_stream_no_tty_demux(self):
         assert res == (self.stdout_data, self.stderr_data)
 
 
+class TCPSocketStreamUpgradeTest(unittest.TestCase):
+    """The daemon may write the first frames of an upgraded stream in the same
+    packet as the response headers. http.client then reads them into the
+    buffered reader it parses the headers with, and they never reach the
+    socket. See https://github.com/docker/docker-py/issues/3332.
+    """
+
+    stdout_data = b'hello\n'
+    stderr_data = b'oh no\n'
+
+    # Long enough for the delayed writes below to be observable, short enough
+    # not to slow the suite down.
+    delay = 0.5
+
+    # Silence longer than the client timeout used in test_stream_quiet.
+    quiet_delay = 2
+
+    @classmethod
+    def setup_class(cls):
+        cls.clients = []
+        cls.server = socketserver.ThreadingTCPServer(
+            ('', 0), cls.get_handler_class())
+        cls.thread = threading.Thread(target=cls.server.serve_forever)
+        cls.thread.daemon = True
+        cls.thread.start()
+        cls.address = f'http://{socket.gethostname()}:{cls.server.server_address[1]}'
+
+    @classmethod
+    def teardown_class(cls):
+        for client in cls.clients:
+            client.close()
+        cls.server.shutdown()
+        cls.server.server_close()
+        cls.thread.join()
+
+    @classmethod
+    def get_handler_class(cls):
+        stdout_data = cls.stdout_data
+        stderr_data = cls.stderr_data
+        delay = cls.delay
+        quiet_delay = cls.quiet_delay
+
+        headers = (
+            b'HTTP/1.1 101 UPGRADED\r\n'
+            b'Content-Type: application/vnd.docker.multiplexed-stream\r\n'
+            b'Connection: Upgrade\r\n'
+            b'Upgrade: tcp\r\n'
+            b'\r\n'
+        )
+
+        def frame(stream, data):
+            return struct.pack('>BxxxL', stream, len(data)) + data
+
+        class Handler(http.server.BaseHTTPRequestHandler):
+            def do_POST(self):
+                path = self.path.split('/')[-1]
+                if path == 'tty':
+                    # One write, so the headers and the payload reach the
+                    # client in a single packet.
+                    self.wfile.write(headers + stdout_data + stderr_data)
+                elif path == 'no-tty':
+                    self.wfile.write(
+                        headers
+                        + frame(1, stdout_data)
+                        + frame(2, stderr_data)
+                    )
+                elif path == 'no-tty-delayed':
+                    # The first frame shares a packet with the headers, the
+                    # second one only shows up later: a caller that streams
+                    # must get the first frame without waiting for it.
+                    self.wfile.write(headers + frame(1, stdout_data))
+                    time.sleep(delay)
+                    self.wfile.write(frame(2, stderr_data))
+                elif path == 'tty-delayed':
+                    self.wfile.write(headers + stdout_data)
+                    time.sleep(delay)
+                    self.wfile.write(stderr_data)
+                elif path == 'no-tty-quiet':
+                    # Nothing is buffered here, and the stream stays quiet
+                    # for longer than the client timeout.
+                    self.wfile.write(headers)
+                    time.sleep(quiet_delay)
+                    self.wfile.write(frame(1, stdout_data))
+                else:
+                    raise Exception(f'Unknown path {path}')
+
+            def log_message(self, fmt, *args):
+                pass
+
+        return Handler
+
+    def request(self, path, stream, tty, demux=False, timeout=None):
+        client = APIClient(
+            base_url=self.address, version=DEFAULT_DOCKER_API_VERSION,
+            timeout=timeout or DEFAULT_TIMEOUT_SECONDS)
+        # The streaming tests read from the connection after this returns, so
+        # the client is closed in teardown rather than here.
+        self.clients.append(client)
+        resp = client._post(client._url(path), stream=True)
+        return client._read_from_socket(
+            resp, stream=stream, tty=tty, demux=demux)
+
+    @staticmethod
+    def with_timeout(fn, timeout=10):
+        """Run fn in a thread and fail if it does not return in time.
+
+        Without the fix this blocks in select/poll on a socket that will never
+        have anything to report, so a plain call would hang the suite.
+        """
+        result = []
+        error = []
+
+        def target():
+            try:
+                result.append(fn())
+            except BaseException as e:
+                error.append(e)
+
+        thread = threading.Thread(target=target, daemon=True)
+        thread.start()
+        thread.join(timeout)
+        if thread.is_alive():
+            raise AssertionError(
+                f'timed out after {timeout}s waiting for the stream')
+        if error:
+            raise error[0]
+        return result[0]
+
+    def test_no_stream_tty(self):
+        res = self.with_timeout(
+            lambda: self.request('/tty', stream=False, tty=True))
+        assert res == self.stdout_data + self.stderr_data
+
+    def test_no_stream_no_tty(self):
+        res = self.with_timeout(
+            lambda: self.request('/no-tty', stream=False, tty=False))
+        assert res == self.stdout_data + self.stderr_data
+
+    def test_no_stream_no_tty_demux(self):
+        res = self.with_timeout(
+            lambda: self.request(
+                '/no-tty', stream=False, tty=False, demux=True))
+        assert res == (self.stdout_data, self.stderr_data)
+
+    def test_stream_no_tty(self):
+        def read_first_frame():
+            gen = self.request('/no-tty-delayed', stream=True, tty=False)
+            return next(gen)
+
+        start = time.monotonic()
+        assert self.with_timeout(read_first_frame) == self.stdout_data
+        # The buffered frame is handed over as soon as it is read, not held
+        # back until the rest of the stream arrives.
+        assert time.monotonic() - start < self.delay
+
+    def test_stream_tty(self):
+        def read_first_chunk():
+            gen = self.request('/tty-delayed', stream=True, tty=True)
+            return next(gen)
+
+        start = time.monotonic()
+        assert self.with_timeout(read_first_chunk) == self.stdout_data
+        assert time.monotonic() - start < self.delay
+
+    def test_stream_quiet(self):
+        # An exec or attach that produces nothing for a while is not cut short
+        # by the client timeout.
+        def read_first_frame():
+            gen = self.request(
+                '/no-tty-quiet', stream=True, tty=False, timeout=1)
+            return next(gen)
+
+        res = self.with_timeout(read_first_frame, self.quiet_delay + 10)
+        assert res == self.stdout_data
+
+
 class UserAgentTest(unittest.TestCase):
     def setUp(self):
         self.patcher = mock.patch.object(