fix key length check

rudderlabs · Sep 6, 2024 · 455b1ad · 455b1ad
1 parent 6408ce2
commit 455b1ad
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 0 deletions.
diff --git a/encrypt/aes_cfb.go b/encrypt/aes_cfb.go
@@ -13,6 +13,10 @@ type encryptionAESCFB struct {
 }
 
 func (e *encryptionAESCFB) Encrypt(src []byte, key string) ([]byte, error) {
+	if len(key) != e.level/8 {
+		return nil, fmt.Errorf("key length must be %d bytes", e.level/8)
+	}
+
 	block, err := aes.NewCipher([]byte(key))
 	if err != nil {
 		return nil, err
@@ -31,6 +35,10 @@ func (e *encryptionAESCFB) Encrypt(src []byte, key string) ([]byte, error) {
 }
 
 func (e *encryptionAESCFB) Decrypt(src []byte, key string) ([]byte, error) {
+	if len(key) != e.level/8 {
+		return nil, fmt.Errorf("key length must be %d bytes", e.level/8)
+	}
+
 	block, err := aes.NewCipher([]byte(key))
 	if err != nil {
 		return nil, err

diff --git a/encrypt/aes_gcm.go b/encrypt/aes_gcm.go
@@ -13,6 +13,10 @@ type encryptionAESGCM struct {
 }
 
 func (e *encryptionAESGCM) Encrypt(src []byte, key string) ([]byte, error) {
+	if len(key) != e.level/8 {
+		return nil, fmt.Errorf("key length must be %d bytes", e.level/8)
+	}
+
 	block, err := aes.NewCipher([]byte(key))
 	if err != nil {
 		return nil, err
@@ -33,6 +37,10 @@ func (e *encryptionAESGCM) Encrypt(src []byte, key string) ([]byte, error) {
 }
 
 func (e *encryptionAESGCM) Decrypt(src []byte, key string) ([]byte, error) {
+	if len(key) != e.level/8 {
+		return nil, fmt.Errorf("key length must be %d bytes", e.level/8)
+	}
+
 	block, err := aes.NewCipher([]byte(key))
 	if err != nil {
 		return nil, err

diff --git a/encrypt/encrypt_test.go b/encrypt/encrypt_test.go
@@ -44,6 +44,9 @@ func Test_EncryptDecrypt(t *testing.T) {
 			require.NoError(t, err)
 
 			plaintext := loremIpsumDolor
+			_, err = encrypter.Encrypt(plaintext, key[:len(key)-1])
+			require.Error(t, err)
+
 			ciphertext, err := encrypter.Encrypt(plaintext, key)
 			require.NoError(t, err)