From 562a19849cee37de6995a505db51d1e8694bd538 Mon Sep 17 00:00:00 2001
From: Chris Gavin <chris@chrisgavin.me>
Date: Thu, 23 Jul 2020 22:38:08 +0100
Subject: [PATCH] Allow specifying arbitrary headers when downloading tools to
 the tool cache.

---
 .../tool-cache/__tests__/tool-cache.test.ts   | 55 +++++++++++++++++++
 packages/tool-cache/src/tool-cache.ts         | 15 +++--
 2 files changed, 64 insertions(+), 6 deletions(-)

diff --git a/packages/tool-cache/__tests__/tool-cache.test.ts b/packages/tool-cache/__tests__/tool-cache.test.ts
index d03c9b392c..d929c1b86c 100644
--- a/packages/tool-cache/__tests__/tool-cache.test.ts
+++ b/packages/tool-cache/__tests__/tool-cache.test.ts
@@ -763,6 +763,61 @@ describe('@actions/tool-cache', function() {
       expect(err.toString()).toContain('404')
     }
   })
+
+  it('supports authorization headers', async function() {
+    nock('http://example.com', {
+      reqheaders: {
+        authorization: 'token abc123'
+      }
+    })
+      .get('/some-file-that-needs-authorization')
+      .reply(200, undefined)
+
+    await tc.downloadTool(
+      'http://example.com/some-file-that-needs-authorization',
+      undefined,
+      'token abc123'
+    )
+  })
+
+  it('supports custom headers', async function() {
+    nock('http://example.com', {
+      reqheaders: {
+        accept: 'application/octet-stream'
+      }
+    })
+      .get('/some-file-that-needs-headers')
+      .reply(200, undefined)
+
+    await tc.downloadTool(
+      'http://example.com/some-file-that-needs-headers',
+      undefined,
+      undefined,
+      {
+        accept: 'application/octet-stream'
+      }
+    )
+  })
+
+  it('supports authorization and custom headers', async function() {
+    nock('http://example.com', {
+      reqheaders: {
+        accept: 'application/octet-stream',
+        authorization: 'token abc123'
+      }
+    })
+      .get('/some-file-that-needs-authorization-and-headers')
+      .reply(200, undefined)
+
+    await tc.downloadTool(
+      'http://example.com/some-file-that-needs-authorization-and-headers',
+      undefined,
+      'token abc123',
+      {
+        accept: 'application/octet-stream'
+      }
+    )
+  })
 })
 
 /**
diff --git a/packages/tool-cache/src/tool-cache.ts b/packages/tool-cache/src/tool-cache.ts
index 3e9ecab409..c271d43050 100644
--- a/packages/tool-cache/src/tool-cache.ts
+++ b/packages/tool-cache/src/tool-cache.ts
@@ -32,12 +32,14 @@ const userAgent = 'actions/tool-cache'
  * @param url       url of tool to download
  * @param dest      path to download tool
  * @param auth      authorization header
+ * @param headers   other headers
  * @returns         path to downloaded tool
  */
 export async function downloadTool(
   url: string,
   dest?: string,
-  auth?: string
+  auth?: string,
+  headers?: IHeaders
 ): Promise<string> {
   dest = dest || path.join(_getTempDirectory(), uuidV4())
   await io.mkdirP(path.dirname(dest))
@@ -56,7 +58,7 @@ export async function downloadTool(
   const retryHelper = new RetryHelper(maxAttempts, minSeconds, maxSeconds)
   return await retryHelper.execute(
     async () => {
-      return await downloadToolAttempt(url, dest || '', auth)
+      return await downloadToolAttempt(url, dest || '', auth, headers)
     },
     (err: Error) => {
       if (err instanceof HTTPError && err.httpStatusCode) {
@@ -79,7 +81,8 @@ export async function downloadTool(
 async function downloadToolAttempt(
   url: string,
   dest: string,
-  auth?: string
+  auth?: string,
+  headers?: IHeaders
 ): Promise<string> {
   if (fs.existsSync(dest)) {
     throw new Error(`Destination file path ${dest} already exists`)
@@ -90,12 +93,12 @@ async function downloadToolAttempt(
     allowRetries: false
   })
 
-  let headers: IHeaders | undefined
   if (auth) {
     core.debug('set auth')
-    headers = {
-      authorization: auth
+    if (headers === undefined) {
+      headers = {}
     }
+    headers.authorization = auth
   }
 
   const response: httpm.HttpClientResponse = await http.get(url, headers)