From a1d18e6f2c8a53485100017b3c69ad83a3f273ed Mon Sep 17 00:00:00 2001
From: Wenqi Mou <wenqimou@gmail.com>
Date: Tue, 17 Sep 2024 18:57:53 -0400
Subject: [PATCH] add checksum and some tests

---
 br/pkg/stream/stream_mgr.go   | 16 +++++++++++++---
 br/tests/br_encryption/run.sh | 20 +++++++++++++++++---
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/br/pkg/stream/stream_mgr.go b/br/pkg/stream/stream_mgr.go
index 4b977b3243450..fa51d68d99fdc 100644
--- a/br/pkg/stream/stream_mgr.go
+++ b/br/pkg/stream/stream_mgr.go
@@ -15,7 +15,9 @@
 package stream
 
 import (
+	"bytes"
 	"context"
+	"crypto/sha256"
 	"strings"
 
 	"github.com/klauspost/compress/zstd"
@@ -195,7 +197,7 @@ func (m *MetadataHelper) decodeCompressedData(data []byte, compressionType backu
 		"failed to decode compressed data: compression type is unimplemented. type id is %d", compressionType)
 }
 
-func (m *MetadataHelper) decryptIfNeeded(ctx context.Context, data []byte, encryptionInfo *encryptionpb.FileEncryptionInfo) ([]byte, error) {
+func (m *MetadataHelper) verifyChecksumAndDecryptIfNeeded(ctx context.Context, data []byte, encryptionInfo *encryptionpb.FileEncryptionInfo) ([]byte, error) {
 	// no need to decrypt
 	if encryptionInfo == nil {
 		return data, nil
@@ -205,6 +207,14 @@ func (m *MetadataHelper) decryptIfNeeded(ctx context.Context, data []byte, encry
 		return data, errors.New("need to decrypt data but encryption manager not set")
 	}
 
+	// Verify checksum before decryption
+	if encryptionInfo.Checksum != nil {
+		actualChecksum := sha256.Sum256(data)
+		if !bytes.Equal(actualChecksum[:], encryptionInfo.Checksum) {
+			return nil, errors.New("checksum mismatch before decryption")
+		}
+	}
+
 	decryptedContent, err := m.encryptionManager.Decrypt(ctx, data, encryptionInfo)
 	if err != nil {
 		return nil, errors.Trace(err)
@@ -235,7 +245,7 @@ func (m *MetadataHelper) ReadFile(
 			return nil, errors.Trace(err)
 		}
 		// decrypt if needed
-		decryptedData, err := m.decryptIfNeeded(ctx, data, encryptionInfo)
+		decryptedData, err := m.verifyChecksumAndDecryptIfNeeded(ctx, data, encryptionInfo)
 		if err != nil {
 			return nil, errors.Trace(err)
 		}
@@ -251,7 +261,7 @@ func (m *MetadataHelper) ReadFile(
 		}
 	}
 	// decrypt if needed
-	decryptedData, err := m.decryptIfNeeded(ctx, cref.data[offset:offset+length], encryptionInfo)
+	decryptedData, err := m.verifyChecksumAndDecryptIfNeeded(ctx, cref.data[offset:offset+length], encryptionInfo)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
diff --git a/br/tests/br_encryption/run.sh b/br/tests/br_encryption/run.sh
index 5252b9a38a251..054efaf617baf 100755
--- a/br/tests/br_encryption/run.sh
+++ b/br/tests/br_encryption/run.sh
@@ -166,12 +166,17 @@ run_backup_restore_test() {
     # Drop databases before restoring
     drop_db || { echo "Failed to drop databases before restore"; exit 1; }
 
-    # Run pitr restore
+    # Run pitr restore and measure the performance
     echo "restore log backup with $full_encryption_args and $log_encryption_args"
+    local start_time=$(date +%s.%N)
     timeout 300 run_br --pd "$PD_ADDR" restore point -s "local://$TEST_DIR/$PREFIX/log" --full-backup-storage "local://$TEST_DIR/$PREFIX/full" $full_encryption_args $log_encryption_args || { 
         echo "Log backup restore failed or timed out after 5 minutes"
         exit 1
     }
+    local end_time=$(date +%s.%N)
+    local duration=$(echo "$end_time - $start_time" | bc | awk '{printf "%.3f", $0}')
+    echo "${encryption_mode} took ${duration} seconds"
+    echo "${encryption_mode},${duration}" >> "$TEST_DIR/performance_results.csv"
 
     # Check data consistency after restore
     echo "check data consistency after restore"
@@ -318,14 +323,23 @@ test_mixed_full_plain_log_encrypted() {
     run_backup_restore_test "mixed_full_plain_log_encrypted" "$full_encryption_args" "$log_encryption_args"
 }
 
+# Initialize performance results file
+echo "Operation,Encryption Mode,Duration (seconds)" > "$TEST_DIR/performance_results.csv"
+
 # Run tests
-#test_plaintext
+test_plaintext
 test_plaintext_data_key
 test_local_master_key
 test_aws_kms
+test_mixed_full_encrypted_log_plain
+test_mixed_full_plain_log_encrypted
 
-# uncomment for manual testing
+# uncomment for manual GCP KMS testing
 #test_gcp_kms
 
 echo "All encryption tests passed successfully"
 
+# Display performance results
+echo "Performance Results:"
+cat "$TEST_DIR/performance_results.csv"
+