From 4b30748efcca66c259ac976201119168878ca884 Mon Sep 17 00:00:00 2001
From: Bradley Ayers <brad@dovetail.com>
Date: Fri, 27 Oct 2023 16:35:44 +1100
Subject: [PATCH] Discard broken S3 downloads

---
 lib/backends/s3.bash | 10 ++++-
 tests/command.bats   | 92 +++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 95 insertions(+), 7 deletions(-)

diff --git a/lib/backends/s3.bash b/lib/backends/s3.bash
index dd34f31..1db2610 100644
--- a/lib/backends/s3.bash
+++ b/lib/backends/s3.bash
@@ -112,7 +112,15 @@ function restore() {
   fi
 
   if [[ ! "${BK_AWS_FOUND}" =~ (false) ]]; then
-    aws s3 cp ${BK_CUSTOM_AWS_ARGS} "s3://${BUCKET}/${TAR_FILE}" .
+    TMP_FILE="$(mktemp)"
+    aws s3 cp ${BK_CUSTOM_AWS_ARGS} "s3://${BUCKET}/${TAR_FILE}" "${TMP_FILE}"  || s3_download_failed=true
+    if ${s3_download_failed:-false}; then
+      echo -e "S3 download failed, soft failing and skipping cache restore..."
+      return 0
+    else
+      mv -f "${TMP_FILE}" "${TAR_FILE}"
+    fi
+
     [ "${BK_CACHE_SAVE_CACHE}" == "true" ] && cp "${TAR_FILE}" "${BK_CACHE_LOCAL_PATH}/${TAR_FILE}"
     tar ${BK_TAR_EXTRACT_ARGS} "${TAR_FILE}" -C .
   else
diff --git a/tests/command.bats b/tests/command.bats
index 93876e0..4fb9851 100644
--- a/tests/command.bats
+++ b/tests/command.bats
@@ -10,10 +10,15 @@ setup() {
 }
 
 @test "Pre-command restores cache with basic key" {
+  stub mktemp \
+   " : echo '/tmp/tempfile'"
+
+  stub mv \
+    "-f /tmp/tempfile v1-cache-key.tar : true"
 
   stub aws \
    "s3api head-object --bucket my-bucket --key 'my-org/my-pipeline/v1-cache-key.tar' --profile my-profile : true" \
-   "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-key.tar . : echo Copied from S3"
+   "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-key.tar /tmp/tempfile : echo Copied from S3"
 
   stub tar \
    "-xf v1-cache-key.tar -C . : echo Extracted tar archive"
@@ -40,6 +45,8 @@ setup() {
 
   unstub aws
   unstub tar
+  unstub mv
+  unstub mktemp
 }
 
 @test "Pre-command restores S3 backed cache using local file" {
@@ -118,9 +125,15 @@ setup() {
 @test "Cache key template evaluation on file" {
   CHECKSUM=355831032f586e782b45744f2ed79316cc830244
 
+  stub mktemp \
+   " : echo '/tmp/tempfile'"
+
+  stub mv \
+    "-f /tmp/tempfile v1-cache-key-${CHECKSUM}.tar : true"
+
   stub aws \
    "s3api head-object --bucket my-bucket --key 'my-org/my-pipeline/v1-cache-key-${CHECKSUM}.tar' --profile my-profile : true" \
-   "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-key-${CHECKSUM}.tar . : echo Copied from S3"
+   "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-key-${CHECKSUM}.tar /tmp/tempfile : echo Copied from S3"
 
   stub tar \
    "-xf v1-cache-key-${CHECKSUM}.tar -C . : echo Extracted tar archive"
@@ -144,6 +157,8 @@ setup() {
   unset BUILDKITE_PIPELINE_SLUG
   unset BUILDKITE_ORGANIZATION_SLUG
 
+  unstub mktemp
+  unstub mv
   unstub aws
   unstub tar
 }
@@ -151,9 +166,15 @@ setup() {
 @test "Cache key template evaluation on dir" {
   CHECKSUM=4cfa4e590847976f26d761074e355e4d95fa8107
 
+  stub mktemp \
+   " : echo '/tmp/tempfile'"
+
+  stub mv \
+    "-f /tmp/tempfile v1-cache-key-${CHECKSUM}.tar : true"
+
   stub aws \
    "s3api head-object --bucket my-bucket --key 'my-org/my-pipeline/v1-cache-key-${CHECKSUM}.tar' --profile my-profile : true" \
-   "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-key-${CHECKSUM}.tar . : echo Copied from S3"
+   "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-key-${CHECKSUM}.tar /tmp/tempfile : echo Copied from S3"
 
   stub tar \
    "-xf v1-cache-key-${CHECKSUM}.tar -C . : echo Extracted tar archive"
@@ -177,6 +198,8 @@ setup() {
   unset BUILDKITE_PIPELINE_SLUG
   unset BUILDKITE_ORGANIZATION_SLUG
 
+  unstub mktemp
+  unstub mv
   unstub aws
   unstub tar
 }
@@ -184,9 +207,15 @@ setup() {
 @test "Cache key multi-template evaluation" {
   CHECKSUMS=355831032f586e782b45744f2ed79316cc830244-241bc31c8ddc004c48e6d88d7fa51ee981b8ce51
 
+  stub mktemp \
+   " : echo '/tmp/tempfile'"
+
+  stub mv \
+    "-f /tmp/tempfile v1-cache-key-${CHECKSUMS}.tar : true"
+
   stub aws \
    "s3api head-object --bucket my-bucket --key 'my-org/my-pipeline/v1-cache-key-${CHECKSUMS}.tar' --profile my-profile : true" \
-   "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-key-${CHECKSUMS}.tar . : echo Copied from S3"
+   "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-key-${CHECKSUMS}.tar /tmp/tempfile : echo Copied from S3"
 
   stub tar \
    "-xf v1-cache-key-${CHECKSUMS}.tar -C . : echo Extracted tar archive"
@@ -210,6 +239,8 @@ setup() {
   unset BUILDKITE_PIPELINE_SLUG
   unset BUILDKITE_ORGANIZATION_SLUG
 
+  unstub mktemp
+  unstub mv
   unstub aws
   unstub tar
 }
@@ -217,9 +248,15 @@ setup() {
 @test "Cache key template evaluation in middle of key" {
   CHECKSUM=355831032f586e782b45744f2ed79316cc830244
 
+  stub mktemp \
+   " : echo '/tmp/tempfile'"
+
+  stub mv \
+    "-f /tmp/tempfile v1-cache-${CHECKSUM}-key.tar : true"
+
   stub aws \
    "s3api head-object --bucket my-bucket --key 'my-org/my-pipeline/v1-cache-$CHECKSUM-key.tar' --profile my-profile : true" \
-   "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-$CHECKSUM-key.tar . : echo Copied from S3"
+   "s3 cp --profile my-profile s3://my-bucket/my-org/my-pipeline/v1-cache-$CHECKSUM-key.tar /tmp/tempfile : echo Copied from S3"
 
   stub tar \
    "-xf v1-cache-$CHECKSUM-key.tar -C . : echo Extracted tar archive"
@@ -243,6 +280,8 @@ setup() {
   unset BUILDKITE_PIPELINE_SLUG
   unset BUILDKITE_ORGANIZATION_SLUG
 
+  unstub mktemp
+  unstub mv
   unstub aws
   unstub tar
 }
@@ -271,9 +310,15 @@ setup() {
 
 @test "S3 arguments are passed through to copy command" {
 
+  stub mktemp \
+   " : echo '/tmp/tempfile'"
+
+  stub mv \
+    "-f /tmp/tempfile v1-cache-key.tar : true"
+
   stub aws \
    "s3api head-object --bucket my-bucket --key 'my-org/my-pipeline/v1-cache-key.tar' --profile my-profile : true" \
-   "s3 cp --profile my-profile --acl bucket-owner-full-control s3://my-bucket/my-org/my-pipeline/v1-cache-key.tar . : echo Copied from S3"
+   "s3 cp --profile my-profile --acl bucket-owner-full-control s3://my-bucket/my-org/my-pipeline/v1-cache-key.tar /tmp/tempfile : echo Copied from S3"
 
   stub tar \
    "-xf v1-cache-key.tar -C . : echo Extracted tar archive"
@@ -299,6 +344,41 @@ setup() {
   unset BUILDKITE_PIPELINE_SLUG
   unset BUILDKITE_ORGANIZATION_SLUG
 
+  unstub mktemp
+  unstub mv
   unstub aws
   unstub tar
 }
+
+@test "S3 download errors cause soft failure" {
+
+  stub mktemp \
+   " : echo '/tmp/tempfile'"
+
+  stub aws \
+   "s3api head-object --bucket my-bucket --key 'my-org/my-pipeline/v1-cache-key.tar' --profile my-profile : true" \
+   "s3 cp --profile my-profile --acl bucket-owner-full-control s3://my-bucket/my-org/my-pipeline/v1-cache-key.tar /tmp/tempfile : false"
+
+  export BUILDKITE_ORGANIZATION_SLUG="my-org"
+  export BUILDKITE_PIPELINE_SLUG="my-pipeline"
+  export BUILDKITE_PLUGIN_CACHE_S3_BUCKET="my-bucket"
+  export BUILDKITE_PLUGIN_CACHE_S3_PROFILE="my-profile"
+  export BUILDKITE_PLUGIN_CACHE_S3_ARGS="--acl bucket-owner-full-control"
+  export BUILDKITE_PLUGIN_CACHE_BACKEND="s3"
+  export BUILDKITE_PLUGIN_CACHE_KEY="v1-cache-key"
+
+  run "$PWD/hooks/pre-command"
+  assert_success
+  assert_output --partial "S3 download failed"
+
+  unset BUILDKITE_PLUGIN_CACHE_KEY
+  unset BUILDKITE_PLUGIN_CACHE_BACKEND
+  unset BUILDKITE_PLUGIN_CACHE_S3_PROFILE
+  unset BUILDKITE_PLUGIN_CACHE_S3_ARGS
+  unset BUILDKITE_PLUGIN_CACHE_S3_BUCKET
+  unset BUILDKITE_PIPELINE_SLUG
+  unset BUILDKITE_ORGANIZATION_SLUG
+
+  unstub mktemp
+  unstub aws
+}
\ No newline at end of file