From 0281a60c3ff54e06fe2bb070fa44dd7e3a36d3c8 Mon Sep 17 00:00:00 2001
From: Stefan Agner <stefan@agner.ch>
Date: Thu, 7 Dec 2023 15:13:37 +0100
Subject: [PATCH] [Python] Keep reference to callback function on timeout

When using a timeout when calling GetConnectedDeviceSync() the callback
function DeviceAvailableCallback() gets potentially GC'ed. Make sure
we hold a reference to the instance.
---
 .../ChipDeviceController-ScriptBinding.cpp    | 11 +++----
 src/controller/python/chip/ChipDeviceCtrl.py  | 29 +++++++++++--------
 2 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/src/controller/python/ChipDeviceController-ScriptBinding.cpp b/src/controller/python/ChipDeviceController-ScriptBinding.cpp
index 2d6857d17a4fe6..7cd85d77443251 100644
--- a/src/controller/python/ChipDeviceController-ScriptBinding.cpp
+++ b/src/controller/python/ChipDeviceController-ScriptBinding.cpp
@@ -207,7 +207,7 @@ const char * pychip_Stack_StatusReportToString(uint32_t profileId, uint16_t stat
 void pychip_Stack_SetLogFunct(LogMessageFunct logFunct);
 
 PyChipError pychip_GetConnectedDeviceByNodeId(chip::Controller::DeviceCommissioner * devCtrl, chip::NodeId nodeId,
-                                              DeviceAvailableFunc callback);
+                                              PyObject * context, DeviceAvailableFunc callback);
 PyChipError pychip_FreeOperationalDeviceProxy(chip::OperationalDeviceProxy * deviceProxy);
 PyChipError pychip_GetLocalSessionId(chip::OperationalDeviceProxy * deviceProxy, uint16_t * localSessionId);
 PyChipError pychip_GetNumSessionsToPeer(chip::OperationalDeviceProxy * deviceProxy, uint32_t * numSessions);
@@ -737,8 +737,8 @@ namespace {
 
 struct GetDeviceCallbacks
 {
-    GetDeviceCallbacks(DeviceAvailableFunc callback) :
-        mOnSuccess(OnDeviceConnectedFn, this), mOnFailure(OnConnectionFailureFn, this), mCallback(callback)
+    GetDeviceCallbacks(PyObject * context, DeviceAvailableFunc callback) :
+        mOnSuccess(OnDeviceConnectedFn, this), mOnFailure(OnConnectionFailureFn, this), mContext(context), mCallback(callback)
     {}
 
     static void OnDeviceConnectedFn(void * context, Messaging::ExchangeManager & exchangeMgr, const SessionHandle & sessionHandle)
@@ -758,15 +758,16 @@ struct GetDeviceCallbacks
 
     Callback::Callback<OnDeviceConnected> mOnSuccess;
     Callback::Callback<OnDeviceConnectionFailure> mOnFailure;
+    PyObject * const mContext;
     DeviceAvailableFunc mCallback;
 };
 } // anonymous namespace
 
 PyChipError pychip_GetConnectedDeviceByNodeId(chip::Controller::DeviceCommissioner * devCtrl, chip::NodeId nodeId,
-                                              DeviceAvailableFunc callback)
+                                              PyObject * context, DeviceAvailableFunc callback)
 {
     VerifyOrReturnError(devCtrl != nullptr, ToPyChipError(CHIP_ERROR_INVALID_ARGUMENT));
-    auto * callbacks = new GetDeviceCallbacks(callback);
+    auto * callbacks = new GetDeviceCallbacks(context, callback);
     return ToPyChipError(devCtrl->GetConnectedDevice(nodeId, &callbacks->mOnSuccess, &callbacks->mOnFailure));
 }
 
diff --git a/src/controller/python/chip/ChipDeviceCtrl.py b/src/controller/python/chip/ChipDeviceCtrl.py
index 4a38348575eccf..0c10636011598e 100644
--- a/src/controller/python/chip/ChipDeviceCtrl.py
+++ b/src/controller/python/chip/ChipDeviceCtrl.py
@@ -76,7 +76,7 @@
 #
 # CHIP_ERROR is actually signed, so using c_uint32 is weird, but everything
 # else seems to do it.
-_DeviceAvailableFunct = CFUNCTYPE(None, c_void_p, PyChipError)
+_DeviceAvailableFunct = CFUNCTYPE(None, py_object, c_void_p, PyChipError)
 
 _IssueNOCChainCallbackPythonCallbackFunct = CFUNCTYPE(
     None, py_object, PyChipError, c_void_p, c_size_t, c_void_p, c_size_t, c_void_p, c_size_t, c_void_p, c_size_t, c_uint64)
@@ -759,16 +759,6 @@ def GetConnectedDeviceSync(self, nodeid, allowPASE=True, timeoutMs: int = None):
         returnErr = None
         deviceAvailableCV = threading.Condition()
 
-        @_DeviceAvailableFunct
-        def DeviceAvailableCallback(device, err):
-            nonlocal returnDevice
-            nonlocal returnErr
-            nonlocal deviceAvailableCV
-            with deviceAvailableCV:
-                returnDevice = c_void_p(device)
-                returnErr = err
-                deviceAvailableCV.notify_all()
-
         if allowPASE:
             res = self._ChipStack.Call(lambda: self._dmLib.pychip_GetDeviceBeingCommissioned(
                 self.devCtrl, nodeid, byref(returnDevice)), timeoutMs)
@@ -776,8 +766,23 @@ def DeviceAvailableCallback(device, err):
                 print('Using PASE connection')
                 return DeviceProxyWrapper(returnDevice)
 
+        class DeviceAvailableClosure():
+            @_DeviceAvailableFunct
+            def DeviceAvailableCallback(self, device, err):
+                nonlocal returnDevice
+                nonlocal returnErr
+                nonlocal deviceAvailableCV
+                with deviceAvailableCV:
+                    returnDevice = c_void_p(device)
+                    returnErr = err
+                    deviceAvailableCV.notify_all()
+                ctypes.pythonapi.Py_DecRef(ctypes.py_object(self))
+
+        closure = DeviceAvailableClosure()
+        ctypes.pythonapi.Py_IncRef(ctypes.py_object(closure))
         self._ChipStack.Call(lambda: self._dmLib.pychip_GetConnectedDeviceByNodeId(
-            self.devCtrl, nodeid, DeviceAvailableCallback), timeoutMs).raise_on_error()
+            self.devCtrl, nodeid, ctypes.py_object(closure), closure.DeviceAvailableCallback),
+            timeoutMs).raise_on_error()
 
         # The callback might have been received synchronously (during self._ChipStack.Call()).
         # Check if the device is already set before waiting for the callback.