Fix UnauthenticatedSession leak (#10754)

* Fix UnauthenticatedSession leak

* Resovle comments

Author: kghost

src/lib/core/Optional.h

@@ -29,6 +29,13 @@
 
 namespace chip {
 
+/// An empty class type used to indicate optional type with uninitialized state.
+struct NullOptionalType
+{
+    explicit NullOptionalType() = default;
+};
+constexpr NullOptionalType NullOptional{};
+
 /**
  * Pairs an object with a boolean value to determine if the object value
  * is actually valid or not.
@@ -38,6 +45,8 @@ class Optional
 {
 public:
     constexpr Optional() : mHasValue(false) {}
+    constexpr Optional(NullOptionalType) : mHasValue(false) {}
+
     ~Optional()
     {
         if (mHasValue)

src/transport/SessionManager.cpp

@@ -189,7 +189,7 @@ CHIP_ERROR SessionManager::SendPreparedMessage(SessionHandle session, const Encr
     else
     {
         auto unauthenticated = session.GetUnauthenticatedSession();
-        mUnauthenticatedSessions.MarkSessionActive(unauthenticated.Get());
+        mUnauthenticatedSessions.MarkSessionActive(unauthenticated);
         destination = &unauthenticated->GetPeerAddress();
 
         ChipLogProgress(Inet,
@@ -339,13 +339,14 @@ void SessionManager::OnMessageReceived(const PeerAddress & peerAddress, System::
 void SessionManager::MessageDispatch(const PacketHeader & packetHeader, const Transport::PeerAddress & peerAddress,
                                      System::PacketBufferHandle && msg)
 {
-    Transport::UnauthenticatedSession * session = mUnauthenticatedSessions.FindOrAllocateEntry(peerAddress);
-    if (session == nullptr)
+    Optional<Transport::UnauthenticatedSessionHandle> optionalSession = mUnauthenticatedSessions.FindOrAllocateEntry(peerAddress);
+    if (!optionalSession.HasValue())
     {
         ChipLogError(Inet, "UnauthenticatedSession exhausted");
         return;
     }
 
+    Transport::UnauthenticatedSessionHandle session      = optionalSession.Value();
     SessionManagerDelegate::DuplicateMessage isDuplicate = SessionManagerDelegate::DuplicateMessage::No;
 
     // Verify message counter
@@ -357,7 +358,7 @@ void SessionManager::MessageDispatch(const PacketHeader & packetHeader, const Tr
     }
     VerifyOrDie(err == CHIP_NO_ERROR);
 
-    mUnauthenticatedSessions.MarkSessionActive(*session);
+    mUnauthenticatedSessions.MarkSessionActive(session);
 
     PayloadHeader payloadHeader;
     ReturnOnFailure(payloadHeader.DecodeAndConsume(msg));
@@ -374,8 +375,7 @@ void SessionManager::MessageDispatch(const PacketHeader & packetHeader, const Tr
 
     if (mCB != nullptr)
     {
-        mCB->OnMessageReceived(packetHeader, payloadHeader, SessionHandle(Transport::UnauthenticatedSessionHandle(*session)),
-                               peerAddress, isDuplicate, std::move(msg));
+        mCB->OnMessageReceived(packetHeader, payloadHeader, SessionHandle(session), peerAddress, isDuplicate, std::move(msg));
     }
 }
 

src/transport/SessionManager.h

@@ -265,11 +265,8 @@ class DLL_EXPORT SessionManager : public TransportMgrDelegate
 
     Optional<SessionHandle> CreateUnauthenticatedSession(const Transport::PeerAddress & peerAddress)
     {
-        Transport::UnauthenticatedSession * session = mUnauthenticatedSessions.FindOrAllocateEntry(peerAddress);
-        if (session == nullptr)
-            return Optional<SessionHandle>::Missing();
-
-        return Optional<SessionHandle>::Value(SessionHandle(Transport::UnauthenticatedSessionHandle(*session)));
+        Optional<Transport::UnauthenticatedSessionHandle> session = mUnauthenticatedSessions.FindOrAllocateEntry(peerAddress);
+        return session.HasValue() ? MakeOptional<SessionHandle>(session.Value()) : NullOptional;
     }
 
 private:

src/transport/UnauthenticatedSessionTable.h

@@ -44,7 +44,7 @@ class UnauthenticatedSessionDeleter
  * @brief
  *   An UnauthenticatedSession stores the binding of TransportAddress, and message counters.
  */
-class UnauthenticatedSession : public ReferenceCounted<UnauthenticatedSession, UnauthenticatedSessionDeleter>
+class UnauthenticatedSession : public ReferenceCounted<UnauthenticatedSession, UnauthenticatedSessionDeleter, 0>
 {
 public:
     UnauthenticatedSession(const PeerAddress & address) : mPeerAddress(address) { mLocalMessageCounter.Init(); }
@@ -82,6 +82,39 @@ template <size_t kMaxConnectionCount, Time::Source kTimeSource = Time::Source::k
 class UnauthenticatedSessionTable
 {
 public:
+    /**
+     * Get a session given the peer address. If the session doesn't exist in the cache, allocate a new entry for it.
+     *
+     * @return the session found or allocated, nullptr if not found and allocation failed.
+     */
+    CHECK_RETURN_VALUE
+    Optional<UnauthenticatedSessionHandle> FindOrAllocateEntry(const PeerAddress & address)
+    {
+        UnauthenticatedSession * result = FindEntry(address);
+        if (result != nullptr)
+            return MakeOptional<UnauthenticatedSessionHandle>(*result);
+
+        CHIP_ERROR err = AllocEntry(address, result);
+        if (err == CHIP_NO_ERROR)
+        {
+            return MakeOptional<UnauthenticatedSessionHandle>(*result);
+        }
+        else
+        {
+            return Optional<UnauthenticatedSessionHandle>::Missing();
+        }
+    }
+
+    /// Mark a session as active
+    void MarkSessionActive(UnauthenticatedSessionHandle session)
+    {
+        session->SetLastActivityTimeMs(mTimeSource.GetCurrentMonotonicTimeMs());
+    }
+
+    /// Allows access to the underlying time source used for keeping track of connection active time
+    Time::TimeSource<kTimeSource> & GetTimeSource() { return mTimeSource; }
+
+private:
     /**
      * Allocates a new session out of the internal resource pool.
      *
@@ -125,36 +158,6 @@ class UnauthenticatedSessionTable
         return result;
     }
 
-    /**
-     * Get a peer given the peer id. If the peer doesn't exist in the cache, allocate a new entry for it.
-     *
-     * @return the peer found or allocated, nullptr if not found and allocate failed.
-     */
-    CHECK_RETURN_VALUE
-    UnauthenticatedSession * FindOrAllocateEntry(const PeerAddress & address)
-    {
-        UnauthenticatedSession * result = FindEntry(address);
-        if (result != nullptr)
-            return result;
-
-        CHIP_ERROR err = AllocEntry(address, result);
-        if (err == CHIP_NO_ERROR)
-        {
-            return result;
-        }
-        else
-        {
-            return nullptr;
-        }
-    }
-
-    /// Mark a session as active
-    void MarkSessionActive(UnauthenticatedSession & entry) { entry.SetLastActivityTimeMs(mTimeSource.GetCurrentMonotonicTimeMs()); }
-
-    /// Allows access to the underlying time source used for keeping track of connection active time
-    Time::TimeSource<kTimeSource> & GetTimeSource() { return mTimeSource; }
-
-private:
     UnauthenticatedSession * FindLeastRecentUsedEntry()
     {
         UnauthenticatedSession * result = nullptr;