From ead0907b1fad1cebab2ebca76f6a76e60edc64f0 Mon Sep 17 00:00:00 2001
From: Terence Hampson <thampson@google.com>
Date: Thu, 4 Aug 2022 17:06:19 +0000
Subject: [PATCH] Address PR comments

---
 .../ota-requestor/DefaultOTARequestor.cpp     |  2 +-
 src/controller/CHIPDeviceController.cpp       |  2 +-
 src/controller/python/chip/ChipDeviceCtrl.py  |  6 ++++++
 src/transport/SessionManager.cpp              | 21 ++++++++++++-------
 src/transport/SessionManager.h                |  6 +++---
 5 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/src/app/clusters/ota-requestor/DefaultOTARequestor.cpp b/src/app/clusters/ota-requestor/DefaultOTARequestor.cpp
index f3bf092bd885c8..98f11b3fee0811 100644
--- a/src/app/clusters/ota-requestor/DefaultOTARequestor.cpp
+++ b/src/app/clusters/ota-requestor/DefaultOTARequestor.cpp
@@ -387,7 +387,7 @@ void DefaultOTARequestor::DisconnectFromProvider()
     }
 
     auto providerNodeId = GetProviderScopedId();
-    mServer->GetSecureSessionManager().MarkSessionAsDefunct(providerNodeId, MakeOptional(Transport::SecureSession::Type::kCASE));
+    mServer->GetSecureSessionManager().MarkSessionsAsDefunct(providerNodeId, MakeOptional(Transport::SecureSession::Type::kCASE));
 }
 
 // Requestor is directed to cancel image update in progress. All the Requestor state is
diff --git a/src/controller/CHIPDeviceController.cpp b/src/controller/CHIPDeviceController.cpp
index e547b49265e856..da29daeac289b9 100644
--- a/src/controller/CHIPDeviceController.cpp
+++ b/src/controller/CHIPDeviceController.cpp
@@ -364,7 +364,7 @@ CHIP_ERROR DeviceController::DisconnectDevice(NodeId nodeId)
 {
     ChipLogProgress(Controller, "Force close session for node 0x%" PRIx64, nodeId);
 
-    if (SessionMgr()->MarkSessionAsDefunct(GetPeerScopedId(nodeId), MakeOptional(Transport::SecureSession::Type::kCASE)))
+    if (SessionMgr()->MarkSessionsAsDefunct(GetPeerScopedId(nodeId), MakeOptional(Transport::SecureSession::Type::kCASE)))
     {
         return CHIP_NO_ERROR;
     }
diff --git a/src/controller/python/chip/ChipDeviceCtrl.py b/src/controller/python/chip/ChipDeviceCtrl.py
index f3dba1a7aa396c..3ec8ced797e066 100644
--- a/src/controller/python/chip/ChipDeviceCtrl.py
+++ b/src/controller/python/chip/ChipDeviceCtrl.py
@@ -118,6 +118,12 @@ class DCState(enum.IntEnum):
 
 
 class DeviceProxyWrapper():
+    ''' Encapsulates a pointer to OperationalDeviceProxy on the c++ side that needs to be
+        freed when DeviceProxyWrapper goes out of scope. There is a potential issue where
+        if this is copied around that a double free will occure, but how this is used today
+        that is not an issue that needs to be accounted for and it will become very apparent
+        if that happens.
+    '''
     def __init__(self, deviceProxy: ctypes.c_void_p, dmLib=None):
         self._deviceProxy = deviceProxy
         self._dmLib = dmLib
diff --git a/src/transport/SessionManager.cpp b/src/transport/SessionManager.cpp
index b8505a0537f05c..5fa99e0e677e91 100644
--- a/src/transport/SessionManager.cpp
+++ b/src/transport/SessionManager.cpp
@@ -445,15 +445,20 @@ void SessionManager::ExpireAllPASESessions()
     });
 }
 
-bool SessionManager::MarkSessionAsDefunct(const ScopedNodeId & node, const Optional<Transport::SecureSession::Type> & type)
+bool SessionManager::MarkSessionsAsDefunct(const ScopedNodeId & node, const Optional<Transport::SecureSession::Type> & type)
 {
-    auto optionalSessionHandle = FindSecureSessionForNode(node, type);
-    if (optionalSessionHandle.HasValue())
-    {
-        optionalSessionHandle.Value()->AsSecureSession()->MarkAsDefunct();
-        return true;
-    }
-    return false;
+    bool found = false;
+    mSecureSessions.ForEachSession([&node, &type, &found](auto session) {
+        if (session->IsActiveSession() && session->GetPeer() == node &&
+            (!type.HasValue() || type.Value() == session->GetSecureSessionType()))
+        {
+            session->AsSecureSession()->MarkAsDefunct();
+            found = true;
+        }
+        return Loop::Continue;
+    });
+
+    return found;
 }
 
 Optional<SessionHandle> SessionManager::AllocateSession(SecureSession::Type secureSessionType,
diff --git a/src/transport/SessionManager.h b/src/transport/SessionManager.h
index 13c6f3af9255d0..2ed3961d8fa1d7 100644
--- a/src/transport/SessionManager.h
+++ b/src/transport/SessionManager.h
@@ -351,14 +351,14 @@ class DLL_EXPORT SessionManager : public TransportMgrDelegate, public FabricTabl
 
     /**
      * @brief
-     *   Mark most recently active session that matches provided arguments as defunct.
+     *   Marks all active session that matches provided arguments as defunct.
      *
-     * @param node    Scoped node ID of the session we should mark as defunct.
+     * @param node    Scoped node ID of the active sessions we should mark as defunct.
      * @param type    Type of session we are looking to mark as defunct. If matching
      *                against all types of sessions is desired, NullOptional should
      *                be passed into type.
      */
-    bool MarkSessionAsDefunct(const ScopedNodeId & node, const Optional<Transport::SecureSession::Type> & type);
+    bool MarkSessionsAsDefunct(const ScopedNodeId & node, const Optional<Transport::SecureSession::Type> & type);
 
     /**
      * @brief