Fix memory leak in UnauthenticatedSessionTable.

src/transport/UnauthenticatedSessionTable.h

@@ -22,6 +22,7 @@
 #include <lib/support/CodeUtils.h>
 #include <lib/support/Pool.h>
 #include <messaging/ReliableMessageProtocolConfig.h>
+#include <system/SystemConfig.h>
 #include <system/TimeSource.h>
 #include <transport/PeerMessageCounter.h>
 #include <transport/Session.h>
@@ -34,8 +35,7 @@ namespace Transport {
  * @brief
  *   An UnauthenticatedSession stores the binding of TransportAddress, and message counters.
  */
-class UnauthenticatedSession : public Session,
-                               public ReferenceCounted<UnauthenticatedSession, NoopDeletor<UnauthenticatedSession>, 0>
+class UnauthenticatedSession : public Session, public ReferenceCounted<UnauthenticatedSession, UnauthenticatedSession, 0>
 {
 public:
     enum class SessionRole
@@ -44,6 +44,7 @@ class UnauthenticatedSession : public Session,
         kResponder,
     };
 
+protected:
     UnauthenticatedSession(SessionRole sessionRole, NodeId ephemeralInitiatorNodeID, const ReliableMessageProtocolConfig & config) :
         mEphemeralInitiatorNodeId(ephemeralInitiatorNodeID), mSessionRole(sessionRole),
         mLastActivityTime(System::SystemClock().GetMonotonicTimestamp()),
@@ -52,6 +53,7 @@ class UnauthenticatedSession : public Session,
     {}
     ~UnauthenticatedSession() override { VerifyOrDie(GetReferenceCount() == 0); }
 
+public:
     UnauthenticatedSession(const UnauthenticatedSession &)             = delete;
     UnauthenticatedSession & operator=(const UnauthenticatedSession &) = delete;
     UnauthenticatedSession(UnauthenticatedSession &&)                  = delete;
@@ -68,8 +70,8 @@ class UnauthenticatedSession : public Session,
 
     Session::SessionType GetSessionType() const override { return Session::SessionType::kUnauthenticated; }
 
-    void Retain() override { ReferenceCounted<UnauthenticatedSession, NoopDeletor<UnauthenticatedSession>, 0>::Retain(); }
-    void Release() override { ReferenceCounted<UnauthenticatedSession, NoopDeletor<UnauthenticatedSession>, 0>::Release(); }
+    void Retain() override { ReferenceCounted<UnauthenticatedSession, UnauthenticatedSession, 0>::Retain(); }
+    void Release() override { ReferenceCounted<UnauthenticatedSession, UnauthenticatedSession, 0>::Release(); }
 
     bool IsActiveSession() const override { return true; }
 
@@ -132,6 +134,23 @@ class UnauthenticatedSession : public Session,
 
     PeerMessageCounter & GetPeerMessageCounter() { return mPeerMessageCounter; }
 
+    static void Release(UnauthenticatedSession * obj)
+    {
+        // When using heap pools, we need to make sure to release ourselves back to
+        // the pool.  When not using heap pools, we don't want the extra cost of the
+        // table pointer here, and the table itself handles entry reuse and cleanup
+        // as needed.
+#if CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+        obj->ReleaseSelfToPool();
+#else
+        // Just do nothing.
+#endif // CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+    }
+
+#if CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+    virtual void ReleaseSelfToPool() = 0;
+#endif // CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+
 private:
     const NodeId mEphemeralInitiatorNodeId;
     const SessionRole mSessionRole;
@@ -142,6 +161,35 @@ class UnauthenticatedSession : public Session,
     PeerMessageCounter mPeerMessageCounter;
 };
 
+template <size_t kMaxSessionCount>
+class UnauthenticatedSessionTable;
+
+namespace detail {
+
+template <size_t kMaxSessionCount>
+class UnauthenticatedSessionPoolEntry : public UnauthenticatedSession
+{
+public:
+    UnauthenticatedSessionPoolEntry(SessionRole sessionRole, NodeId ephemeralInitiatorNodeID,
+                                    const ReliableMessageProtocolConfig & config,
+                                    UnauthenticatedSessionTable<kMaxSessionCount> & sessionTable) :
+        UnauthenticatedSession(sessionRole, ephemeralInitiatorNodeID, config)
+#if CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+        ,
+        mSessionTable(sessionTable)
+#endif // CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+    {}
+
+private:
+#if CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+    virtual void ReleaseSelfToPool();
+
+    UnauthenticatedSessionTable<kMaxSessionCount> & mSessionTable;
+#endif // CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+};
+
+} // namespace detail
+
 /*
  * @brief
  *   An table which manages UnauthenticatedSessions
@@ -153,7 +201,17 @@ template <size_t kMaxSessionCount>
 class UnauthenticatedSessionTable
 {
 public:
-    ~UnauthenticatedSessionTable() { mEntries.ReleaseAll(); }
+    ~UnauthenticatedSessionTable()
+    {
+#if !CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+        // When not using heap pools, our entries never actually get released
+        // back to the pool (which lets us make the entries 4 bytes smaller by
+        // not storing a reference to the table in them) and we LRU reuse ones
+        // that have 0 refcount.  But we should release them all here, to ensure
+        // that we don't hit fatal asserts in our pool destructor.
+        mEntries.ReleaseAll();
+#endif // CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+    }
 
     /**
      * Get a responder session with the given ephemeralInitiatorNodeID. If the session doesn't exist in the cache, allocate a new
@@ -203,6 +261,9 @@ class UnauthenticatedSessionTable
     }
 
 private:
+    using EntryType = detail::UnauthenticatedSessionPoolEntry<kMaxSessionCount>;
+    friend EntryType;
+
     /**
      * Allocates a new session out of the internal resource pool.
      *
@@ -213,17 +274,23 @@ class UnauthenticatedSessionTable
     CHIP_ERROR AllocEntry(UnauthenticatedSession::SessionRole sessionRole, NodeId ephemeralInitiatorNodeID,
                           const ReliableMessageProtocolConfig & config, UnauthenticatedSession *& entry)
     {
-        entry = mEntries.CreateObject(sessionRole, ephemeralInitiatorNodeID, config);
-        if (entry != nullptr)
+        auto entryToUse = mEntries.CreateObject(sessionRole, ephemeralInitiatorNodeID, config, *this);
+        if (entryToUse != nullptr)
+        {
+            entry = entryToUse;
             return CHIP_NO_ERROR;
+        }
 
-        entry = FindLeastRecentUsedEntry();
-        if (entry == nullptr)
+#if !CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+        entryToUse = FindLeastRecentUsedEntry();
+#endif // CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+        if (entryToUse == nullptr)
         {
             return CHIP_ERROR_NO_MEMORY;
         }
 
-        mEntries.ResetObject(entry, sessionRole, ephemeralInitiatorNodeID, config);
+        mEntries.ResetObject(entryToUse, sessionRole, ephemeralInitiatorNodeID, config, *this);
+        entry = entryToUse;
         return CHIP_NO_ERROR;
     }
 
@@ -242,12 +309,12 @@ class UnauthenticatedSessionTable
         return result;
     }
 
-    UnauthenticatedSession * FindLeastRecentUsedEntry()
+    EntryType * FindLeastRecentUsedEntry()
     {
-        UnauthenticatedSession * result     = nullptr;
+        EntryType * result                  = nullptr;
         System::Clock::Timestamp oldestTime = System::Clock::Timestamp(std::numeric_limits<System::Clock::Timestamp::rep>::max());
 
-        mEntries.ForEachActiveObject([&](UnauthenticatedSession * entry) {
+        mEntries.ForEachActiveObject([&](EntryType * entry) {
             if (entry->GetReferenceCount() == 0 && entry->GetLastActivityTime() < oldestTime)
             {
                 result     = entry;
@@ -259,8 +326,18 @@ class UnauthenticatedSessionTable
         return result;
     }
 
-    ObjectPool<UnauthenticatedSession, kMaxSessionCount> mEntries;
+    void ReleaseEntry(EntryType * entry) { mEntries.ReleaseObject(entry); }
+
+    ObjectPool<EntryType, kMaxSessionCount> mEntries;
 };
 
+#if CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+template <size_t kMaxSessionCount>
+void detail::UnauthenticatedSessionPoolEntry<kMaxSessionCount>::ReleaseSelfToPool()
+{
+    mSessionTable.ReleaseEntry(this);
+}
+#endif // CHIP_SYSTEM_CONFIG_POOL_USE_HEAP
+
 } // namespace Transport
 } // namespace chip