From 145671090a41350ad431a3a8f6dc89a4faf335ea Mon Sep 17 00:00:00 2001
From: Zang MingJie <mingjiez@google.com>
Date: Thu, 14 Jul 2022 21:20:54 +0800
Subject: [PATCH] Fix OT UDP receive and UDP close racing (#20536)

* Fix OT UDP receive and UDP close racing

* Add failure log
---
 src/inet/UDPEndPoint.cpp               |  2 +-
 src/inet/UDPEndPointImplLwIP.cpp       | 17 +++++++++++++++++
 src/inet/UDPEndPointImplOpenThread.cpp | 16 ++++++++++++++++
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/src/inet/UDPEndPoint.cpp b/src/inet/UDPEndPoint.cpp
index e0907d951a9e65..b0b6f3f861e989 100644
--- a/src/inet/UDPEndPoint.cpp
+++ b/src/inet/UDPEndPoint.cpp
@@ -125,8 +125,8 @@ void UDPEndPoint::Close()
 {
     if (mState != State::kClosed)
     {
-        CloseImpl();
         mState = State::kClosed;
+        CloseImpl();
     }
 }
 
diff --git a/src/inet/UDPEndPointImplLwIP.cpp b/src/inet/UDPEndPointImplLwIP.cpp
index 32b10c62c5d720..9e6b98f780a330 100644
--- a/src/inet/UDPEndPointImplLwIP.cpp
+++ b/src/inet/UDPEndPointImplLwIP.cpp
@@ -241,6 +241,19 @@ void UDPEndPointImplLwIP::CloseImpl()
         udp_remove(mUDP);
         mUDP              = nullptr;
         mLwIPEndPointType = LwIPEndPointType::Unknown;
+
+        // In case that there is a UDPEndPointImplLwIP::LwIPReceiveUDPMessage
+        // event pending in the event queue (SystemLayer::ScheduleLambda), we
+        // schedule a release call to the end of the queue, to ensure that the
+        // queued pointer to UDPEndPointImplLwIP is not dangling.
+        Retain();
+        CHIP_ERROR err = GetSystemLayer().ScheduleLambda([this] { Release(); });
+        if (err != CHIP_NO_ERROR)
+        {
+            ChipLogError(Inet, "Unable scedule lambda: %" CHIP_ERROR_FORMAT, err.Format());
+            // There is nothing we can do here, accept the chance of racing
+            Release();
+        }
     }
 
     // Unlock LwIP stack
@@ -342,6 +355,10 @@ void UDPEndPointImplLwIP::LwIPReceiveUDPMessage(void * arg, struct udp_pcb * pcb
     UDPEndPointImplLwIP * ep       = static_cast<UDPEndPointImplLwIP *>(arg);
     IPPacketInfo * pktInfo         = nullptr;
     System::PacketBufferHandle buf = System::PacketBufferHandle::Adopt(p);
+
+    if (ep->mState == State::kClosed)
+        return;
+
     if (buf->HasChainedBuffer())
     {
         // Try the simple expedient of flattening in-place.
diff --git a/src/inet/UDPEndPointImplOpenThread.cpp b/src/inet/UDPEndPointImplOpenThread.cpp
index 2f78e80a4253ef..073f633bd8a38c 100644
--- a/src/inet/UDPEndPointImplOpenThread.cpp
+++ b/src/inet/UDPEndPointImplOpenThread.cpp
@@ -54,6 +54,9 @@ void UDPEndPointImplOT::handleUdpReceive(void * aContext, otMessage * aMessage,
     char destStr[Inet::IPAddress::kMaxStringLength];
 #endif
 
+    if (ep->mState == State::kClosed)
+        return;
+
     if (msgLen > System::PacketBuffer::kMaxSizeWithoutReserve)
     {
         ChipLogError(Inet, "UDP message too long, discarding. Size received %d", msgLen);
@@ -258,6 +261,19 @@ void UDPEndPointImplOT::CloseImpl()
     if (otUdpIsOpen(mOTInstance, &mSocket))
     {
         otUdpClose(mOTInstance, &mSocket);
+
+        // In case that there is a UDPEndPointImplOT::handleUdpReceive event
+        // pending in the event queue (SystemLayer::ScheduleLambda), we
+        // schedule a release call to the end of the queue, to ensure that the
+        // queued pointer to UDPEndPointImplOT is not dangling.
+        Retain();
+        CHIP_ERROR err = GetSystemLayer().ScheduleLambda([this] { Release(); });
+        if (err != CHIP_NO_ERROR)
+        {
+            ChipLogError(Inet, "Unable scedule lambda: %" CHIP_ERROR_FORMAT, err.Format());
+            // There is nothing we can do here, accept the chance of racing
+            Release();
+        }
     }
     UnlockOpenThread();
 }