From 89ea321d96e8727b00bbf6ea632a40586b88a502 Mon Sep 17 00:00:00 2001
From: Jerry Yu <jerry.h.yu@arm.com>
Date: Thu, 9 Sep 2021 14:31:24 +0800
Subject: [PATCH] tls13: add key_schedule_stage_early_data

Signed-off-by: Jerry Yu <jerry.h.yu@arm.com>
---
 library/ssl_misc.h       |  7 +++++++
 library/ssl_tls13_keys.c | 27 +++++++++++++++++++++++++++
 library/ssl_tls13_keys.h | 25 +++++++++++++++++++++++++
 3 files changed, 59 insertions(+)

diff --git a/library/ssl_misc.h b/library/ssl_misc.h
index 604976f5f..5afdc4c5f 100644
--- a/library/ssl_misc.h
+++ b/library/ssl_misc.h
@@ -675,6 +675,13 @@ struct mbedtls_ssl_handshake_params
     int extensions_present;             /*!< extension presence; Each bitfield
                                              represents an extension and defined
                                              as \c MBEDTLS_SSL_EXT_XXX */
+
+    union
+    {
+        unsigned char early    [MBEDTLS_MD_MAX_SIZE];
+        unsigned char handshake[MBEDTLS_MD_MAX_SIZE];
+        unsigned char app      [MBEDTLS_MD_MAX_SIZE];
+    } tls13_master_secrets;
 #endif /* MBEDTLS_SSL_PROTO_TLS1_3_EXPERIMENTAL */
 
 #if defined(MBEDTLS_SSL_SESSION_TICKETS)
diff --git a/library/ssl_tls13_keys.c b/library/ssl_tls13_keys.c
index 7aec21dff..b7beb125c 100644
--- a/library/ssl_tls13_keys.c
+++ b/library/ssl_tls13_keys.c
@@ -820,4 +820,31 @@ int mbedtls_ssl_tls13_populate_transform( mbedtls_ssl_transform *transform,
     return( 0 );
 }
 
+int mbedtls_ssl_tls13_key_schedule_stage_early_data( mbedtls_ssl_context *ssl )
+{
+    int ret = 0;
+
+    if( ssl->handshake->ciphersuite_info == NULL )
+    {
+        MBEDTLS_SSL_DEBUG_MSG( 1, ( "cipher suite info not found" ) );
+        return( MBEDTLS_ERR_SSL_INTERNAL_ERROR );
+    }
+    mbedtls_md_type_t const md_type = ssl->handshake->ciphersuite_info->mac;
+    const unsigned char *input = NULL;
+    size_t input_len = 0;
+#if defined(MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED)
+    input = ssl->handshake->psk;
+    input_len = ssl->handshake->psk_len;
+#endif
+    ret = mbedtls_ssl_tls1_3_evolve_secret( md_type, NULL, input, input_len,
+                                            ssl->handshake->tls13_master_secrets.early );
+    if( ret != 0 )
+    {
+        MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_tls1_3_evolve_secret", ret );
+        return( ret );
+    }
+
+    return( 0 );
+}
+
 #endif /* MBEDTLS_SSL_PROTO_TLS1_3_EXPERIMENTAL */
diff --git a/library/ssl_tls13_keys.h b/library/ssl_tls13_keys.h
index ca892b166..592ba12a0 100644
--- a/library/ssl_tls13_keys.h
+++ b/library/ssl_tls13_keys.h
@@ -531,4 +531,29 @@ int mbedtls_ssl_tls13_populate_transform( mbedtls_ssl_transform *transform,
                                           mbedtls_ssl_key_set const *traffic_keys,
                                           mbedtls_ssl_context *ssl );
 
+/*
+ * TLS 1.3 key schedule evolutions
+ *
+ *   Early Data -> Handshake -> Application
+ *
+ * Small wrappers around mbedtls_ssl_tls1_3_evolve_secret().
+ */
+
+/**
+ * \brief Begin TLS 1.3 key schedule by calculating early secret
+ *        from chosen PSK.
+ *
+ *        The TLS 1.3 key schedule can be viewed as a simple state machine
+ *        with states Initial -> Early -> Handshake -> Application, and
+ *        this function represents the Initial -> Early transition.
+ *
+ *        In the early stage, mbedtls_ssl_tls1_3_generate_early_data_keys()
+ *        can be used to derive the 0-RTT traffic keys.
+ *
+ * \param ssl  The SSL context to operate on.
+ *
+ * \returns    \c 0 on success.
+ * \returns    A negative error code on failure.
+ */
+int mbedtls_ssl_tls13_key_schedule_stage_early_data( mbedtls_ssl_context *ssl );
 #endif /* MBEDTLS_SSL_TLS1_3_KEYS_H */