Wire up get_device() instead of emulating it

This way it is automatically cancellable, and we get the performance
benefit of not having to wait for all backends to start.

frida/__init__.py

@@ -70,56 +70,28 @@ def inject_library_blob(target, blob, entrypoint, data, **kwargs):
     return get_local_device().inject_library_blob(target, blob, entrypoint, data, **kwargs)
 
 
-def enumerate_devices(**kwargs):
-    return get_device_manager().enumerate_devices(**kwargs)
-
-
 def get_local_device(**kwargs):
-    return get_device_matching(lambda device: device.type == 'local', timeout=0, **kwargs)
+    return get_device_matching(lambda d: d.type == 'local', timeout=0, **kwargs)
 
 
 def get_remote_device(**kwargs):
-    return get_device_matching(lambda device: device.type == 'remote', timeout=0, **kwargs)
+    return get_device_matching(lambda d: d.type == 'remote', timeout=0, **kwargs)
 
 
 def get_usb_device(timeout=0, **kwargs):
-    return get_device_matching(lambda device: device.type == 'usb', timeout, **kwargs)
+    return get_device_matching(lambda d: d.type == 'usb', timeout, **kwargs)
 
 
 def get_device(id, timeout=0, **kwargs):
-    return get_device_matching(lambda device: device.id == id, timeout, **kwargs)
+    return get_device_manager().get_device(id, timeout, **kwargs)
 
 
 def get_device_matching(predicate, timeout=0, **kwargs):
-    matches = []
-    lock = threading.Lock()
-    done = threading.Event()
+    return get_device_manager().get_device_matching(predicate, timeout, **kwargs)
 
-    def on_device_added(device):
-        if predicate(device):
-            with lock:
-                matches.append(device)
-            done.set()
 
-    manager = get_device_manager()
-    manager.on('added', on_device_added)
-    try:
-        initial_matches = [device for device in manager.enumerate_devices(**kwargs) if predicate(device)]
-        if len(initial_matches) > 0:
-            return initial_matches[0]
-
-        done.wait(timeout)
-
-        with lock:
-            if len(matches) == 0:
-                if timeout == 0:
-                    raise InvalidArgumentError("device not found")
-                else:
-                    raise TimedOutError("timed out while waiting for device to appear")
-
-            return matches[0]
-    finally:
-        manager.off('added', on_device_added)
+def enumerate_devices(**kwargs):
+    return get_device_manager().enumerate_devices(**kwargs)
 
 
 def shutdown(**kwargs):

frida/core.py

@@ -34,6 +34,28 @@ class DeviceManager(object):
     def __repr__(self):
         return repr(self._impl)
 
+    def get_local_device(self, **kwargs):
+        return self.get_device_matching(lambda d: d.type == 'local', timeout=0, **kwargs)
+
+    def get_remote_device(self, **kwargs):
+        return self.get_device_matching(lambda d: d.type == 'remote', timeout=0, **kwargs)
+
+    def get_usb_device(self, timeout=0, **kwargs):
+        return self.get_device_matching(lambda d: d.type == 'usb', timeout, **kwargs)
+
+    def get_device(self, id, timeout=0, **kwargs):
+        return self.get_device_matching(lambda d: d.id == id, timeout, **kwargs)
+
+    @cancellable
+    def get_device_matching(self, predicate, timeout=0):
+        if timeout < 0:
+            raw_timeout = -1
+        elif timeout == 0:
+            raw_timeout = 0
+        else:
+            raw_timeout = int(timeout * 1000.0)
+        return Device(self._impl.get_device_matching(lambda d: predicate(Device(d)), raw_timeout))
+
     @cancellable
     def enumerate_devices(self):
         return [Device(device) for device in self._impl.enumerate_devices()]
@@ -46,16 +68,6 @@ class DeviceManager(object):
     def remove_remote_device(self, host):
         self._impl.remove_remote_device(host)
 
-    @cancellable
-    def get_device(self, device_id):
-        devices = self._impl.enumerate_devices()
-        if device_id is None:
-            return Device(devices[0])
-        for device in devices:
-            if device.id == device_id:
-                return Device(device)
-        raise _frida.InvalidArgumentError("unable to find device with id %s" % device_id)
-
     def on(self, signal, callback):
         self._impl.on(signal, callback)
 

src/_frida.c

@@ -260,6 +260,8 @@ static PyObject * PyGObject_marshal_object (gpointer handle, GType type);
 static int PyDeviceManager_init (PyDeviceManager * self, PyObject * args, PyObject * kwds);
 static void PyDeviceManager_dealloc (PyDeviceManager * self);
 static PyObject * PyDeviceManager_close (PyDeviceManager * self);
+static PyObject * PyDeviceManager_get_device_matching (PyDeviceManager * self, PyObject * args);
+static gboolean PyDeviceManager_is_matching_device (FridaDevice * device, PyObject * predicate);
 static PyObject * PyDeviceManager_enumerate_devices (PyDeviceManager * self);
 static PyObject * PyDeviceManager_add_remote_device (PyDeviceManager * self, PyObject * args);
 static PyObject * PyDeviceManager_remove_remote_device (PyDeviceManager * self, PyObject * args);
@@ -376,6 +378,7 @@ static PyMethodDef PyGObject_methods[] =
 static PyMethodDef PyDeviceManager_methods[] =
 {
   { "close", (PyCFunction) PyDeviceManager_close, METH_NOARGS, "Close the device manager." },
+  { "get_device_matching", (PyCFunction) PyDeviceManager_get_device_matching, METH_VARARGS, "Get device matching predicate." },
   { "enumerate_devices", (PyCFunction) PyDeviceManager_enumerate_devices, METH_NOARGS, "Enumerate devices." },
   { "add_remote_device", (PyCFunction) PyDeviceManager_add_remote_device, METH_VARARGS, "Add a remote device." },
   { "remove_remote_device", (PyCFunction) PyDeviceManager_remove_remote_device, METH_VARARGS, "Remove a remote device." },
@@ -1830,6 +1833,66 @@ PyDeviceManager_close (PyDeviceManager * self)
   Py_RETURN_NONE;
 }
 
+static PyObject *
+PyDeviceManager_get_device_matching (PyDeviceManager * self, PyObject * args)
+{
+  PyObject * predicate;
+  gint timeout;
+  GError * error = NULL;
+  FridaDevice * result;
+
+  if (!PyArg_ParseTuple (args, "Oi", &predicate, &timeout))
+    return NULL;
+
+  if (!PyCallable_Check (predicate))
+    goto not_callable;
+
+  Py_BEGIN_ALLOW_THREADS
+  result = frida_device_manager_get_device_sync (PY_GOBJECT_HANDLE (self), (FridaDeviceManagerPredicate) PyDeviceManager_is_matching_device,
+      predicate, timeout, g_cancellable_get_current (), &error);
+  Py_END_ALLOW_THREADS
+  if (error != NULL)
+    return PyFrida_raise (error);
+
+  return PyDevice_new_take_handle (result);
+
+not_callable:
+  {
+    PyErr_SetString (PyExc_TypeError, "object must be callable");
+    return NULL;
+  }
+}
+
+static gboolean
+PyDeviceManager_is_matching_device (FridaDevice * device, PyObject * predicate)
+{
+  gboolean is_matching = FALSE;
+  PyGILState_STATE gstate;
+  PyObject * device_object, * result;
+
+  gstate = PyGILState_Ensure ();
+
+  device_object = PyDevice_new_take_handle (g_object_ref (device));
+
+  result = PyObject_CallFunction (predicate, "O", device_object);
+  if (result != NULL)
+  {
+    is_matching = result == Py_True;
+
+    Py_DECREF (result);
+  }
+  else
+  {
+    PyErr_Print ();
+  }
+
+  Py_DECREF (device_object);
+
+  PyGILState_Release (gstate);
+
+  return is_matching;
+}
+
 static PyObject *
 PyDeviceManager_enumerate_devices (PyDeviceManager * self)
 {

tests/test_core.py

@@ -1,5 +1,8 @@
 # -*- coding: utf-8 -*-
 
+import sys
+import threading
+import time
 try:
     import unittest2 as unittest
 except:
@@ -13,6 +16,40 @@ class TestCore(unittest.TestCase):
         devices = frida.get_device_manager().enumerate_devices()
         self.assertTrue(len(devices) > 0)
 
+    def test_get_existing_device(self):
+        device = frida.get_device_matching(lambda d: d.id == 'local')
+        self.assertEqual(device.name, "Local System")
+
+        device = frida.get_device_manager().get_device_matching(lambda d: d.id == 'local')
+        self.assertEqual(device.name, "Local System")
+
+    def test_get_nonexistent_device(self):
+        def get_nonexistent():
+            frida.get_device_manager().get_device_matching(lambda device: device.type == 'lol')
+        self.assertRaisesMatching(frida.InvalidArgumentError, "device not found", get_nonexistent)
+
+    def test_wait_for_nonexistent_device(self):
+        def wait_for_nonexistent():
+            frida.get_device_manager().get_device_matching(lambda device: device.type == 'lol', timeout=0.1)
+        self.assertRaisesMatching(frida.InvalidArgumentError, "device not found", wait_for_nonexistent)
+
+    def test_cancel_wait_for_nonexistent_device(self):
+        cancellable = frida.Cancellable()
+
+        def wait_for_nonexistent():
+            frida.get_device_manager().get_device_matching(lambda device: device.type == 'lol', timeout=-1, cancellable=cancellable)
+
+        def cancel_after_100ms():
+            time.sleep(0.1)
+            cancellable.cancel()
+
+        threading.Thread(target=cancel_after_100ms).start()
+        self.assertRaisesMatching(frida.OperationCancelledError, "operation was cancelled", wait_for_nonexistent)
+
+    def assertRaisesMatching(self, exception, regex, operation):
+        m = self.assertRaisesRegex if sys.version_info[0] >= 3 else self.assertRaisesRegexp
+        m(exception, regex, operation)
+
 
 if __name__ == '__main__':
     unittest.main()