From 0155b8177100b07a39765b80e128487f07c8f125 Mon Sep 17 00:00:00 2001
From: Steve Bazyl <sbazyl@google.com>
Date: Fri, 22 Apr 2016 14:30:30 -0700
Subject: [PATCH] Rewind output stream if requesting partial content but get
 the full media body instead

---
 lib/google/apis/core/download.rb       | 11 ++++++++++-
 spec/google/apis/core/download_spec.rb | 14 +++++++++++++-
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/lib/google/apis/core/download.rb b/lib/google/apis/core/download.rb
index 6ddbc3984..db55509a7 100644
--- a/lib/google/apis/core/download.rb
+++ b/lib/google/apis/core/download.rb
@@ -67,12 +67,21 @@ module Google
         def execute_once(client, &block)
           client.get(@download_url || url) do |req|
             apply_request_options(req)
+            check_if_rewind_needed = false
             if @offset > 0
               logger.debug { sprintf('Resuming download from offset %d', @offset) }
               req.header[RANGE_HEADER] = sprintf('bytes=%d-', @offset)
+              check_if_rewind_needed = true
             end
-            req.on_body(200, 201) do |res, chunk|
+            req.on_body(200, 201, 206) do |res, chunk|
               check_status(res.status_code, chunk) unless res.status_code.nil?
+              if check_if_rewind_needed && res.status_code != 206
+                # Oh no! Requested a chunk, but received the entire content
+                # Attempt to rewind the stream
+                @download_io.rewind
+                check_if_rewind_needed = false
+              end
+
               logger.debug { sprintf('Writing chunk (%d bytes)', chunk.length) }
               @offset += chunk.length
               @download_io.write(chunk)
diff --git a/spec/google/apis/core/download_spec.rb b/spec/google/apis/core/download_spec.rb
index b7516ce36..9f56beb0f 100644
--- a/spec/google/apis/core/download_spec.rb
+++ b/spec/google/apis/core/download_spec.rb
@@ -52,7 +52,19 @@ RSpec.describe Google::Apis::Core::DownloadCommand do
       before(:example) do
         stub_request(:get, 'https://www.googleapis.com/zoo/animals')
           .to_return(body: ['Hello ', Timeout::Error])
-          .to_return(body: 'world')
+          .to_return(status: [206, 'Partial content'], body: 'world')
+      end
+
+      it 'should receive entire content' do
+        expect(received).to eql('Hello world')
+      end
+    end
+
+    context 'with disconnects and no partial response' do
+      before(:example) do
+        stub_request(:get, 'https://www.googleapis.com/zoo/animals')
+            .to_return(body: ['Hello ', Timeout::Error])
+            .to_return(body: 'Hello world')
       end
 
       it 'should receive entire content' do