diff --git a/src/crypto/PSASpake2p.cpp b/src/crypto/PSASpake2p.cpp
index 281a9266e0..55a0b1d906 100644
--- a/src/crypto/PSASpake2p.cpp
+++ b/src/crypto/PSASpake2p.cpp
@@ -30,7 +30,7 @@ CHIP_ERROR PSASpake2p_P256_SHA256_HKDF_HMAC::Init(const uint8_t * context, size_
 {
     Clear();
 
-    VerifyOrReturnError(context_len <= sizeof(mContext), CHIP_ERROR_INVALID_ARGUMENT);
+    VerifyOrReturnError(context_len <= sizeof(mContext), CHIP_ERROR_BUFFER_TOO_SMALL);
 
     psa_pake_cipher_suite_t cs = PSA_PAKE_CIPHER_SUITE_INIT;
     psa_pake_cs_set_algorithm(&cs, PSA_ALG_SPAKE2P);
@@ -137,6 +137,8 @@ CHIP_ERROR PSASpake2p_P256_SHA256_HKDF_HMAC::BeginProver(const uint8_t * my_iden
 
 CHIP_ERROR PSASpake2p_P256_SHA256_HKDF_HMAC::ComputeRoundOne(const uint8_t * pab, size_t pab_len, uint8_t * out, size_t * out_len)
 {
+    VerifyOrReturnError(out_len != nullptr, CHIP_ERROR_INVALID_ARGUMENT);
+
     psa_status_t status;
 
     if (mRole == PSA_PAKE_ROLE_SERVER)
@@ -153,6 +155,8 @@ CHIP_ERROR PSASpake2p_P256_SHA256_HKDF_HMAC::ComputeRoundOne(const uint8_t * pab
 
 CHIP_ERROR PSASpake2p_P256_SHA256_HKDF_HMAC::ComputeRoundTwo(const uint8_t * in, size_t in_len, uint8_t * out, size_t * out_len)
 {
+    VerifyOrReturnError(out_len != nullptr, CHIP_ERROR_INVALID_ARGUMENT);
+
     psa_status_t status;
 
     if (mRole == PSA_PAKE_ROLE_CLIENT)
@@ -177,6 +181,9 @@ CHIP_ERROR PSASpake2p_P256_SHA256_HKDF_HMAC::KeyConfirm(const uint8_t * in, size
 
 CHIP_ERROR PSASpake2p_P256_SHA256_HKDF_HMAC::GetKeys(uint8_t * out, size_t * out_len)
 {
+    VerifyOrReturnError(out != nullptr, CHIP_ERROR_INVALID_ARGUMENT);
+    VerifyOrReturnError(out_len != nullptr, CHIP_ERROR_INVALID_ARGUMENT);
+
     /*
      * TODO: either:
      * - use psa_pake_shared_secret() proposed in https://github.com/ARM-software/psa-api/issues/86
@@ -184,7 +191,6 @@ CHIP_ERROR PSASpake2p_P256_SHA256_HKDF_HMAC::GetKeys(uint8_t * out, size_t * out
      */
     oberon_spake2p_operation_t & oberonCtx = mOperation.MBEDTLS_PRIVATE(ctx).oberon_spake2p_ctx;
 
-    VerifyOrReturnError(out_len != nullptr, CHIP_ERROR_INVALID_ARGUMENT);
     VerifyOrReturnError((oberonCtx.hash_len / 2) <= *out_len, CHIP_ERROR_BUFFER_TOO_SMALL);
 
     memcpy(out, oberonCtx.shared, oberonCtx.hash_len / 2);
diff --git a/src/crypto/PSASpake2p.h b/src/crypto/PSASpake2p.h
index e407b7e3d9..b416fa002b 100644
--- a/src/crypto/PSASpake2p.h
+++ b/src/crypto/PSASpake2p.h
@@ -107,6 +107,11 @@ class PSASpake2p_P256_SHA256_HKDF_HMAC
      * @param out     The output first round Spake2+ contribution.
      * @param out_len The output first round Spake2+ contribution length.
      *
+     * The out_len parameter is expected to point to an integer that holds
+     * the size of the buffer to put the first round Spake2+ contribution.
+     * After successful execution of this method, the variable is set to the
+     * actual size of the generated output.
+     *
      * @return Returns a CHIP_ERROR on error, CHIP_NO_ERROR otherwise
      **/
     CHIP_ERROR ComputeRoundOne(const uint8_t * pab, size_t pab_len, uint8_t * out, size_t * out_len);
@@ -119,6 +124,11 @@ class PSASpake2p_P256_SHA256_HKDF_HMAC
      * @param out     The output second round Spake2+ contribution.
      * @param out_len The output second round Spake2+ contribution length.
      *
+     * The out_len parameter is expected to point to an integer that holds
+     * the size of the buffer to put the second round Spake2+ contribution.
+     * After successful execution of this method, the variable is set to the
+     * actual size of the generated output.
+     *
      * @return Returns a CHIP_ERROR on error, CHIP_NO_ERROR otherwise
      **/
     CHIP_ERROR ComputeRoundTwo(const uint8_t * in, size_t in_len, uint8_t * out, size_t * out_len);