Better fix for empty password / salt

Signed-off-by: Paul Elliott <paul.elliott@arm.com>
diff --git a/library/pkcs12.c b/library/pkcs12.c
index 15be702..f44ac82 100644
--- a/library/pkcs12.c
+++ b/library/pkcs12.c
@@ -199,8 +199,6 @@
             data_len -= use_len;
         }
     }
-    else
-        memset( data, 0, data_len );
 }
 
 int mbedtls_pkcs12_derivation( unsigned char *data, size_t datalen,
@@ -213,9 +211,12 @@
 
     unsigned char diversifier[128];
     unsigned char salt_block[128], pwd_block[128], hash_block[128];
+    unsigned char empty_string[2] = { 0, 0 };
     unsigned char hash_output[MBEDTLS_MD_MAX_SIZE];
     unsigned char *p;
     unsigned char c;
+    int           use_password = 0;
+    int           use_salt = 0;
 
     size_t hlen, use_len, v, i;
 
@@ -229,6 +230,12 @@
     if( pwd == NULL && pwdlen != 0 )
         return( MBEDTLS_ERR_PKCS12_BAD_INPUT_DATA );
 
+    if( salt == NULL && saltlen != 0 )
+        return( MBEDTLS_ERR_PKCS12_BAD_INPUT_DATA );
+
+    use_password = ( pwd && pwdlen != 0 );
+    use_salt = ( salt && saltlen != 0 );
+
     md_info = mbedtls_md_info_from_type( md_type );
     if( md_info == NULL )
         return( MBEDTLS_ERR_PKCS12_FEATURE_UNAVAILABLE );
@@ -246,8 +253,15 @@
 
     memset( diversifier, (unsigned char) id, v );
 
-    pkcs12_fill_buffer( salt_block, v, salt, saltlen );
-    pkcs12_fill_buffer( pwd_block,  v, pwd,  pwdlen  );
+    if( use_salt != 0 )
+    {
+        pkcs12_fill_buffer( salt_block, v, salt, saltlen );
+    }
+
+    if( use_password != 0 )
+    {
+        pkcs12_fill_buffer( pwd_block,  v, pwd,  pwdlen  );
+    }
 
     p = data;
     while( datalen > 0 )
@@ -259,11 +273,29 @@
         if( ( ret = mbedtls_md_update( &md_ctx, diversifier, v ) ) != 0 )
             goto exit;
 
-        if( ( ret = mbedtls_md_update( &md_ctx, salt_block, v ) ) != 0 )
-            goto exit;
+        if( use_salt != 0 )
+        {
+            if( ( ret = mbedtls_md_update( &md_ctx, salt_block, v )) != 0 )
+                goto exit;
+        }
+        else
+        {
+            if( ( ret = mbedtls_md_update( &md_ctx, empty_string,
+                                           sizeof( empty_string ) )) != 0 )
+                goto exit;
+        }
 
-        if( ( ret = mbedtls_md_update( &md_ctx, pwd_block, v ) ) != 0 )
-            goto exit;
+        if( use_password != 0)
+        {
+            if( ( ret = mbedtls_md_update( &md_ctx, pwd_block, v )) != 0 )
+                goto exit;
+        }
+        else
+        {
+            if( ( ret = mbedtls_md_update( &md_ctx, empty_string,
+                                           sizeof( empty_string ) )) != 0 )
+                goto exit;
+        }
 
         if( ( ret = mbedtls_md_finish( &md_ctx, hash_output ) ) != 0 )
             goto exit;
@@ -291,22 +323,28 @@
             if( ++hash_block[i - 1] != 0 )
                 break;
 
-        // salt_block += B
-        c = 0;
-        for( i = v; i > 0; i-- )
+        if( use_salt != 0 )
         {
-            j = salt_block[i - 1] + hash_block[i - 1] + c;
-            c = MBEDTLS_BYTE_1( j );
-            salt_block[i - 1] = MBEDTLS_BYTE_0( j );
+            // salt_block += B
+            c = 0;
+            for( i = v; i > 0; i-- )
+            {
+                j = salt_block[i - 1] + hash_block[i - 1] + c;
+                c = MBEDTLS_BYTE_1( j );
+                salt_block[i - 1] = MBEDTLS_BYTE_0( j );
+            }
         }
 
-        // pwd_block  += B
-        c = 0;
-        for( i = v; i > 0; i-- )
+        if( use_password != 0 )
         {
-            j = pwd_block[i - 1] + hash_block[i - 1] + c;
-            c = MBEDTLS_BYTE_1( j );
-            pwd_block[i - 1] = MBEDTLS_BYTE_0( j );
+            // pwd_block  += B
+            c = 0;
+            for( i = v; i > 0; i-- )
+            {
+                j = pwd_block[i - 1] + hash_block[i - 1] + c;
+                c = MBEDTLS_BYTE_1( j );
+                pwd_block[i - 1] = MBEDTLS_BYTE_0( j );
+            }
         }
     }