Fixed some overflow bugs in TransferManager and reliable download setup in BlobURL

Author: rickle-msft

ChangeLog.txt

@@ -4,6 +4,7 @@ XXXX.XX.XX Version XX.X.X
 * Added CopyFromURL, which will do a synchronous server-side copy, meaning the service will not return an HTTP response until it has completed the copy.
 * Added support for IProgressReceiver in TransferManager operations. This parameter was previously ignored but is now supported.
 * Removed internal dependency on javafx to be compatible with openjdk.
+* Fixed a bug that would cause downloading large files with the TransferManager to fail and a bug in BlobURL.download() logic for setting up reliable download.
 
 2018.09.11 Version 10.1.0
 * Interfaces for helper types updated to be more consistent throughout the library. All types, with the exception of the options for pipeline factories, use a fluent pattern.

src/main/java/com/microsoft/azure/storage/blob/BlobURL.java

@@ -322,26 +322,6 @@ public Single<DownloadResponse> download() {
         return this.download(null, null, false, null);
     }
 
-    /**
-     * Reads a range of bytes from a blob. The response also includes the blob's properties and metadata. For more
-     * information, see the <a href="https://docs.microsoft.com/rest/api/storageservices/get-blob">Azure Docs</a>.
-     * <p>
-     * Note that the response body has reliable download functionality built in, meaning that a failed download stream
-     * will be automatically retried. This behavior may be configured with {@link ReliableDownloadOptions}.
-     *
-     * @param range
-     *         {@link BlobRange}
-     *
-     * @return Emits the successful response.
-     *
-     * @apiNote ## Sample Code \n
-     * [!code-java[Sample_Code](../azure-storage-java/src/test/java/com/microsoft/azure/storage/Samples.java?name=upload_download "Sample code for BlobURL.download")] \n
-     * For more samples, please see the [Samples file](%https://github.com/Azure/azure-storage-java/blob/master/src/test/java/com/microsoft/azure/storage/Samples.java)
-     */
-    public Single<DownloadResponse> download(BlobRange range) {
-        return this.download(range, null, false, null);
-    }
-
     /**
      * Reads a range of bytes from a blob. The response also includes the blob's properties and metadata. For more
      * information, see the <a href="https://docs.microsoft.com/rest/api/storageservices/get-blob">Azure Docs</a>.
@@ -374,7 +354,7 @@ public Single<DownloadResponse> download(BlobRange range, BlobAccessConditions a
         range = range == null ? BlobRange.DEFAULT : range;
         accessConditions = accessConditions == null ? BlobAccessConditions.NONE : accessConditions;
         HTTPGetterInfo info = new HTTPGetterInfo()
-                .withCount(range.offset())
+                .withOffset(range.offset())
                 .withCount(range.count())
                 .withETag(accessConditions.modifiedAccessConditions().ifMatch());
 

src/main/java/com/microsoft/azure/storage/blob/DownloadResponse.java

@@ -73,21 +73,12 @@ public Flowable<ByteBuffer> body(ReliableDownloadOptions options) {
             return this.rawResponse.body();
         }
 
-        return this.rawResponse.body()
-                /*
-                Update how much data we have received in case we need to retry and propagate to the user the data we
-                have received.
-                 */
-                .doOnNext(buffer -> {
-                    this.info.withOffset(this.info.offset() + buffer.remaining());
-                    if (info.count() != null) {
-                        this.info.withCount(this.info.count() - buffer.remaining());
-                    }
-                })
-                .onErrorResumeNext(throwable -> {
-                    // So far we have tried once but retried zero times.
-                    return tryContinueFlowable(throwable, 0, optionsReal);
-                });
+        /*
+        We pass -1 for currentRetryCount because we want tryContinueFlowable to receive a value of 0 for number of
+        retries as we have not actually retried yet, only made the initial try. Because applyReliableDownload() will
+        add 1 before calling into tryContinueFlowable, we set the initial value to -1.
+         */
+        return this.applyReliableDownload(this.rawResponse.body(), -1, optionsReal);
     }
 
     private Flowable<ByteBuffer> tryContinueFlowable(Throwable t, int retryCount, ReliableDownloadOptions options) {
@@ -98,30 +89,38 @@ private Flowable<ByteBuffer> tryContinueFlowable(Throwable t, int retryCount, Re
             try {
                 // Get a new response and try reading from it.
                 return getter.apply(this.info)
-                        .flatMapPublisher(response -> {
-                            // Do not compound the number of retries; just get the raw body.
-                            ReliableDownloadOptions newOptions = new ReliableDownloadOptions()
-                                    .withMaxRetryRequests(0);
-
-                            return response.body(newOptions)
-                                    .doOnNext(buffer -> {
-                                        this.info.withOffset(this.info.offset() + buffer.remaining());
-                                        if (info.count() != null) {
-                                            this.info.withCount(this.info.count() - buffer.remaining());
-                                        }
-                                    })
-                                    .onErrorResumeNext(t2 -> {
-                                        // Increment the retry count and try again with the new exception.
-                                        return tryContinueFlowable(t2, retryCount + 1, options);
-                                    });
-                        });
+                        .flatMapPublisher(response ->
+                            /*
+                            Do not compound the number of retries by passing in another set of downloadOptions; just get
+                            the raw body.
+                             */
+                            this.applyReliableDownload(this.rawResponse.body(), retryCount, options));
             } catch (Exception e) {
                 // If the getter fails, return the getter failure to the user.
                 return Flowable.error(e);
             }
         }
     }
 
+    private Flowable<ByteBuffer> applyReliableDownload(Flowable<ByteBuffer> data,
+            int currentRetryCount, ReliableDownloadOptions options) {
+        return data
+                .doOnNext(buffer -> {
+                    /*
+                    Update how much data we have received in case we need to retry and propagate to the user the data we
+                    have received.
+                     */
+                    this.info.withOffset(this.info.offset() + buffer.remaining());
+                    if (this.info.count() != null) {
+                        this.info.withCount(this.info.count() - buffer.remaining());
+                    }
+                })
+                .onErrorResumeNext(t2 -> {
+                    // Increment the retry count and try again with the new exception.
+                    return tryContinueFlowable(t2, currentRetryCount + 1, options);
+                });
+    }
+
     public int statusCode() {
         return this.rawResponse.statusCode();
     }

src/main/java/com/microsoft/azure/storage/blob/TransferManager.java

@@ -106,8 +106,10 @@ public static Single<CommonRestResponse> uploadFileToBlockBlob(
                 the list of Ids later. Eager ensures parallelism but may require some internal buffering.
                  */
                 .concatMapEager(i -> {
-                    int count = Math.min(blockLength, (int) (file.size() - i * blockLength));
-                    Flowable<ByteBuffer> data = FlowableUtil.readFile(file, i * blockLength, count);
+                    // The max number of bytes for a block is currently 100MB, so the final result must be an int.
+                    int count = (int) Math.min((long)blockLength, (file.size() - i * (long)blockLength));
+                    // i * blockLength could be a long, so we need a cast to prevent overflow.
+                    Flowable<ByteBuffer> data = FlowableUtil.readFile(file, i * (long)blockLength, count);
 
                     // Report progress as necessary.
                     data = ProgressReporter.addParallelProgressReporting(data, optionsReal.progressReceiver(),

src/test/java/com/microsoft/azure/storage/APISpec.groovy

@@ -229,7 +229,7 @@ class APISpec extends Specification {
         }
     }
 
-    static ByteBuffer getRandomData(long size) {
+    static ByteBuffer getRandomData(int size) {
         Random rand = new Random(getRandomSeed())
         byte[] data = new byte[size]
         rand.nextBytes(data)
@@ -240,7 +240,14 @@ class APISpec extends Specification {
         File file = File.createTempFile(UUID.randomUUID().toString(), ".txt")
         file.deleteOnExit()
         FileOutputStream fos = new FileOutputStream(file)
-        fos.write(getRandomData(size).array())
+        Random rand = new Random(getRandomSeed())
+        while (size > 0) {
+            int sizeToWrite = (int) Math.min((long)(Integer.MAX_VALUE.longValue()/10L), size)
+            byte[] data = new byte[sizeToWrite]
+            rand.nextBytes(data)
+            fos.write(data)
+            size -= sizeToWrite
+        }
         fos.close()
         return file
     }
@@ -431,7 +438,6 @@ class APISpec extends Specification {
     to play too nicely with mocked objects and the complex reflection stuff on both ends made it more difficult to work
     with than was worth it.
      */
-
     def getStubResponse(int code, Class responseHeadersType) {
         return new HttpResponse() {
 
@@ -477,6 +483,58 @@ class APISpec extends Specification {
         }
     }
 
+    /*
+    This is for stubbing responses that will actually go through the pipeline and autorest code. Autorest does not seem
+    to play too nicely with mocked objects and the complex reflection stuff on both ends made it more difficult to work
+    with than was worth it. Because this type is just for BlobDownload, we don't need to accept a header type.
+     */
+    def getStubResponseForBlobDownload(int code, Flowable<ByteBuffer> body, String etag) {
+        return new HttpResponse() {
+
+            @Override
+            int statusCode() {
+                return code
+            }
+
+            @Override
+            String headerValue(String s) {
+                return null
+            }
+
+            @Override
+            HttpHeaders headers() {
+                return new HttpHeaders()
+            }
+
+            @Override
+            Flowable<ByteBuffer> body() {
+                return body
+            }
+
+            @Override
+            Single<byte[]> bodyAsByteArray() {
+                return null
+            }
+
+            @Override
+            Single<String> bodyAsString() {
+                return null
+            }
+
+            @Override
+            Object deserializedHeaders() {
+                def headers = new BlobDownloadHeaders()
+                headers.withETag(etag)
+                return headers
+            }
+
+            @Override
+            boolean isDecoded() {
+                return true
+            }
+        }
+    }
+
     def getContextStubPolicy(int successCode, Class responseHeadersType) {
         return Mock(RequestPolicy) {
             sendAsync(_) >> { HttpRequest request ->

src/test/java/com/microsoft/azure/storage/BlobAPITest.groovy

@@ -18,8 +18,11 @@ package com.microsoft.azure.storage
 import com.microsoft.azure.storage.blob.*
 import com.microsoft.azure.storage.blob.models.*
 import com.microsoft.rest.v2.http.HttpPipeline
+import com.microsoft.rest.v2.http.HttpRequest
+import com.microsoft.rest.v2.policy.RequestPolicy
 import com.microsoft.rest.v2.util.FlowableUtil
 import io.reactivex.Flowable
+import io.reactivex.Single
 import spock.lang.Unroll
 
 import java.nio.ByteBuffer
@@ -70,6 +73,48 @@ class BlobAPITest extends APISpec {
         headers.blobContentMD5() != null
     }
 
+    /*
+    This is to test the appropriate integration of DownloadResponse, including setting the correct range values on
+    HTTPGetterInfo.
+     */
+    def "Download with retry range"() {
+        /*
+        We are going to make a request for some range on a blob. The Flowable returned will throw an exception, forcing
+        a retry per the ReliableDownloadOptions. The next request should have the same range header, which was generated
+        from the count and offset values in HTTPGetterInfo that was constructed on the initial call to download. We
+        don't need to check the data here, but we want to ensure that the correct range is set each time. This will
+        test the correction of a bug that was found which caused HTTPGetterInfo to have an incorrect offset when it was
+        constructed in BlobURL.download().
+         */
+        setup:
+        def mockPolicy = Mock(RequestPolicy) {
+            sendAsync(_) >> { HttpRequest request ->
+                if (request.headers().value("x-ms-range") != "bytes=2-") {
+                    return Single.error(new IllegalArgumentException("The range header was not set correctly on retry."))
+                }
+                else {
+                    // ETag can be a dummy value. It's not validated, but DownloadResponse requires one
+                    return Single.just(getStubResponseForBlobDownload(206, Flowable.error(new IOException()), "etag"))
+                }
+            }
+        }
+        def pipeline = HttpPipeline.build(getStubFactory(mockPolicy))
+        bu = bu.withPipeline(pipeline)
+
+        when:
+        def range = new BlobRange().withOffset(2)
+        bu.download(range, null, false, null).blockingGet().body(new ReliableDownloadOptions().withMaxRetryRequests(3))
+                .blockingSubscribe()
+
+        then:
+        /*
+        Because the dummy Flowable always throws an error. This will also validate that an IllegalArgumentException is
+        NOT thrown because the types would not match.
+         */
+        def e = thrown(RuntimeException)
+        e.getCause() instanceof IOException
+    }
+
     def "Download min"() {
         expect:
         FlowableUtil.collectBytesInBuffer(bu.download().blockingGet().body(null)).blockingGet() == defaultData

src/test/java/com/microsoft/azure/storage/TransferManagerTest.groovy

@@ -426,7 +426,7 @@ class TransferManagerTest extends APISpec {
         getRandomFile(16 * 1024 * 1024)        | _ // medium file in several chunks
         getRandomFile(8L * 1026 * 1024 + 10)   | _ // medium file not aligned to block
         getRandomFile(0)                       | _ // empty file
-        getRandomFile(5L * 1024 * 1024 * 1024) | _ // file size exceeds max int
+        // Files larger than 2GB to test no integer overflow are left to stress/perf tests to keep test passes short.
     }
 
     def compareFiles(AsynchronousFileChannel channel1, long offset, long count, AsynchronousFileChannel channel2) {