From 7d8c3fe12c9db11e31aa8ee305f454ae018402d1 Mon Sep 17 00:00:00 2001
From: Jerry Yu <jerry.h.yu@arm.com>
Date: Mon, 12 Dec 2022 12:59:44 +0800
Subject: [PATCH] Add wait flight2 state.

The state is come from RFC8446 section A.2

Signed-off-by: Jerry Yu <jerry.h.yu@arm.com>
---
 include/mbedtls/ssl.h      |  2 +-
 library/ssl_tls13_server.c | 45 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/include/mbedtls/ssl.h b/include/mbedtls/ssl.h
index 3c2696fe4..2bca21a2f 100644
--- a/include/mbedtls/ssl.h
+++ b/include/mbedtls/ssl.h
@@ -687,12 +687,12 @@ typedef enum {
     MBEDTLS_SSL_SERVER_FINISHED,
     MBEDTLS_SSL_FLUSH_BUFFERS,
     MBEDTLS_SSL_HANDSHAKE_WRAPUP,
-
     MBEDTLS_SSL_NEW_SESSION_TICKET,
     MBEDTLS_SSL_SERVER_HELLO_VERIFY_REQUEST_SENT,
     MBEDTLS_SSL_HELLO_RETRY_REQUEST,
     MBEDTLS_SSL_ENCRYPTED_EXTENSIONS,
     MBEDTLS_SSL_END_OF_EARLY_DATA,
+    MBEDTLS_SSL_WAIT_FLIGHT2,
     MBEDTLS_SSL_CLIENT_CERTIFICATE_VERIFY,
     MBEDTLS_SSL_CLIENT_CCS_AFTER_SERVER_FINISHED,
     MBEDTLS_SSL_CLIENT_CCS_BEFORE_2ND_CLIENT_HELLO,
diff --git a/library/ssl_tls13_server.c b/library/ssl_tls13_server.c
index e8967e6c2..40d51d806 100644
--- a/library/ssl_tls13_server.c
+++ b/library/ssl_tls13_server.c
@@ -2782,6 +2782,30 @@ static int ssl_tls13_write_server_finished(mbedtls_ssl_context *ssl)
     MBEDTLS_SSL_DEBUG_MSG(1, ("Switch to handshake keys for inbound traffic"));
     mbedtls_ssl_set_inbound_transform(ssl, ssl->handshake->transform_handshake);
 
+    mbedtls_ssl_handshake_set_state(ssl, MBEDTLS_SSL_WAIT_FLIGHT2);
+
+    return 0;
+}
+
+/*
+ * Handler for MBEDTLS_SSL_WAIT_FLIGHT2
+ *
+ * RFC 8446 section A.2
+ *
+ *                      WAIT_FLIGHT2
+ *                            |
+ *                   +--------+--------+
+ *           No auth |                 | Client auth
+ *                   |                 |
+ *                   |                 v
+ *                   |             WAIT_CERT
+ *                   |        Recv |       | Recv Certificate
+ */
+MBEDTLS_CHECK_RETURN_CRITICAL
+static int ssl_tls13_process_wait_flight2(mbedtls_ssl_context *ssl)
+{
+    MBEDTLS_SSL_DEBUG_MSG(2, ("=> ssl_tls13_process_wait_flight2"));
+
     if (ssl->handshake->certificate_request_sent) {
         mbedtls_ssl_handshake_set_state(ssl, MBEDTLS_SSL_CLIENT_CERTIFICATE);
     } else {
@@ -2790,6 +2814,7 @@ static int ssl_tls13_write_server_finished(mbedtls_ssl_context *ssl)
         mbedtls_ssl_handshake_set_state(ssl, MBEDTLS_SSL_CLIENT_FINISHED);
     }
 
+    MBEDTLS_SSL_DEBUG_MSG(2, ("<= ssl_tls13_process_wait_flight2"));
     return 0;
 }
 
@@ -3213,10 +3238,30 @@ int mbedtls_ssl_tls13_handshake_server_step(mbedtls_ssl_context *ssl)
             break;
 #endif /* MBEDTLS_SSL_TLS1_3_COMPATIBILITY_MODE */
 
+        /* RFC 8446 section A.2
+         *
+         *                                 | Send Finished ( SERVER_FINISHED )
+         *                                 | K_send = application
+         *                        +--------+--------+
+         *               No 0-RTT |                 | 0-RTT
+         *                        |                 |
+         *    K_recv = handshake  |                 | K_recv = early data
+         *  [Skip decrypt errors] |    +------> WAIT_EOED -+
+         *                        |    |       Recv |      | Recv EndOfEarlyData
+         *                        |    | early data |      | K_recv = handshake
+         *                        |    +------------+      |
+         *                        |                        |
+         *                        +> WAIT_FLIGHT2 <--------+
+         *                                 |
+         */
         case MBEDTLS_SSL_SERVER_FINISHED:
             ret = ssl_tls13_write_server_finished(ssl);
             break;
 
+        case MBEDTLS_SSL_WAIT_FLIGHT2:
+            ret = ssl_tls13_process_wait_flight2(ssl);
+            break;
+
         case MBEDTLS_SSL_CLIENT_FINISHED:
             ret = ssl_tls13_process_client_finished(ssl);
             break;