From b2f59f733a0f03ce9359e16eb8a47cb7120449b6 Mon Sep 17 00:00:00 2001
From: Juliusz Sosinowicz <juliusz@wolfssl.com>
Date: Wed, 28 Aug 2024 12:58:50 +0200
Subject: [PATCH] DTLS 1.3: check size including headers

---
 .github/workflows/os-check.yml |  2 ++
 src/dtls.c                     |  7 ++++++-
 src/tls13.c                    | 23 +++++++++++++----------
 3 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/.github/workflows/os-check.yml b/.github/workflows/os-check.yml
index 68557a31ab..fcad812100 100644
--- a/.github/workflows/os-check.yml
+++ b/.github/workflows/os-check.yml
@@ -33,6 +33,8 @@ jobs:
           '--enable-dtls --enable-dtls13 --enable-earlydata 
              --enable-session-ticket --enable-psk 
              CPPFLAGS=''-DWOLFSSL_DTLS13_NO_HRR_ON_RESUME'' ',
+          '--enable-experimental --enable-kyber --enable-dtls --enable-dtls13
+             --enable-dtls-frag-ch',
         ]
     name: make check
     runs-on: ${{ matrix.os }}
diff --git a/src/dtls.c b/src/dtls.c
index 9961ac21a4..1bdb7ce464 100644
--- a/src/dtls.c
+++ b/src/dtls.c
@@ -953,8 +953,13 @@ int DoClientHelloStateless(WOLFSSL* ssl, const byte* input, word32 helloSz,
             int tlsxFound;
             ret = FindExtByType(&ch.cookieExt, TLSX_COOKIE, ch.extension,
                                  &tlsxFound);
-            if (ret != 0)
+            if (ret != 0) {
+                if (isFirstCHFrag) {
+                    WOLFSSL_MSG("\t\tCookie probably missing from first "
+                                "fragment. Dropping.");
+                }
                 return ret;
+            }
         }
     }
 #endif
diff --git a/src/tls13.c b/src/tls13.c
index 04e6cc5585..a27f34fdf2 100644
--- a/src/tls13.c
+++ b/src/tls13.c
@@ -4455,8 +4455,17 @@ int SendTls13ClientHello(WOLFSSL* ssl)
         if (ret != 0)
             return ret;
 
+        /* Total message size. */
+        args->sendSz =
+                (int)(args->length + HANDSHAKE_HEADER_SZ + RECORD_HEADER_SZ);
+
+#ifdef WOLFSSL_DTLS13
+        if (ssl->options.dtls)
+            args->sendSz += DTLS_RECORD_EXTRA + DTLS_HANDSHAKE_EXTRA;
+#endif /* WOLFSSL_DTLS13 */
+
 #ifdef WOLFSSL_DTLS_CH_FRAG
-        if (ssl->options.dtls && args->length > maxFrag &&
+        if (ssl->options.dtls && args->sendSz > maxFrag &&
                 TLSX_Find(ssl->extensions, TLSX_COOKIE) == NULL) {
             /* Try again with an empty key share if we would be fragmenting
              * without a cookie */
@@ -4467,7 +4476,9 @@ int SendTls13ClientHello(WOLFSSL* ssl)
             ret = TLSX_GetRequestSize(ssl, client_hello, &args->length);
             if (ret != 0)
                 return ret;
-            if (args->length > maxFrag) {
+            args->sendSz = (int)(args->length +
+                    DTLS_HANDSHAKE_HEADER_SZ + DTLS_RECORD_HEADER_SZ);
+            if (args->sendSz > maxFrag) {
                 WOLFSSL_MSG("Can't fit first CH in one fragment.");
                 return BUFFER_ERROR;
             }
@@ -4476,14 +4487,6 @@ int SendTls13ClientHello(WOLFSSL* ssl)
 #endif
     }
 
-    /* Total message size. */
-    args->sendSz = (int)(args->length + HANDSHAKE_HEADER_SZ + RECORD_HEADER_SZ);
-
-#ifdef WOLFSSL_DTLS13
-    if (ssl->options.dtls)
-        args->sendSz += DTLS_RECORD_EXTRA + DTLS_HANDSHAKE_EXTRA;
-#endif /* WOLFSSL_DTLS13 */
-
     /* Check buffers are big enough and grow if needed. */
     if ((ret = CheckAvailableSize(ssl, args->sendSz)) != 0)
         return ret;