From 692a01238e0218daea233cb45be613bfea4065e1 Mon Sep 17 00:00:00 2001
From: Marco Oliverio <marco@wolfssl.com>
Date: Fri, 6 May 2022 10:09:53 +0200
Subject: [PATCH] tests: support udp in test_server_nofail()

---
 tests/api.c    | 30 ++++++++++++++++++++++++++++--
 wolfssl/test.h |  1 +
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/tests/api.c b/tests/api.c
index 7edb136f5..75bbadc2a 100644
--- a/tests/api.c
+++ b/tests/api.c
@@ -4518,6 +4518,10 @@ static THREAD_RETURN WOLFSSL_THREAD test_server_nofail(void* args)
     int idx;
     int ret, err = 0;
     int sharedCtx = 0;
+    int doUdp = 0;
+    SOCKADDR_IN_T cliAddr;
+    socklen_t     cliLen;
+
 #if defined(OPENSSL_ALL) || defined(WOLFSSL_HAPROXY) || defined(WOLFSSL_WPAS)
     size_t msg_len = 0;
 #endif
@@ -4540,6 +4544,7 @@ static THREAD_RETURN WOLFSSL_THREAD test_server_nofail(void* args)
         WOLFSSL_METHOD* method = NULL;
         if (cbf != NULL && cbf->method != NULL) {
             method = cbf->method();
+
         }
         else {
             method = wolfSSLv23_server_method();
@@ -4574,9 +4579,24 @@ static THREAD_RETURN WOLFSSL_THREAD test_server_nofail(void* args)
     port = wolfSSLPort;
 #endif
 
+    if (cbf != NULL)
+        doUdp = cbf->doUdp;
+
     /* do it here to detect failure */
-    tcp_accept(&sockfd, &clientfd, opts, port, 0, 0, 0, 0, 1, 0, 0);
-    CloseSocket(sockfd);
+    tcp_accept(
+        &sockfd, &clientfd, opts, port, 0, doUdp, 0, 0, 1, 0, 0);
+
+    if (doUdp) {
+        cliLen = sizeof(cliAddr);
+
+        idx = (int)recvfrom(sockfd, input, sizeof(input), MSG_PEEK,
+                  (struct sockaddr*)&cliAddr, &cliLen);
+
+        AssertIntGT(idx, 0);
+    }
+    else {
+        CloseSocket(sockfd);
+    }
 
     wolfSSL_CTX_set_verify(ctx,
                   WOLFSSL_VERIFY_PEER | WOLFSSL_VERIFY_FAIL_IF_NO_PEER_CERT, 0);
@@ -4623,6 +4643,12 @@ static THREAD_RETURN WOLFSSL_THREAD test_server_nofail(void* args)
         goto done;
     }
 
+    if (doUdp) {
+        err = wolfSSL_dtls_set_peer(ssl, &cliAddr, cliLen);
+        if (err != WOLFSSL_SUCCESS)
+            goto done;
+    }
+
 #ifdef WOLFSSL_SESSION_EXPORT
     /* only add in more complex nonblocking case with session export tests */
     if (args && opts->argc > 0) {
diff --git a/wolfssl/test.h b/wolfssl/test.h
index a075643e6..c2ad11e4b 100644
--- a/wolfssl/test.h
+++ b/wolfssl/test.h
@@ -515,6 +515,7 @@ typedef struct callback_functions {
     unsigned char isSharedCtx:1;
     unsigned char loadToSSL:1;
     unsigned char ticNoInit:1;
+    unsigned char doUdp:1;
 } callback_functions;
 
 #if defined(WOLFSSL_SRTP) && !defined(SINGLE_THREADED) && defined(_POSIX_THREADS)