initial

Signed-off-by: Wenqi Mou <wenqimou@gmail.com>

br/cmd/br/restore.go

@@ -74,14 +74,14 @@ func runRestoreCommand(command *cobra.Command, cmdName string) error {
 
 	if err := task.RunRestore(GetDefaultContext(), tidbGlue, cmdName, &cfg); err != nil {
 		log.Error("failed to restore", zap.Error(err))
-		printWorkaroundOnFullRestoreError(command, err)
+		printWorkaroundOnFullRestoreError(err)
 		return errors.Trace(err)
 	}
 	return nil
 }
 
 // print workaround when we met not fresh or incompatible cluster error on full cluster restore
-func printWorkaroundOnFullRestoreError(command *cobra.Command, err error) {
+func printWorkaroundOnFullRestoreError(err error) {
 	if !errors.ErrorEqual(err, berrors.ErrRestoreNotFreshCluster) &&
 		!errors.ErrorEqual(err, berrors.ErrRestoreIncompatibleSys) {
 		return

br/pkg/checkpoint/checkpoint.go

@@ -724,7 +724,7 @@ func walkCheckpointFile[K KeyType, V ValueType](
 				pastDureTime = checkpointData.DureTime
 			}
 			for _, meta := range checkpointData.RangeGroupMetas {
-				decryptContent, err := metautil.Decrypt(meta.RangeGroupsEncriptedData, cipher, meta.CipherIv)
+				decryptContent, err := utils.Decrypt(meta.RangeGroupsEncriptedData, cipher, meta.CipherIv)
 				if err != nil {
 					return errors.Trace(err)
 				}

br/pkg/conn/conn.go

@@ -296,8 +296,8 @@ func (mgr *Mgr) Close() {
 	mgr.PdController.Close()
 }
 
-// GetTS gets current ts from pd.
-func (mgr *Mgr) GetTS(ctx context.Context) (uint64, error) {
+// GetCurrentTsFromPd gets current ts from pd.
+func (mgr *Mgr) GetCurrentTsFromPd(ctx context.Context) (uint64, error) {
 	p, l, err := mgr.GetPDClient().GetTS(ctx)
 	if err != nil {
 		return 0, errors.Trace(err)

br/pkg/encryption/BUILD.bazel

@@ -0,0 +1,15 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "encryption",
+    srcs = ["manager.go"],
+    importpath = "github.com/pingcap/tidb/br/pkg/encryption",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//br/pkg/encryption/master_key",
+        "//br/pkg/utils",
+        "@com_github_pingcap_errors//:errors",
+        "@com_github_pingcap_kvproto//pkg/brpb",
+        "@com_github_pingcap_kvproto//pkg/encryptionpb",
+    ],
+)

br/pkg/encryption/manager.go

@@ -0,0 +1,82 @@
+package encryption
+
+import (
+	"context"
+
+	"github.com/pingcap/errors"
+	backuppb "github.com/pingcap/kvproto/pkg/brpb"
+	"github.com/pingcap/kvproto/pkg/encryptionpb"
+	encryption "github.com/pingcap/tidb/br/pkg/encryption/master_key"
+	"github.com/pingcap/tidb/br/pkg/utils"
+)
+
+type Manager struct {
+	cipherInfo        *backuppb.CipherInfo
+	masterKeyBackends *encryption.MultiMasterKeyBackend
+	encryptionMethod  *encryptionpb.EncryptionMethod
+}
+
+func NewManager(cipherInfo *backuppb.CipherInfo, masterKeyConfigs *backuppb.MasterKeyConfig) (*Manager, error) {
+	// should never happen since config has default
+	if cipherInfo == nil || masterKeyConfigs == nil {
+		return nil, errors.New("cipherInfo or masterKeyConfigs is nil")
+	}
+
+	if cipherInfo.CipherType != encryptionpb.EncryptionMethod_PLAINTEXT {
+		return &Manager{
+			cipherInfo:        cipherInfo,
+			masterKeyBackends: nil,
+			encryptionMethod:  nil,
+		}, nil
+	}
+
+	if masterKeyConfigs.EncryptionType != encryptionpb.EncryptionMethod_PLAINTEXT {
+		masterKeyBackends, err := encryption.NewMultiMasterKeyBackend(masterKeyConfigs.GetMasterKeys())
+		if err != nil {
+			return nil, errors.Trace(err)
+		}
+		return &Manager{
+			cipherInfo:        nil,
+			masterKeyBackends: masterKeyBackends,
+			encryptionMethod:  &masterKeyConfigs.EncryptionType,
+		}, nil
+	}
+	return nil, nil
+}
+
+func (m *Manager) Decrypt(ctx context.Context, content []byte, fileEncryptionInfo *encryptionpb.FileEncryptionInfo) ([]byte, error) {
+	switch mode := fileEncryptionInfo.Mode.(type) {
+	case *encryptionpb.FileEncryptionInfo_PlainTextDataKey:
+		if m.cipherInfo == nil {
+			return nil, errors.New("plaintext data key info is required but not set")
+		}
+		decryptedContent, err := utils.Decrypt(content, m.cipherInfo, fileEncryptionInfo.FileIv)
+		if err != nil {
+			return nil, errors.Annotate(err, "failed to decrypt content using plaintext data key")
+		}
+		return decryptedContent, nil
+	case *encryptionpb.FileEncryptionInfo_MasterKeyBased:
+		encryptedContents := fileEncryptionInfo.GetMasterKeyBased().DataKeyEncryptedContent
+		if encryptedContents == nil || len(encryptedContents) == 0 {
+			return nil, errors.New("should contain at least one encrypted data key")
+		}
+		// pick first one, the list is for future expansion of multiple encrypted data keys by different master key backend
+		encryptedContent := encryptedContents[0]
+		decryptedDataKey, err := m.masterKeyBackends.Decrypt(ctx, encryptedContent)
+		if err != nil {
+			return nil, errors.Annotate(err, "failed to decrypt data key using master key")
+		}
+
+		cipherInfo := backuppb.CipherInfo{
+			CipherType: fileEncryptionInfo.EncryptionMethod,
+			CipherKey:  decryptedDataKey,
+		}
+		decryptedContent, err := utils.Decrypt(content, &cipherInfo, fileEncryptionInfo.FileIv)
+		if err != nil {
+			return nil, errors.Annotate(err, "failed to decrypt content using decrypted data key")
+		}
+		return decryptedContent, nil
+	default:
+		return nil, errors.Errorf("internal error: unsupported encryption mode type %T", mode)
+	}
+}

br/pkg/encryption/master_key/BUILD.bazel

@@ -0,0 +1,41 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "master_key",
+    srcs = [
+        "common.go",
+        "file_backend.go",
+        "kms_backend.go",
+        "master_key.go",
+        "mem_backend.go",
+        "multi_master_key_backend.go",
+    ],
+    importpath = "github.com/pingcap/tidb/br/pkg/encryption/master_key",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//br/pkg/kms:aws",
+        "//br/pkg/utils",
+        "@com_github_pingcap_errors//:errors",
+        "@com_github_pingcap_kvproto//pkg/encryptionpb",
+        "@com_github_pingcap_log//:log",
+        "@org_uber_go_zap//:zap",
+    ],
+)
+
+go_test(
+    name = "master_key_test",
+    srcs = [
+        "file_backend_test.go",
+        "kms_backend_test.go",
+        "mem_backend_test.go",
+        "multi_master_key_backend_test.go",
+    ],
+    embed = [":master_key"],
+    deps = [
+        "@com_github_pingcap_errors//:errors",
+        "@com_github_pingcap_kvproto//pkg/encryptionpb",
+        "@com_github_stretchr_testify//assert",
+        "@com_github_stretchr_testify//mock",
+        "@com_github_stretchr_testify//require",
+    ],
+)

br/pkg/encryption/master_key/common.go

@@ -0,0 +1,29 @@
+package encryption
+
+import (
+	"crypto/rand"
+	"encoding/binary"
+	"time"
+)
+
+// must keep it same with the constants in TiKV implementation
+const (
+	MetadataKeyMethod           string = "method"
+	MetadataKeyIv               string = "iv"
+	MetadataKeyAesGcmTag        string = "aes_gcm_tag"
+	MetadataKeyKmsVendor        string = "kms_vendor"
+	MetadataKeyKmsCiphertextKey string = "kms_ciphertext_key"
+	MetadataMethodAes256Gcm     string = "aes256-gcm"
+)
+
+type IV [12]byte
+
+func NewIV() IV {
+	var iv IV
+	binary.BigEndian.PutUint64(iv[:8], uint64(time.Now().UnixNano()))
+	// Fill the remaining 4 bytes with random data
+	if _, err := rand.Read(iv[8:]); err != nil {
+		panic(err) // Handle this error appropriately in production code
+	}
+	return iv
+}

br/pkg/encryption/master_key/file_backend.go

@@ -0,0 +1,58 @@
+package encryption
+
+import (
+	"context"
+	"encoding/hex"
+	"os"
+
+	"github.com/pingcap/errors"
+	"github.com/pingcap/kvproto/pkg/encryptionpb"
+)
+
+const AesGcmKeyLen = 32 // AES-256 key length
+
+type FileBackend struct {
+	memCache *MemAesGcmBackend
+}
+
+func createFileBackend(keyPath string) (*FileBackend, error) {
+	// FileBackend uses AES-256-GCM
+	keyLen := AesGcmKeyLen
+
+	content, err := os.ReadFile(keyPath)
+	if err != nil {
+		return nil, errors.Annotate(err, "failed to read master key file from disk")
+	}
+
+	fileLen := len(content)
+	expectedLen := keyLen*2 + 1 // hex-encoded key + newline
+
+	if fileLen != expectedLen {
+		return nil, errors.Errorf("mismatch master key file size, expected %d, actual %d", expectedLen, fileLen)
+	}
+
+	if content[fileLen-1] != '\n' {
+		return nil, errors.Errorf("master key file should end with newline")
+	}
+
+	key, err := hex.DecodeString(string(content[:fileLen-1]))
+	if err != nil {
+		return nil, errors.Annotate(err, "failed to decode master key from file")
+	}
+
+	backend, err := NewMemAesGcmBackend(key)
+	if err != nil {
+		return nil, errors.Annotate(err, "failed to create MemAesGcmBackend")
+	}
+
+	return &FileBackend{memCache: backend}, nil
+}
+
+func (f *FileBackend) Encrypt(ctx context.Context, plaintext []byte) (*encryptionpb.EncryptedContent, error) {
+	iv := NewIV()
+	return f.memCache.EncryptContent(ctx, plaintext, iv)
+}
+
+func (f *FileBackend) Decrypt(ctx context.Context, content *encryptionpb.EncryptedContent) ([]byte, error) {
+	return f.memCache.DecryptContent(ctx, content)
+}

br/pkg/encryption/master_key/file_backend_test.go

@@ -0,0 +1,103 @@
+package encryption
+
+import (
+	"context"
+	"encoding/hex"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TempKeyFile represents a temporary key file for testing
+type TempKeyFile struct {
+	Path string
+	file *os.File
+}
+
+// Cleanup closes and removes the temporary file
+func (tkf *TempKeyFile) Cleanup() {
+	if tkf.file != nil {
+		tkf.file.Close()
+	}
+	os.Remove(tkf.Path)
+}
+
+// createMasterKeyFile creates a temporary master key file for testing
+func createMasterKeyFile() (*TempKeyFile, error) {
+	tempFile, err := os.CreateTemp("", "test_key_*.txt")
+	if err != nil {
+		return nil, err
+	}
+
+	_, err = tempFile.WriteString("c3d99825f2181f4808acd2068eac7441a65bd428f14d2aab43fefc0129091139\n")
+	if err != nil {
+		tempFile.Close()
+		os.Remove(tempFile.Name())
+		return nil, err
+	}
+
+	return &TempKeyFile{
+		Path: tempFile.Name(),
+		file: tempFile,
+	}, nil
+}
+
+func TestFileBackendAes256Gcm(t *testing.T) {
+	pt, err := hex.DecodeString("25431587e9ecffc7c37f8d6d52a9bc3310651d46fb0e3bad2726c8f2db653749")
+	require.NoError(t, err)
+	ct, err := hex.DecodeString("84e5f23f95648fa247cb28eef53abec947dbf05ac953734618111583840bd980")
+	require.NoError(t, err)
+	iv, err := hex.DecodeString("cafabd9672ca6c79a2fbdc22")
+	require.NoError(t, err)
+
+	tempKeyFile, err := createMasterKeyFile()
+	require.NoError(t, err)
+	defer tempKeyFile.Cleanup()
+
+	backend, err := createFileBackend(tempKeyFile.Path)
+	require.NoError(t, err)
+
+	ctx := context.Background()
+	encryptedContent, err := backend.memCache.EncryptContent(ctx, pt, IV(iv))
+	require.NoError(t, err)
+	assert.Equal(t, ct, encryptedContent.Content)
+
+	plaintext, err := backend.Decrypt(ctx, encryptedContent)
+	require.NoError(t, err)
+	assert.Equal(t, pt, plaintext)
+}
+
+func TestFileBackendAuthenticate(t *testing.T) {
+	pt := []byte{1, 2, 3}
+
+	tempKeyFile, err := createMasterKeyFile()
+	require.NoError(t, err)
+	defer tempKeyFile.Cleanup()
+
+	backend, err := createFileBackend(tempKeyFile.Path)
+	require.NoError(t, err)
+
+	ctx := context.Background()
+	encryptedContent, err := backend.Encrypt(ctx, pt)
+	require.NoError(t, err)
+
+	plaintext, err := backend.Decrypt(ctx, encryptedContent)
+	require.NoError(t, err)
+	assert.Equal(t, pt, plaintext)
+
+	// Test checksum mismatch
+	encryptedContent1 := *encryptedContent
+	encryptedContent1.Metadata[MetadataKeyAesGcmTag][0] ^= 0xFF
+	_, err = backend.Decrypt(ctx, &encryptedContent1)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), wrongMasterKey)
+
+	// Test checksum not found
+	encryptedContent2 := *encryptedContent
+	delete(encryptedContent2.Metadata, MetadataKeyAesGcmTag)
+	_, err = backend.Decrypt(ctx, &encryptedContent2)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), gcmTagNotFound)
+}