feat(backup_service): support backup encryption

Author: ulya-sidorina

cmd/ydbcp/main.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"flag"
 	"fmt"
-	"github.com/ydb-platform/ydb-go-sdk/v3/log"
 	"net/http"
 	_ "net/http/pprof"
 	"os"
@@ -19,6 +18,7 @@ import (
 	"ydbcp/internal/connectors/db/yql/queries"
 	"ydbcp/internal/connectors/s3"
 	"ydbcp/internal/handlers"
+	"ydbcp/internal/kms"
 	"ydbcp/internal/metrics"
 	"ydbcp/internal/processor"
 	"ydbcp/internal/server"
@@ -32,6 +32,9 @@ import (
 	"ydbcp/internal/watchers/schedule_watcher"
 	"ydbcp/internal/watchers/ttl_watcher"
 	ap "ydbcp/pkg/plugins/auth"
+	kp "ydbcp/pkg/plugins/kms"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/log"
 
 	"github.com/jonboulle/clockwork"
 
@@ -114,6 +117,24 @@ func main() {
 		}
 	}()
 	xlog.Info(ctx, "Initialized AuthProvider")
+
+	var kmsProvider kp.KmsProvider
+	if len(configInstance.KMS.PluginPath) == 0 {
+		kmsProvider, err = kms.NewDummyKmsProvider(ctx)
+	} else {
+		kmsProvider, err = kms.NewKmsProvider(ctx, configInstance.KMS)
+	}
+	if err != nil {
+		xlog.Error(ctx, "Error init KmsProvider", zap.Error(err))
+		os.Exit(1)
+	}
+	defer func() {
+		if err := kmsProvider.Close(ctx); err != nil {
+			xlog.Error(ctx, "Error close kms provider", zap.Error(err))
+		}
+	}()
+	xlog.Info(ctx, "Initialized KmsProvider")
+
 	metrics.InitializeMetricsRegistry(ctx, &wg, &configInstance.MetricsServer, clockwork.NewRealClock())
 	xlog.Info(ctx, "Initialized metrics registry")
 	audit.EventsDestination = configInstance.Audit.EventsDestination
@@ -144,6 +165,7 @@ func main() {
 		dbConnector,
 		clientConnector,
 		authProvider,
+		kmsProvider,
 		*configInstance,
 	).Register(server)
 	operation.NewOperationService(dbConnector, authProvider).Register(server)
@@ -192,6 +214,7 @@ func main() {
 			queries.NewWriteTableQuery,
 			clockwork.NewRealClock(),
 			*configInstance,
+			kmsProvider,
 		),
 	); err != nil {
 		xlog.Error(ctx, "failed to register TBWR handler", zap.Error(err))
@@ -210,7 +233,7 @@ func main() {
 		xlog.Info(ctx, "Created TtlWatcher")
 	}
 
-	backupScheduleHandler := handlers.NewBackupScheduleHandler(queries.NewWriteTableQuery, clockwork.NewRealClock())
+	backupScheduleHandler := handlers.NewBackupScheduleHandler(queries.NewWriteTableQuery, clockwork.NewRealClock(), configInstance.FeatureFlags)
 
 	schedule_watcher.NewScheduleWatcher(
 		ctx, &wg, configInstance.OperationProcessor.ProcessorIntervalSeconds, dbConnector,

internal/backup_operations/make_backup.go

@@ -2,10 +2,9 @@ package backup_operations
 
 import (
 	"context"
+	"crypto/rand"
 	"errors"
 	"fmt"
-	"github.com/jonboulle/clockwork"
-	"github.com/ydb-platform/ydb-go-sdk/v3"
 	"path"
 	"regexp"
 	"strings"
@@ -14,8 +13,12 @@ import (
 	"ydbcp/internal/connectors/client"
 	"ydbcp/internal/types"
 	"ydbcp/internal/util/xlog"
+	kp "ydbcp/pkg/plugins/kms"
 	pb "ydbcp/pkg/proto/ydbcp/v1alpha1"
 
+	"github.com/jonboulle/clockwork"
+	"github.com/ydb-platform/ydb-go-sdk/v3"
+
 	"go.uber.org/zap"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -36,6 +39,7 @@ type MakeBackupInternalRequest struct {
 	ScheduleID           *string
 	Ttl                  *time.Duration
 	ParentOperationID    *string
+	EncryptionSettings   *pb.EncryptionSettings
 }
 
 func FromBackupSchedule(schedule *types.BackupSchedule) MakeBackupInternalRequest {
@@ -65,6 +69,7 @@ func FromTBWROperation(tbwr *types.TakeBackupWithRetryOperation) MakeBackupInter
 		ScheduleID:           tbwr.ScheduleID,
 		Ttl:                  tbwr.Ttl,
 		ParentOperationID:    &tbwr.ID,
+		EncryptionSettings:   tbwr.EncryptionSettings,
 	}
 }
 
@@ -282,6 +287,31 @@ func IsEmptyBackup(backup *types.Backup) bool {
 	return backup.Size == 0 && backup.S3Endpoint == ""
 }
 
+func GetEncryptionParams(settings *pb.EncryptionSettings) ([]byte, string, error) {
+	var algorithm string
+	var length int
+
+	switch settings.Algorithm {
+	case pb.EncryptionSettings_UNSPECIFIED:
+	case pb.EncryptionSettings_AES_128_GCM:
+		algorithm = "AES-128-GCM"
+		length = 16
+	case pb.EncryptionSettings_AES_256_GCM:
+		algorithm = "AES-256-GCM"
+		length = 32
+	case pb.EncryptionSettings_CHACHA20_POLY1305:
+		algorithm = "ChaCha20-Poly1305"
+		length = 32
+	}
+
+	dek := make([]byte, length)
+	_, err := rand.Read(dek)
+	if err != nil {
+		return nil, "", err
+	}
+	return dek, algorithm, nil
+}
+
 func MakeBackup(
 	ctx context.Context,
 	clientConn client.ClientConnector,
@@ -292,6 +322,7 @@ func MakeBackup(
 	subject string,
 	clock clockwork.Clock,
 	featureFlags config.FeatureFlagsConfig,
+	kmsProvider kp.KmsProvider,
 ) (*types.Backup, *types.TakeBackupOperation, error) {
 	if req.ScheduleID != nil {
 		ctx = xlog.With(ctx, zap.String("ScheduleID", *req.ScheduleID))
@@ -359,6 +390,37 @@ func MakeBackup(
 		S3ForcePathStyle:  s3.S3ForcePathStyle,
 	}
 
+	if req.EncryptionSettings != nil && featureFlags.EnableBackupEncryption {
+		dek, algorithm, err := GetEncryptionParams(req.EncryptionSettings)
+		if err != nil {
+			return nil, nil, err
+		}
+
+		s3Settings.EncryptionKey = dek
+		s3Settings.EncryptionAlgorithm = algorithm
+
+		kmsKey := req.EncryptionSettings.GetKmsKey()
+		if kmsKey == nil {
+			xlog.Error(ctx, "kms key is not specified")
+			return nil, nil, status.Errorf(codes.InvalidArgument, "kms key is not specified")
+		}
+
+		_, err = kmsProvider.Encrypt(
+			ctx,
+			&kp.EncryptRequest{
+				KeyID:     kmsKey.GetKeyId(),
+				Plaintext: dek,
+			},
+		)
+
+		if err != nil {
+			xlog.Error(ctx, "can't encrypt data encryption key", zap.Error(err))
+			return nil, nil, err
+		}
+
+		// TODO: save encrypted key to s3
+	}
+
 	clientOperationID, err := clientConn.ExportToS3(ctx, client, s3Settings, featureFlags)
 	if err != nil {
 		xlog.Error(ctx, "can't start export operation", zap.Error(err))
@@ -388,9 +450,10 @@ func MakeBackup(
 			CreatedAt: now,
 			Creator:   subject,
 		},
-		ScheduleID:  req.ScheduleID,
-		ExpireAt:    expireAt,
-		SourcePaths: pathsForExport,
+		ScheduleID:         req.ScheduleID,
+		ExpireAt:           expireAt,
+		SourcePaths:        pathsForExport,
+		EncryptionSettings: req.EncryptionSettings,
 	}
 
 	op := &types.TakeBackupOperation{
@@ -409,9 +472,10 @@ func MakeBackup(
 			CreatedAt: now,
 			Creator:   subject,
 		},
-		YdbOperationId:    clientOperationID,
-		UpdatedAt:         now,
-		ParentOperationID: req.ParentOperationID,
+		YdbOperationId:     clientOperationID,
+		UpdatedAt:          now,
+		ParentOperationID:  req.ParentOperationID,
+		EncryptionSettings: req.EncryptionSettings,
 	}
 
 	return backup, op, nil

internal/config/config.go

@@ -65,8 +65,9 @@ type MetricsServerConfig struct {
 }
 
 type FeatureFlagsConfig struct {
-	DisableTTLDeletion   bool `yaml:"disable_ttl_deletion" default:"false"`
-	EnableNewPathsFormat bool `yaml:"enable_new_paths_format" default:"false"`
+	DisableTTLDeletion     bool `yaml:"disable_ttl_deletion" default:"false"`
+	EnableNewPathsFormat   bool `yaml:"enable_new_paths_format" default:"false"`
+	EnableBackupEncryption bool `yaml:"enable_backup_encryption" default:"false"`
 }
 
 type LogConfig struct {
@@ -96,6 +97,7 @@ type Config struct {
 	ClientConnection   ClientConnectionConfig   `yaml:"client_connection"`
 	S3                 S3Config                 `yaml:"s3"`
 	Auth               PluginConfig             `yaml:"auth"`
+	KMS                PluginConfig             `yaml:"kms"`
 	GRPCServer         GRPCServerConfig         `yaml:"grpc_server"`
 	MetricsServer      MetricsServerConfig      `yaml:"metrics_server"`
 	OperationProcessor OperationProcessorConfig `yaml:"operation_processor"`

internal/connectors/client/connector.go

@@ -272,6 +272,17 @@ func (d *ClientYdbConnector) ExportToS3(
 		exportRequest.Settings.DestinationPrefix = s3Settings.DestinationPrefix
 	}
 
+	if featureFlags.EnableBackupEncryption && len(s3Settings.EncryptionKey) > 0 {
+		exportRequest.Settings.EncryptionSettings = &Ydb_Export.EncryptionSettings{
+			EncryptionAlgorithm: s3Settings.EncryptionAlgorithm,
+			Key: &Ydb_Export.EncryptionSettings_SymmetricKey_{
+				SymmetricKey: &Ydb_Export.EncryptionSettings_SymmetricKey{
+					Key: s3Settings.EncryptionKey,
+				},
+			},
+		}
+	}
+
 	response, err := exportClient.ExportToS3(ctx, exportRequest)
 
 	if err != nil {
@@ -425,6 +436,17 @@ func (d *ClientYdbConnector) ImportFromS3(
 		importRequest.Settings.DestinationPath = path.Join(clientDb.Name(), s3Settings.DestinationPath)
 	}
 
+	if len(s3Settings.EncryptionKey) > 0 {
+		importRequest.Settings.EncryptionSettings = &Ydb_Export.EncryptionSettings{
+			EncryptionAlgorithm: s3Settings.EncryptionAlgorithm,
+			Key: &Ydb_Export.EncryptionSettings_SymmetricKey_{
+				SymmetricKey: &Ydb_Export.EncryptionSettings_SymmetricKey{
+					Key: s3Settings.EncryptionKey,
+				},
+			},
+		}
+	}
+
 	response, err := importClient.ImportFromS3(ctx, importRequest)
 
 	if err != nil {

internal/connectors/db/yql/queries/write.go

@@ -132,6 +132,16 @@ func BuildCreateOperationQuery(operation types.Operation, index int) WriteSingle
 		if tb.ParentOperationID != nil {
 			d.AddValueParam("$parent_operation_id", table_types.StringValueFromString(*tb.ParentOperationID))
 		}
+		if tb.EncryptionSettings != nil {
+			d.AddValueParam(
+				"$encryption_algorithm",
+				table_types.StringValueFromString(tb.EncryptionSettings.GetAlgorithm().String()),
+			)
+			d.AddValueParam(
+				"$kms_key_id",
+				table_types.StringValueFromString(tb.EncryptionSettings.GetKmsKey().GetKeyId()),
+			)
+		}
 	} else if operation.GetType() == types.OperationTypeTBWR {
 		tbwr, ok := operation.(*types.TakeBackupWithRetryOperation)
 		if !ok {
@@ -184,6 +194,16 @@ func BuildCreateOperationQuery(operation types.Operation, index int) WriteSingle
 		if tbwr.Ttl != nil {
 			d.AddValueParam("$ttl", table_types.IntervalValueFromDuration(*tbwr.Ttl))
 		}
+		if tbwr.EncryptionSettings != nil {
+			d.AddValueParam(
+				"$encryption_algorithm",
+				table_types.StringValueFromString(tbwr.EncryptionSettings.GetAlgorithm().String()),
+			)
+			d.AddValueParam(
+				"$kms_key_id",
+				table_types.StringValueFromString(tbwr.EncryptionSettings.GetKmsKey().GetKeyId()),
+			)
+		}
 	} else if operation.GetType() == types.OperationTypeRB {
 		rb, ok := operation.(*types.RestoreBackupOperation)
 		if !ok {
@@ -336,6 +356,17 @@ func BuildCreateBackupQuery(b types.Backup, index int) WriteSingleTableQueryImpl
 	if b.ExpireAt != nil {
 		d.AddValueParam("$expire_at", table_types.TimestampValueFromTime(*b.ExpireAt))
 	}
+
+	if b.EncryptionSettings != nil {
+		d.AddValueParam(
+			"$encryption_algorithm",
+			table_types.StringValueFromString(b.EncryptionSettings.GetAlgorithm().String()),
+		)
+		d.AddValueParam(
+			"$kms_key_id",
+			table_types.StringValueFromString(b.EncryptionSettings.GetKmsKey().GetKeyId()),
+		)
+	}
 	return d
 }
 
@@ -389,6 +420,17 @@ func BuildCreateBackupScheduleQuery(schedule types.BackupSchedule, index int) Wr
 			"$recovery_point_objective",
 			table_types.IntervalValueFromDuration(schedule.ScheduleSettings.RecoveryPointObjective.AsDuration()),
 		)
+
+		if schedule.ScheduleSettings.EncryptionSettings != nil {
+			d.AddValueParam(
+				"$encryption_algorithm",
+				table_types.StringValueFromString(schedule.ScheduleSettings.EncryptionSettings.GetAlgorithm().String()),
+			)
+			d.AddValueParam(
+				"$kms_key_id",
+				table_types.StringValueFromString(schedule.ScheduleSettings.EncryptionSettings.GetKmsKey().GetKeyId()),
+			)
+		}
 	}
 	if schedule.NextLaunch != nil {
 		d.AddValueParam("$next_launch", table_types.TimestampValueFromTime(*schedule.NextLaunch))

internal/handlers/schedule_backup.go

@@ -8,6 +8,7 @@ import (
 	"google.golang.org/protobuf/types/known/durationpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 	"ydbcp/internal/audit"
+	"ydbcp/internal/config"
 	"ydbcp/internal/connectors/db"
 	"ydbcp/internal/connectors/db/yql/queries"
 	"ydbcp/internal/types"
@@ -20,11 +21,13 @@ type BackupScheduleHandlerType func(context.Context, db.DBConnector, *types.Back
 func NewBackupScheduleHandler(
 	queryBuilderFactory queries.WriteQueryBuilderFactory,
 	clock clockwork.Clock,
+	featureFlags config.FeatureFlagsConfig,
 ) BackupScheduleHandlerType {
 	return func(ctx context.Context, driver db.DBConnector, schedule *types.BackupSchedule) error {
 		return BackupScheduleHandler(
 			ctx, driver, schedule,
 			queryBuilderFactory, clock,
+			featureFlags,
 		)
 	}
 }
@@ -46,6 +49,7 @@ func BackupScheduleHandler(
 	schedule *types.BackupSchedule,
 	queryBuilderFactory queries.WriteQueryBuilderFactory,
 	clock clockwork.Clock,
+	featureFlags config.FeatureFlagsConfig,
 ) error {
 	if schedule.Status != types.BackupScheduleStateActive {
 		xlog.Error(ctx, "backup schedule is not active", zap.String("ScheduleID", schedule.ID))
@@ -86,6 +90,10 @@ func BackupScheduleHandler(
 				d := schedule.ScheduleSettings.Ttl.AsDuration()
 				tbwr.Ttl = &d
 			}
+
+			if schedule.ScheduleSettings.EncryptionSettings != nil && featureFlags.EnableBackupEncryption {
+				tbwr.EncryptionSettings = schedule.ScheduleSettings.EncryptionSettings
+			}
 		}
 
 		xlog.Info(