diff --git a/go.mod b/go.mod index 6cbe67e..8b38de2 100644 --- a/go.mod +++ b/go.mod @@ -23,6 +23,7 @@ require ( gorm.io/driver/postgres v1.6.0 gorm.io/driver/sqlite v1.6.0 gorm.io/gorm v1.31.1 + github.com/disintegration/imaging v1.6.2 ) require ( @@ -34,7 +35,6 @@ require ( github.com/cespare/xxhash/v2 v2.2.0 // indirect github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect - github.com/disintegration/imaging v1.6.2 // indirect github.com/dustin/go-humanize v1.0.1 // indirect github.com/felixge/httpsnoop v1.0.3 // indirect github.com/fsnotify/fsnotify v1.7.0 // indirect diff --git a/internal/config/config.go b/internal/config/config.go index 84d2cd7..9d0e739 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -144,6 +144,10 @@ func Load(configPath string) (*Config, error) { viper.SetDefault("casbin.enable_cache", true) viper.SetDefault("casbin.cache_ttl", 300) viper.SetDefault("casbin.skip_routes", []string{"/health", "/api/v1/auth/login", "/api/v1/auth/register"}) + // Encryption 默认值 + viper.SetDefault("encryption.enabled", false) + viper.SetDefault("encryption.key", "") + viper.SetDefault("encryption.key_version", 1) if err := viper.ReadInConfig(); err != nil { return nil, fmt.Errorf("failed to read config: %w", err) @@ -231,6 +235,10 @@ func Load(configPath string) (*Config, error) { cfg.Casbin.AutoLoad, _ = strconv.ParseBool(getEnvOrDefault("APP_CASBIN_AUTO_LOAD", fmt.Sprintf("%t", cfg.Casbin.AutoLoad))) cfg.Casbin.EnableCache, _ = strconv.ParseBool(getEnvOrDefault("APP_CASBIN_ENABLE_CACHE", fmt.Sprintf("%t", cfg.Casbin.EnableCache))) cfg.Casbin.CacheTTL, _ = strconv.Atoi(getEnvOrDefault("APP_CASBIN_CACHE_TTL", fmt.Sprintf("%d", cfg.Casbin.CacheTTL))) + // Encryption 环境变量覆盖 + cfg.Encryption.Enabled, _ = strconv.ParseBool(getEnvOrDefault("APP_ENCRYPTION_ENABLED", fmt.Sprintf("%t", cfg.Encryption.Enabled))) + cfg.Encryption.Key = getEnvOrDefault("APP_ENCRYPTION_KEY", cfg.Encryption.Key) + cfg.Encryption.KeyVersion, _ = strconv.Atoi(getEnvOrDefault("APP_ENCRYPTION_KEY_VERSION", fmt.Sprintf("%d", cfg.Encryption.KeyVersion))) return &cfg, nil } diff --git a/internal/model/message.go b/internal/model/message.go index cfec4a0..6aa34ea 100644 --- a/internal/model/message.go +++ b/internal/model/message.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "strconv" + "sync" "time" "gorm.io/gorm" @@ -198,8 +199,17 @@ func (m *Message) BeforeUpdate(tx *gorm.DB) error { return nil } -// AfterFind 查询后自动解密消息内容 +// AfterFind 查询后标记需要解密（实际解密延迟到批量处理时） +// 这样可以支持批量并行解密优化 func (m *Message) AfterFind(tx *gorm.DB) error { + // 标记为未解密状态，等待批量解密 + // 如果是单条查询，AfterFind 后会立即使用，需要解密 + m.decrypted = false + return nil +} + +// Decrypt 解密单条消息（显式调用） +func (m *Message) Decrypt() error { return m.decryptSegments() } @@ -310,3 +320,137 @@ func (m *Message) IsInteractionNotification() bool { return false } } + +// BatchDecryptMessages 批量并行解密消息 +// 用于批量查询后提高解密性能，比逐条调用 AfterFind 更高效 +// workerCount 指定并行工作数，如果 <= 0 则使用默认值 +func BatchDecryptMessages(messages []*Message, workerCount int) error { + if len(messages) == 0 { + return nil + } + + encryptor := crypto.GetMessageEncryptor() + if encryptor == nil { + // 加密器未初始化，直接解析 JSON（兼容模式） + for _, m := range messages { + if m.decrypted || m.SegmentsEncrypted == "" { + m.decrypted = true + continue + } + if err := json.Unmarshal([]byte(m.SegmentsEncrypted), &m.Segments); err != nil { + return err + } + m.decrypted = true + } + return nil + } + + // 小批量直接串行处理 + if len(messages) <= 10 { + for _, m := range messages { + if err := m.decryptSegments(); err != nil { + return err + } + } + return nil + } + + // 并行解密 + // 先收集需要解密的密文 + ciphertexts := make([]string, len(messages)) + for i, m := range messages { + ciphertexts[i] = m.SegmentsEncrypted + } + + // 批量并行解密 + plaintexts := encryptor.BatchDecrypt(ciphertexts, workerCount) + + // 解析结果 + for i, m := range messages { + if m.decrypted { + continue + } + if m.SegmentsEncrypted == "" { + m.decrypted = true + continue + } + + plaintext := plaintexts[i] + if plaintext == nil { + // 解密失败，尝试直接解析（兼容旧数据） + if err := json.Unmarshal([]byte(m.SegmentsEncrypted), &m.Segments); err != nil { + continue // 忽略单条错误 + } + } else { + if err := json.Unmarshal(plaintext, &m.Segments); err != nil { + continue // 忽略单条错误 + } + } + m.decrypted = true + } + + return nil +} + +// BatchDecryptMessagesParallel 使用 worker pool 并行解密 +// 这是一个更高效的版本，适用于大量消息 +func BatchDecryptMessagesParallel(messages []*Message) { + if len(messages) == 0 { + return + } + + encryptor := crypto.GetMessageEncryptor() + if encryptor == nil { + // 加密器未初始化，串行解析 + for _, m := range messages { + if !m.decrypted && m.SegmentsEncrypted != "" { + _ = json.Unmarshal([]byte(m.SegmentsEncrypted), &m.Segments) + m.decrypted = true + } + } + return + } + + // 确定并行度 + workerCount := 4 + if len(messages) < 20 { + workerCount = 2 + } else if len(messages) > 100 { + workerCount = 8 + } + + jobs := make(chan int, len(messages)) + var wg sync.WaitGroup + + // 启动 workers + for w := 0; w < workerCount; w++ { + wg.Add(1) + go func() { + defer wg.Done() + for i := range jobs { + m := messages[i] + if m.decrypted || m.SegmentsEncrypted == "" { + m.decrypted = true + continue + } + + plaintext, err := encryptor.Decrypt(m.SegmentsEncrypted) + if err != nil { + // 尝试直接解析 + _ = json.Unmarshal([]byte(m.SegmentsEncrypted), &m.Segments) + } else { + _ = json.Unmarshal(plaintext, &m.Segments) + } + m.decrypted = true + } + }() + } + + // 分发任务 + for i := range messages { + jobs <- i + } + close(jobs) + + wg.Wait() +} diff --git a/internal/pkg/crypto/batch_decrypt.go b/internal/pkg/crypto/batch_decrypt.go new file mode 100644 index 0000000..61fcc5f --- /dev/null +++ b/internal/pkg/crypto/batch_decrypt.go @@ -0,0 +1,109 @@ +package crypto + +import ( + "encoding/json" + "sync" +) + +// BatchDecryptResult 批量解密结果 +type BatchDecryptResult struct { + Index int + Plaintext []byte + Error error +} + +// BatchDecrypt 批量并行解密 +// workerCount 指定并行工作数，如果 <= 0 则使用 CPU 核数 +func (e *MessageEncryptor) BatchDecrypt(ciphertexts []string, workerCount int) [][]byte { + if e == nil || len(ciphertexts) == 0 { + return nil + } + + results := make([][]byte, len(ciphertexts)) + + // 小批量直接串行处理 + if len(ciphertexts) <= 10 { + for i, ct := range ciphertexts { + if ct == "" { + continue + } + plaintext, err := e.Decrypt(ct) + if err == nil { + results[i] = plaintext + } + } + return results + } + + // 并行处理 + jobs := make(chan int, len(ciphertexts)) + var wg sync.WaitGroup + + // 确定工作数 + if workerCount <= 0 { + workerCount = 4 // 默认4个并行 + } + if workerCount > len(ciphertexts) { + workerCount = len(ciphertexts) + } + + // 启动 worker + for w := 0; w < workerCount; w++ { + wg.Add(1) + go func() { + defer wg.Done() + for i := range jobs { + ct := ciphertexts[i] + if ct == "" { + continue + } + plaintext, err := e.Decrypt(ct) + if err == nil { + results[i] = plaintext + } + } + }() + } + + // 分发任务 + for i := range ciphertexts { + jobs <- i + } + close(jobs) + + wg.Wait() + return results +} + +// DecryptMessageSegments 解密消息段 JSON +// 这是一个便捷方法，直接返回解析后的 MessageSegments +func (e *MessageEncryptor) DecryptMessageSegments(ciphertext string) ([]byte, error) { + if e == nil || ciphertext == "" { + return nil, nil + } + return e.Decrypt(ciphertext) +} + +// TryParseMessageSegments 尝试解析消息段 +// 先尝试解密，如果失败则尝试直接解析（兼容未加密的旧数据） +func TryParseMessageSegments(ciphertext string, segments interface{}) error { + if ciphertext == "" { + return nil + } + + encryptor := GetMessageEncryptor() + if encryptor == nil { + // 加密器未初始化，直接解析 + return json.Unmarshal([]byte(ciphertext), segments) + } + + // 尝试解密 + plaintext, err := encryptor.Decrypt(ciphertext) + if err != nil { + // 解密失败，可能是未加密的旧数据，尝试直接解析 + return json.Unmarshal([]byte(ciphertext), segments) + } + + // 解析解密后的数据 + return json.Unmarshal(plaintext, segments) +} diff --git a/internal/pkg/crypto/batch_decrypt_test.go b/internal/pkg/crypto/batch_decrypt_test.go new file mode 100644 index 0000000..65bd009 --- /dev/null +++ b/internal/pkg/crypto/batch_decrypt_test.go @@ -0,0 +1,137 @@ +package crypto + +import ( + "encoding/json" + "fmt" + "testing" +) + +// BenchmarkEncryptDecryptSingle 单条加密解密性能基准 +func BenchmarkEncryptDecryptSingle(b *testing.B) { + key := "12345678901234567890123456789012" + _ = InitMessageEncryptor(key, 1) + encryptor := GetMessageEncryptor() + + plaintext := []byte(`{"type":"text","data":{"content":"这是一条测试消息"}}`) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + ciphertext, _ := encryptor.Encrypt(plaintext) + _, _ = encryptor.Decrypt(ciphertext) + } +} + +// BenchmarkBatchDecrypt 批量并行解密性能基准 +func BenchmarkBatchDecrypt(b *testing.B) { + key := "12345678901234567890123456789012" + _ = InitMessageEncryptor(key, 1) + encryptor := GetMessageEncryptor() + + // 准备测试数据 + sizes := []int{10, 50, 100, 500} + + for _, size := range sizes { + b.Run(fmt.Sprintf("size_%d", size), func(b *testing.B) { + // 生成加密数据 + ciphertexts := make([]string, size) + for i := 0; i < size; i++ { + msg := map[string]interface{}{ + "type": "text", + "data": map[string]string{"content": fmt.Sprintf("消息内容 %d", i)}, + } + plaintext, _ := json.Marshal(msg) + ciphertexts[i], _ = encryptor.Encrypt(plaintext) + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _ = encryptor.BatchDecrypt(ciphertexts, 4) + } + }) + } +} + +// BenchmarkSerialDecrypt 串行解密性能基准（对比用） +func BenchmarkSerialDecrypt(b *testing.B) { + key := "12345678901234567890123456789012" + _ = InitMessageEncryptor(key, 1) + encryptor := GetMessageEncryptor() + + sizes := []int{10, 50, 100, 500} + + for _, size := range sizes { + b.Run(fmt.Sprintf("size_%d", size), func(b *testing.B) { + ciphertexts := make([]string, size) + for i := 0; i < size; i++ { + msg := map[string]interface{}{ + "type": "text", + "data": map[string]string{"content": fmt.Sprintf("消息内容 %d", i)}, + } + plaintext, _ := json.Marshal(msg) + ciphertexts[i], _ = encryptor.Encrypt(plaintext) + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + for _, ct := range ciphertexts { + _, _ = encryptor.Decrypt(ct) + } + } + }) + } +} + +func TestBatchDecrypt(t *testing.T) { + key := "12345678901234567890123456789012" + _ = InitMessageEncryptor(key, 1) + encryptor := GetMessageEncryptor() + + // 准备测试数据 + count := 100 + ciphertexts := make([]string, count) + expectedContents := make([]string, count) + + for i := 0; i < count; i++ { + content := fmt.Sprintf("测试消息 %d", i) + expectedContents[i] = content + msg := fmt.Sprintf(`{"type":"text","data":{"content":"%s"}}`, content) + ciphertext, err := encryptor.Encrypt([]byte(msg)) + if err != nil { + t.Fatalf("Encrypt failed: %v", err) + } + ciphertexts[i] = ciphertext + } + + // 批量解密 + results := encryptor.BatchDecrypt(ciphertexts, 4) + + // 验证结果 + for i, result := range results { + if result == nil { + t.Errorf("Decrypt result[%d] is nil", i) + continue + } + + var msg map[string]interface{} + if err := json.Unmarshal(result, &msg); err != nil { + t.Errorf("Unmarshal result[%d] failed: %v", i, err) + continue + } + + data, ok := msg["data"].(map[string]interface{}) + if !ok { + t.Errorf("result[%d] data is not map", i) + continue + } + + content, ok := data["content"].(string) + if !ok { + t.Errorf("result[%d] content is not string", i) + continue + } + + if content != expectedContents[i] { + t.Errorf("result[%d] content = %s, want %s", i, content, expectedContents[i]) + } + } +} diff --git a/internal/repository/message_repo.go b/internal/repository/message_repo.go index aefc332..17e373d 100644 --- a/internal/repository/message_repo.go +++ b/internal/repository/message_repo.go @@ -131,6 +131,10 @@ func (r *MessageRepository) GetMessages(conversationID string, page, pageSize in Order("seq DESC"). Find(&messages).Error + if err == nil && len(messages) > 0 { + model.BatchDecryptMessagesParallel(messages) + } + return messages, total, err } @@ -141,6 +145,11 @@ func (r *MessageRepository) GetMessagesAfterSeq(conversationID string, afterSeq Order("seq ASC"). Limit(limit). Find(&messages).Error + + if err == nil && len(messages) > 0 { + model.BatchDecryptMessagesParallel(messages) + } + return messages, err } @@ -155,6 +164,11 @@ func (r *MessageRepository) GetMessagesBeforeSeq(conversationID string, beforeSe for i, j := 0, len(messages)-1; i < j; i, j = i+1, j-1 { messages[i], messages[j] = messages[j], messages[i] } + + if err == nil && len(messages) > 0 { + model.BatchDecryptMessagesParallel(messages) + } + return messages, err } @@ -350,6 +364,8 @@ func (r *MessageRepository) GetMessageByID(messageID string) (*model.Message, er if err != nil { return nil, err } + // 单条消息直接解密 + _ = message.Decrypt() return &message, nil } diff --git a/internal/wire/infrastructure.go b/internal/wire/infrastructure.go index 9c22d4f..0c37cba 100644 --- a/internal/wire/infrastructure.go +++ b/internal/wire/infrastructure.go @@ -6,6 +6,7 @@ import ( "carrot_bbs/internal/cache" "carrot_bbs/internal/config" "carrot_bbs/internal/model" + "carrot_bbs/internal/pkg/crypto" "carrot_bbs/internal/pkg/email" "carrot_bbs/internal/pkg/gorse" "carrot_bbs/internal/pkg/openai" @@ -45,6 +46,9 @@ var InfrastructureSet = wire.NewSet( ProvideOpenAIClient, ProvideEmailClient, ProvideS3Client, + + // 消息加密器 + ProvideMessageEncryptor, ) // ProvideConfig 加载配置 @@ -125,6 +129,31 @@ func ProvideTransactionManager(db *gorm.DB) repository.TransactionManager { return repository.NewTransactionManager(db) } +// ProvideMessageEncryptor 提供消息加密器 +func ProvideMessageEncryptor(cfg *config.Config) error { + if !cfg.Encryption.Enabled { + zap.L().Info("Message encryption is disabled") + return nil + } + + if len(cfg.Encryption.Key) != 32 { + zap.L().Error("Encryption key must be 32 bytes for AES-256", + zap.Int("actual_length", len(cfg.Encryption.Key)), + ) + return nil + } + + if err := crypto.InitMessageEncryptor(cfg.Encryption.Key, cfg.Encryption.KeyVersion); err != nil { + zap.L().Error("Failed to initialize message encryptor", zap.Error(err)) + return err + } + + zap.L().Info("Message encryption initialized successfully", + zap.Int("key_version", cfg.Encryption.KeyVersion), + ) + return nil +} + // ProvideLogger 提供 Zap Logger func ProvideLogger() *zap.Logger { logger, err := zap.NewProduction()