Avoid over-large allocation in EncryptBytes (#1901)

mindersec · Dec 11, 2023 · c4e382a · c4e382a
1 parent 52b5da8
commit c4e382a
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/internal/crypto/crypto.go b/internal/crypto/crypto.go
@@ -142,6 +142,9 @@ func VerifyCertChain(certIn []byte, roots *x509.CertPool) (bool, error) {
 
 // EncryptBytes encrypts a row of data using AES-CFB.
 func EncryptBytes(key string, data []byte) ([]byte, error) {
+	if len(data) > 1024*1024*32 {
+		return nil, status.Errorf(codes.InvalidArgument, "data is too large (>32MB)")
+	}
 	block, err := aes.NewCipher(deriveKey(key))
 	if err != nil {
 		return nil, status.Errorf(codes.Unknown, "failed to create cipher: %s", err)

diff --git a/internal/crypto/crypto_test.go b/internal/crypto/crypto_test.go
@@ -22,6 +22,8 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 func TestGetCert(t *testing.T) {
@@ -71,6 +73,14 @@ func TestEncryptDecryptBytes(t *testing.T) {
 	assert.Equal(t, "test", string(decrypted))
 }
 
+func TestEncryptTooLarge(t *testing.T) {
+	t.Parallel()
+
+	large := make([]byte, 34000000) // More than 32 MB
+	_, err := EncryptBytes("test", large)
+	assert.ErrorIs(t, err, status.Error(codes.InvalidArgument, "data is too large (>32MB)"))
+}
+
 func TestGenerateNonce(t *testing.T) {
 	t.Parallel()