fix: only store received content when the header is within success ra…

…nge (#7970)
googleapis · Jan 10, 2025 · 6216964 · 6216964
1 parent bef0f90
commit 6216964
Show file tree

Hide file tree

Showing 2 changed files with 101 additions and 1 deletion.
diff --git a/Storage/src/Connection/Rest.php b/Storage/src/Connection/Rest.php
@@ -324,7 +324,11 @@ public function downloadObject(array $args = [])
             &$attempt,
         ) {
             // if the exception has a response for us to use
-            if ($e instanceof RequestException && $e->hasResponse()) {
+            if ($e instanceof RequestException
+                && $e->hasResponse()
+                && $e->getResponse()->getStatusCode() >= 200
+                && $e->getResponse()->getStatusCode() < 300
+            ) {
                 $msg = (string) $e->getResponse()->getBody();
 
                 $fetchedStream = Utils::streamFor($msg);

diff --git a/Storage/tests/Unit/Connection/RestTest.php b/Storage/tests/Unit/Connection/RestTest.php
@@ -17,6 +17,7 @@
 
 namespace Google\Cloud\Storage\Tests\Unit\Connection;
 
+use Google\Auth\HttpHandler\Guzzle7HttpHandler;
 use Google\Cloud\Core\RequestBuilder;
 use Google\Cloud\Core\RequestWrapper;
 use Google\Cloud\Core\Retry;
@@ -25,6 +26,9 @@
 use Google\Cloud\Core\Upload\ResumableUploader;
 use Google\Cloud\Core\Upload\StreamableUploader;
 use Google\Cloud\Storage\Connection\Rest;
+use Google\Cloud\Storage\Connection\RetryTrait;
+use GuzzleHttp\Client;
+use GuzzleHttp\Exception\RequestException;
 use GuzzleHttp\Promise\Create;
 use GuzzleHttp\Promise\PromiseInterface;
 use GuzzleHttp\Psr7\Request;
@@ -273,6 +277,98 @@ function ($args) use (&$actualRequest, $response) {
         );
     }
 
+    public function downloadObjectWithResumeProvider()
+    {
+        $part1 = 'first part';
+        $part2 = 'second part';
+        $full = $part1 . $part2;
+        /** @see RetryTrait::$httpRetryCodes for list of status codes that are allowed to retry */
+        return [
+            [0, '', 200, $full, $full],
+            [200, $part1, 200, $part2, $full, ['Range' => 'bytes=10-']],
+            [408, 'Server Message', 200, $full, $full],
+            [429, 'Server Message', 200, $full, $full],
+            [500, 'Server Error', 200, $full, $full],
+            [502, 'Server Error', 200, $full, $full],
+            [503, 'Server Error', 200, $full, $full],
+            [504, 'Server Error', 200, $full, $full],
+        ];
+    }
+
+    /**
+     * @dataProvider downloadObjectWithResumeProvider
+     */
+    public function testDownloadObjectWithResume(
+        int $status1,
+        string $body1,
+        int $status2,
+        string $body2,
+        string $expectedResult,
+        array $expectedSecondRequestHeaders = []
+    ) {
+        /** @var Request[] $actualRequests */
+        $actualRequests = [];
+        $requestHeaders = [];
+
+        $mockClient = $this->prophesize(Client::class);
+        $requestIndex = 0;
+        $mockClient->send(
+            Argument::type(RequestInterface::class),
+            Argument::type('array')
+        )->will(
+            function ($args) use (
+                &$actualRequests,
+                &$requestHeaders,
+                $status1,
+                $body1,
+                $status2,
+                $body2,
+                &$requestIndex
+            ) {
+                $actualRequests[$requestIndex] = $args[0];
+                $requestHeaders[$requestIndex] = $args[1]['headers'] ?? [];
+                if ($requestIndex++ === 0) {
+                    throw new RequestException("Server error", $args[0], new Response($status1, [], $body1));
+                }
+                return new Response($status2, [], $body2);
+            }
+        );
+        $requestWrapper = new RequestWrapper([
+            'httpHandler' => new Guzzle7HttpHandler($mockClient->reveal()),
+            'accessToken' => 'Fake token',
+            'retries' => 3,
+        ]);
+
+        $rest = new Rest();
+        $rest->setRequestWrapper($requestWrapper);
+
+        $options = self::$downloadOptions;
+        $options['retries'] = 3;
+        // ensure the tests execute quickly (no reason to wait for the delay)
+        $options['restDelayFunction'] = function () {
+        };
+        $actualBody = (string) $rest->downloadObject($options);
+        $actualUri1 = (string) $actualRequests[0]->getUri();
+        $actualUri2 = (string) $actualRequests[1]->getUri();
+
+        $expectedUri = sprintf(
+            '%s/storage/v1/b/bigbucket/o/myfile.txt?generation=100&alt=media&userProject=myProject',
+            Rest::DEFAULT_API_ENDPOINT
+        );
+
+        $this->assertEquals(
+            $expectedUri,
+            $actualUri1
+        );
+        $this->assertEquals([], $requestHeaders[0]);
+        $this->assertEquals(
+            $expectedUri,
+            $actualUri2
+        );
+        $this->assertEquals($expectedSecondRequestHeaders, $requestHeaders[1]);
+        $this->assertEquals($expectedResult, $actualBody);
+    }
+
     /**
      * @dataProvider insertObjectProvider
      */