encryptx/aesx/aesx_test.go

package aesx_test

import (
	"bytes"
	"errors"
	"testing"

	"code.yun.ink/pkg/encryptx/aesx"
	"code.yun.ink/pkg/encryptx/cipherx"
)

func testKey(t *testing.T, size int) []byte {
	t.Helper()
	key, err := aesx.GenerateKey(size)
	if err != nil {
		t.Fatalf("GenerateKey(%d) 失败: %v", size, err)
	}
	if len(key) != size {
		t.Fatalf("GenerateKey(%d) 返回长度 %d", size, len(key))
	}
	return key
}

func TestGenerateKey(t *testing.T) {
	for _, size := range []int{aesx.KeySize128, aesx.KeySize192, aesx.KeySize256} {
		key := testKey(t, size)
		if !aesx.ValidKeySize(len(key)) {
			t.Errorf("ValidKeySize(%d) = false, 期望 true", len(key))
		}
	}

	// 非法长度必须被拒绝
	for _, size := range []int{-1, 0, 1, 15, 17, 31, 33, 64} {
		if _, err := aesx.GenerateKey(size); err == nil {
			t.Errorf("GenerateKey(%d) 期望错误, 得到 nil", size)
		}
		if aesx.ValidKeySize(size) {
			t.Errorf("ValidKeySize(%d) = true, 期望 false", size)
		}
	}

	// 两次生成的密钥不得相同
	k1, _ := aesx.GenerateKey(aesx.KeySize256)
	k2, _ := aesx.GenerateKey(aesx.KeySize256)
	if bytes.Equal(k1, k2) {
		t.Error("两次 GenerateKey 产生了相同密钥")
	}

	hexKey, err := aesx.GenerateKeyHex(aesx.KeySize256)
	if err != nil {
		t.Fatalf("GenerateKeyHex 失败: %v", err)
	}
	if len(hexKey) != 64 {
		t.Errorf("GenerateKeyHex(32) 长度 = %d, 期望 64", len(hexKey))
	}
	if _, err := aesx.GenerateKeyHex(7); err == nil {
		t.Error("GenerateKeyHex(7) 期望错误")
	}
}

// TestKeySizeValidation 是 P0-2 的回归测试。
//
// 旧实现的 GenerateKey 返回十六进制字符串（32 字节 → 64 字符），而加密函数把
// 入参直接当原始密钥字节用，导致「生成密钥 → 加密」这条最基本的路径必然失败
// （invalid key size 64）。现在密钥统一为 []byte，两者天然一致。
func TestKeySizeValidation(t *testing.T) {
	plaintext := []byte("hello world")

	for _, wrongLen := range []int{5, 15, 17, 31, 33, 64} {
		key := bytes.Repeat([]byte{0x01}, wrongLen)
		if _, err := aesx.EncryptCBC(plaintext, key); !errors.Is(err, cipherx.ErrInvalidKeySize) {
			t.Errorf("密钥长度 %d: 期望 ErrInvalidKeySize, 得到 %v", wrongLen, err)
		}
	}

	// GenerateKey 的返回值必须能直接用于加密（旧实现做不到）
	key, err := aesx.GenerateKey(aesx.KeySize256)
	if err != nil {
		t.Fatalf("GenerateKey 失败: %v", err)
	}
	ciphertext, err := aesx.EncryptCBC(plaintext, key)
	if err != nil {
		t.Fatalf("用 GenerateKey 生成的密钥加密失败: %v", err)
	}
	decrypted, err := aesx.DecryptCBC(ciphertext, key)
	if err != nil {
		t.Fatalf("解密失败: %v", err)
	}
	if !bytes.Equal(decrypted, plaintext) {
		t.Errorf("往返不一致: 得到 %q, 期望 %q", decrypted, plaintext)
	}
}

func TestAllModesRoundTrip(t *testing.T) {
	key := testKey(t, aesx.KeySize256)

	encryptors := map[string]func([]byte, []byte) ([]byte, error){
		"ECB": aesx.EncryptECB, "CBC": aesx.EncryptCBC, "CFB": aesx.EncryptCFB,
		"OFB": aesx.EncryptOFB, "CTR": aesx.EncryptCTR,
	}
	decryptors := map[string]func([]byte, []byte) ([]byte, error){
		"ECB": aesx.DecryptECB, "CBC": aesx.DecryptCBC, "CFB": aesx.DecryptCFB,
		"OFB": aesx.DecryptOFB, "CTR": aesx.DecryptCTR,
	}

	plaintexts := [][]byte{
		{}, []byte("a"), []byte("0123456789abcde"), []byte("0123456789abcdef"),
		[]byte("0123456789abcdefg"), bytes.Repeat([]byte("AES"), 500),
	}

	for name, encrypt := range encryptors {
		for i, plaintext := range plaintexts {
			ciphertext, err := encrypt(plaintext, key)
			if err != nil {
				t.Fatalf("%s 明文#%d 加密失败: %v", name, i, err)
			}
			decrypted, err := decryptors[name](ciphertext, key)
			if err != nil {
				t.Fatalf("%s 明文#%d 解密失败: %v", name, i, err)
			}
			if !bytes.Equal(decrypted, plaintext) {
				t.Errorf("%s 明文#%d 往返不一致", name, i)
			}
		}
	}
}

func TestHexAPIRoundTrip(t *testing.T) {
	key := testKey(t, aesx.KeySize128)
	plaintext := []byte("十六进制 API 往返测试")

	cases := []struct {
		name    string
		encrypt func([]byte, []byte) (string, error)
		decrypt func(string, []byte) ([]byte, error)
	}{
		{"ECB", aesx.EncryptECBHex, aesx.DecryptECBHex},
		{"CBC", aesx.EncryptCBCHex, aesx.DecryptCBCHex},
		{"CFB", aesx.EncryptCFBHex, aesx.DecryptCFBHex},
		{"OFB", aesx.EncryptOFBHex, aesx.DecryptOFBHex},
		{"CTR", aesx.EncryptCTRHex, aesx.DecryptCTRHex},
	}

	for _, c := range cases {
		hexStr, err := c.encrypt(plaintext, key)
		if err != nil {
			t.Fatalf("%s Hex 加密失败: %v", c.name, err)
		}
		// 输出必须是大写十六进制，且不含换行符（旧实现会带上 \n）
		if bytes.ContainsAny([]byte(hexStr), "\n\r \t") {
			t.Errorf("%s Hex 输出含空白字符: %q", c.name, hexStr)
		}
		for _, r := range hexStr {
			if (r < '0' || r > '9') && (r < 'A' || r > 'F') {
				t.Fatalf("%s Hex 输出含非大写十六进制字符: %q", c.name, hexStr)
			}
		}

		decrypted, err := c.decrypt(hexStr, key)
		if err != nil {
			t.Fatalf("%s Hex 解密失败: %v", c.name, err)
		}
		if !bytes.Equal(decrypted, plaintext) {
			t.Errorf("%s Hex 往返不一致", c.name)
		}
	}
}

func TestWithIVRoundTrip(t *testing.T) {
	key := testKey(t, aesx.KeySize256)
	iv := bytes.Repeat([]byte{0x5A}, aesx.BlockSize)
	plaintext := []byte("fixed iv interop")

	cases := []struct {
		name    string
		encrypt func([]byte, []byte, []byte) ([]byte, error)
		decrypt func([]byte, []byte, []byte) ([]byte, error)
	}{
		{"CBC", aesx.EncryptCBCWithIV, aesx.DecryptCBCWithIV},
		{"CFB", aesx.EncryptCFBWithIV, aesx.DecryptCFBWithIV},
		{"OFB", aesx.EncryptOFBWithIV, aesx.DecryptOFBWithIV},
		{"CTR", aesx.EncryptCTRWithIV, aesx.DecryptCTRWithIV},
	}
	for _, c := range cases {
		ciphertext, err := c.encrypt(plaintext, key, iv)
		if err != nil {
			t.Fatalf("%s WithIV 加密失败: %v", c.name, err)
		}
		decrypted, err := c.decrypt(ciphertext, key, iv)
		if err != nil {
			t.Fatalf("%s WithIV 解密失败: %v", c.name, err)
		}
		if !bytes.Equal(decrypted, plaintext) {
			t.Errorf("%s WithIV 往返不一致", c.name)
		}
		// 错误 IV 下，CBC 通常会因填充校验失败报错，流式模式则会静默返回垃圾
		// 明文（这正是这些模式不提供完整性保护的体现）。无论哪种情况，都绝不
		// 允许还原出正确明文。
		wrongIV, err := c.decrypt(ciphertext, key, bytes.Repeat([]byte{0x00}, aesx.BlockSize))
		if err == nil && bytes.Equal(wrongIV, plaintext) {
			t.Errorf("%s 使用错误 IV 却解出了正确明文", c.name)
		}
	}
}

// TestECBIsDeterministic 记录 ECB 的固有弱点（相同明文块 → 相同密文块），
// 与 CBC 对比以证明文档中的安全警告属实。
func TestECBIsDeterministic(t *testing.T) {
	key := testKey(t, aesx.KeySize256)
	plaintext := bytes.Repeat([]byte("0123456789abcdef"), 4)

	ecb1, err := aesx.EncryptECB(plaintext, key)
	if err != nil {
		t.Fatalf("ECB 加密失败: %v", err)
	}
	ecb2, err := aesx.EncryptECB(plaintext, key)
	if err != nil {
		t.Fatalf("ECB 加密失败: %v", err)
	}
	if !bytes.Equal(ecb1, ecb2) {
		t.Error("ECB 应当是确定性的（这正是它的弱点）")
	}

	cbc1, err := aesx.EncryptCBC(plaintext, key)
	if err != nil {
		t.Fatalf("CBC 加密失败: %v", err)
	}
	cbc2, err := aesx.EncryptCBC(plaintext, key)
	if err != nil {
		t.Fatalf("CBC 加密失败: %v", err)
	}
	if bytes.Equal(cbc1, cbc2) {
		t.Error("CBC 使用了随机 IV，两次加密结果不应相同")
	}
}

func TestMalformedCiphertextNeverPanics(t *testing.T) {
	key := testKey(t, aesx.KeySize256)
	decryptors := map[string]func([]byte, []byte) ([]byte, error){
		"ECB": aesx.DecryptECB, "CBC": aesx.DecryptCBC, "CFB": aesx.DecryptCFB,
		"OFB": aesx.DecryptOFB, "CTR": aesx.DecryptCTR,
	}
	for _, decrypt := range decryptors {
		for n := 0; n <= 40; n++ {
			// 不应 panic
			_, _ = decrypt(bytes.Repeat([]byte{0x7F}, n), key)
		}
	}

	// ECB / CBC 的非对齐密文必须返回 ErrNotBlockAligned
	nonAligned := bytes.Repeat([]byte{0xAB}, 20)
	if _, err := aesx.DecryptECB(nonAligned, key); !errors.Is(err, cipherx.ErrNotBlockAligned) {
		t.Errorf("ECB 非对齐: 期望 ErrNotBlockAligned, 得到 %v", err)
	}
	if _, err := aesx.DecryptCBC(nonAligned, key); !errors.Is(err, cipherx.ErrNotBlockAligned) {
		t.Errorf("CBC 非对齐: 期望 ErrNotBlockAligned, 得到 %v", err)
	}
}

// ---- GCM ----

func TestGCMRoundTrip(t *testing.T) {
	key := testKey(t, aesx.KeySize256)
	plaintexts := [][]byte{
		{}, []byte("a"), []byte("GCM 往返测试"), bytes.Repeat([]byte("x"), 100000),
	}
	for i, plaintext := range plaintexts {
		ciphertext, nonce, err := aesx.EncryptGCM(plaintext, key, nil)
		if err != nil {
			t.Fatalf("GCM 加密#%d 失败: %v", i, err)
		}
		if len(nonce) != aesx.GCMNonceSize {
			t.Fatalf("nonce 长度 = %d, 期望 %d", len(nonce), aesx.GCMNonceSize)
		}
		// 密文应等于明文长度 + 认证标签
		if len(ciphertext) != len(plaintext)+aesx.GCMTagSize {
			t.Errorf("密文长度 = %d, 期望 %d", len(ciphertext), len(plaintext)+aesx.GCMTagSize)
		}

		decrypted, err := aesx.DecryptGCM(ciphertext, key, nonce, nil)
		if err != nil {
			t.Fatalf("GCM 解密#%d 失败: %v", i, err)
		}
		if !bytes.Equal(decrypted, plaintext) {
			t.Errorf("GCM 往返#%d 不一致", i)
		}
	}
}

func TestGCMNonceIsRandom(t *testing.T) {
	key := testKey(t, aesx.KeySize256)
	seen := make(map[string]bool)
	for i := 0; i < 64; i++ {
		_, nonce, err := aesx.EncryptGCM([]byte("same"), key, nil)
		if err != nil {
			t.Fatalf("GCM 加密失败: %v", err)
		}
		if seen[string(nonce)] {
			t.Fatal("GCM nonce 重复，会导致灾难性的密钥恢复攻击")
		}
		seen[string(nonce)] = true
	}
}

func TestGCMBindsAAD(t *testing.T) {
	key := testKey(t, aesx.KeySize256)
	plaintext := []byte("authenticated data test")
	aad := []byte("v1|user=42")

	ciphertext, nonce, err := aesx.EncryptGCM(plaintext, key, aad)
	if err != nil {
		t.Fatalf("GCM 加密失败: %v", err)
	}

	// 正确的 AAD 可解密
	if _, err := aesx.DecryptGCM(ciphertext, key, nonce, aad); err != nil {
		t.Fatalf("正确 AAD 解密失败: %v", err)
	}
	// 不同的 AAD 必须失败
	if _, err := aesx.DecryptGCM(ciphertext, key, nonce, []byte("v1|user=43")); !errors.Is(err, aesx.ErrAuthenticationFailed) {
		t.Errorf("错误 AAD: 期望 ErrAuthenticationFailed, 得到 %v", err)
	}
	// 缺少 AAD 必须失败
	if _, err := aesx.DecryptGCM(ciphertext, key, nonce, nil); !errors.Is(err, aesx.ErrAuthenticationFailed) {
		t.Errorf("缺失 AAD: 期望 ErrAuthenticationFailed, 得到 %v", err)
	}
}

// TestGCMTamperDetection 验证任何一处篡改都会导致解密失败，且不返回明文。
func TestGCMTamperDetection(t *testing.T) {
	key := testKey(t, aesx.KeySize256)
	plaintext := []byte("tamper detection")
	ciphertext, nonce, err := aesx.EncryptGCM(plaintext, key, nil)
	if err != nil {
		t.Fatalf("GCM 加密失败: %v", err)
	}

	// 翻转密文（含标签）中每一个字节
	for i := range ciphertext {
		tampered := bytes.Clone(ciphertext)
		tampered[i] ^= 0x01
		got, err := aesx.DecryptGCM(tampered, key, nonce, nil)
		if !errors.Is(err, aesx.ErrAuthenticationFailed) {
			t.Errorf("篡改密文第 %d 字节: 期望认证失败, 得到 %v", i, err)
		}
		if got != nil {
			t.Errorf("篡改密文第 %d 字节: 认证失败却返回了明文 %q", i, got)
		}
	}

	// 翻转 nonce 中每一个字节
	for i := range nonce {
		tampered := bytes.Clone(nonce)
		tampered[i] ^= 0x01
		if _, err := aesx.DecryptGCM(ciphertext, key, tampered, nil); !errors.Is(err, aesx.ErrAuthenticationFailed) {
			t.Errorf("篡改 nonce 第 %d 字节: 期望认证失败, 得到 %v", i, err)
		}
	}

	// 错误密钥
	wrongKey := testKey(t, aesx.KeySize256)
	if _, err := aesx.DecryptGCM(ciphertext, wrongKey, nonce, nil); !errors.Is(err, aesx.ErrAuthenticationFailed) {
		t.Errorf("错误密钥: 期望认证失败, 得到 %v", err)
	}
}

func TestGCMInvalidInputs(t *testing.T) {
	key := testKey(t, aesx.KeySize256)
	ciphertext, nonce, err := aesx.EncryptGCM([]byte("data"), key, nil)
	if err != nil {
		t.Fatalf("GCM 加密失败: %v", err)
	}

	// nonce 长度错误
	for _, badNonce := range [][]byte{nil, {}, bytes.Repeat([]byte{1}, 8), bytes.Repeat([]byte{1}, 16)} {
		if _, err := aesx.DecryptGCM(ciphertext, key, badNonce, nil); !errors.Is(err, aesx.ErrInvalidNonceSize) {
			t.Errorf("nonce 长度 %d: 期望 ErrInvalidNonceSize, 得到 %v", len(badNonce), err)
		}
	}

	// 密文短于认证标签
	if _, err := aesx.DecryptGCM(bytes.Repeat([]byte{1}, aesx.GCMTagSize-1), key, nonce, nil); !errors.Is(err, aesx.ErrCiphertextTooShort) {
		t.Errorf("超短密文: 期望 ErrCiphertextTooShort, 得到 %v", err)
	}
	// 空密文
	if _, err := aesx.DecryptGCM(nil, key, nonce, nil); !errors.Is(err, aesx.ErrCiphertextTooShort) {
		t.Errorf("空密文: 期望 ErrCiphertextTooShort, 得到 %v", err)
	}
	// 密钥长度错误
	if _, err := aesx.DecryptGCM(ciphertext, []byte("short"), nonce, nil); !errors.Is(err, cipherx.ErrInvalidKeySize) {
		t.Errorf("错误密钥长度: 期望 ErrInvalidKeySize, 得到 %v", err)
	}
}

func TestSealOpenGCM(t *testing.T) {
	key := testKey(t, aesx.KeySize128)
	plaintext := []byte("single field storage")
	aad := []byte("ctx")

	sealed, err := aesx.SealGCM(plaintext, key, aad)
	if err != nil {
		t.Fatalf("SealGCM 失败: %v", err)
	}
	if len(sealed) != aesx.GCMNonceSize+len(plaintext)+aesx.GCMTagSize {
		t.Errorf("SealGCM 长度 = %d, 期望 %d", len(sealed), aesx.GCMNonceSize+len(plaintext)+aesx.GCMTagSize)
	}

	opened, err := aesx.OpenGCM(sealed, key, aad)
	if err != nil {
		t.Fatalf("OpenGCM 失败: %v", err)
	}
	if !bytes.Equal(opened, plaintext) {
		t.Errorf("Seal/Open 往返不一致")
	}

	// AAD 不匹配
	if _, err := aesx.OpenGCM(sealed, key, []byte("other")); !errors.Is(err, aesx.ErrAuthenticationFailed) {
		t.Errorf("错误 AAD: 期望认证失败, 得到 %v", err)
	}
	// 数据过短
	for n := 0; n < aesx.GCMNonceSize+aesx.GCMTagSize; n++ {
		if _, err := aesx.OpenGCM(bytes.Repeat([]byte{0x01}, n), key, aad); err == nil {
			t.Errorf("长度 %d 的数据: 期望错误, 得到 nil", n)
		}
	}
	// 篡改
	tampered := bytes.Clone(sealed)
	tampered[len(tampered)-1] ^= 0x01
	if _, err := aesx.OpenGCM(tampered, key, aad); !errors.Is(err, aesx.ErrAuthenticationFailed) {
		t.Errorf("篡改数据: 期望认证失败, 得到 %v", err)
	}
}

func TestGCMHexRoundTrip(t *testing.T) {
	key := testKey(t, aesx.KeySize256)
	plaintext := []byte("hex gcm")

	ciphertextHex, nonceHex, err := aesx.EncryptGCMHex(plaintext, key, nil)
	if err != nil {
		t.Fatalf("EncryptGCMHex 失败: %v", err)
	}
	decrypted, err := aesx.DecryptGCMHex(ciphertextHex, nonceHex, key, nil)
	if err != nil {
		t.Fatalf("DecryptGCMHex 失败: %v", err)
	}
	if !bytes.Equal(decrypted, plaintext) {
		t.Error("GCM Hex 往返不一致")
	}

	if _, err := aesx.DecryptGCMHex("not hex", nonceHex, key, nil); err == nil {
		t.Error("非法密文十六进制: 期望错误")
	}
	if _, err := aesx.DecryptGCMHex(ciphertextHex, "not hex", key, nil); err == nil {
		t.Error("非法 nonce 十六进制: 期望错误")
	}
}

func TestGCMWithNonceInterop(t *testing.T) {
	key := testKey(t, aesx.KeySize256)
	nonce := bytes.Repeat([]byte{0x42}, aesx.GCMNonceSize)
	plaintext := []byte("interop")

	ciphertext, err := aesx.EncryptGCMWithNonce(plaintext, key, nonce, nil)
	if err != nil {
		t.Fatalf("EncryptGCMWithNonce 失败: %v", err)
	}
	decrypted, err := aesx.DecryptGCM(ciphertext, key, nonce, nil)
	if err != nil {
		t.Fatalf("解密失败: %v", err)
	}
	if !bytes.Equal(decrypted, plaintext) {
		t.Error("往返不一致")
	}

	if _, err := aesx.EncryptGCMWithNonce(plaintext, key, []byte("bad"), nil); !errors.Is(err, aesx.ErrInvalidNonceSize) {
		t.Errorf("非法 nonce 长度: 期望 ErrInvalidNonceSize, 得到 %v", err)
	}
}