From 6a971fd61a3fc0d52f4c31379dc3a2e0651b60f9 Mon Sep 17 00:00:00 2001
From: Waleed Elmelegy <waleed.elmelegy@arm.com>
Date: Thu, 28 Dec 2023 17:48:16 +0000
Subject: [PATCH] Refactor and improve Record size limit handling

Signed-off-by: Waleed Elmelegy <waleed.elmelegy@arm.com>
---
 library/ssl_tls.c          | 81 +++++++++++++++++++-------------------
 library/ssl_tls13_client.c |  8 ++--
 2 files changed, 44 insertions(+), 45 deletions(-)

diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 452970ebe8..250adfb28a 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -3387,6 +3387,31 @@ const char *mbedtls_ssl_get_version(const mbedtls_ssl_context *ssl)
     }
 }
 
+#if defined(MBEDTLS_SSL_RECORD_SIZE_LIMIT)
+
+size_t mbedtls_ssl_get_output_record_size_limit(const mbedtls_ssl_context *ssl)
+{
+    const size_t max_len = MBEDTLS_SSL_OUT_CONTENT_LEN;
+    size_t record_size_limit = max_len;
+
+    if (ssl->session != NULL &&
+        ssl->session->record_size_limit >= MBEDTLS_SSL_RECORD_SIZE_LIMIT_MIN &&
+        ssl->session->record_size_limit < max_len) {
+        record_size_limit = ssl->session->record_size_limit;
+    }
+
+    // TODO: this is currently untested
+    /* During a handshake, use the value being negotiated */
+    if (ssl->session_negotiate != NULL &&
+        ssl->session_negotiate->record_size_limit >= MBEDTLS_SSL_RECORD_SIZE_LIMIT_MIN &&
+        ssl->session_negotiate->record_size_limit < max_len) {
+        record_size_limit = ssl->session_negotiate->record_size_limit;
+    }
+
+    return record_size_limit;
+}
+#endif /* MBEDTLS_SSL_RECORD_SIZE_LIMIT */
+
 #if defined(MBEDTLS_SSL_MAX_FRAGMENT_LENGTH)
 size_t mbedtls_ssl_get_input_max_frag_len(const mbedtls_ssl_context *ssl)
 {
@@ -3420,31 +3445,6 @@ size_t mbedtls_ssl_get_input_max_frag_len(const mbedtls_ssl_context *ssl)
     return max_len;
 }
 
-#if defined(MBEDTLS_SSL_RECORD_SIZE_LIMIT)
-
-size_t mbedtls_ssl_get_output_record_size_limit(const mbedtls_ssl_context *ssl)
-{
-    const size_t max_len = MBEDTLS_SSL_OUT_CONTENT_LEN;
-    size_t record_size_limit = max_len;
-
-    if (ssl->session != NULL &&
-        ssl->session->record_size_limit >= MBEDTLS_SSL_RECORD_SIZE_LIMIT_MIN &&
-        ssl->session->record_size_limit < max_len) {
-        record_size_limit = ssl->session->record_size_limit;
-    }
-
-    // TODO: this is currently untested
-    /* During a handshake, use the value being negotiated */
-    if (ssl->session_negotiate != NULL &&
-        ssl->session_negotiate->record_size_limit >= MBEDTLS_SSL_RECORD_SIZE_LIMIT_MIN &&
-        ssl->session_negotiate->record_size_limit < max_len) {
-        record_size_limit = ssl->session_negotiate->record_size_limit;
-    }
-
-    return record_size_limit;
-}
-#endif /* MBEDTLS_SSL_RECORD_SIZE_LIMIT */
-
 size_t mbedtls_ssl_get_output_max_frag_len(const mbedtls_ssl_context *ssl)
 {
     size_t max_len;
@@ -3516,24 +3516,25 @@ int mbedtls_ssl_get_max_out_record_payload(const mbedtls_ssl_context *ssl)
 
     if (max_len > record_size_limit) {
         max_len = record_size_limit;
-        if (ssl->transform_out != NULL &&
-            ssl->transform_out->tls_version == MBEDTLS_SSL_VERSION_TLS1_3) {
-            /* RFC 8449, section 4:
-             *
-             * This value [record_size_limit] is the length of the plaintext
-             * of a protected record.
-             * The value includes the content type and padding added in TLS 1.3
-             * (that is, the complete length of TLSInnerPlaintext).
-             *
-             * Thus, round down to a multiple of MBEDTLS_SSL_CID_TLS1_3_PADDING_GRANULARITY
-             * and subtract 1 (for the content type that will be added later)
-             */
-            max_len = ((max_len / MBEDTLS_SSL_CID_TLS1_3_PADDING_GRANULARITY) *
-                       MBEDTLS_SSL_CID_TLS1_3_PADDING_GRANULARITY) - 1;
-        }
     }
 #endif
 
+    if (ssl->transform_out != NULL &&
+        ssl->transform_out->tls_version == MBEDTLS_SSL_VERSION_TLS1_3) {
+        /* RFC 8449, section 4:
+         *
+         * This value [record_size_limit] is the length of the plaintext
+         * of a protected record.
+         * The value includes the content type and padding added in TLS 1.3
+         * (that is, the complete length of TLSInnerPlaintext).
+         *
+         * Thus, round down to a multiple of MBEDTLS_SSL_CID_TLS1_3_PADDING_GRANULARITY
+         * and subtract 1 (for the content type that will be added later)
+         */
+        max_len = ((max_len / MBEDTLS_SSL_CID_TLS1_3_PADDING_GRANULARITY) *
+                   MBEDTLS_SSL_CID_TLS1_3_PADDING_GRANULARITY) - 1;
+    }
+
 #if defined(MBEDTLS_SSL_PROTO_DTLS)
     if (mbedtls_ssl_get_current_mtu(ssl) != 0) {
         const size_t mtu = mbedtls_ssl_get_current_mtu(ssl);
diff --git a/library/ssl_tls13_client.c b/library/ssl_tls13_client.c
index 503db5862a..5775a3ea5f 100644
--- a/library/ssl_tls13_client.c
+++ b/library/ssl_tls13_client.c
@@ -2133,11 +2133,9 @@ static int ssl_tls13_parse_encrypted_extensions(mbedtls_ssl_context *ssl,
 
     if ((handshake->received_extensions & MBEDTLS_SSL_EXT_MASK(RECORD_SIZE_LIMIT)) &&
         (handshake->received_extensions & MBEDTLS_SSL_EXT_MASK(MAX_FRAGMENT_LENGTH))) {
-        mbedtls_debug_print_msg(ssl,
-                                3,
-                                __FILE__,
-                                __LINE__,
-                                "Record size limit extension cannot be used with max fragment length extension");
+        MBEDTLS_SSL_DEBUG_MSG(3,
+                              (
+                                  "Record size limit extension cannot be used with max fragment length extension"));
         MBEDTLS_SSL_PEND_FATAL_ALERT(
             MBEDTLS_SSL_ALERT_MSG_ILLEGAL_PARAMETER,
             MBEDTLS_ERR_SSL_ILLEGAL_PARAMETER);