From 3fac706d86f54327b0bcaf84d82dff11063a31b1 Mon Sep 17 00:00:00 2001
From: Stephen Seo <seo.disparate@gmail.com>
Date: Wed, 19 Apr 2023 18:26:45 +0900
Subject: [PATCH] Add UDPC_free_PacketInfo_ptr(...)

Also add unit test for this function.
---
 src/UDPC.h            | 31 +++++++++++++++++++++++++------
 src/UDPConnection.cpp | 10 ++++++++--
 src/test/TestUDPC.cpp | 10 ++++++++++
 3 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/src/UDPC.h b/src/UDPC.h
index 06b995f..25026de 100644
--- a/src/UDPC.h
+++ b/src/UDPC.h
@@ -161,8 +161,9 @@ typedef struct UDPC_EXPORT UDPC_ConnectionId {
  *
  * If \ref data is NULL or \ref dataSize is 0, then this packet is invalid.
  *
- * \warning This struct must be free'd with a call to UDPC_free_PacketInfo to
- * avoid a memory leak.
+ * \warning This struct must be free'd with a call to
+ * \ref UDPC_free_PacketInfo_ptr or \ref UDPC_free_PacketInfo to avoid a memory
+ * leak.
  */
 typedef struct UDPC_EXPORT UDPC_PacketInfo {
     /*!
@@ -657,19 +658,37 @@ UDPC_EXPORT UDPC_Event UDPC_get_event(UDPC_HContext ctx, unsigned long *remainin
  * \brief Get a received packet from a given UDPC context.
  *
  * \warning The received packet (if valid) must be free'd with a call to
- * \ref UDPC_free_PacketInfo() to avoid a memory leak.
+ * \ref UDPC_free_PacketInfo_ptr or \ref UDPC_free_PacketInfo to avoid a memory
+ * leak.
  */
 UDPC_EXPORT UDPC_PacketInfo UDPC_get_received(UDPC_HContext ctx, unsigned long *remaining);
 
 /*!
  * \brief Frees a UDPC_PacketInfo.
  *
- * Internally, the member variable \ref UDPC_PacketInfo::data will be free'd and
- * set to NULL and \ref UDPC_PacketInfo::dataSize will be set to 0 if the given
- * packet is valid.
+ * Internally, the member variable \ref UDPC_PacketInfo::data will be free'd.
+ * \ref UDPC_free_PacketInfo_ptr is safer to use than this function, as it
+ * also zeros out the relevant data to avoid double frees.
  */
 UDPC_EXPORT void UDPC_free_PacketInfo(UDPC_PacketInfo pInfo);
 
+/*!
+ * \brief Frees a UDPC_PacketInfo.
+ *
+ * This is a safer alternative to \ref UDPC_free_PacketInfo because it
+ * internally zeroes out the internal pointer and size variables, making it
+ * safe to pass the same ptr multiple times to this function as it avoids a
+ * double free.
+ *
+ * Usage:
+ * \code{.c}
+ * UDPC_PacketInfo pinfo = UDPC_get_received(ctx, NULL);
+ * UDPC_free_PacketInfo_ptr(&pinfo);
+ * UDPC_free_PacketInfo_ptr(&pinfo); // This is safe, no double free.
+ * \endcode
+ */
+UDPC_EXPORT void UDPC_free_PacketInfo_ptr(UDPC_PacketInfo *pInfoPtr);
+
 /*!
  * \brief Sets public/private keys used for packet verification
  *
diff --git a/src/UDPConnection.cpp b/src/UDPConnection.cpp
index d420400..b8b0f22 100644
--- a/src/UDPConnection.cpp
+++ b/src/UDPConnection.cpp
@@ -2496,8 +2496,14 @@ UDPC_PacketInfo UDPC_get_received(UDPC_HContext ctx, unsigned long *remaining) {
 void UDPC_free_PacketInfo(UDPC_PacketInfo pInfo) {
     if(pInfo.data && pInfo.dataSize > 0) {
         std::free(pInfo.data);
-        pInfo.data = 0;
-        pInfo.dataSize = 0;
+    }
+}
+
+void UDPC_free_PacketInfo_ptr(UDPC_PacketInfo *pInfoPtr) {
+    if (pInfoPtr->data && pInfoPtr->dataSize > 0) {
+        std::free(pInfoPtr->data);
+        pInfoPtr->data = nullptr;
+        pInfoPtr->dataSize = 0;
     }
 }
 
diff --git a/src/test/TestUDPC.cpp b/src/test/TestUDPC.cpp
index e0e0f76..75ba976 100644
--- a/src/test/TestUDPC.cpp
+++ b/src/test/TestUDPC.cpp
@@ -4,6 +4,7 @@
 #include <UDPC_Defines.hpp>
 
 #include <cstdio>
+#include <cstdlib>
 #include <cstring>
 #include <future>
 
@@ -328,3 +329,12 @@ TEST(UDPC, a4toa6) {
     a4 = htonl(0x01020304);
     EXPECT_EQ(UDPC_a4toa6(a4), a6);
 }
+
+TEST(UDPC, free_packet_ptr) {
+    UDPC_PacketInfo pinfo;
+    pinfo.dataSize = 8;
+    pinfo.data = (char*)std::malloc(pinfo.dataSize);
+
+    UDPC_free_PacketInfo_ptr(&pinfo);
+    UDPC_free_PacketInfo_ptr(&pinfo);
+}