feat(ydbcp): implement kms plugin (#178)

Author: ulya-sidorina

internal/auth/auth.go

@@ -4,9 +4,9 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"plugin"
 
 	"ydbcp/internal/config"
+	"ydbcp/internal/plugin"
 	"ydbcp/internal/util/xlog"
 	"ydbcp/pkg/plugins/auth"
 
@@ -27,30 +27,8 @@ var (
 	ErrGetAuthToken = errors.New("can't get auth token")
 )
 
-func NewAuthProvider(ctx context.Context, cfg config.AuthConfig) (auth.AuthProvider, error) {
-	xlog.Info(ctx, "Loading auth provider plugin", zap.String("PluginPath", cfg.PluginPath))
-
-	plug, err := plugin.Open(cfg.PluginPath)
-	if err != nil {
-		return nil, fmt.Errorf("can't load auth provider plugin, path %s: %w", cfg.PluginPath, err)
-	}
-	symbol, err := plug.Lookup("AuthProvider")
-	if err != nil {
-		return nil, fmt.Errorf("can't lookup AuthProvider symbol, plugin path %s: %w", cfg.PluginPath, err)
-	}
-	var instance auth.AuthProvider
-	instance, ok := symbol.(auth.AuthProvider)
-	if !ok {
-		return nil, fmt.Errorf("can't cast AuthProvider symbol, plugin path %s", cfg.PluginPath)
-	}
-	pluginConfig, err := cfg.ConfigurationString()
-	if err != nil {
-		return nil, fmt.Errorf("can't get auth provider configuration: %w", err)
-	}
-	if err = instance.Init(ctx, pluginConfig); err != nil {
-		return nil, fmt.Errorf("can't initialize auth provider plugin: %w", err)
-	}
-	return instance, nil
+func NewAuthProvider(ctx context.Context, cfg config.PluginConfig) (auth.AuthProvider, error) {
+	return plugin.Load[auth.AuthProvider](ctx, cfg, "AuthProvider", "auth")
 }
 
 func Authenticate(ctx context.Context, provider auth.AuthProvider) (string, error) {

internal/config/config.go

@@ -45,7 +45,7 @@ type ClientConnectionConfig struct {
 	AllowInsecureEndpoint  bool     `yaml:"allow_insecure_endpoint" default:"false"`
 }
 
-type AuthConfig struct {
+type PluginConfig struct {
 	PluginPath    string      `yaml:"plugin_path"`
 	Configuration interface{} `yaml:"configuration"`
 }
@@ -95,7 +95,7 @@ type Config struct {
 	DBConnection       YDBConnectionConfig      `yaml:"db_connection"`
 	ClientConnection   ClientConnectionConfig   `yaml:"client_connection"`
 	S3                 S3Config                 `yaml:"s3"`
-	Auth               AuthConfig               `yaml:"auth"`
+	Auth               PluginConfig             `yaml:"auth"`
 	GRPCServer         GRPCServerConfig         `yaml:"grpc_server"`
 	MetricsServer      MetricsServerConfig      `yaml:"metrics_server"`
 	OperationProcessor OperationProcessorConfig `yaml:"operation_processor"`
@@ -223,7 +223,7 @@ func (c *S3Config) SecretKey() (string, error) {
 
 }
 
-func (c *AuthConfig) ConfigurationString() (string, error) {
+func (c *PluginConfig) ConfigurationString() (string, error) {
 	txt, err := yaml.Marshal(c.Configuration)
 	if err != nil {
 		return "", fmt.Errorf("can't marshal Auth.Configuration to YAML: %w", err)

internal/kms/kms.go

@@ -0,0 +1,13 @@
+package kms
+
+import (
+	"context"
+
+	"ydbcp/internal/config"
+	"ydbcp/internal/plugin"
+	"ydbcp/pkg/plugins/kms"
+)
+
+func NewKmsProvider(ctx context.Context, cfg config.PluginConfig) (kms.KmsProvider, error) {
+	return plugin.Load[kms.KmsProvider](ctx, cfg, "KmsProvider", "kms")
+}

internal/kms/mock.go

@@ -0,0 +1,84 @@
+package kms
+
+import (
+	"context"
+	"fmt"
+
+	"ydbcp/pkg/plugins/kms"
+)
+
+type MockKmsProvider struct {
+	keys map[string][]byte
+}
+
+func NewMockKmsProvider(keys map[string][]byte) *MockKmsProvider {
+	return &MockKmsProvider{
+		keys: keys,
+	}
+}
+
+func (p *MockKmsProvider) Init(_ context.Context, _ string) error {
+	if p.keys == nil {
+		p.keys = make(map[string][]byte)
+	}
+	return nil
+}
+
+func (p *MockKmsProvider) Close(_ context.Context) error {
+	p.keys = nil
+	return nil
+}
+
+func xor(data []byte, key []byte) []byte {
+	out := make([]byte, len(data))
+	for i := range data {
+		out[i] = data[i] ^ key[i%len(key)]
+	}
+	return out
+}
+
+func (p *MockKmsProvider) Encrypt(
+	_ context.Context,
+	req *kms.EncryptRequest,
+) (*kms.EncryptResponse, error) {
+	if req == nil {
+		return nil, fmt.Errorf("mock kms: encrypt request is nil")
+	}
+	if len(req.Plaintext) == 0 {
+		return &kms.EncryptResponse{KeyID: req.KeyID}, nil
+	}
+
+	key, ok := p.keys[req.KeyID]
+	if !ok {
+		return nil, fmt.Errorf("mock kms: key not found")
+	}
+
+	ciphertext := xor(req.Plaintext, key)
+	return &kms.EncryptResponse{
+		KeyID:      req.KeyID,
+		Ciphertext: ciphertext,
+	}, nil
+}
+
+func (p *MockKmsProvider) Decrypt(
+	_ context.Context,
+	req *kms.DecryptRequest,
+) (*kms.DecryptResponse, error) {
+	if req == nil {
+		return nil, fmt.Errorf("mock kms: decrypt request is nil")
+	}
+	if len(req.Ciphertext) == 0 {
+		return &kms.DecryptResponse{KeyID: req.KeyID}, nil
+	}
+
+	key, ok := p.keys[req.KeyID]
+	if !ok {
+		return nil, fmt.Errorf("mock kms: key not found")
+	}
+
+	plaintext := xor(req.Ciphertext, key)
+	return &kms.DecryptResponse{
+		KeyID:     req.KeyID,
+		Plaintext: plaintext,
+	}, nil
+}

internal/plugin/plugin.go

@@ -0,0 +1,55 @@
+package plugin
+
+import (
+	"context"
+	"fmt"
+	"plugin"
+
+	"ydbcp/internal/config"
+	"ydbcp/internal/util/xlog"
+
+	"go.uber.org/zap"
+)
+
+type Plugin interface {
+	Init(ctx context.Context, config string) error
+}
+
+// Load loads a Go plugin, looks up the provided symbol name and initializes it with the
+// plugin configuration. The type parameter must be an interface that embeds the Init method.
+func Load[T Plugin](
+	ctx context.Context,
+	cfg config.PluginConfig,
+	symbolName string,
+	pluginKind string,
+) (T, error) {
+	var zero T
+
+	xlog.Info(ctx, fmt.Sprintf("Loading %s plugin", pluginKind), zap.String("plugin_path", cfg.PluginPath))
+
+	plug, err := plugin.Open(cfg.PluginPath)
+	if err != nil {
+		return zero, fmt.Errorf("can't load %s plugin, path %s: %w", pluginKind, cfg.PluginPath, err)
+	}
+
+	symbol, err := plug.Lookup(symbolName)
+	if err != nil {
+		return zero, fmt.Errorf("can't lookup %s symbol, plugin path %s: %w", symbolName, cfg.PluginPath, err)
+	}
+
+	instance, ok := symbol.(T)
+	if !ok {
+		return zero, fmt.Errorf("can't cast %s symbol, plugin path %s", symbolName, cfg.PluginPath)
+	}
+
+	pluginConfig, err := cfg.ConfigurationString()
+	if err != nil {
+		return zero, fmt.Errorf("can't get %s plugin configuration: %w", pluginKind, err)
+	}
+
+	if err = instance.Init(ctx, pluginConfig); err != nil {
+		return zero, fmt.Errorf("can't initialize %s plugin: %w", pluginKind, err)
+	}
+
+	return instance, nil
+}

internal/util/tls_setup/tls_setup.go

@@ -5,17 +5,64 @@ import (
 	"crypto/x509"
 	"errors"
 	"fmt"
+	"os"
+	"strings"
+
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
-	"os"
 )
 
-func LoadTLSCredentials(RootCAPath *string, withInsecure bool) (grpc.DialOption, error) {
+func ParseEndpoint(e string) (string, bool) {
+	if strings.HasPrefix(e, "grpcs://") {
+		return e[8:], true
+	}
+	if strings.HasPrefix(e, "grpc://") {
+		return e[7:], false
+	}
+	return e, false
+}
+
+func LoadTLSCredentials(rootCAPath *string, withInsecure bool) (grpc.DialOption, error) {
+	tlsConfig, err := buildTLSConfig(rootCAPath, withInsecure)
+	if err != nil {
+		return nil, err
+	}
+	return grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig)), nil
+}
+
+func LoadMTLSCredentials(
+	rootCAPath *string,
+	clientCertPath *string,
+	clientKeyPath *string,
+	withInsecure bool,
+) (grpc.DialOption, error) {
+	if clientCertPath == nil || len(*clientCertPath) == 0 {
+		return nil, fmt.Errorf("client certificate path is required for mTLS")
+	}
+	if clientKeyPath == nil || len(*clientKeyPath) == 0 {
+		return nil, fmt.Errorf("client key path is required for mTLS")
+	}
+
+	tlsConfig, err := buildTLSConfig(rootCAPath, withInsecure)
+	if err != nil {
+		return nil, err
+	}
+
+	cert, err := tls.LoadX509KeyPair(*clientCertPath, *clientKeyPath)
+	if err != nil {
+		return nil, fmt.Errorf("failed to load client certificate/key: %w", err)
+	}
+	tlsConfig.Certificates = []tls.Certificate{cert}
+
+	return grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig)), nil
+}
+
+func buildTLSConfig(rootCAPath *string, withInsecure bool) (*tls.Config, error) {
 	var certPool *x509.CertPool
-	if RootCAPath != nil && len(*RootCAPath) > 0 {
-		caBundle, err := os.ReadFile(*RootCAPath)
+	if rootCAPath != nil && len(*rootCAPath) > 0 {
+		caBundle, err := os.ReadFile(*rootCAPath)
 		if err != nil {
-			return nil, fmt.Errorf("unable to read root ca bundle from file %s: %w", *RootCAPath, err)
+			return nil, fmt.Errorf("unable to read root ca bundle from file %s: %w", *rootCAPath, err)
 		}
 		certPool = x509.NewCertPool()
 		if ok := certPool.AppendCertsFromPEM(caBundle); !ok {
@@ -36,5 +83,5 @@ func LoadTLSCredentials(RootCAPath *string, withInsecure bool) (grpc.DialOption,
 	if withInsecure {
 		tlsConfig.InsecureSkipVerify = true
 	}
-	return grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig)), nil
+	return tlsConfig, nil
 }

pkg/plugins/kms/kms.go

@@ -0,0 +1,32 @@
+package kms
+
+import "context"
+
+type EncryptRequest struct {
+	KeyID     string
+	Plaintext []byte
+}
+
+type EncryptResponse struct {
+	KeyID      string
+	Ciphertext []byte
+}
+
+type DecryptRequest struct {
+	KeyID      string
+	Ciphertext []byte
+}
+
+type DecryptResponse struct {
+	KeyID     string
+	Plaintext []byte
+}
+
+// KmsProvider is an interface that KMS plugins must implement.
+type KmsProvider interface {
+	Init(ctx context.Context, config string) error
+	Close(ctx context.Context) error
+
+	Encrypt(ctx context.Context, req *EncryptRequest) (*EncryptResponse, error)
+	Decrypt(ctx context.Context, req *DecryptRequest) (*DecryptResponse, error)
+}

plugins/auth_nebius/auth_nebius.go

@@ -29,16 +29,6 @@ type pluginConfig struct {
 	RootCAPath            string `yaml:"root_ca_path"`
 }
 
-func parseEndpoint(e string) (string, bool) {
-	if strings.HasPrefix(e, "grpcs://") {
-		return e[8:], true
-	}
-	if strings.HasPrefix(e, "grpc://") {
-		return e[7:], false
-	}
-	return e, false
-}
-
 func (p *authProviderNebius) loadTLSCredentials() (grpc.DialOption, error) {
 	if !p.tls {
 		return grpc.WithTransportCredentials(insecure.NewCredentials()), nil
@@ -53,7 +43,7 @@ func (p *authProviderNebius) Init(ctx context.Context, config string) error {
 		xlog.Error(ctx, "Unable to parse configuration file", zap.Error(err))
 		return fmt.Errorf("unable to parse configuration %w", err)
 	}
-	p.endpoint, p.tls = parseEndpoint(p.config.AccessServiceEndpoint)
+	p.endpoint, p.tls = tls_setup.ParseEndpoint(p.config.AccessServiceEndpoint)
 	return nil
 }