pgcryptocipherccl: don't panic when passed invalid padding size

pkg/ccl/pgcryptoccl/pgcryptocipherccl/cipher.go

@@ -68,13 +68,17 @@ func Decrypt(data []byte, key []byte, iv []byte, cipherType string) ([]byte, err
 func newCipher(method cipherMethod, key []byte) (cipher.Block, error) {
 	switch a := method.algorithm; a {
 	case aesCipher:
+		var err error
 		switch l := len(key); {
 		case l >= 32:
-			key = zeroPadOrTruncate(key, 32)
+			key, err = zeroPadOrTruncate(key, 32)
 		case l >= 24:
-			key = zeroPadOrTruncate(key, 24)
+			key, err = zeroPadOrTruncate(key, 24)
 		default:
-			key = zeroPadOrTruncate(key, 16)
+			key, err = zeroPadOrTruncate(key, 16)
+		}
+		if err != nil {
+			return nil, err
 		}
 		return aes.NewCipher(key)
 	default:
@@ -120,8 +124,12 @@ func validateDataLength(data []byte, blockSize int) error {
 func encrypt(method cipherMethod, block cipher.Block, iv []byte, data []byte) ([]byte, error) {
 	switch m := method.mode; m {
 	case cbcMode:
+		var err error
 		ret := make([]byte, len(data))
-		iv = zeroPadOrTruncate(iv, block.BlockSize())
+		iv, err = zeroPadOrTruncate(iv, block.BlockSize())
+		if err != nil {
+			return nil, err
+		}
 		mode := cipher.NewCBCEncrypter(block, iv)
 		mode.CryptBlocks(ret, data)
 		return ret, nil
@@ -133,21 +141,16 @@ func encrypt(method cipherMethod, block cipher.Block, iv []byte, data []byte) ([
 func decrypt(method cipherMethod, block cipher.Block, iv []byte, data []byte) ([]byte, error) {
 	switch m := method.mode; m {
 	case cbcMode:
+		var err error
 		ret := make([]byte, len(data))
-		iv = zeroPadOrTruncate(iv, block.BlockSize())
+		iv, err = zeroPadOrTruncate(iv, block.BlockSize())
+		if err != nil {
+			return nil, err
+		}
 		mode := cipher.NewCBCDecrypter(block, iv)
 		mode.CryptBlocks(ret, data)
 		return ret, nil
 	default:
 		return nil, errors.Newf("cannot encrypt for unknown mode: %d", m)
 	}
 }
-
-func zeroPadOrTruncate(data []byte, size int) []byte {
-	if len(data) >= size {
-		return data[:size]
-	}
-	paddedData := make([]byte, size)
-	copy(paddedData, data)
-	return paddedData
-}

pkg/ccl/pgcryptoccl/pgcryptocipherccl/padding.go

@@ -51,3 +51,17 @@ func pkcsUnpad(data []byte) ([]byte, error) {
 
 	return data[:len(data)-int(paddingLen)], nil
 }
+
+// zeroPadOrTruncate pads a slice of bytes with zeroes if its length is smaller
+// than size and truncates the slice to length size otherwise.
+func zeroPadOrTruncate(data []byte, size int) ([]byte, error) {
+	if size < 0 {
+		return nil, errors.AssertionFailedf("cannot zero pad or truncate to negative size")
+	}
+	if len(data) >= size {
+		return data[:size], nil
+	}
+	paddedData := make([]byte, size)
+	copy(paddedData, data)
+	return paddedData, nil
+}

pkg/ccl/pgcryptoccl/pgcryptocipherccl/padding_test.go

@@ -110,3 +110,49 @@ func TestPKCSUnpad(t *testing.T) {
 		})
 	}
 }
+
+func TestZeroPadOrTruncate(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	for name, tc := range map[string]struct {
+		data        []byte
+		size        int
+		expected    []byte
+		expectedErr string
+	}{
+		"data length less than size": {
+			data:     []byte{1, 2},
+			size:     3,
+			expected: []byte{1, 2, 0},
+		},
+		"data length equal to size": {
+			data:     []byte{1, 2, 3},
+			size:     3,
+			expected: []byte{1, 2, 3},
+		},
+		"data length greater than size": {
+			data:     []byte{1, 2, 3, 4},
+			size:     3,
+			expected: []byte{1, 2, 3},
+		},
+		"empty data": {
+			data:     nil,
+			size:     3,
+			expected: []byte{0, 0, 0},
+		},
+		"negative size": {
+			data:        []byte{1, 2, 3},
+			size:        -1,
+			expectedErr: "cannot zero pad or truncate to negative size",
+		},
+	} {
+		t.Run(name, func(t *testing.T) {
+			actual, err := zeroPadOrTruncate(tc.data, tc.size)
+			if tc.expectedErr != "" {
+				require.EqualError(t, err, tc.expectedErr)
+				return
+			}
+			require.Equal(t, tc.expected, actual)
+		})
+	}
+}