Add AES 256 and update AES test

This commit is contained in:
2025-07-05 23:59:59 -05:00
parent 26371cdb6a
commit bc05835dae
2 changed files with 277 additions and 33 deletions

View File

@@ -8,7 +8,57 @@ import (
"io"
)
const (
// AES key sizes in bytes
AES128KeySize = 16 // 128 bits
AES192KeySize = 24 // 192 bits
AES256KeySize = 32 // 256 bits
// Default to AES-256
DefaultKeySize = AES256KeySize
)
// GenerateAESKey generates a random AES key of the specified size
func GenerateAESKey(keySize int) ([]byte, error) {
if keySize != AES128KeySize && keySize != AES192KeySize && keySize != AES256KeySize {
return nil, errors.New("invalid key size: must be 16, 24, or 32 bytes")
}
key := make([]byte, keySize)
if _, err := io.ReadFull(rand.Reader, key); err != nil {
return nil, err
}
return key, nil
}
// GenerateAES256Key generates a random AES-256 key (default)
func GenerateAES256Key() ([]byte, error) {
return GenerateAESKey(AES256KeySize)
}
// EncryptAES256CBC encrypts data using AES-256 in CBC mode
func EncryptAES256CBC(key, plaintext []byte) ([]byte, error) {
if len(key) != AES256KeySize {
return nil, errors.New("key must be 32 bytes for AES-256")
}
return EncryptAESCBC(key, plaintext)
}
// DecryptAES256CBC decrypts data using AES-256 in CBC mode
func DecryptAES256CBC(key, ciphertext []byte) ([]byte, error) {
if len(key) != AES256KeySize {
return nil, errors.New("key must be 32 bytes for AES-256")
}
return DecryptAESCBC(key, ciphertext)
}
// EncryptAESCBC encrypts data using AES in CBC mode (accepts any valid AES key size)
func EncryptAESCBC(key, plaintext []byte) ([]byte, error) {
// Validate key size
if len(key) != AES128KeySize && len(key) != AES192KeySize && len(key) != AES256KeySize {
return nil, errors.New("invalid key size: must be 16, 24, or 32 bytes")
}
block, err := aes.NewCipher(key)
if err != nil {
return nil, err
@@ -36,7 +86,13 @@ func EncryptAESCBC(key, plaintext []byte) ([]byte, error) {
return append(iv, ciphertext...), nil
}
// DecryptAESCBC decrypts data using AES in CBC mode (accepts any valid AES key size)
func DecryptAESCBC(key, ciphertext []byte) ([]byte, error) {
// Validate key size
if len(key) != AES128KeySize && len(key) != AES192KeySize && len(key) != AES256KeySize {
return nil, errors.New("invalid key size: must be 16, 24, or 32 bytes")
}
block, err := aes.NewCipher(key)
if err != nil {
return nil, err
@@ -58,10 +114,22 @@ func DecryptAESCBC(key, ciphertext []byte) ([]byte, error) {
mode.CryptBlocks(plaintext, ciphertext)
// Remove PKCS7 padding
if len(plaintext) == 0 {
return nil, errors.New("invalid padding: empty plaintext")
}
padding := int(plaintext[len(plaintext)-1])
if padding == 0 || padding > len(plaintext) {
if padding == 0 || padding > aes.BlockSize || padding > len(plaintext) {
return nil, errors.New("invalid PKCS7 padding")
}
// TODO: Add check to ensure all padding bytes are correct?
// Verify all padding bytes are correct
for i := len(plaintext) - padding; i < len(plaintext); i++ {
if plaintext[i] != byte(padding) {
return nil, errors.New("invalid PKCS7 padding")
}
}
return plaintext[:len(plaintext)-padding], nil
}

View File

@@ -4,14 +4,49 @@ import (
"bytes"
"crypto/aes"
"crypto/rand"
"fmt"
"testing"
)
func TestAESCBCEncryptionDecryption(t *testing.T) {
// Generate a random key (AES-256)
key := make([]byte, 32)
if _, err := rand.Read(key); err != nil {
t.Fatalf("Failed to generate random key: %v", err)
func TestGenerateAESKeys(t *testing.T) {
t.Run("GenerateAES256Key", func(t *testing.T) {
key, err := GenerateAES256Key()
if err != nil {
t.Fatalf("GenerateAES256Key failed: %v", err)
}
if len(key) != AES256KeySize {
t.Errorf("Expected key size %d, got %d", AES256KeySize, len(key))
}
})
t.Run("GenerateAESKey_AllSizes", func(t *testing.T) {
sizes := []int{AES128KeySize, AES192KeySize, AES256KeySize}
for _, size := range sizes {
key, err := GenerateAESKey(size)
if err != nil {
t.Fatalf("GenerateAESKey(%d) failed: %v", size, err)
}
if len(key) != size {
t.Errorf("Expected key size %d, got %d", size, len(key))
}
}
})
t.Run("GenerateAESKey_InvalidSize", func(t *testing.T) {
invalidSizes := []int{8, 15, 17, 23, 25, 31, 33, 64}
for _, size := range invalidSizes {
_, err := GenerateAESKey(size)
if err == nil {
t.Errorf("GenerateAESKey(%d) should have failed but didn't", size)
}
}
})
}
func TestAES256CBCEncryptionDecryption(t *testing.T) {
key, err := GenerateAES256Key()
if err != nil {
t.Fatalf("Failed to generate AES-256 key: %v", err)
}
testCases := []struct {
@@ -20,67 +55,208 @@ func TestAESCBCEncryptionDecryption(t *testing.T) {
}{
{"ShortMessage", []byte("Hello")},
{"BlockSizeMessage", []byte("This is 16 bytes")},
{"LongMessage", []byte("This is a longer message that spans multiple AES blocks.")},
{"LongMessage", []byte("This is a longer message that spans multiple AES blocks and tests the padding.")},
{"EmptyMessage", []byte("")},
{"SingleByte", []byte("A")},
{"ExactlyTwoBlocks", []byte("This is exactly 32 bytes long!!!")},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
ciphertext, err := EncryptAESCBC(key, tc.plaintext)
ciphertext, err := EncryptAES256CBC(key, tc.plaintext)
if err != nil {
t.Fatalf("EncryptAESCBC failed: %v", err)
t.Fatalf("EncryptAES256CBC failed: %v", err)
}
decrypted, err := DecryptAESCBC(key, ciphertext)
decrypted, err := DecryptAES256CBC(key, ciphertext)
if err != nil {
t.Fatalf("DecryptAESCBC failed: %v", err)
t.Fatalf("DecryptAES256CBC failed: %v", err)
}
if !bytes.Equal(tc.plaintext, decrypted) {
t.Errorf("Decrypted text does not match original plaintext. Got %q, want %q", decrypted, tc.plaintext)
t.Errorf("Decrypted text does not match original plaintext.\nGot: %q (%x)\nWant: %q (%x)",
decrypted, decrypted, tc.plaintext, tc.plaintext)
}
})
}
}
func TestAES256CBC_InvalidKeySize(t *testing.T) {
plaintext := []byte("test message")
invalidKeys := [][]byte{
make([]byte, 16), // AES-128
make([]byte, 24), // AES-192
make([]byte, 15), // Too short
make([]byte, 33), // Too long
nil, // Nil key
}
for i, key := range invalidKeys {
t.Run(fmt.Sprintf("InvalidKey_%d", i), func(t *testing.T) {
_, err := EncryptAES256CBC(key, plaintext)
if err == nil {
t.Error("EncryptAES256CBC should have failed with invalid key size")
}
// Test with some dummy ciphertext
dummyCiphertext := make([]byte, 32) // Just enough for IV + one block
rand.Read(dummyCiphertext)
_, err = DecryptAES256CBC(key, dummyCiphertext)
if err == nil {
t.Error("DecryptAES256CBC should have failed with invalid key size")
}
})
}
}
func TestAESCBCEncryptionDecryption(t *testing.T) {
keySizes := []int{AES128KeySize, AES192KeySize, AES256KeySize}
for _, keySize := range keySizes {
t.Run(fmt.Sprintf("AES_%d", keySize*8), func(t *testing.T) {
key, err := GenerateAESKey(keySize)
if err != nil {
t.Fatalf("Failed to generate AES-%d key: %v", keySize*8, err)
}
testCases := []struct {
name string
plaintext []byte
}{
{"ShortMessage", []byte("Hello")},
{"BlockSizeMessage", []byte("This is 16 bytes")},
{"LongMessage", []byte("This is a longer message that spans multiple AES blocks.")},
{"EmptyMessage", []byte("")},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
ciphertext, err := EncryptAESCBC(key, tc.plaintext)
if err != nil {
t.Fatalf("EncryptAESCBC failed: %v", err)
}
decrypted, err := DecryptAESCBC(key, ciphertext)
if err != nil {
t.Fatalf("DecryptAESCBC failed: %v", err)
}
if !bytes.Equal(tc.plaintext, decrypted) {
t.Errorf("Decrypted text does not match original plaintext. Got %q, want %q", decrypted, tc.plaintext)
}
})
}
})
}
}
func TestDecryptAESCBCErrorCases(t *testing.T) {
key := make([]byte, 32)
_, _ = rand.Read(key)
key, err := GenerateAES256Key()
if err != nil {
t.Fatalf("Failed to generate key: %v", err)
}
t.Run("CiphertextTooShort", func(t *testing.T) {
shortCiphertext := []byte{0x01, 0x02, 0x03} // Less than AES block size
_, err := DecryptAESCBC(key, shortCiphertext)
if err == nil {
t.Error("DecryptAESCBC should have failed for ciphertext shorter than block size, but it didn't")
t.Error("DecryptAESCBC should have failed for ciphertext shorter than block size")
}
})
t.Run("InvalidPadding", func(t *testing.T) {
// Encrypt something valid first
plaintext := []byte("valid data")
ciphertext, _ := EncryptAESCBC(key, plaintext)
// Tamper with the ciphertext (specifically the part that would affect padding)
if len(ciphertext) > aes.BlockSize {
ciphertext[len(ciphertext)-1] = ^ciphertext[len(ciphertext)-1] // Flip bits of last byte
}
_, err := DecryptAESCBC(key, ciphertext)
t.Run("InvalidKeySize", func(t *testing.T) {
invalidKey := make([]byte, 17) // Invalid key size
validCiphertext := make([]byte, 32) // IV + one block
rand.Read(validCiphertext)
_, err := DecryptAESCBC(invalidKey, validCiphertext)
if err == nil {
// Note: Depending on the padding implementation and the nature of the tampering,
// CBC decryption might not always error out on bad padding. It might return garbage data.
// A more robust test might check the decrypted content, but error checking is a start.
t.Logf("DecryptAESCBC did not error on potentially invalid padding (this might be expected)")
t.Error("DecryptAESCBC should have failed for invalid key size")
}
})
t.Run("CiphertextNotMultipleOfBlockSize", func(t *testing.T) {
iv := make([]byte, aes.BlockSize)
_, _ = rand.Read(iv)
rand.Read(iv)
invalidCiphertext := append(iv, []byte{0x01, 0x02, 0x03}...) // IV + data not multiple of block size
_, err := DecryptAESCBC(key, invalidCiphertext)
if err == nil {
t.Error("DecryptAESCBC should have failed for ciphertext not multiple of block size, but it didn't")
t.Error("DecryptAESCBC should have failed for ciphertext not multiple of block size")
}
})
t.Run("InvalidPadding", func(t *testing.T) {
// Create a valid ciphertext first
plaintext := []byte("valid data")
ciphertext, err := EncryptAESCBC(key, plaintext)
if err != nil {
t.Fatalf("Failed to create test ciphertext: %v", err)
}
// Corrupt the last byte (which affects padding)
corruptedCiphertext := make([]byte, len(ciphertext))
copy(corruptedCiphertext, ciphertext)
corruptedCiphertext[len(corruptedCiphertext)-1] ^= 0xFF
_, err = DecryptAESCBC(key, corruptedCiphertext)
if err == nil {
t.Error("DecryptAESCBC should have failed for corrupted padding")
}
})
t.Run("EmptyPlaintext", func(t *testing.T) {
// Create a ciphertext that would result in empty plaintext
invalidCiphertext := make([]byte, aes.BlockSize) // Only IV, no data
_, err := DecryptAESCBC(key, invalidCiphertext)
if err == nil {
t.Error("DecryptAESCBC should have failed for empty ciphertext data")
}
})
}
func TestConstants(t *testing.T) {
if AES128KeySize != 16 {
t.Errorf("AES128KeySize should be 16, got %d", AES128KeySize)
}
if AES192KeySize != 24 {
t.Errorf("AES192KeySize should be 24, got %d", AES192KeySize)
}
if AES256KeySize != 32 {
t.Errorf("AES256KeySize should be 32, got %d", AES256KeySize)
}
if DefaultKeySize != AES256KeySize {
t.Errorf("DefaultKeySize should be AES256KeySize (%d), got %d", AES256KeySize, DefaultKeySize)
}
}
func BenchmarkAES256CBC(b *testing.B) {
key, err := GenerateAES256Key()
if err != nil {
b.Fatalf("Failed to generate key: %v", err)
}
data := make([]byte, 1024) // 1KB of data
rand.Read(data)
b.Run("Encrypt", func(b *testing.B) {
b.ResetTimer()
for i := 0; i < b.N; i++ {
_, err := EncryptAES256CBC(key, data)
if err != nil {
b.Fatal(err)
}
}
})
ciphertext, _ := EncryptAES256CBC(key, data)
b.Run("Decrypt", func(b *testing.B) {
b.ResetTimer()
for i := 0; i < b.N; i++ {
_, err := DecryptAES256CBC(key, ciphertext)
if err != nil {
b.Fatal(err)
}
}
})
}