From 0bd675340f42e0b4783836c30d8a4ea1027931cf Mon Sep 17 00:00:00 2001
From: dudaodong <lanliddd.2007@163.com>
Date: Thu, 16 Feb 2023 11:45:08 +0800
Subject: [PATCH] fix: fix AesEcbEncrypt failed with key lenght is 24 or 32

---
 cryptor/aes.go             | 20 +++++++++++++++++---
 cryptor/crypto_internal.go | 14 +++++++-------
 2 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/cryptor/aes.go b/cryptor/aes.go
index 501098f..2b9879d 100644
--- a/cryptor/aes.go
+++ b/cryptor/aes.go
@@ -17,14 +17,23 @@ import (
 // AesEcbEncrypt encrypt data with key use AES ECB algorithm
 // len(key) should be 16, 24 or 32
 func AesEcbEncrypt(data, key []byte) []byte {
-	cipher, _ := aes.NewCipher(generateAesKey(key))
+	size := len(key)
+	if size != 16 && size != 24 && size != 32 {
+		panic("key length shoud be 16 or 24 or 32")
+	}
+
+	cipher, _ := aes.NewCipher(generateAesKey(key, size))
 	length := (len(data) + aes.BlockSize) / aes.BlockSize
+
 	plain := make([]byte, length*aes.BlockSize)
+
 	copy(plain, data)
+
 	pad := byte(len(plain) - len(data))
 	for i := len(data); i < len(plain); i++ {
 		plain[i] = pad
 	}
+
 	encrypted := make([]byte, len(plain))
 	for bs, be := 0, cipher.BlockSize(); bs <= len(data); bs, be = bs+cipher.BlockSize(), be+cipher.BlockSize() {
 		cipher.Encrypt(encrypted[bs:be], plain[bs:be])
@@ -36,9 +45,14 @@ func AesEcbEncrypt(data, key []byte) []byte {
 // AesEcbDecrypt decrypt data with key use AES ECB algorithm
 // len(key) should be 16, 24 or 32
 func AesEcbDecrypt(encrypted, key []byte) []byte {
-	cipher, _ := aes.NewCipher(generateAesKey(key))
+	size := len(key)
+	if size != 16 && size != 24 && size != 32 {
+		panic("key length shoud be 16 or 24 or 32")
+	}
+
+	cipher, _ := aes.NewCipher(generateAesKey(key, size))
 	decrypted := make([]byte, len(encrypted))
-	//
+
 	for bs, be := 0, cipher.BlockSize(); bs < len(encrypted); bs, be = bs+cipher.BlockSize(), be+cipher.BlockSize() {
 		cipher.Decrypt(decrypted[bs:be], encrypted[bs:be])
 	}
diff --git a/cryptor/crypto_internal.go b/cryptor/crypto_internal.go
index 6005346..0729542 100644
--- a/cryptor/crypto_internal.go
+++ b/cryptor/crypto_internal.go
@@ -2,15 +2,15 @@ package cryptor
 
 import "bytes"
 
-func generateAesKey(key []byte) []byte {
-	genKey := make([]byte, 16)
-	copy(genKey, key)
-	for i := 16; i < len(key); {
-		for j := 0; j < 16 && i < len(key); j, i = j+1, i+1 {
-			genKey[j] ^= key[i]
+func generateAesKey(key []byte, size int) []byte {
+	aesKey := make([]byte, size)
+	copy(aesKey, key)
+	for i := size; i < len(key); {
+		for j := 0; j < size && i < len(key); j, i = j+1, i+1 {
+			aesKey[j] ^= key[i]
 		}
 	}
-	return genKey
+	return aesKey
 }
 
 func generateDesKey(key []byte) []byte {