From 2445fe844a7758ae509d2a5b7bf275ae22f8b45f Mon Sep 17 00:00:00 2001
From: JacobBarthelmeh <jacob@wolfssl.com>
Date: Fri, 31 May 2024 16:45:50 -0600
Subject: [PATCH] rework get max fragment length

---
 src/internal.c     |  3 ---
 src/ssl_sess.c     |  2 +-
 src/tls.c          |  2 ++
 tests/api.c        | 34 +++++++++++++++++++++++++++-------
 wolfssl/internal.h |  3 ++-
 5 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/src/internal.c b/src/internal.c
index 8a2a584b2..56e72376f 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -7605,9 +7605,6 @@ int InitSSL(WOLFSSL* ssl, WOLFSSL_CTX* ctx, int writeDup)
                        "err = %d", ret);
         return MEMORY_E;
     }
-#ifdef HAVE_MAX_FRAGMENT
-    ssl->session->maxFragmentSz = ssl->max_fragment;
-#endif /* HAVE_MAX_FRAGMENT */
 #ifdef HAVE_SESSION_TICKET
     ssl->options.noTicketTls12 = ctx->noTicketTls12;
 #endif
diff --git a/src/ssl_sess.c b/src/ssl_sess.c
index 695eb0677..f2f4baaf1 100644
--- a/src/ssl_sess.c
+++ b/src/ssl_sess.c
@@ -754,7 +754,7 @@ uint8_t wolfSSL_SESSION_get_max_fragment_length(WOLFSSL_SESSION* session)
         return 0;
     }
 
-    return session->maxFragmentSz;
+    return session->mfl;
 }
 
 
diff --git a/src/tls.c b/src/tls.c
index 1a347a757..03df15674 100644
--- a/src/tls.c
+++ b/src/tls.c
@@ -2988,6 +2988,8 @@ static int TLSX_MFL_Parse(WOLFSSL* ssl, const byte* input, word16 length,
             WOLFSSL_ERROR_VERBOSE(UNKNOWN_MAX_FRAG_LEN_E);
             return UNKNOWN_MAX_FRAG_LEN_E;
     }
+    if (ssl->session != NULL)
+        ssl->session->mfl = *input;
 
 #ifndef NO_WOLFSSL_SERVER
     if (isRequest) {
diff --git a/tests/api.c b/tests/api.c
index 1f9fb52cb..4ecc5fa6f 100644
--- a/tests/api.c
+++ b/tests/api.c
@@ -11307,6 +11307,31 @@ static int test_wolfSSL_UseMaxFragment(void)
 
     wolfSSL_free(ssl);
     wolfSSL_CTX_free(ctx);
+
+#if defined(OPENSSL_EXTRA) && defined(HAVE_MANUAL_MEMIO_TESTS_DEPENDENCIES)
+    /* check negotiated max fragment size */
+    {
+        WOLFSSL *ssl_c = NULL;
+        WOLFSSL *ssl_s = NULL;
+        struct test_memio_ctx test_ctx;
+        WOLFSSL_CTX *ctx_c = NULL;
+        WOLFSSL_CTX *ctx_s = NULL;
+
+        XMEMSET(&test_ctx, 0, sizeof(test_ctx));
+        ExpectIntEQ(test_memio_setup(&test_ctx, &ctx_c, &ctx_s, &ssl_c, &ssl_s,
+            wolfTLSv1_2_client_method, wolfTLSv1_2_server_method), 0);
+        ExpectIntEQ(wolfSSL_UseMaxFragment(ssl_c, WOLFSSL_MFL_2_8),
+            WOLFSSL_SUCCESS);
+        ExpectIntEQ(test_memio_do_handshake(ssl_c, ssl_s, 10, NULL), 0);
+        ExpectIntEQ(SSL_SESSION_get_max_fragment_length(
+            wolfSSL_get_session(ssl_c)), WOLFSSL_MFL_2_8);
+
+        wolfSSL_free(ssl_c);
+        wolfSSL_free(ssl_s);
+        wolfSSL_CTX_free(ctx_c);
+        wolfSSL_CTX_free(ctx_s);
+    }
+#endif
 #endif /* !NO_WOLFSSL_CLIENT || !NO_WOLFSSL_SERVER */
 #endif
     return EXPECT_RESULT();
@@ -47946,6 +47971,7 @@ static int test_wolfSSL_CTX_sess_set_remove_cb(void)
     /* Both should have been allocated */
     ExpectIntEQ(clientSessRemCountMalloc, 1);
     ExpectIntEQ(serverSessRemCountMalloc, 1);
+
     /* This should not be called yet. Session wasn't evicted from cache yet. */
     ExpectIntEQ(clientSessRemCountFree, 0);
 #if (defined(WOLFSSL_TLS13) && defined(HAVE_SESSION_TICKET)) || \
@@ -47972,13 +47998,6 @@ static int test_wolfSSL_CTX_sess_set_remove_cb(void)
     ExpectIntEQ(SSL_CTX_remove_session(serverSessCtx, serverSess), 0);
     ExpectNull(SSL_SESSION_get_ex_data(serverSess, serverSessRemIdx));
     ExpectIntEQ(serverSessRemCountFree, 1);
-
-    /* check on the max fragment size */
-#ifdef HAVE_MAX_FRAGMENT
-    ExpectIntEQ(SSL_SESSION_get_max_fragment_length(serverSess),
-        MAX_RECORD_SIZE);
-#endif
-
     /* Need to free the references that we kept */
     SSL_CTX_free(serverSessCtx);
     SSL_SESSION_free(serverSess);
@@ -67062,6 +67081,7 @@ static int test_wolfSSL_dtls_stateless_maxfrag(void)
     /* CH without cookie shouldn't change state */
     ExpectIntEQ(ssl_s->max_fragment, max_fragment);
     ExpectIntNE(test_ctx.c_len, 0);
+
     /* consume HRR from buffer */
     test_ctx.c_len = 0;
     ExpectIntEQ(test_memio_do_handshake(ssl_c, ssl_s, 10, NULL), 0);
diff --git a/wolfssl/internal.h b/wolfssl/internal.h
index f9160b693..7213016ea 100644
--- a/wolfssl/internal.h
+++ b/wolfssl/internal.h
@@ -4434,7 +4434,8 @@ struct WOLFSSL_SESSION {
     WOLFSSL_CRYPTO_EX_DATA ex_data;
 #endif
 #ifdef OPENSSL_EXTRA
-    word32 maxFragmentSz;
+    byte               mfl; /* max fragment length negotiated i.e.
+                             * WOLFSSL_MFL_2_8  (6) */
 #endif
     byte               isSetup:1;
 };