Merge pull request #9328 from valeriosetti/psasim-improvements

psasim: some improvements to the core
diff --git a/tests/psa-client-server/psasim/src/psa_ff_server.c b/tests/psa-client-server/psasim/src/psa_ff_server.c
index 7f97b9b..b106092 100644
--- a/tests/psa-client-server/psasim/src/psa_ff_server.c
+++ b/tests/psa-client-server/psasim/src/psa_ff_server.c
@@ -26,8 +26,6 @@
 #define MAX_CLIENTS 128
 #define MAX_MESSAGES 32
 
-#define SLEEP_US        1
-
 struct connection {
     uint32_t client;
     void *rhandle;
@@ -104,9 +102,6 @@
     uint32_t requested_version;
     ssize_t len;
     int idx;
-#if !defined(PSASIM_USE_USLEEP)
-    const struct timespec ts_delay = { .tv_sec = 0, .tv_nsec = SLEEP_US * 1000 };
-#endif
 
     if (timeout == PSA_POLL) {
         INFO("psa_wait: Called in polling mode");
@@ -261,11 +256,6 @@
             break;
         } else {
             /* There is no 'select' function in SysV to block on multiple queues, so busy-wait :( */
-#if defined(PSASIM_USE_USLEEP)
-            usleep(SLEEP_US);
-#else /* PSASIM_USE_USLEEP */
-            nanosleep(&ts_delay, NULL);
-#endif /* PSASIM_USE_USLEEP */
         }
     } while (timeout == PSA_BLOCK);
 
@@ -474,7 +464,7 @@
 
     while (sofar < num_bytes) {
         size_t sending = (num_bytes - sofar);
-        if (sending >= MAX_FRAGMENT_SIZE) {
+        if (sending > (MAX_FRAGMENT_SIZE - (sizeof(size_t) * 2))) {
             sending = MAX_FRAGMENT_SIZE - (sizeof(size_t) * 2);
         }
 
diff --git a/tests/psa-client-server/psasim/src/psa_sim_crypto_client.c b/tests/psa-client-server/psasim/src/psa_sim_crypto_client.c
index 28dff38..4200f6c 100644
--- a/tests/psa-client-server/psasim/src/psa_sim_crypto_client.c
+++ b/tests/psa-client-server/psasim/src/psa_sim_crypto_client.c
@@ -44,7 +44,7 @@
     invec.base = in_params;
     invec.len = in_params_len;
 
-    size_t max_receive = 8192;
+    size_t max_receive = 24576;
     uint8_t *receive = malloc(max_receive);
     if (receive == NULL) {
         fprintf(stderr, "FAILED to allocate %u bytes\n", (unsigned) max_receive);
@@ -119,6 +119,11 @@
 
 void mbedtls_psa_crypto_free(void)
 {
+    /* Do not try to close a connection that was never started.*/
+    if (handle == -1) {
+        return;
+    }
+
     CLIENT_PRINT("Closing handle");
     psa_close(handle);
     handle = -1;
diff --git a/tests/psa-client-server/psasim/src/psa_sim_crypto_server.c b/tests/psa-client-server/psasim/src/psa_sim_crypto_server.c
index 5259751..cab32c4 100644
--- a/tests/psa-client-server/psasim/src/psa_sim_crypto_server.c
+++ b/tests/psa-client-server/psasim/src/psa_sim_crypto_server.c
@@ -131,7 +131,7 @@
 
     ok = psasim_server_serialise_psa_aead_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 1);
     if (!ok) {
         goto fail;
     }
@@ -375,7 +375,7 @@
 
     ok = psasim_server_serialise_psa_aead_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -619,7 +619,7 @@
 
     ok = psasim_server_serialise_psa_aead_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -736,7 +736,7 @@
 
     ok = psasim_server_serialise_psa_aead_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 1);
     if (!ok) {
         goto fail;
     }
@@ -866,7 +866,7 @@
 
     ok = psasim_server_serialise_psa_aead_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -977,7 +977,7 @@
 
     ok = psasim_server_serialise_psa_aead_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -1062,7 +1062,7 @@
 
     ok = psasim_server_serialise_psa_aead_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -1172,7 +1172,7 @@
 
     ok = psasim_server_serialise_psa_aead_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -1277,7 +1277,7 @@
 
     ok = psasim_server_serialise_psa_aead_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -1387,7 +1387,7 @@
 
     ok = psasim_server_serialise_psa_aead_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -1760,7 +1760,7 @@
 
     ok = psasim_server_serialise_psa_cipher_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 1);
     if (!ok) {
         goto fail;
     }
@@ -1980,7 +1980,7 @@
 
     ok = psasim_server_serialise_psa_cipher_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -2200,7 +2200,7 @@
 
     ok = psasim_server_serialise_psa_cipher_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -2296,7 +2296,7 @@
 
     ok = psasim_server_serialise_psa_cipher_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 1);
     if (!ok) {
         goto fail;
     }
@@ -2410,7 +2410,7 @@
 
     ok = psasim_server_serialise_psa_cipher_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -2513,7 +2513,7 @@
 
     ok = psasim_server_serialise_psa_cipher_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -2623,7 +2623,7 @@
 
     ok = psasim_server_serialise_psa_cipher_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -3436,7 +3436,7 @@
 
     ok = psasim_server_serialise_psa_hash_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 1);
     if (!ok) {
         goto fail;
     }
@@ -3520,7 +3520,7 @@
 
     ok = psasim_server_serialise_psa_hash_operation_t(
         &rpos, &rremain,
-        target_operation);
+        target_operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -3827,7 +3827,7 @@
 
     ok = psasim_server_serialise_psa_hash_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 1);
     if (!ok) {
         goto fail;
     }
@@ -3929,7 +3929,7 @@
 
     ok = psasim_server_serialise_psa_hash_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -4014,7 +4014,7 @@
 
     ok = psasim_server_serialise_psa_hash_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -4103,7 +4103,7 @@
 
     ok = psasim_server_serialise_psa_hash_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 1);
     if (!ok) {
         goto fail;
     }
@@ -4389,7 +4389,7 @@
 
     ok = psasim_server_serialise_psa_key_derivation_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 1);
     if (!ok) {
         goto fail;
     }
@@ -4567,7 +4567,7 @@
 
     ok = psasim_server_serialise_psa_key_derivation_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -4664,7 +4664,7 @@
 
     ok = psasim_server_serialise_psa_key_derivation_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -4757,7 +4757,7 @@
 
     ok = psasim_server_serialise_psa_key_derivation_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -4860,7 +4860,7 @@
 
     ok = psasim_server_serialise_psa_key_derivation_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -4950,7 +4950,7 @@
 
     ok = psasim_server_serialise_psa_key_derivation_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -5055,7 +5055,7 @@
 
     ok = psasim_server_serialise_psa_key_derivation_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -5166,7 +5166,7 @@
 
     ok = psasim_server_serialise_psa_key_derivation_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -5261,7 +5261,7 @@
 
     ok = psasim_server_serialise_psa_key_derivation_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -5345,7 +5345,7 @@
 
     ok = psasim_server_serialise_psa_key_derivation_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -5420,7 +5420,7 @@
 
     ok = psasim_server_serialise_psa_mac_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 1);
     if (!ok) {
         goto fail;
     }
@@ -5643,7 +5643,7 @@
 
     ok = psasim_server_serialise_psa_mac_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 1);
     if (!ok) {
         goto fail;
     }
@@ -5754,7 +5754,7 @@
 
     ok = psasim_server_serialise_psa_mac_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -5839,7 +5839,7 @@
 
     ok = psasim_server_serialise_psa_mac_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -6030,7 +6030,7 @@
 
     ok = psasim_server_serialise_psa_mac_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 1);
     if (!ok) {
         goto fail;
     }
@@ -6127,7 +6127,7 @@
 
     ok = psasim_server_serialise_psa_mac_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -6589,7 +6589,7 @@
 
     ok = psasim_server_serialise_psa_sign_hash_interruptible_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 1);
     if (!ok) {
         goto fail;
     }
@@ -6685,7 +6685,7 @@
 
     ok = psasim_server_serialise_psa_sign_hash_interruptible_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -6873,7 +6873,7 @@
 
     ok = psasim_server_serialise_psa_sign_hash_interruptible_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -7181,7 +7181,7 @@
 
     ok = psasim_server_serialise_psa_verify_hash_interruptible_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 1);
     if (!ok) {
         goto fail;
     }
@@ -7256,7 +7256,7 @@
 
     ok = psasim_server_serialise_psa_verify_hash_interruptible_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
@@ -7436,7 +7436,7 @@
 
     ok = psasim_server_serialise_psa_verify_hash_interruptible_operation_t(
         &rpos, &rremain,
-        operation);
+        operation, 0);
     if (!ok) {
         goto fail;
     }
diff --git a/tests/psa-client-server/psasim/src/psa_sim_generate.pl b/tests/psa-client-server/psasim/src/psa_sim_generate.pl
index ac23807..dd2fe9e 100755
--- a/tests/psa-client-server/psasim/src/psa_sim_generate.pl
+++ b/tests/psa-client-server/psasim/src/psa_sim_generate.pl
@@ -349,7 +349,7 @@
     invec.base = in_params;
     invec.len = in_params_len;
 
-    size_t max_receive = 8192;
+    size_t max_receive = 24576;
     uint8_t *receive = malloc(max_receive);
     if (receive == NULL) {
         fprintf(stderr, "FAILED to allocate %u bytes\n", (unsigned) max_receive);
@@ -424,6 +424,11 @@
 
 void mbedtls_psa_crypto_free(void)
 {
+    /* Do not try to close a connection that was never started.*/
+    if (handle == -1) {
+        return;
+    }
+
     CLIENT_PRINT("Closing handle");
     psa_close(handle);
     handle = -1;
@@ -745,11 +750,19 @@
 
             my $server_specific = ($argtype =~ /^psa_\w+_operation_t/) ? "server_" : "";
 
+            my $completed = ""; # Only needed on server serialise calls
+            if (length($server_specific)) {
+                # On server serialisation, which is only for operation types,
+                # we need to mark the operation as completed (variously called
+                # terminated or inactive in psa/crypto.h) on certain calls.
+                $completed = ($name =~ /_(abort|finish|hash_verify)$/) ? ", 1" : ", 0";
+            }
+
             print $fh <<EOF;
 
     ok = psasim_${server_specific}serialise_${argtype}(
         &rpos, &rremain,
-        $argname);
+        $argname$completed);
     if (!ok) {
         goto fail;
     }
diff --git a/tests/psa-client-server/psasim/src/psa_sim_serialise.c b/tests/psa-client-server/psasim/src/psa_sim_serialise.c
index e655e07..92ecdd2 100644
--- a/tests/psa-client-server/psasim/src/psa_sim_serialise.c
+++ b/tests/psa-client-server/psasim/src/psa_sim_serialise.c
@@ -936,7 +936,8 @@
 
 int psasim_server_serialise_psa_hash_operation_t(uint8_t **pos,
                                                  size_t *remaining,
-                                                 psa_hash_operation_t *operation)
+                                                 psa_hash_operation_t *operation,
+                                                 int completed)
 {
     psasim_operation_t client_operation;
 
@@ -946,6 +947,13 @@
 
     ssize_t slot = operation - hash_operations;
 
+    if (completed) {
+        memset(&hash_operations[slot],
+               0,
+               sizeof(psa_hash_operation_t));
+        hash_operation_handles[slot] = 0;
+    }
+
     client_operation.handle = hash_operation_handles[slot];
 
     memcpy(*pos, &client_operation, sizeof(client_operation));
@@ -1031,7 +1039,8 @@
 
 int psasim_server_serialise_psa_aead_operation_t(uint8_t **pos,
                                                  size_t *remaining,
-                                                 psa_aead_operation_t *operation)
+                                                 psa_aead_operation_t *operation,
+                                                 int completed)
 {
     psasim_operation_t client_operation;
 
@@ -1041,6 +1050,13 @@
 
     ssize_t slot = operation - aead_operations;
 
+    if (completed) {
+        memset(&aead_operations[slot],
+               0,
+               sizeof(psa_aead_operation_t));
+        aead_operation_handles[slot] = 0;
+    }
+
     client_operation.handle = aead_operation_handles[slot];
 
     memcpy(*pos, &client_operation, sizeof(client_operation));
@@ -1162,7 +1178,8 @@
 
 int psasim_server_serialise_psa_mac_operation_t(uint8_t **pos,
                                                 size_t *remaining,
-                                                psa_mac_operation_t *operation)
+                                                psa_mac_operation_t *operation,
+                                                int completed)
 {
     psasim_operation_t client_operation;
 
@@ -1172,6 +1189,13 @@
 
     ssize_t slot = operation - mac_operations;
 
+    if (completed) {
+        memset(&mac_operations[slot],
+               0,
+               sizeof(psa_mac_operation_t));
+        mac_operation_handles[slot] = 0;
+    }
+
     client_operation.handle = mac_operation_handles[slot];
 
     memcpy(*pos, &client_operation, sizeof(client_operation));
@@ -1257,7 +1281,8 @@
 
 int psasim_server_serialise_psa_cipher_operation_t(uint8_t **pos,
                                                    size_t *remaining,
-                                                   psa_cipher_operation_t *operation)
+                                                   psa_cipher_operation_t *operation,
+                                                   int completed)
 {
     psasim_operation_t client_operation;
 
@@ -1267,6 +1292,13 @@
 
     ssize_t slot = operation - cipher_operations;
 
+    if (completed) {
+        memset(&cipher_operations[slot],
+               0,
+               sizeof(psa_cipher_operation_t));
+        cipher_operation_handles[slot] = 0;
+    }
+
     client_operation.handle = cipher_operation_handles[slot];
 
     memcpy(*pos, &client_operation, sizeof(client_operation));
@@ -1352,7 +1384,8 @@
 
 int psasim_server_serialise_psa_key_derivation_operation_t(uint8_t **pos,
                                                            size_t *remaining,
-                                                           psa_key_derivation_operation_t *operation)
+                                                           psa_key_derivation_operation_t *operation,
+                                                           int completed)
 {
     psasim_operation_t client_operation;
 
@@ -1362,6 +1395,13 @@
 
     ssize_t slot = operation - key_derivation_operations;
 
+    if (completed) {
+        memset(&key_derivation_operations[slot],
+               0,
+               sizeof(psa_key_derivation_operation_t));
+        key_derivation_operation_handles[slot] = 0;
+    }
+
     client_operation.handle = key_derivation_operation_handles[slot];
 
     memcpy(*pos, &client_operation, sizeof(client_operation));
@@ -1447,7 +1487,8 @@
 
 int psasim_server_serialise_psa_sign_hash_interruptible_operation_t(uint8_t **pos,
                                                                     size_t *remaining,
-                                                                    psa_sign_hash_interruptible_operation_t *operation)
+                                                                    psa_sign_hash_interruptible_operation_t *operation,
+                                                                    int completed)
 {
     psasim_operation_t client_operation;
 
@@ -1457,6 +1498,13 @@
 
     ssize_t slot = operation - sign_hash_interruptible_operations;
 
+    if (completed) {
+        memset(&sign_hash_interruptible_operations[slot],
+               0,
+               sizeof(psa_sign_hash_interruptible_operation_t));
+        sign_hash_interruptible_operation_handles[slot] = 0;
+    }
+
     client_operation.handle = sign_hash_interruptible_operation_handles[slot];
 
     memcpy(*pos, &client_operation, sizeof(client_operation));
@@ -1542,7 +1590,8 @@
 
 int psasim_server_serialise_psa_verify_hash_interruptible_operation_t(uint8_t **pos,
                                                                       size_t *remaining,
-                                                                      psa_verify_hash_interruptible_operation_t *operation)
+                                                                      psa_verify_hash_interruptible_operation_t *operation,
+                                                                      int completed)
 {
     psasim_operation_t client_operation;
 
@@ -1552,6 +1601,13 @@
 
     ssize_t slot = operation - verify_hash_interruptible_operations;
 
+    if (completed) {
+        memset(&verify_hash_interruptible_operations[slot],
+               0,
+               sizeof(psa_verify_hash_interruptible_operation_t));
+        verify_hash_interruptible_operation_handles[slot] = 0;
+    }
+
     client_operation.handle = verify_hash_interruptible_operation_handles[slot];
 
     memcpy(*pos, &client_operation, sizeof(client_operation));
diff --git a/tests/psa-client-server/psasim/src/psa_sim_serialise.h b/tests/psa-client-server/psasim/src/psa_sim_serialise.h
index f85faad..f60e371 100644
--- a/tests/psa-client-server/psasim/src/psa_sim_serialise.h
+++ b/tests/psa-client-server/psasim/src/psa_sim_serialise.h
@@ -664,12 +664,15 @@
  * \param remaining[in,out]  Pointer to a `size_t` holding number of bytes
  *                           remaining in the buffer.
  * \param value              The value to serialise into the buffer.
+ * \param completed          Non-zero if the operation is now completed (set by
+ *                           finish and abort calls).
  *
  * \return                   \c 1 on success ("okay"), \c 0 on error.
  */
 int psasim_server_serialise_psa_hash_operation_t(uint8_t **pos,
                                                  size_t *remaining,
-                                                 psa_hash_operation_t *value);
+                                                 psa_hash_operation_t *value,
+                                                 int completed);
 
 /** Deserialise a `psa_hash_operation_t` from a buffer on the server side.
  *
@@ -750,12 +753,15 @@
  * \param remaining[in,out]  Pointer to a `size_t` holding number of bytes
  *                           remaining in the buffer.
  * \param value              The value to serialise into the buffer.
+ * \param completed          Non-zero if the operation is now completed (set by
+ *                           finish and abort calls).
  *
  * \return                   \c 1 on success ("okay"), \c 0 on error.
  */
 int psasim_server_serialise_psa_aead_operation_t(uint8_t **pos,
                                                  size_t *remaining,
-                                                 psa_aead_operation_t *value);
+                                                 psa_aead_operation_t *value,
+                                                 int completed);
 
 /** Deserialise a `psa_aead_operation_t` from a buffer on the server side.
  *
@@ -879,12 +885,15 @@
  * \param remaining[in,out]  Pointer to a `size_t` holding number of bytes
  *                           remaining in the buffer.
  * \param value              The value to serialise into the buffer.
+ * \param completed          Non-zero if the operation is now completed (set by
+ *                           finish and abort calls).
  *
  * \return                   \c 1 on success ("okay"), \c 0 on error.
  */
 int psasim_server_serialise_psa_mac_operation_t(uint8_t **pos,
                                                 size_t *remaining,
-                                                psa_mac_operation_t *value);
+                                                psa_mac_operation_t *value,
+                                                int completed);
 
 /** Deserialise a `psa_mac_operation_t` from a buffer on the server side.
  *
@@ -965,12 +974,15 @@
  * \param remaining[in,out]  Pointer to a `size_t` holding number of bytes
  *                           remaining in the buffer.
  * \param value              The value to serialise into the buffer.
+ * \param completed          Non-zero if the operation is now completed (set by
+ *                           finish and abort calls).
  *
  * \return                   \c 1 on success ("okay"), \c 0 on error.
  */
 int psasim_server_serialise_psa_cipher_operation_t(uint8_t **pos,
                                                    size_t *remaining,
-                                                   psa_cipher_operation_t *value);
+                                                   psa_cipher_operation_t *value,
+                                                   int completed);
 
 /** Deserialise a `psa_cipher_operation_t` from a buffer on the server side.
  *
@@ -1051,12 +1063,15 @@
  * \param remaining[in,out]  Pointer to a `size_t` holding number of bytes
  *                           remaining in the buffer.
  * \param value              The value to serialise into the buffer.
+ * \param completed          Non-zero if the operation is now completed (set by
+ *                           finish and abort calls).
  *
  * \return                   \c 1 on success ("okay"), \c 0 on error.
  */
 int psasim_server_serialise_psa_key_derivation_operation_t(uint8_t **pos,
                                                            size_t *remaining,
-                                                           psa_key_derivation_operation_t *value);
+                                                           psa_key_derivation_operation_t *value,
+                                                           int completed);
 
 /** Deserialise a `psa_key_derivation_operation_t` from a buffer on the server side.
  *
@@ -1137,12 +1152,15 @@
  * \param remaining[in,out]  Pointer to a `size_t` holding number of bytes
  *                           remaining in the buffer.
  * \param value              The value to serialise into the buffer.
+ * \param completed          Non-zero if the operation is now completed (set by
+ *                           finish and abort calls).
  *
  * \return                   \c 1 on success ("okay"), \c 0 on error.
  */
 int psasim_server_serialise_psa_sign_hash_interruptible_operation_t(uint8_t **pos,
                                                                     size_t *remaining,
-                                                                    psa_sign_hash_interruptible_operation_t *value);
+                                                                    psa_sign_hash_interruptible_operation_t *value,
+                                                                    int completed);
 
 /** Deserialise a `psa_sign_hash_interruptible_operation_t` from a buffer on the server side.
  *
@@ -1223,12 +1241,15 @@
  * \param remaining[in,out]  Pointer to a `size_t` holding number of bytes
  *                           remaining in the buffer.
  * \param value              The value to serialise into the buffer.
+ * \param completed          Non-zero if the operation is now completed (set by
+ *                           finish and abort calls).
  *
  * \return                   \c 1 on success ("okay"), \c 0 on error.
  */
 int psasim_server_serialise_psa_verify_hash_interruptible_operation_t(uint8_t **pos,
                                                                       size_t *remaining,
-                                                                      psa_verify_hash_interruptible_operation_t *value);
+                                                                      psa_verify_hash_interruptible_operation_t *value,
+                                                                      int completed);
 
 /** Deserialise a `psa_verify_hash_interruptible_operation_t` from a buffer on the server side.
  *
diff --git a/tests/psa-client-server/psasim/src/psa_sim_serialise.pl b/tests/psa-client-server/psasim/src/psa_sim_serialise.pl
index 81808ca..75e6cd0 100755
--- a/tests/psa-client-server/psasim/src/psa_sim_serialise.pl
+++ b/tests/psa-client-server/psasim/src/psa_sim_serialise.pl
@@ -158,11 +158,15 @@
     my $type_d = $type;
     $type_d =~ s/ /_/g;
 
+    if (length($server) && $type !~ /^psa_(\w+)_operation_t$/) {
+        die("$0: declare_server_serialise: $type: not supported\n");
+    }
+
     my $server_side = (length($server)) ? " on the server side" : "";
 
     my $ptr = (length($server)) ? "*" : "";
 
-    return align_declaration(<<EOF);
+    my $code = <<EOF;
 
 /** Serialise $an `$type` into a buffer${server_side}.
  *
@@ -171,13 +175,29 @@
  * \\param remaining[in,out]  Pointer to a `size_t` holding number of bytes
  *                           remaining in the buffer.
  * \\param value              The value to serialise into the buffer.
+EOF
+
+    $code .= <<EOF if length($server);
+ * \\param completed          Non-zero if the operation is now completed (set by
+ *                           finish and abort calls).
+EOF
+
+    my $value_sep = (length($server)) ? "," : ");";
+
+    $code .= <<EOF;
  *
  * \\return                   \\c 1 on success ("okay"), \\c 0 on error.
  */
 int psasim_${server}serialise_$type_d(uint8_t **pos,
                              size_t *remaining,
-                             $type ${ptr}value);
+                             $type ${ptr}value$value_sep
 EOF
+
+    $code .= <<EOF if length($server);
+                             int completed);
+EOF
+
+    return align_declaration($code);
 }
 
 sub declare_deserialise
@@ -543,7 +563,8 @@
 
 int psasim_server_serialise_$type_d(uint8_t **pos,
                              size_t *remaining,
-                             $type *operation)
+                             $type *operation,
+                             int completed)
 {
     psasim_operation_t client_operation;
 
@@ -553,6 +574,13 @@
 
     ssize_t slot = operation - ${t}_operations;
 
+    if (completed) {
+        memset(&${t}_operations[slot],
+               0,
+               sizeof($type_d));
+        ${t}_operation_handles[slot] = 0;
+    }
+
     client_operation.handle = ${t}_operation_handles[slot];
 
     memcpy(*pos, &client_operation, sizeof(client_operation));
@@ -619,7 +647,7 @@
     if ($type =~ /^psa_(\w+)_operation_t$/) {
         $t = $1;
     } else {
-        die("$0: define_server_serialise: $type: not supported\n");
+        die("$0: define_server_deserialise: $type: not supported\n");
     }
 
     my $type_d = $type;
@@ -1123,8 +1151,8 @@
 EOF
 }
 
-# Horrible way to align first, second and third lines of function signature to
-# appease uncrustify (these are the 2nd-4th lines of code, indices 1, 2 and 3)
+# Horrible way to align first few lines of function signature to appease
+# uncrustify (these are usually the 2nd-4th lines of code, indices 1, 2 and 3)
 #
 sub align_signature
 {
@@ -1132,13 +1160,17 @@
 
     my @code = split(/\n/, $code);
 
+    my $i = 1;
     # Find where the ( is
-    my $idx = index($code[1], "(");
+    my $idx = index($code[$i], "(");
     die("can't find (") if $idx < 0;
 
     my $indent = " " x ($idx + 1);
-    $code[2] =~ s/^\s+/$indent/;
-    $code[3] =~ s/^\s+/$indent/;
+
+    do {
+        # Indent each line up until the one that ends with )
+        $code[++$i] =~ s/^\s+/$indent/;
+    } while $code[$i] !~ /\)$/;
 
     return join("\n", @code) . "\n";
 }
@@ -1163,8 +1195,10 @@
     die("can't find (") if $idx < 0;
 
     my $indent = " " x ($idx + 1);
-    $code[$i + 1] =~ s/^\s+/$indent/;
-    $code[$i + 2] =~ s/^\s+/$indent/;
+    do {
+        # Indent each line up until the one with the ; on it
+        $code[++$i] =~ s/^\s+/$indent/;
+    } while ($code[$i] !~ /;/);
 
     return join("\n", @code) . "\n";
 }