backend/internal/pkg/crypto/message_encryptor.go

package crypto

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"io"
	"sync"
)

var (
	ErrInvalidKey        = errors.New("invalid encryption key: must be 32 bytes")
	ErrInvalidCiphertext = errors.New("invalid ciphertext: too short")
	ErrDecryptFailed     = errors.New("decryption failed")
)

// MessageEncryptor 消息加密器
// 使用 AES-256-GCM 算法进行加密
type MessageEncryptor struct {
	key        []byte
	gcm        cipher.AEAD
	keyVersion int
	mu         sync.RWMutex
}

// globalEncryptor 全局加密器实例
var globalEncryptor *MessageEncryptor
var encryptorOnce sync.Once

// InitMessageEncryptor 初始化全局消息加密器
// key 必须是32字节（256位）的密钥
func InitMessageEncryptor(key string, keyVersion int) error {
	keyBytes := []byte(key)
	if len(keyBytes) != 32 {
		return ErrInvalidKey
	}

	block, err := aes.NewCipher(keyBytes)
	if err != nil {
		return err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return err
	}

	encryptorOnce.Do(func() {
		globalEncryptor = &MessageEncryptor{
			key:        keyBytes,
			gcm:        gcm,
			keyVersion: keyVersion,
		}
	})

	return nil
}

// GetMessageEncryptor 获取全局加密器实例
func GetMessageEncryptor() *MessageEncryptor {
	return globalEncryptor
}

// Encrypt 加密数据
// 返回 base64 编码的密文（格式：nonce + ciphertext + tag）
func (e *MessageEncryptor) Encrypt(plaintext []byte) (string, error) {
	if e == nil {
		return "", errors.New("encryptor not initialized")
	}

	if len(plaintext) == 0 {
		return "", nil
	}

	e.mu.RLock()
	defer e.mu.RUnlock()

	// 生成随机nonce（12字节）
	nonce := make([]byte, e.gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return "", err
	}

	// 加密数据，nonce附加在密文前面
	ciphertext := e.gcm.Seal(nonce, nonce, plaintext, nil)

	// 返回base64编码的密文
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

// Decrypt 解密数据
// 输入 base64 编码的密文
func (e *MessageEncryptor) Decrypt(ciphertextBase64 string) ([]byte, error) {
	if e == nil {
		return nil, errors.New("encryptor not initialized")
	}

	if ciphertextBase64 == "" {
		return nil, nil
	}

	e.mu.RLock()
	defer e.mu.RUnlock()

	// base64解码
	ciphertext, err := base64.StdEncoding.DecodeString(ciphertextBase64)
	if err != nil {
		return nil, err
	}

	// 检查密文长度
	nonceSize := e.gcm.NonceSize()
	if len(ciphertext) < nonceSize {
		return nil, ErrInvalidCiphertext
	}

	// 提取nonce和实际密文
	nonce := ciphertext[:nonceSize]
	actualCiphertext := ciphertext[nonceSize:]

	// 解密
	plaintext, err := e.gcm.Open(nil, nonce, actualCiphertext, nil)
	if err != nil {
		return nil, ErrDecryptFailed
	}

	return plaintext, nil
}

// GetKeyVersion 获取当前密钥版本
func (e *MessageEncryptor) GetKeyVersion() int {
	if e == nil {
		return 0
	}
	e.mu.RLock()
	defer e.mu.RUnlock()
	return e.keyVersion
}

// BatchDecrypt 批量解密多条密文，返回与输入等长的明文字节切片数组。
// 并发度由 workers 指定（<=0 时按密文数量自适应）。
// 解密失败或空密文的位置返回 nil，调用方可据此判断。
// 这是 Decrypt 的无锁读优化版：AEAD 的 Open 不修改内部状态，
// 但仍走 RLock 以兼容未来可能的密钥轮换语义。
func (e *MessageEncryptor) BatchDecrypt(ciphertexts []string, workers int) [][]byte {
	results := make([][]byte, len(ciphertexts))
	if e == nil || len(ciphertexts) == 0 {
		return results
	}

	// 自适应并发度：与历史实现保持一致
	if workers <= 0 {
		switch {
		case len(ciphertexts) < 20:
			workers = 2
		case len(ciphertexts) > 100:
			workers = 8
		default:
			workers = 4
		}
	}
	if workers > len(ciphertexts) {
		workers = len(ciphertexts)
	}

	// 预取一次 AEAD，避免 worker 内重复进入 RLock
	e.mu.RLock()
	gcm := e.gcm
	nonceSize := e.gcm.NonceSize()
	e.mu.RUnlock()

	jobs := make(chan int, len(ciphertexts))
	var wg sync.WaitGroup

	for w := 0; w < workers; w++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for i := range jobs {
				ct := ciphertexts[i]
				if ct == "" {
					continue
				}
				raw, err := base64.StdEncoding.DecodeString(ct)
				if err != nil || len(raw) < nonceSize {
					continue
				}
				nonce := raw[:nonceSize]
				actual := raw[nonceSize:]
				plaintext, err := gcm.Open(nil, nonce, actual, nil)
				if err != nil {
					continue
				}
				results[i] = plaintext
			}
		}()
	}

	for i := range ciphertexts {
		jobs <- i
	}
	close(jobs)
	wg.Wait()

	return results
}