diff --git a/library/psa_crypto_pake.c b/library/psa_crypto_pake.c
index b8a08a1..7f147b7 100644
--- a/library/psa_crypto_pake.c
+++ b/library/psa_crypto_pake.c
@@ -328,6 +328,9 @@
         step != PSA_PAKE_STEP_ZK_PROOF )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
+    if( output == NULL || output_size == 0 || output_length == NULL )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_ECJPAKE)
     if( operation->alg == PSA_ALG_JPAKE )
     {
@@ -508,6 +511,9 @@
         step != PSA_PAKE_STEP_ZK_PROOF )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
+    if( input == NULL || input_length == 0 )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_ECJPAKE)
     if( operation->alg == PSA_ALG_JPAKE )
     {