Add testing for extra calls during a round

Test that extra calls to psa_pake_input() and psa_pake_output() during a
round return the correct error.

Signed-off-by: David Horstmann <david.horstmann@arm.com>
diff --git a/tests/suites/test_suite_psa_crypto_pake.data b/tests/suites/test_suite_psa_crypto_pake.data
index 9e1cc63..89f1562 100644
--- a/tests/suites/test_suite_psa_crypto_pake.data
+++ b/tests/suites/test_suite_psa_crypto_pake.data
@@ -210,6 +210,14 @@
 depends_on:MBEDTLS_PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_LEGACY:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
 ecjpake_rounds_inject:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:0:"abcdef":ERR_INJECT_ROUND2_SERVER_ZK_PROOF:PSA_ERROR_DATA_INVALID
 
+PSA PAKE: inject ERR_INJECT_EXTRA_OUTPUT
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_rounds_inject:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:0:"abcdef":ERR_INJECT_EXTRA_OUTPUT:PSA_ERROR_BAD_STATE
+
+PSA PAKE: inject ERR_INJECT_EXTRA_INPUT
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
+ecjpake_rounds_inject:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:1:"abcdef":ERR_INJECT_EXTRA_INPUT:PSA_ERROR_BAD_STATE
+
 PSA PAKE: ecjpake size macros
 depends_on:MBEDTLS_PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_LEGACY:PSA_WANT_ECC_SECP_R1_256
 ecjpake_size_macros:
diff --git a/tests/suites/test_suite_psa_crypto_pake.function b/tests/suites/test_suite_psa_crypto_pake.function
index 52380de..87c40f5 100644
--- a/tests/suites/test_suite_psa_crypto_pake.function
+++ b/tests/suites/test_suite_psa_crypto_pake.function
@@ -2,6 +2,7 @@
 #include <stdint.h>
 
 #include "psa/crypto.h"
+#include "psa/crypto_extra.h"
 
 typedef enum {
     ERR_NONE = 0,
@@ -39,6 +40,8 @@
     ERR_INJECT_ROUND2_SERVER_KEY_SHARE,
     ERR_INJECT_ROUND2_SERVER_ZK_PUBLIC,
     ERR_INJECT_ROUND2_SERVER_ZK_PROOF,
+    ERR_INJECT_EXTRA_OUTPUT,
+    ERR_INJECT_EXTRA_INPUT,
     /* erros issued from the .data file */
     ERR_IN_SETUP,
     ERR_IN_SET_USER,
@@ -69,6 +72,13 @@
         *(buf + 7) ^= 1;                           \
     }
 
+#define DO_ROUND_CONDITIONAL_CHECK_FAILURE(this_stage, function) \
+    if (this_stage == err_stage)                                 \
+    {                                                            \
+        TEST_EQUAL(function, expected_error_arg);                \
+        break;                                                   \
+    }
+
 #define DO_ROUND_UPDATE_OFFSETS(main_buf_offset, step_offset, step_size) \
     {                                       \
         step_offset = main_buf_offset;      \
@@ -185,6 +195,12 @@
                 buffer0 + buffer0_off);
             DO_ROUND_UPDATE_OFFSETS(buffer0_off, s_x2_pr_off, s_x2_pr_len);
 
+            size_t extra_output_len;
+            DO_ROUND_CONDITIONAL_CHECK_FAILURE(
+                ERR_INJECT_EXTRA_OUTPUT,
+                psa_pake_output(server, PSA_PAKE_STEP_KEY_SHARE,
+                                buffer0 + s_g2_off, 512 - s_g2_off, &extra_output_len));
+            (void) extra_output_len;
             /*
              * When injecting errors in inputs, the implementation is
              * free to detect it right away of with a delay.
@@ -223,6 +239,12 @@
                                         s_x2_pr_len);
                 DO_ROUND_CHECK_FAILURE();
 
+                /* Note: Must have client_input_first == 1 to inject extra input */
+                DO_ROUND_CONDITIONAL_CHECK_FAILURE(
+                    ERR_INJECT_EXTRA_INPUT,
+                    psa_pake_input(client, PSA_PAKE_STEP_KEY_SHARE,
+                                   buffer0 + s_g2_off, s_g2_len));
+
                 /* Error didn't trigger, make test fail */
                 if ((err_stage >= ERR_INJECT_ROUND1_SERVER_KEY_SHARE_PART1) &&
                     (err_stage <= ERR_INJECT_ROUND1_SERVER_ZK_PROOF_PART2)) {