From 8b2a57849f89113fb4b7174d124a9b582746e34a Mon Sep 17 00:00:00 2001
From: Dave Hadka <dhadka@github.com>
Date: Wed, 22 Apr 2020 18:23:41 -0400
Subject: [PATCH] Adds socket timeout and validate file size

---
 src/cacheHttpClient.ts | 28 +++++++++++++++++++++++++++-
 src/constants.ts       |  2 ++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/src/cacheHttpClient.ts b/src/cacheHttpClient.ts
index c83e307..fc8f577 100644
--- a/src/cacheHttpClient.ts
+++ b/src/cacheHttpClient.ts
@@ -9,7 +9,7 @@ import {
 import * as crypto from "crypto";
 import * as fs from "fs";
 
-import { Inputs } from "./constants";
+import { Inputs, SocketTimeout } from "./constants";
 import {
     ArtifactCacheEntry,
     CommitCacheRequest,
@@ -144,7 +144,33 @@ export async function downloadCache(
     const stream = fs.createWriteStream(archivePath);
     const httpClient = new HttpClient("actions/cache");
     const downloadResponse = await httpClient.get(archiveLocation);
+
+    // Abort download if no traffic received over the socket.
+    downloadResponse.message.socket.setTimeout(SocketTimeout, () => {
+        downloadResponse.message.destroy();
+        core.debug(
+            `Aborting download, socket timed out after ${SocketTimeout} ms`
+        );
+    });
+
     await pipeResponseToStream(downloadResponse, stream);
+
+    // Validate download size.
+    var contentLengthHeader =
+        downloadResponse.message.headers["content-length"];
+
+    if (contentLengthHeader) {
+        const expectedLength = parseInt(contentLengthHeader);
+        const actualLength = utils.getArchiveFileSize(archivePath);
+
+        if (actualLength != expectedLength) {
+            throw new Error(
+                `Incomplete download. Expected file size: ${expectedLength}, actual file size: ${actualLength}`
+            );
+        }
+    } else {
+        core.debug("Unable to validate download, no Content-Length header");
+    }
 }
 
 // Reserve Cache
diff --git a/src/constants.ts b/src/constants.ts
index 2b78f62..320b15c 100644
--- a/src/constants.ts
+++ b/src/constants.ts
@@ -20,3 +20,5 @@ export enum Events {
 }
 
 export const CacheFilename = "cache.tgz";
+
+export const SocketTimeout = 5000;