From 21df158961b1d51f1adbb3bcf81c5d57532e9cbc Mon Sep 17 00:00:00 2001
From: Gilles Peskine <Gilles.Peskine@arm.com>
Date: Wed, 3 Jun 2026 15:43:18 +0200
Subject: [PATCH] Add some hooks in mbedtls_test_ssl_perform_handshake()

Let test code customize the kitchen sink function
`mbedtls_test_ssl_perform_handshake()` at a few points.

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
---
 tests/include/test/ssl_helpers.h     | 32 ++++++++++++++++++++++++++++
 tests/src/test_helpers/ssl_helpers.c | 17 +++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/tests/include/test/ssl_helpers.h b/tests/include/test/ssl_helpers.h
index a653c9643e..467a27f342 100644
--- a/tests/include/test/ssl_helpers.h
+++ b/tests/include/test/ssl_helpers.h
@@ -117,6 +117,9 @@ enum {
 #define MBEDTLS_TEST_MAX_ALPN_LIST_SIZE 10
 #endif
 
+/* Forward declaration. Defined below. */
+struct mbedtls_test_ssl_endpoint;
+
 typedef struct mbedtls_test_ssl_log_pattern {
     const char *pattern;
     size_t counter;
@@ -147,6 +150,35 @@ typedef struct mbedtls_test_handshake_test_options {
     int expected_srv_fragments;
     int renegotiate;
     int legacy_renegotiation;
+    /** Hook that mbedtls_test_ssl_perform_handshake() runs just before
+     * the initial handshake. */
+    void (*pre_handshake_fun)(struct mbedtls_test_ssl_endpoint *client,
+                              struct mbedtls_test_ssl_endpoint *server,
+                              void *param);
+    /** Value passed to ::pre_handshake_fun. */
+    void *pre_handshake_param;
+    /** Hook that mbedtls_test_ssl_perform_handshake() runs after
+     * the initial handshake succeeds. */
+    void (*post_handshake_fun)(struct mbedtls_test_ssl_endpoint *client,
+                               struct mbedtls_test_ssl_endpoint *server,
+                               void *param);
+    /** Value passed to ::post_handshake_fun. */
+    void *post_handshake_param;
+    /** Hook that mbedtls_test_ssl_perform_handshake() runs after
+     * exchanging some data, before testing additional features such as
+     * serialization and renegotiation. */
+    void (*post_data_fun)(struct mbedtls_test_ssl_endpoint *client,
+                          struct mbedtls_test_ssl_endpoint *server,
+                          void *param);
+    /** Value passed to ::post_data_fun. */
+    void *post_data_param;
+    /** Hook that mbedtls_test_ssl_perform_handshake() runs after a successful
+     * connection, just before closing down. */
+    void (*pre_shutdown_fun)(struct mbedtls_test_ssl_endpoint *client,
+                             struct mbedtls_test_ssl_endpoint *server,
+                             void *param);
+    /** Value passed to ::pre_shutdown_fun. */
+    void *pre_shutdown_param;
     void *srv_log_obj;
     void *cli_log_obj;
     void (*srv_log_fun)(void *, int, const char *, int, const char *);
diff --git a/tests/src/test_helpers/ssl_helpers.c b/tests/src/test_helpers/ssl_helpers.c
index db29413de6..08cfe252af 100644
--- a/tests/src/test_helpers/ssl_helpers.c
+++ b/tests/src/test_helpers/ssl_helpers.c
@@ -2282,6 +2282,10 @@ int mbedtls_test_ssl_perform_connection(
         expected_handshake_result = MBEDTLS_ERR_SSL_BAD_PROTOCOL_VERSION;
     }
 
+    if (options->pre_handshake_fun != NULL) {
+        options->pre_handshake_fun(client, server, options->pre_handshake_param);
+    }
+
     TEST_EQUAL(mbedtls_test_move_handshake_to_state(&(client->ssl),
                                                     &(server->ssl),
                                                     MBEDTLS_SSL_HANDSHAKE_OVER),
@@ -2313,6 +2317,10 @@ int mbedtls_test_ssl_perform_connection(
                    options->expected_ciphersuite);
     }
 
+    if (options->post_handshake_fun != NULL) {
+        options->post_handshake_fun(client, server, options->post_handshake_param);
+    }
+
 #if defined(MBEDTLS_SSL_VARIABLE_BUFFER_LENGTH)
     if (options->resize_buffers != 0) {
         /* A server, when using DTLS, might delay a buffer resize to happen
@@ -2339,6 +2347,11 @@ int mbedtls_test_ssl_perform_connection(
                        options->expected_srv_fragments),
                    0);
     }
+
+    if (options->post_data_fun != NULL) {
+        options->post_data_fun(client, server, options->post_data_param);
+    }
+
 #if defined(MBEDTLS_SSL_CONTEXT_SERIALIZATION)
     if (options->serialize == 1) {
         TEST_ASSERT(test_serialization(options, client, server));
@@ -2396,6 +2409,10 @@ void mbedtls_test_ssl_perform_handshake(
     TEST_ASSERT(mbedtls_ssl_conf_get_user_data_p(&server->conf) == server);
     TEST_ASSERT(mbedtls_ssl_get_user_data_p(&server->ssl) == server);
 
+    if (options->pre_shutdown_fun != NULL) {
+        options->pre_shutdown_fun(client, server, options->pre_shutdown_param);
+    }
+
 exit:
     mbedtls_test_ssl_endpoint_free(client);
     mbedtls_test_ssl_endpoint_free(server);