Add new mbedtls_pkcs12_pbe_ext function to replace old function

Add new mbedtls_pkcs12_pbe_ext function to replace
old mbedtls_pkcs12_pbe function that have security
issues.

Signed-off-by: Waleed Elmelegy <waleed.elmelegy@arm.com>
diff --git a/library/pkcs12.c b/library/pkcs12.c
index 1f45f45..4670afc 100644
--- a/library/pkcs12.c
+++ b/library/pkcs12.c
@@ -49,7 +49,7 @@
                                    mbedtls_asn1_buf *salt, int *iterations)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    unsigned char **p = &params->p;
+    unsigned char *p = params->p;
     const unsigned char *end = params->p + params->len;
 
     /*
@@ -64,18 +64,18 @@
                                  MBEDTLS_ERR_ASN1_UNEXPECTED_TAG);
     }
 
-    if ((ret = mbedtls_asn1_get_tag(p, end, &salt->len, MBEDTLS_ASN1_OCTET_STRING)) != 0) {
+    if ((ret = mbedtls_asn1_get_tag(&p, end, &salt->len, MBEDTLS_ASN1_OCTET_STRING)) != 0) {
         return MBEDTLS_ERROR_ADD(MBEDTLS_ERR_PKCS12_PBE_INVALID_FORMAT, ret);
     }
 
-    salt->p = *p;
-    *p += salt->len;
+    salt->p = p;
+    p += salt->len;
 
-    if ((ret = mbedtls_asn1_get_int(p, end, iterations)) != 0) {
+    if ((ret = mbedtls_asn1_get_int(&p, end, iterations)) != 0) {
         return MBEDTLS_ERROR_ADD(MBEDTLS_ERR_PKCS12_PBE_INVALID_FORMAT, ret);
     }
 
-    if (*p != end) {
+    if (p != end) {
         return MBEDTLS_ERROR_ADD(MBEDTLS_ERR_PKCS12_PBE_INVALID_FORMAT,
                                  MBEDTLS_ERR_ASN1_LENGTH_MISMATCH);
     }
@@ -172,18 +172,46 @@
 #endif /* MBEDTLS_ARC4_C */
 }
 
+#if !defined(MBEDTLS_CIPHER_PADDING_PKCS7)
+int mbedtls_pkcs12_pbe_ext(mbedtls_asn1_buf *pbe_params, int mode,
+                           mbedtls_cipher_type_t cipher_type, mbedtls_md_type_t md_type,
+                           const unsigned char *pwd,  size_t pwdlen,
+                           const unsigned char *data, size_t len,
+                           unsigned char *output, size_t output_size,
+                           size_t *output_len);
+#endif
+
 int mbedtls_pkcs12_pbe(mbedtls_asn1_buf *pbe_params, int mode,
                        mbedtls_cipher_type_t cipher_type, mbedtls_md_type_t md_type,
                        const unsigned char *pwd,  size_t pwdlen,
                        const unsigned char *data, size_t len,
                        unsigned char *output)
 {
+    size_t output_len = 0;
+
+    /* We assume caller of the function is providing a big enough output buffer
+     * so we pass output_size as SIZE_MAX to pass checks, However, no guarantees
+     * for the output size actually being correct.
+     */
+    return mbedtls_pkcs12_pbe_ext(pbe_params, mode, cipher_type, md_type,
+                                  pwd, pwdlen, data, len, output, SIZE_MAX,
+                                  &output_len);
+}
+
+int mbedtls_pkcs12_pbe_ext(mbedtls_asn1_buf *pbe_params, int mode,
+                           mbedtls_cipher_type_t cipher_type, mbedtls_md_type_t md_type,
+                           const unsigned char *pwd,  size_t pwdlen,
+                           const unsigned char *data, size_t len,
+                           unsigned char *output, size_t output_size,
+                           size_t *output_len)
+{
     int ret, keylen = 0;
     unsigned char key[32];
     unsigned char iv[16];
     const mbedtls_cipher_info_t *cipher_info;
     mbedtls_cipher_context_t cipher_ctx;
-    size_t olen = 0;
+    size_t finish_olen = 0;
+    unsigned int padlen = 0;
 
     if (pwd == NULL && pwdlen != 0) {
         return MBEDTLS_ERR_PKCS12_BAD_INPUT_DATA;
@@ -196,6 +224,19 @@
 
     keylen = cipher_info->key_bitlen / 8;
 
+    if (mode == MBEDTLS_PKCS12_PBE_DECRYPT) {
+        if (output_size < len) {
+            return MBEDTLS_ERR_ASN1_BUF_TOO_SMALL;
+        }
+    }
+
+    if (mode == MBEDTLS_PKCS12_PBE_ENCRYPT) {
+        padlen = cipher_info->block_size - (len % cipher_info->block_size);
+        if (output_size < (len + padlen)) {
+            return MBEDTLS_ERR_ASN1_BUF_TOO_SMALL;
+        }
+    }
+
     if ((ret = pkcs12_pbe_derive_key_iv(pbe_params, md_type, pwd, pwdlen,
                                         key, keylen,
                                         iv, cipher_info->iv_size)) != 0) {
@@ -242,14 +283,16 @@
     }
 
     if ((ret = mbedtls_cipher_update(&cipher_ctx, data, len,
-                                     output, &olen)) != 0) {
+                                     output, output_len)) != 0) {
         goto exit;
     }
 
-    if ((ret = mbedtls_cipher_finish(&cipher_ctx, output + olen, &olen)) != 0) {
+    if ((ret = mbedtls_cipher_finish(&cipher_ctx, output + (*output_len), &finish_olen)) != 0) {
         ret = MBEDTLS_ERR_PKCS12_PASSWORD_MISMATCH;
     }
 
+    *output_len += finish_olen;
+
 exit:
     mbedtls_platform_zeroize(key, sizeof(key));
     mbedtls_platform_zeroize(iv,  sizeof(iv));