backend/internal/pkg/crypto/message_encryptor_test.go

package crypto

import (
	"encoding/json"
	"sync"
	"testing"
)

func TestMessageEncryptor_EncryptDecrypt(t *testing.T) {
	// 初始化加密器（32字节密钥）
	key := "12345678901234567890123456789012" // 32 bytes
	err := InitMessageEncryptor(key, 1)
	if err != nil {
		t.Fatalf("InitMessageEncryptor failed: %v", err)
	}

	encryptor := GetMessageEncryptor()
	if encryptor == nil {
		t.Fatal("GetMessageEncryptor returned nil")
	}

	tests := []struct {
		name      string
		plaintext []byte
	}{
		{
			name:      "simple text",
			plaintext: []byte("Hello, World!"),
		},
		{
			name:      "empty",
			plaintext: []byte(""),
		},
		{
			name:      "json data",
			plaintext: []byte(`{"type":"text","data":{"content":"这是一条测试消息"}}`),
		},
		{
			name: "long message",
			plaintext: []byte("这是一条很长的消息，用于测试加密长文本的性能和正确性。" +
				"加密后的数据应该能够正确解密，并且每次加密的结果应该不同（因为有随机nonce）。"),
		},
		{
			name: "complex message segments",
			plaintext: mustMarshal([]map[string]interface{}{
				{"type": "text", "data": map[string]interface{}{"content": "你好"}},
				{"type": "image", "data": map[string]interface{}{"url": "https://example.com/image.png"}},
			}),
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// 空字符串特殊处理
			if len(tt.plaintext) == 0 {
				ciphertext, err := encryptor.Encrypt(tt.plaintext)
				if err != nil {
					t.Errorf("Encrypt() error = %v", err)
				}
				if ciphertext != "" {
					t.Errorf("Encrypt() for empty input should return empty string, got %s", ciphertext)
				}
				return
			}

			// 加密
			ciphertext, err := encryptor.Encrypt(tt.plaintext)
			if err != nil {
				t.Errorf("Encrypt() error = %v", err)
				return
			}

			// 解密
			decrypted, err := encryptor.Decrypt(ciphertext)
			if err != nil {
				t.Errorf("Decrypt() error = %v", err)
				return
			}

			// 比较
			if string(decrypted) != string(tt.plaintext) {
				t.Errorf("Decrypt() = %s, want %s", decrypted, tt.plaintext)
			}
		})
	}
}

func TestMessageEncryptor_DifferentNonce(t *testing.T) {
	key := "12345678901234567890123456789012"
	_ = InitMessageEncryptor(key, 1)
	encryptor := GetMessageEncryptor()

	plaintext := []byte("same message")

	// 加密两次相同的内容，密文应该不同（因为随机nonce）
	ciphertext1, _ := encryptor.Encrypt(plaintext)
	ciphertext2, _ := encryptor.Encrypt(plaintext)

	if ciphertext1 == ciphertext2 {
		t.Error("Same plaintext should produce different ciphertext due to random nonce")
	}

	// 但两个都应该能正确解密
	decrypted1, _ := encryptor.Decrypt(ciphertext1)
	decrypted2, _ := encryptor.Decrypt(ciphertext2)

	if string(decrypted1) != string(plaintext) || string(decrypted2) != string(plaintext) {
		t.Error("Both ciphertexts should decrypt to the same plaintext")
	}
}

func TestMessageEncryptor_InvalidKey(t *testing.T) {
	tests := []struct {
		name    string
		key     string
		wantErr bool
	}{
		{"valid 32 bytes", "12345678901234567890123456789012", false},
		{"invalid 16 bytes", "1234567890123456", true},
		{"invalid 24 bytes", "123456789012345678901234", true},
		{"invalid 31 bytes", "1234567890123456789012345678901", true},
		{"invalid 33 bytes", "123456789012345678901234567890123", true},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// 重置全局加密器
			globalEncryptor = nil
			encryptorOnce = sync.Once{}

			err := InitMessageEncryptor(tt.key, 1)
			if (err != nil) != tt.wantErr {
				t.Errorf("InitMessageEncryptor() error = %v, wantErr %v", err, tt.wantErr)
			}
		})
	}
}

func TestMessageEncryptor_KeyVersion(t *testing.T) {
	// 重置
	globalEncryptor = nil
	encryptorOnce = sync.Once{}

	err := InitMessageEncryptor("12345678901234567890123456789012", 1)
	if err != nil {
		t.Fatal(err)
	}

	encryptor := GetMessageEncryptor()
	if encryptor.GetKeyVersion() != 1 {
		t.Errorf("GetKeyVersion() = %d, want 1", encryptor.GetKeyVersion())
	}

	// 测试密钥轮换
	err = encryptor.RotateKey("abcdefghijklmnopqrstuvwxyz123456", 2)
	if err != nil {
		t.Fatal(err)
	}

	if encryptor.GetKeyVersion() != 2 {
		t.Errorf("GetKeyVersion() after rotation = %d, want 2", encryptor.GetKeyVersion())
	}
}

func mustMarshal(v interface{}) []byte {
	data, _ := json.Marshal(v)
	return data
}