diff --git a/CMakeLists.txt b/CMakeLists.txt
index 2f65549a85c..5713558d349 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -197,7 +197,12 @@ if(ASAN)
     target_link_options(${PROJECT_NAME} PUBLIC -fsanitize=address)
 endif()
 
-if(TSAN OR ASAN)
+if (UBSAN)
+    target_compile_options(${PROJECT_NAME} PUBLIC -fsanitize=undefined -fno-sanitize-recover=all)
+    target_link_options(${PROJECT_NAME} PUBLIC -fsanitize=undefined -fno-sanitize-recover=all)
+endif()
+
+if(TSAN OR ASAN OR UBSAN)
     # no-omit-frame-pointer and no-optimize-sibling-calls provide better stack traces
     target_compile_options(${PROJECT_NAME} PUBLIC -fno-omit-frame-pointer -fno-optimize-sibling-calls)
 endif()
diff --git a/tests/testlib/s2n_key_schedule_testlib.c b/tests/testlib/s2n_key_schedule_testlib.c
index 4041081bd9e..273f8b8b3cc 100644
--- a/tests/testlib/s2n_key_schedule_testlib.c
+++ b/tests/testlib/s2n_key_schedule_testlib.c
@@ -18,10 +18,7 @@
 S2N_RESULT s2n_connection_set_test_transcript_hash(struct s2n_connection *conn,
         message_type_t message_type, const struct s2n_blob *digest)
 {
-    conn->handshake.handshake_type = conn->handshake.handshake_type & NEGOTIATED;
-    while (s2n_conn_get_current_message_type(conn) != message_type) {
-        conn->handshake.message_number++;
-    }
+    RESULT_GUARD(s2n_connection_set_test_message_type(conn, message_type));
     RESULT_CHECKED_MEMCPY(conn->handshake.hashes->transcript_hash_digest,
             digest->data, digest->size);
     return S2N_RESULT_OK;
@@ -59,3 +56,23 @@ S2N_RESULT s2n_connection_set_test_master_secret(struct s2n_connection *conn,
     conn->secrets.extract_secret_type = S2N_MASTER_SECRET;
     return S2N_RESULT_OK;
 }
+
+/* This function will iterate over all rows and columns of the handshake state 
+ * machine until it finds a valid (handshake_type, handshake_number) such that 
+ * the active message is `expected_message_type`. If callers need to depend on a
+ * specific `message_number` or `handshake_type` this function should not be 
+ * used.
+ */
+S2N_RESULT s2n_connection_set_test_message_type(struct s2n_connection *conn, message_type_t expected_message_type)
+{
+    for (uint32_t handshake = 0; handshake < S2N_HANDSHAKES_COUNT; handshake++) {
+        for (int message = 0; message < S2N_MAX_HANDSHAKE_LENGTH; message++) {
+            conn->handshake.handshake_type = handshake;
+            conn->handshake.message_number = message;
+            if (s2n_conn_get_current_message_type(conn) == expected_message_type) {
+                return S2N_RESULT_OK;
+            }
+        }
+    }
+    RESULT_BAIL(S2N_ERR_HANDSHAKE_UNREACHABLE);
+}
diff --git a/tests/testlib/s2n_testlib.h b/tests/testlib/s2n_testlib.h
index 8df4d499466..113cd2e1b5b 100644
--- a/tests/testlib/s2n_testlib.h
+++ b/tests/testlib/s2n_testlib.h
@@ -76,6 +76,7 @@ int s2n_connection_allow_all_response_extensions(struct s2n_connection *conn);
 int s2n_connection_set_all_protocol_versions(struct s2n_connection *conn, uint8_t version);
 S2N_RESULT s2n_set_all_mutually_supported_groups(struct s2n_connection *conn);
 S2N_RESULT s2n_skip_handshake(struct s2n_connection *conn);
+S2N_RESULT s2n_connection_set_test_message_type(struct s2n_connection *conn, message_type_t expected_message_type);
 
 S2N_RESULT s2n_connection_set_secrets(struct s2n_connection *conn);
 
diff --git a/tests/unit/s2n_tls13_secrets_test.c b/tests/unit/s2n_tls13_secrets_test.c
index 32779e94cb4..427acb7a8f1 100644
--- a/tests/unit/s2n_tls13_secrets_test.c
+++ b/tests/unit/s2n_tls13_secrets_test.c
@@ -410,9 +410,7 @@ int main(int argc, char **argv)
             EXPECT_BYTEARRAY_EQUAL(conn->secrets.version.tls13.server_handshake_secret,
                     empty_secret, sizeof(empty_secret));
 
-            while (s2n_conn_get_current_message_type(conn) != SERVER_HELLO) {
-                conn->handshake.message_number++;
-            }
+            EXPECT_OK(s2n_connection_set_test_message_type(conn, SERVER_HELLO));
             EXPECT_OK(s2n_tls13_secrets_update(conn));
 
             EXPECT_BYTEARRAY_NOT_EQUAL(conn->secrets.version.tls13.client_handshake_secret,
@@ -434,9 +432,7 @@ int main(int argc, char **argv)
             EXPECT_BYTEARRAY_EQUAL(conn->handshake.server_finished,
                     empty_secret, sizeof(empty_secret));
 
-            while (s2n_conn_get_current_message_type(conn) != SERVER_HELLO) {
-                conn->handshake.message_number++;
-            }
+            EXPECT_OK(s2n_connection_set_test_message_type(conn, SERVER_HELLO));
             EXPECT_OK(s2n_tls13_secrets_update(conn));
 
             uint8_t expected_len = 0;
@@ -460,9 +456,7 @@ int main(int argc, char **argv)
             EXPECT_BYTEARRAY_EQUAL(conn->secrets.version.tls13.server_app_secret,
                     empty_secret, sizeof(empty_secret));
 
-            while (s2n_conn_get_current_message_type(conn) != SERVER_FINISHED) {
-                conn->handshake.message_number++;
-            }
+            EXPECT_OK(s2n_connection_set_test_message_type(conn, SERVER_FINISHED));
             EXPECT_OK(s2n_tls13_secrets_update(conn));
 
             EXPECT_BYTEARRAY_NOT_EQUAL(conn->secrets.version.tls13.client_app_secret,
@@ -481,9 +475,7 @@ int main(int argc, char **argv)
             EXPECT_BYTEARRAY_EQUAL(conn->secrets.version.tls13.resumption_master_secret,
                     empty_secret, sizeof(empty_secret));
 
-            while (s2n_conn_get_current_message_type(conn) != CLIENT_FINISHED) {
-                conn->handshake.message_number++;
-            }
+            EXPECT_OK(s2n_connection_set_test_message_type(conn, CLIENT_FINISHED));
             EXPECT_OK(s2n_tls13_secrets_update(conn));
 
             EXPECT_BYTEARRAY_NOT_EQUAL(conn->secrets.version.tls13.resumption_master_secret,
@@ -503,9 +495,7 @@ int main(int argc, char **argv)
             EXPECT_BYTEARRAY_EQUAL(conn->secrets.version.tls13.exporter_master_secret,
                     empty_secret, sizeof(empty_secret));
 
-            while (s2n_conn_get_current_message_type(conn) != SERVER_FINISHED) {
-                conn->handshake.message_number++;
-            }
+            EXPECT_OK(s2n_connection_set_test_message_type(conn, SERVER_FINISHED));
             EXPECT_OK(s2n_tls13_secrets_update(conn));
 
             EXPECT_BYTEARRAY_NOT_EQUAL(conn->secrets.version.tls13.exporter_master_secret,
diff --git a/tls/s2n_handshake.h b/tls/s2n_handshake.h
index 6ff1664cf48..864ba9d7d54 100644
--- a/tls/s2n_handshake.h
+++ b/tls/s2n_handshake.h
@@ -47,6 +47,9 @@
 #define TLS_NPN                       67
 #define TLS_MESSAGE_HASH              254
 
+/* Maximum number of messages in a handshake */
+#define S2N_MAX_HANDSHAKE_LENGTH 32
+
 /* This is the list of message types that we support */
 typedef enum {
     CLIENT_HELLO = 0,
diff --git a/tls/s2n_handshake_io.c b/tls/s2n_handshake_io.c
index 8d8845c8a7d..2b16ac5c3f0 100644
--- a/tls/s2n_handshake_io.c
+++ b/tls/s2n_handshake_io.c
@@ -133,9 +133,6 @@ static const char *message_names[] = {
     MESSAGE_NAME_ENTRY(CLIENT_NPN),
 };
 
-/* Maximum number of messages in a handshake */
-#define S2N_MAX_HANDSHAKE_LENGTH    32
-
 /* We support different ordering of TLS Handshake messages, depending on what is being negotiated. There's also a dummy "INITIAL" handshake
  * that everything starts out as until we know better.
  */