From 1ee0c3a7fd89d17266296a4f449761b7a767af67 Mon Sep 17 00:00:00 2001
From: Juliusz Sosinowicz <juliusz@wolfssl.com>
Date: Mon, 14 Jun 2021 15:51:04 +0200
Subject: [PATCH] Refactor `dtls_expected_peer_handshake_number` handling

Moving the `dtls_expected_peer_handshake_number` value along has been moved to one location. It has also been changed to not keep state before a cookie exchange has been completed.
---
 src/internal.c | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/src/internal.c b/src/internal.c
index 1aa36de28..6932a222e 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -13556,6 +13556,20 @@ static int DoHandShakeMsgType(WOLFSSL* ssl, byte* input, word32* inOutIdx,
     }
 #endif /* WOLFSSL_ASYNC_CRYPT || WOLFSSL_NONBLOCK_OCSP */
 
+#ifdef WOLFSSL_DTLS
+    if (ret == 0) {
+        if (type == client_hello) {
+            /* Advance expected number only if cookie exchange complete */
+            if (ssl->msgsReceived.got_client_hello)
+                ssl->keys.dtls_expected_peer_handshake_number =
+                    ssl->keys.dtls_peer_handshake_number + 1;
+        }
+        else if (type != finished) {
+            ssl->keys.dtls_expected_peer_handshake_number++;
+        }
+    }
+#endif
+
     WOLFSSL_LEAVE("DoHandShakeMsgType()", ret);
     return ret;
 }
@@ -13922,8 +13936,6 @@ static int DtlsMsgDrain(WOLFSSL* ssl)
         word32 idx = 0;
         if ((ret = DoHandShakeMsgType(ssl, item->msg, &idx, item->type,
                                       item->sz, item->sz)) == 0) {
-            if (item->type != finished)
-                ssl->keys.dtls_expected_peer_handshake_number++;
             DtlsTxMsgListClean(ssl);
         }
     #ifdef WOLFSSL_ASYNC_CRYPT
@@ -14012,8 +14024,8 @@ static int DoDtlsHandShakeMsg(WOLFSSL* ssl, byte* input, word32* inOutIdx,
          * number. (If the cookie changes multiple times in quick succession,
          * the client could be sending multiple new client hello messages
          * with newer and newer cookies.) */
-        WOLFSSL_MSG("Current message is out of order");
         if (type != client_hello) {
+            WOLFSSL_MSG("Current message is out of order");
             if (ssl->dtls_rx_msg_list_sz < DTLS_POOL_SZ) {
                 DtlsMsgStore(ssl, ssl->keys.curEpoch,
                              ssl->keys.dtls_peer_handshake_number,
@@ -14047,10 +14059,6 @@ static int DoDtlsHandShakeMsg(WOLFSSL* ssl, byte* input, word32* inOutIdx,
         }
         else {
             ret = DoHandShakeMsgType(ssl, input, inOutIdx, type, size, totalSz);
-            if (ret == 0) {
-                ssl->keys.dtls_expected_peer_handshake_number =
-                    ssl->keys.dtls_peer_handshake_number + 1;
-            }
         }
     }
     else if (ssl->keys.dtls_peer_handshake_number <
@@ -14147,8 +14155,6 @@ static int DoDtlsHandShakeMsg(WOLFSSL* ssl, byte* input, word32* inOutIdx,
         ret = DoHandShakeMsgType(ssl, input, inOutIdx, type, size, totalSz);
         if (ret == 0) {
             DtlsTxMsgListClean(ssl);
-            if (type != finished)
-                ssl->keys.dtls_expected_peer_handshake_number++;
             if (ssl->dtls_rx_msg_list != NULL) {
                 ret = DtlsMsgDrain(ssl);
             }