Merge pull request #1510 from guzzle/stream-read-content-length

Only read up to Content-Length in stream wrapper
guzzle · Jul 1, 2016 · 10a49d5 · 10a49d5
2 parents 90a08e8 + 5b3279a
commit 10a49d5
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 5 deletions.
diff --git a/src/Handler/StreamHandler.php b/src/Handler/StreamHandler.php
@@ -105,7 +105,12 @@ private function createResponse(
         $headers = \GuzzleHttp\headers_from_lines($hdrs);
         list ($stream, $headers) = $this->checkDecode($options, $headers, $stream);
         $stream = Psr7\stream_for($stream);
-        $sink = $this->createSink($stream, $options);
+        $sink = $stream;
+
+        if (strcasecmp('HEAD', $request->getMethod())) {
+            $sink = $this->createSink($stream, $options);
+        }
+
         $response = new Psr7\Response($status, $headers, $sink, $ver, $reason);
 
         if (isset($options['on_headers'])) {
@@ -118,8 +123,14 @@ private function createResponse(
             }
         }
 
+        // Do not drain when the request is a HEAD request because they have
+        // no body.
         if ($sink !== $stream) {
-            $this->drain($stream, $sink);
+            $this->drain(
+                $stream,
+                $sink,
+                $response->getHeaderLine('Content-Length')
+            );
         }
 
         $this->invokeStats($options, $request, $startTime, $response, null);
@@ -181,13 +192,27 @@ private function checkDecode(array $options, array $headers, $stream)
      *
      * @param StreamInterface $source
      * @param StreamInterface $sink
+     * @param string          $contentLength Header specifying the amount of
+     *                                       data to read.
      *
      * @return StreamInterface
      * @throws \RuntimeException when the sink option is invalid.
      */
-    private function drain(StreamInterface $source, StreamInterface $sink)
-    {
-        Psr7\copy_to_stream($source, $sink);
+    private function drain(
+        StreamInterface $source,
+        StreamInterface $sink,
+        $contentLength
+    ) {
+        // If a content-length header is provided, then stop reading once
+        // that number of bytes has been read. This can prevent infinitely
+        // reading from a stream when dealing with servers that do not honor
+        // Connection: Close headers.
+        Psr7\copy_to_stream(
+            $source,
+            $sink,
+            strlen($contentLength) > 0 ? (int) $contentLength : -1
+        );
+
         $sink->seek(0);
         $source->close();
 

diff --git a/tests/Handler/StreamHandlerTest.php b/tests/Handler/StreamHandlerTest.php
@@ -142,6 +142,43 @@ public function testDrainsResponseIntoSaveToBodyAtNonExistentPath()
         unlink($tmpfname);
     }
 
+    public function testDrainsResponseAndReadsOnlyContentLengthBytes()
+    {
+        Server::flush();
+        Server::enqueue([
+            new Response(200, [
+                'Foo' => 'Bar',
+                'Content-Length' => 8,
+            ], 'hi there... This has way too much data!')
+        ]);
+        $handler = new StreamHandler();
+        $request = new Request('GET', Server::$url);
+        $response = $handler($request, [])->wait();
+        $body = $response->getBody();
+        $stream = $body->detach();
+        $this->assertEquals('hi there', stream_get_contents($stream));
+        fclose($stream);
+    }
+
+    public function testDoesNotDrainWhenHeadRequest()
+    {
+        Server::flush();
+        // Say the content-length is 8, but return no response.
+        Server::enqueue([
+            new Response(200, [
+                'Foo' => 'Bar',
+                'Content-Length' => 8,
+            ], '')
+        ]);
+        $handler = new StreamHandler();
+        $request = new Request('HEAD', Server::$url);
+        $response = $handler($request, [])->wait();
+        $body = $response->getBody();
+        $stream = $body->detach();
+        $this->assertEquals('', stream_get_contents($stream));
+        fclose($stream);
+    }
+
     public function testAutomaticallyDecompressGzip()
     {
         Server::flush();