feat(message): add message encryption with batch decryption optimization

- Add encryption configuration with environment variable support
- Implement batch parallel decryption for improved query performance
- Integrate message encryptor via wire dependency injection
- Add explicit Decrypt() method for single message decryption
- Optimize AfterFind hook to defer decryption for batch processing
This commit is contained in:
lafay
2026-03-17 10:10:49 +08:00
parent 5d6c982c9c
commit 8a2b0e0a5e
7 changed files with 445 additions and 2 deletions

View File

@@ -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)
}

View File

@@ -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])
}
}
}