feat: add encryption key management

Signed-off-by: chohee <[email protected]>

Author: ChoHee15

cmd/dependency/dependency.go

@@ -57,6 +57,7 @@ import (
 	"d7y.io/dragonfly/v2/pkg/types"
 	"d7y.io/dragonfly/v2/pkg/unit"
 	"d7y.io/dragonfly/v2/version"
+	manager_cfg "d7y.io/dragonfly/v2/manager/config"
 )
 
 // InitCommandAndConfig initializes flags binding and common sub cmds.
@@ -307,7 +308,8 @@ func initDecoderConfig(dc *mapstructure.DecoderConfig) {
 			reflect.TypeOf(config.URL{}),
 			reflect.TypeOf(net.IP{}),
 			reflect.TypeOf(config.CertPool{}),
-			reflect.TypeOf(config.Regexp{}):
+			reflect.TypeOf(config.Regexp{}),
+			reflect.TypeOf(manager_cfg.EncryptionKey{}):  
 
 			b, _ := yaml.Marshal(v)
 			p := reflect.New(to)

cmd/manager/cmd/root.go

@@ -101,6 +101,10 @@ func init() {
 // initDfpath initializes dfpath.
 func initDfpath(cfg *config.ServerConfig) (dfpath.Dfpath, error) {
 	var options []dfpath.Option
+	if cfg.WorkHome != "" {
+		options = append(options, dfpath.WithWorkHome(cfg.WorkHome))
+	}
+
 	if cfg.LogDir != "" {
 		options = append(options, dfpath.WithLogDir(cfg.LogDir))
 	}
@@ -113,6 +117,10 @@ func initDfpath(cfg *config.ServerConfig) (dfpath.Dfpath, error) {
 		options = append(options, dfpath.WithPluginDir(cfg.PluginDir))
 	}
 
+	if cfg.DataDir != "" {
+		options = append(options, dfpath.WithDataDir(cfg.DataDir))
+	}
+
 	return dfpath.New(options...)
 }
 

cmd/scheduler/cmd/root.go

@@ -100,6 +100,10 @@ func init() {
 
 func initDfpath(cfg *config.ServerConfig) (dfpath.Dfpath, error) {
 	var options []dfpath.Option
+	if cfg.WorkHome != "" {
+		options = append(options, dfpath.WithWorkHome(cfg.WorkHome))
+	}
+
 	if cfg.LogDir != "" {
 		options = append(options, dfpath.WithLogDir(cfg.LogDir))
 	}

go.mod

@@ -280,3 +280,6 @@ require (
 	gorm.io/driver/sqlserver v1.4.1 // indirect
 	gorm.io/plugin/dbresolver v1.3.0 // indirect
 )
+
+// use local api repo
+replace d7y.io/api/v2 => ../api_fork

go.sum

@@ -63,8 +63,6 @@ cloud.google.com/go/storage v1.50.0 h1:3TbVkzTooBvnZsk7WaAQfOsNrdoM8QHusXA1cpk6Q
 cloud.google.com/go/storage v1.50.0/go.mod h1:l7XeiD//vx5lfqE3RavfmU9yvk5Pp0Zhcv482poyafY=
 cloud.google.com/go/trace v1.11.6 h1:2O2zjPzqPYAHrn3OKl029qlqG6W8ZdYaOWRyr8NgMT4=
 cloud.google.com/go/trace v1.11.6/go.mod h1:GA855OeDEBiBMzcckLPE2kDunIpC72N+Pq8WFieFjnI=
-d7y.io/api/v2 v2.1.60 h1:Imk1mw30wmTJNdqKFF5VO/eqI/t1H5PfIPPK4ryKX2U=
-d7y.io/api/v2 v2.1.60/go.mod h1:WzakPywEgs27gr/TwrlRarPRRK2o5nN0YrTzHDdFii8=
 dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU=
 filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA=
 filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4=

manager/config/config.go

@@ -17,12 +17,14 @@
 package config
 
 import (
+	"encoding/base64"
 	"errors"
 	"fmt"
 	"net"
 	"time"
 
 	"d7y.io/dragonfly/v2/cmd/dependency/base"
+	logger "d7y.io/dragonfly/v2/internal/dflog"
 	"d7y.io/dragonfly/v2/pkg/net/ip"
 	"d7y.io/dragonfly/v2/pkg/objectstorage"
 	"d7y.io/dragonfly/v2/pkg/slices"
@@ -56,12 +58,18 @@ type Config struct {
 
 	// Network configuration.
 	Network NetworkConfig `yaml:"network" mapstructure:"network"`
+
+	// Encryption configuration
+	Encryption EncryptionConfig `yaml:"encryption" mapstructure:"encryption"`
 }
 
 type ServerConfig struct {
 	// Server name.
 	Name string `yaml:"name" mapstructure:"name"`
 
+	// Server work home directory.
+	WorkHome string `yaml:"workHome" mapstructure:"workHome"`
+
 	// Server dynamic config cache directory.
 	CacheDir string `yaml:"cacheDir" mapstructure:"cacheDir"`
 
@@ -83,6 +91,9 @@ type ServerConfig struct {
 	// Server plugin directory.
 	PluginDir string `yaml:"pluginDir" mapstructure:"pluginDir"`
 
+	// Server data directory.
+	DataDir string `yaml:"dataDir" mapstructure:"dataDir"`
+
 	// GRPC server configuration.
 	GRPC GRPCConfig `yaml:"grpc" mapstructure:"grpc"`
 
@@ -398,6 +409,34 @@ type NetworkConfig struct {
 	EnableIPv6 bool `mapstructure:"enableIPv6" yaml:"enableIPv6"`
 }
 
+type EncryptionKey [32]byte
+type EncryptionConfig struct {
+	Enable bool `mapstructure:"enable" yaml:"enable"`
+	// AES256 base64, optional
+	Key *EncryptionKey `mapstructure:"key" yaml:"key"`
+}
+
+// UnmarshalText Base64
+func (e *EncryptionKey) UnmarshalText(text []byte) error {
+	logger.Infof("Base64 str: %s", string(text))
+	keyBytes, err := base64.StdEncoding.DecodeString(string(text))
+	if err != nil {
+		return fmt.Errorf("invalid base64 key: %v", err)
+	}
+
+	if len(keyBytes) != 32 {
+		return fmt.Errorf("key must be 32 bytes, got %d", len(keyBytes))
+	}
+
+	copy(e[:], keyBytes)
+	return nil
+}
+
+// MarshalText Base64
+func (e EncryptionKey) MarshalText() ([]byte, error) {
+	return []byte(base64.StdEncoding.EncodeToString(e[:])), nil
+}
+
 // New config instance.
 func New() *Config {
 	return &Config{
@@ -489,6 +528,9 @@ func New() *Config {
 		Network: NetworkConfig{
 			EnableIPv6: DefaultNetworkEnableIPv6,
 		},
+		Encryption: EncryptionConfig{
+			Enable: false,
+		},
 	}
 }
 
@@ -686,6 +728,8 @@ func (cfg *Config) Validate() error {
 		}
 	}
 
+	// TODO: validate key
+
 	return nil
 }
 

manager/database/database.go

@@ -115,6 +115,7 @@ func migrate(db *gorm.DB) error {
 		&models.Application{},
 		&models.PersonalAccessToken{},
 		&models.Peer{},
+		&models.EncryptionKey{},
 	)
 }
 

manager/manager.go

@@ -18,7 +18,11 @@ package manager
 
 import (
 	"context"
+	"crypto/rand"
 	"embed"
+	"encoding/base64"
+	"encoding/hex"
+	"errors"
 	"fmt"
 	"io/fs"
 	"net"
@@ -29,6 +33,8 @@ import (
 	"google.golang.org/grpc"
 	"gorm.io/gorm"
 
+	"bytes"
+
 	logger "d7y.io/dragonfly/v2/internal/dflog"
 	"d7y.io/dragonfly/v2/internal/ratelimiter"
 	"d7y.io/dragonfly/v2/manager/cache"
@@ -37,6 +43,7 @@ import (
 	managergc "d7y.io/dragonfly/v2/manager/gc"
 	"d7y.io/dragonfly/v2/manager/job"
 	"d7y.io/dragonfly/v2/manager/metrics"
+	"d7y.io/dragonfly/v2/manager/models"
 	"d7y.io/dragonfly/v2/manager/permission/rbac"
 	"d7y.io/dragonfly/v2/manager/router"
 	"d7y.io/dragonfly/v2/manager/rpcserver"
@@ -48,6 +55,7 @@ import (
 	"d7y.io/dragonfly/v2/pkg/objectstorage"
 	"d7y.io/dragonfly/v2/pkg/redis"
 	"d7y.io/dragonfly/v2/pkg/rpc"
+	"gorm.io/plugin/soft_delete"
 )
 
 const (
@@ -122,6 +130,16 @@ func New(cfg *config.Config, d dfpath.Dfpath) (*Server, error) {
 		return nil, err
 	}
 
+	// Initialize encryption key
+	if cfg.Encryption.Enable {
+		logger.Infof("Encryption enabled")
+		if err := initializeEncryptionKey(cfg, db.DB); err != nil {
+			return nil, err
+		}
+	} else {
+		logger.Infof("Encryption disabled")
+	}
+
 	// Initialize enforcer.
 	enforcer, err := rbac.NewEnforcer(db.DB)
 	if err != nil {
@@ -250,6 +268,104 @@ func registerGCTasks(gc pkggc.GC, db *gorm.DB) error {
 	return nil
 }
 
+// initializeEncryptionKey
+func initializeEncryptionKey(cfg *config.Config, db *gorm.DB) error {
+	// db.Delete(&models.EncryptionKey{}, "1 = 1")
+
+	var existingKey models.EncryptionKey
+	hasDBKey := false
+	if err := db.First(&existingKey).Error; err == nil {
+		hasDBKey = true
+	} else if !errors.Is(err, gorm.ErrRecordNotFound) {
+		return fmt.Errorf("failed to check encryption key: %v", err)
+	}
+
+	if cfg.Encryption.Key != nil {
+		configKey := cfg.Encryption.Key
+		keyBytes := configKey[:]
+		if hasDBKey {
+			// compare key in config with key in db
+			if bytes.Equal(existingKey.Key, keyBytes) {
+				logger.Infof(
+					"encryption key in config file is the same as in database, key(hex): %s, key(base64): %s",
+					hex.EncodeToString(keyBytes),
+					base64.StdEncoding.EncodeToString(keyBytes),
+				)
+				return nil
+			}
+			// key in config is different from key in db, overwrite db
+			oldKeyHex := hex.EncodeToString(existingKey.Key)
+			oldKeyBase64 := base64.StdEncoding.EncodeToString(existingKey.Key)
+			newKeyHex := hex.EncodeToString(keyBytes)
+			newKeyBase64 := base64.StdEncoding.EncodeToString(keyBytes)
+
+			if err := db.Model(&existingKey).Update("key", keyBytes).Error; err != nil {
+				return fmt.Errorf("failed to update encryption key in database: %v", err)
+			}
+
+			logger.Infof(
+				"encryption key in database is overwritten by config file, old key(hex): %s, old key(base64): %s, new key(hex): %s, new key(base64): %s",
+				oldKeyHex, oldKeyBase64, newKeyHex, newKeyBase64,
+			)
+			return nil
+		} else {
+			// config has key, but db has no key, write it into db
+			// check soft delete
+			var oldKey models.EncryptionKey
+			if err := db.Unscoped().Where("`key` = ?", keyBytes).First(&oldKey).Error; err == nil {
+				if oldKey.IsDel != soft_delete.DeletedAt(soft_delete.FlagActived) {
+					// restore the key soft deleted
+					db.Unscoped().Model(&oldKey).Update("is_del", soft_delete.FlagActived)
+					logger.Infof("Restore the key which was soft deleted before")
+				} else {
+					logger.Fatalf("key should be soft deleted in this situation")
+				}
+			} else if errors.Is(err, gorm.ErrRecordNotFound) {
+				// insert new key
+				if err := db.Create(&models.EncryptionKey{Key: keyBytes}).Error; err != nil {
+					return fmt.Errorf("failed to save encryption key to database: %v", err)
+				}
+			} else {
+				// return fmt.Errorf("unknow failed when update encryption key in database: %v", err)
+				logger.Fatalf("unknow failed when update encryption key in database: %v", err)
+				// panic(err)
+			}
+
+			logger.Infof(
+				"encryption key from config file is saved to database, key(hex): %s, key(base64): %s",
+				hex.EncodeToString(keyBytes),
+				base64.StdEncoding.EncodeToString(keyBytes),
+			)
+			return nil
+		}
+	}
+
+	// config has no key and db has key
+	if hasDBKey {
+		logger.Infof(
+			"encryption key loaded from database, key(hex): %s, key(base64): %s",
+			hex.EncodeToString(existingKey.Key),
+			base64.StdEncoding.EncodeToString(existingKey.Key),
+		)
+		return nil
+	}
+
+	// config and db both have no key, generate one
+	keyBytes := make([]byte, 32)
+	if _, err := rand.Read(keyBytes); err != nil {
+		return fmt.Errorf("failed to generate random encryption key: %v", err)
+	}
+	if err := db.Create(&models.EncryptionKey{Key: keyBytes}).Error; err != nil {
+		return fmt.Errorf("failed to save random encryption key to database: %v", err)
+	}
+	logger.Infof(
+		"generated random encryption key and saved to database, key(hex): %s, key(base64): %s",
+		hex.EncodeToString(keyBytes),
+		base64.StdEncoding.EncodeToString(keyBytes),
+	)
+	return nil
+}
+
 // Serve starts the manager server.
 func (s *Server) Serve() error {
 	// Started REST server.

manager/models/encryption.go

@@ -0,0 +1,6 @@
+package models
+
+type EncryptionKey struct {
+	BaseModel
+	Key []byte `gorm:"type:binary(32);not null;unique" json:"key"`
+}

manager/rpcserver/manager_server_v2.go

@@ -976,3 +976,19 @@ func (s *managerServerV2) KeepAlive(stream managerv2.Manager_KeepAliveServer) er
 		}
 	}
 }
+
+// RequestEncryptionKey implements manager.ManagerServer.
+func (s *managerServerV2) RequestEncryptionKey(ctx context.Context, req *managerv2.RequestEncryptionKeyRequest) (*managerv2.RequestEncryptionKeyResponse, error) {
+	log := logger.WithHostnameAndIP(req.Hostname, req.Ip)
+
+	// Get key from db
+	var encKey models.EncryptionKey
+	if err := s.db.WithContext(ctx).First(&encKey).Error; err != nil {
+		log.Errorf("failed to get encryption key: %v", err)
+		return nil, status.Error(codes.Internal, "failed to get encryption key")
+	}
+
+	return &managerv2.RequestEncryptionKeyResponse{
+		EncryptionKey: encKey.Key,
+	}, nil
+}