From 56106454d34489e684cccfc436ce6d84098ec542 Mon Sep 17 00:00:00 2001
From: Naman Gupta <naman.gupta@gmail.com>
Date: Mon, 27 Nov 2023 02:35:52 -0800
Subject: [PATCH] BR: Support encryption for restored ebs volumes (#48900)

close pingcap/tidb#48899
---
 br/pkg/aws/ebs.go               | 3 ++-
 br/pkg/task/restore.go          | 4 ++++
 br/pkg/task/restore_ebs_meta.go | 5 ++++-
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/br/pkg/aws/ebs.go b/br/pkg/aws/ebs.go
index 2e0cfa36f6e65..18f4f5b484a1b 100644
--- a/br/pkg/aws/ebs.go
+++ b/br/pkg/aws/ebs.go
@@ -550,7 +550,7 @@ func fetchTargetSnapshots(meta *config.EBSBasedBRMeta, specifiedAZ string) map[s
 // CreateVolumes create volumes from snapshots
 // if err happens in the middle, return half-done result
 // returned map: store id -> old volume id -> new volume id
-func (e *EC2Session) CreateVolumes(meta *config.EBSBasedBRMeta, volumeType string, iops, throughput int64, targetAZ string) (map[string]string, error) {
+func (e *EC2Session) CreateVolumes(meta *config.EBSBasedBRMeta, volumeType string, iops, throughput int64, encrypted bool, targetAZ string) (map[string]string, error) {
 	template := ec2.CreateVolumeInput{
 		VolumeType: &volumeType,
 	}
@@ -560,6 +560,7 @@ func (e *EC2Session) CreateVolumes(meta *config.EBSBasedBRMeta, volumeType strin
 	if throughput > 0 {
 		template.SetThroughput(throughput)
 	}
+	template.Encrypted = &encrypted
 
 	newVolumeIDMap := make(map[string]string)
 	var mutex sync.Mutex
diff --git a/br/pkg/task/restore.go b/br/pkg/task/restore.go
index f51f3cacd8a81..a890ce746f93e 100644
--- a/br/pkg/task/restore.go
+++ b/br/pkg/task/restore.go
@@ -217,6 +217,7 @@ type RestoreConfig struct {
 	VolumeType          pconfig.EBSVolumeType `json:"volume-type" toml:"volume-type"`
 	VolumeIOPS          int64                 `json:"volume-iops" toml:"volume-iops"`
 	VolumeThroughput    int64                 `json:"volume-throughput" toml:"volume-throughput"`
+	VolumeEncrypted     bool                  `json:"volume-encrypted" toml:"volume-encrypted"`
 	ProgressFile        string                `json:"progress-file" toml:"progress-file"`
 	TargetAZ            string                `json:"target-az" toml:"target-az"`
 	UseFSR              bool                  `json:"use-fsr" toml:"use-fsr"`
@@ -382,6 +383,9 @@ func (cfg *RestoreConfig) ParseFromFlags(flags *pflag.FlagSet) error {
 		if cfg.VolumeThroughput, err = flags.GetInt64(flagVolumeThroughput); err != nil {
 			return errors.Trace(err)
 		}
+		if cfg.VolumeEncrypted, err = flags.GetBool(flagVolumeEncrypted); err != nil {
+			return errors.Trace(err)
+		}
 
 		cfg.ProgressFile, err = flags.GetString(flagProgressFile)
 		if err != nil {
diff --git a/br/pkg/task/restore_ebs_meta.go b/br/pkg/task/restore_ebs_meta.go
index cbb5d509f6133..2835e5e5a64f3 100644
--- a/br/pkg/task/restore_ebs_meta.go
+++ b/br/pkg/task/restore_ebs_meta.go
@@ -41,6 +41,7 @@ const (
 	flagVolumeType       = "volume-type"
 	flagVolumeIOPS       = "volume-iops"
 	flagVolumeThroughput = "volume-throughput"
+	flagVolumeEncrypted  = "volume-encrypted"
 	flagTargetAZ         = "target-az"
 )
 
@@ -54,6 +55,7 @@ func DefineRestoreSnapshotFlags(command *cobra.Command) {
 	command.Flags().String(flagVolumeType, string(config.GP3Volume), "volume type: gp3, io1, io2")
 	command.Flags().Int64(flagVolumeIOPS, 0, "volume iops(0 means default for that volume type)")
 	command.Flags().Int64(flagVolumeThroughput, 0, "volume throughout in MiB/s(0 means default for that volume type)")
+	command.Flags().Bool(flagVolumeEncrypted, false, "whether encryption is enabled for the volume")
 	command.Flags().String(flagProgressFile, "progress.txt", "the file name of progress file")
 	command.Flags().String(flagTargetAZ, "", "the target AZ for restored volumes")
 
@@ -65,6 +67,7 @@ func DefineRestoreSnapshotFlags(command *cobra.Command) {
 	_ = command.Flags().MarkHidden(flagVolumeType)
 	_ = command.Flags().MarkHidden(flagVolumeIOPS)
 	_ = command.Flags().MarkHidden(flagVolumeThroughput)
+	_ = command.Flags().MarkHidden(flagVolumeEncrypted)
 	_ = command.Flags().MarkHidden(flagProgressFile)
 	_ = command.Flags().MarkHidden(flagTargetAZ)
 }
@@ -256,7 +259,7 @@ func (h *restoreEBSMetaHelper) restoreVolumes(progress glue.Progress) (map[strin
 	}
 
 	volumeIDMap, err = ec2Session.CreateVolumes(h.metaInfo,
-		string(h.cfg.VolumeType), h.cfg.VolumeIOPS, h.cfg.VolumeThroughput, h.cfg.TargetAZ)
+		string(h.cfg.VolumeType), h.cfg.VolumeIOPS, h.cfg.VolumeThroughput, h.cfg.VolumeEncrypted, h.cfg.TargetAZ)
 	if err != nil {
 		return nil, 0, errors.Trace(err)
 	}