From 2a8918bb11bbff6d5418f714a605d44dba5234bf Mon Sep 17 00:00:00 2001
From: Yanchen Ma <yanchen@whist.com>
Date: Tue, 8 Mar 2022 11:34:55 -0500
Subject: [PATCH] tests

---
 .../host-service/mandelbox/user_configs.go    |  2 +-
 .../mandelbox/user_configs_test.go            | 54 +++++++++++++++++++
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/backend/services/host-service/mandelbox/user_configs.go b/backend/services/host-service/mandelbox/user_configs.go
index f98af15fc54..4a76674c820 100644
--- a/backend/services/host-service/mandelbox/user_configs.go
+++ b/backend/services/host-service/mandelbox/user_configs.go
@@ -316,11 +316,11 @@ func (mandelbox *mandelboxData) downloadUserConfig(s3Client *s3.Client, key stri
 		validateIntegrity = false
 	}
 
-	// Allow up to 3 retries on download if failure is because of checksum mismatch
 	var downloadedFile []byte
 	var downloadSuccessful bool
 	var numBytesSuccessfullyDownloaded int64
 
+	// Allow up to 3 retries on download if failure is because of checksum mismatch
 	for i := 0; i < 3; i++ {
 		// Download file into a pre-allocated in-memory buffer
 		// This should be okay as we don't expect configs to be very large
diff --git a/backend/services/host-service/mandelbox/user_configs_test.go b/backend/services/host-service/mandelbox/user_configs_test.go
index 5bc1050a329..3b7a68ca143 100644
--- a/backend/services/host-service/mandelbox/user_configs_test.go
+++ b/backend/services/host-service/mandelbox/user_configs_test.go
@@ -1,10 +1,12 @@
 package mandelbox
 
 import (
+	"bytes"
 	"context"
 	"os"
 	"os/exec"
 	"path"
+	"strings"
 	"sync"
 	"testing"
 
@@ -154,6 +156,58 @@ func TestUserConfigIntegration(t *testing.T) {
 	t.Run("new non-repeated token", downloadTestNotNewToken(false, true))
 }
 
+// TestUserConfigChecksum tests the checksum verification functonality of
+// the upload and download user config code.
+func TestUserConfigChecksum(t *testing.T) {
+	testFile1 := []byte("this is a test file")
+	testFile1Hash := configutils.GetMD5Hash(testFile1)
+
+	// Test upload
+	s3Client, err := configutils.NewS3Client("us-east-1")
+	if err != nil {
+		t.Fatalf("failed to create s3 client: %v", err)
+	}
+
+	_, err = configutils.UploadFileToBucket(s3Client, UserConfigS3Bucket, "user_config_test_user/checksum_test", testFile1)
+	if err != nil {
+		t.Fatalf("failed to upload test file: %v", err)
+	}
+
+	// Check that checksum is present and correct
+	head, err := configutils.GetHeadObject(s3Client, UserConfigS3Bucket, "user_config_test_user/checksum_test")
+	if err != nil {
+		t.Fatalf("failed to get head object: %v", err)
+	}
+
+	metadata := head.Metadata
+	downloadHash, ok := metadata["md5"]
+	if !ok {
+		t.Fatalf("md5 metadata not present on head object")
+	}
+
+	if downloadHash != testFile1Hash {
+		t.Fatalf("md5 metadata on head object does not match file hash")
+	}
+
+	// Test download happy path
+	mandelbox := mandelboxData{ID: types.MandelboxID(utils.PlaceholderTestUUID())}
+	data, err := mandelbox.downloadUserConfig(s3Client, "user_config_test_user/checksum_test", head)
+	if err != nil {
+		t.Errorf("failed to download file: %v", err)
+	}
+
+	if !bytes.Equal(data, testFile1) {
+		t.Errorf("downloaded file does not match original file")
+	}
+
+	// Test download with bad checksum
+	head.Metadata["md5"] = "bad_hash"
+	_, err = mandelbox.downloadUserConfig(s3Client, "user_config_test_user/checksum_test", head)
+	if err == nil || !strings.HasPrefix(err.Error(), "Could not download object (due to MD5 mismatch)") {
+		t.Errorf("download with bad checksum did not fail")
+	}
+}
+
 // cleanupTestDirs removes the created directories created by the integration
 // test. This allows the test to be runnable multiple times without
 // creation errors.