From 520206f818a347616533bee5d4c0b226ee097ebf Mon Sep 17 00:00:00 2001
From: Chris Sidi <hashtagchris@github.com>
Date: Wed, 25 Nov 2020 23:04:07 -0500
Subject: [PATCH] Retry artifact download when response stream is truncated

---
 packages/artifact/__tests__/download.test.ts  | 71 +++++++++++++++++--
 .../src/internal/download-http-client.ts      |  6 ++
 2 files changed, 71 insertions(+), 6 deletions(-)

diff --git a/packages/artifact/__tests__/download.test.ts b/packages/artifact/__tests__/download.test.ts
index 305bfa8a..81759d46 100644
--- a/packages/artifact/__tests__/download.test.ts
+++ b/packages/artifact/__tests__/download.test.ts
@@ -124,7 +124,7 @@ describe('Download Tests', () => {
     )
     const targetPath = path.join(root, 'FileA.txt')
 
-    setupDownloadItemResponse(true, 200, fileContents)
+    setupDownloadItemResponse(true, 200, false, fileContents)
     const downloadHttpClient = new DownloadHttpClient()
 
     const items: DownloadItem[] = []
@@ -147,7 +147,7 @@ describe('Download Tests', () => {
     )
     const targetPath = path.join(root, 'FileB.txt')
 
-    setupDownloadItemResponse(false, 200, fileContents)
+    setupDownloadItemResponse(false, 200, false, fileContents)
     const downloadHttpClient = new DownloadHttpClient()
 
     const items: DownloadItem[] = []
@@ -171,7 +171,7 @@ describe('Download Tests', () => {
       const fileContents = Buffer.from('try, try again\n', defaultEncoding)
       const targetPath = path.join(root, `FileC-${statusCode}.txt`)
 
-      setupDownloadItemResponse(false, statusCode, fileContents)
+      setupDownloadItemResponse(false, statusCode, false, fileContents)
       const downloadHttpClient = new DownloadHttpClient()
 
       const items: DownloadItem[] = []
@@ -188,6 +188,52 @@ describe('Download Tests', () => {
     }
   })
 
+  it('Test retry on truncated response with gzip', async () => {
+    const fileContents = Buffer.from(
+      'Sometimes gunzip fails on the first try\n',
+      defaultEncoding
+    )
+    const targetPath = path.join(root, 'FileD.txt')
+
+    setupDownloadItemResponse(true, 200, true, fileContents)
+    const downloadHttpClient = new DownloadHttpClient()
+
+    const items: DownloadItem[] = []
+    items.push({
+      sourceLocation: `${configVariables.getRuntimeUrl()}_apis/resources/Containers/13?itemPath=my-artifact%2FFileD.txt`,
+      targetPath
+    })
+
+    await expect(
+      downloadHttpClient.downloadSingleArtifact(items)
+    ).resolves.not.toThrow()
+
+    await checkDestinationFile(targetPath, fileContents)
+  })
+
+  it('Test retry on truncated response without gzip', async () => {
+    const fileContents = Buffer.from(
+      'You have to inspect the content-length header to know if you got everything\n',
+      defaultEncoding
+    )
+    const targetPath = path.join(root, 'FileE.txt')
+
+    setupDownloadItemResponse(false, 200, true, fileContents)
+    const downloadHttpClient = new DownloadHttpClient()
+
+    const items: DownloadItem[] = []
+    items.push({
+      sourceLocation: `${configVariables.getRuntimeUrl()}_apis/resources/Containers/13?itemPath=my-artifact%2FFileD.txt`,
+      targetPath
+    })
+
+    await expect(
+      downloadHttpClient.downloadSingleArtifact(items)
+    ).resolves.not.toThrow()
+
+    await checkDestinationFile(targetPath, fileContents)
+  })
+
   /**
    * Helper used to setup mocking for the HttpClient
    */
@@ -253,17 +299,24 @@ describe('Download Tests', () => {
   function setupDownloadItemResponse(
     isGzip: boolean,
     firstHttpResponseCode: number,
+    truncateFirstResponse: boolean,
     fileContents: Buffer
   ): void {
     const spyInstance = jest
       .spyOn(HttpClient.prototype, 'get')
       .mockImplementationOnce(async () => {
         if (firstHttpResponseCode === 200) {
+          const fullResponse = await constructResponse(isGzip, fileContents)
+          const actualResponse = truncateFirstResponse
+            ? fullResponse.subarray(0, 3)
+            : fullResponse
+
           return {
             message: getDownloadResponseMessage(
               firstHttpResponseCode,
               isGzip,
-              await constructResponse(isGzip, fileContents)
+              fullResponse.length,
+              actualResponse
             ),
             readBody: emptyMockReadBody
           }
@@ -272,6 +325,7 @@ describe('Download Tests', () => {
             message: getDownloadResponseMessage(
               firstHttpResponseCode,
               false,
+              0,
               null
             ),
             readBody: emptyMockReadBody
@@ -283,11 +337,13 @@ describe('Download Tests', () => {
     if (firstHttpResponseCode !== 200) {
       spyInstance.mockImplementationOnce(async () => {
         // chained response, if the HTTP GET function gets called again, return a successful response
+        const fullResponse = await constructResponse(isGzip, fileContents)
         return {
           message: getDownloadResponseMessage(
             200,
             isGzip,
-            await constructResponse(isGzip, fileContents)
+            fullResponse.length,
+            fullResponse
           ),
           readBody: emptyMockReadBody
         }
@@ -311,6 +367,7 @@ describe('Download Tests', () => {
   function getDownloadResponseMessage(
     httpResponseCode: number,
     isGzip: boolean,
+    contentLength: number,
     response: Buffer | null
   ): http.IncomingMessage {
     let readCallCount = 0
@@ -335,7 +392,9 @@ describe('Download Tests', () => {
     })
 
     mockMessage.statusCode = httpResponseCode
-    mockMessage.headers = {}
+    mockMessage.headers = {
+      'content-length': contentLength.toString()
+    }
 
     if (isGzip) {
       mockMessage.headers['content-encoding'] = 'gzip'
diff --git a/packages/artifact/src/internal/download-http-client.ts b/packages/artifact/src/internal/download-http-client.ts
index 70ea4e44..3b00872e 100644
--- a/packages/artifact/src/internal/download-http-client.ts
+++ b/packages/artifact/src/internal/download-http-client.ts
@@ -264,6 +264,12 @@ export class DownloadHttpClient {
         const gunzip = zlib.createGunzip()
         response.message
           .pipe(gunzip)
+          .on('error', error => {
+            core.error(
+              `An error has been encountered while attempting to decompress a file`
+            )
+            reject(error)
+          })
           .pipe(destinationStream)
           .on('close', () => {
             resolve()