From 8fb48464e3202034453008fbaa258651ebd151a3 Mon Sep 17 00:00:00 2001
From: Juliusz Sosinowicz <juliusz@wolfssl.com>
Date: Wed, 22 Jun 2022 20:11:08 +0200
Subject: [PATCH] Add callback when we parse a verified ClientHello

---
 src/ssl.c          | 34 ++++++++++++++++++++++++++++++++++
 src/tls13.c        | 11 +++++++++++
 wolfssl/internal.h |  5 +++++
 wolfssl/ssl.h      |  4 ++++
 4 files changed, 54 insertions(+)

diff --git a/src/ssl.c b/src/ssl.c
index 3f19647b9..82d7c9162 100644
--- a/src/ssl.c
+++ b/src/ssl.c
@@ -1183,6 +1183,14 @@ int wolfSSL_dtls_set_peer(WOLFSSL* ssl, void* peer, unsigned int peerSz)
     if (ssl == NULL)
         return WOLFSSL_FAILURE;
 
+    if (peer == NULL || peerSz == 0) {
+        if (ssl->buffers.dtlsCtx.peer.sa != NULL)
+            XFREE(ssl->buffers.dtlsCtx.peer.sa,ssl->heap,DYNAMIC_TYPE_SOCKADDR);
+        ssl->buffers.dtlsCtx.peer.sa = NULL;
+        ssl->buffers.dtlsCtx.peer.sz = 0;
+        return WOLFSSL_SUCCESS;
+    }
+
     sa = (void*)XMALLOC(peerSz, ssl->heap, DYNAMIC_TYPE_SOCKADDR);
     if (sa != NULL) {
         if (ssl->buffers.dtlsCtx.peer.sa != NULL) {
@@ -12527,6 +12535,18 @@ int wolfSSL_DTLS_SetCookieSecret(WOLFSSL* ssl,
                 return wolfSSL_accept_TLSv13(ssl);
             }
 #endif
+
+#ifdef WOLFSSL_DTLS
+            if (ssl->chGoodCb != NULL && !IsSCR(ssl)) {
+                int cbret = ssl->chGoodCb(ssl, ssl->chGoodCtx);
+                if (cbret < 0) {
+                    ssl->error = cbret;
+                    WOLFSSL_MSG("ClientHello Good Cb don't continue error");
+                    return WOLFSSL_FATAL_ERROR;
+                }
+            }
+#endif
+
             ssl->options.acceptState = ACCEPT_FIRST_REPLY_DONE;
             WOLFSSL_MSG("accept state ACCEPT_FIRST_REPLY_DONE");
             FALL_THROUGH;
@@ -12748,6 +12768,20 @@ int wolfSSL_DTLS_SetCookieSecret(WOLFSSL* ssl,
 
 #endif /* NO_WOLFSSL_SERVER */
 
+#ifdef WOLFSSL_DTLS
+int wolfSSL_SetChGoodCb(WOLFSSL* ssl, ClientHelloGoodCb cb, void* user_ctx)
+{
+    WOLFSSL_ENTER("wolfSSL_SetChGoodCb");
+
+    if (ssl == NULL)
+        return BAD_FUNC_ARG;
+
+    ssl->chGoodCb  = cb;
+    ssl->chGoodCtx = user_ctx;
+
+    return WOLFSSL_SUCCESS;
+}
+#endif
 
 #ifndef NO_HANDSHAKE_DONE_CB
 
diff --git a/src/tls13.c b/src/tls13.c
index 431c2e666..ae3e32c5b 100644
--- a/src/tls13.c
+++ b/src/tls13.c
@@ -10798,6 +10798,17 @@ int wolfSSL_accept_TLSv13(WOLFSSL* ssl)
                 }
             }
 
+#ifdef WOLFSSL_DTLS
+            if (ssl->chGoodCb != NULL) {
+                int cbret = ssl->chGoodCb(ssl, ssl->chGoodCtx);
+                if (cbret < 0) {
+                    ssl->error = cbret;
+                    WOLFSSL_MSG("ClientHello Good Cb don't continue error");
+                    return WOLFSSL_FATAL_ERROR;
+                }
+            }
+#endif
+
             ssl->options.acceptState = TLS13_ACCEPT_SECOND_REPLY_DONE;
             WOLFSSL_MSG("accept state ACCEPT_SECOND_REPLY_DONE");
             FALL_THROUGH;
diff --git a/wolfssl/internal.h b/wolfssl/internal.h
index 66d8b1041..efaccc767 100644
--- a/wolfssl/internal.h
+++ b/wolfssl/internal.h
@@ -4478,6 +4478,11 @@ struct WOLFSSL {
 #ifdef WOLFSSL_STATIC_MEMORY
     WOLFSSL_HEAP_HINT heap_hint;
 #endif
+#ifdef WOLFSSL_DTLS
+    ClientHelloGoodCb chGoodCb;        /*  notify user we parsed a verified
+                                        *  ClientHello */
+    void*             chGoodCtx;       /*  user ClientHello cb context  */
+#endif
 #ifndef NO_HANDSHAKE_DONE_CB
     HandShakeDoneCb hsDoneCb;          /*  notify user handshake done */
     void*           hsDoneCtx;         /*  user handshake cb context  */
diff --git a/wolfssl/ssl.h b/wolfssl/ssl.h
index 829fc4745..56be14ad0 100644
--- a/wolfssl/ssl.h
+++ b/wolfssl/ssl.h
@@ -3940,6 +3940,10 @@ WOLFSSL_API int wolfSSL_CTX_DisableExtendedMasterSecret(WOLFSSL_CTX* ctx);
 #define WOLFSSL_CRL_START_MON 0x02   /* start monitoring flag */
 
 
+/* notify user we parsed a verified ClientHello is done. This only has an effect
+ * on the server end. */
+typedef int (*ClientHelloGoodCb)(WOLFSSL* ssl, void*);
+WOLFSSL_API int wolfSSL_SetChGoodCb(WOLFSSL* ssl, ClientHelloGoodCb cb, void* user_ctx);
 /* notify user the handshake is done */
 typedef int (*HandShakeDoneCb)(WOLFSSL* ssl, void*);
 WOLFSSL_API int wolfSSL_SetHsDoneCb(WOLFSSL* ssl, HandShakeDoneCb cb, void* user_ctx);