diff --git a/proton/vpn/backend/networkmanager/killswitch/default/nmclient.py b/proton/vpn/backend/networkmanager/killswitch/default/nmclient.py index ac7bb28..33a6d68 100644 --- a/proton/vpn/backend/networkmanager/killswitch/default/nmclient.py +++ b/proton/vpn/backend/networkmanager/killswitch/default/nmclient.py @@ -28,6 +28,7 @@ from packaging.version import Version import gi + gi.require_version("NM", "1.0") from gi.repository import NM, GLib, Gio, GObject # pylint: disable=C0413 # noqa: E402 @@ -130,7 +131,7 @@ def __init__(self): self.initialize_nm_client_singleton() def add_connection_async( - self, connection: NM.Connection, save_to_disk: bool = False + self, connection: NM.Connection, save_to_disk: bool = False ) -> Future: """ Adds a new connection asynchronously. @@ -228,11 +229,21 @@ def _on_connection_removed(connection, result, _user_data): f"Error removing KS connection: {connection=}, {result=}" ).with_traceback(exc.__traceback__) ) + # If the connection to remove has no live device bound to it will not emit a + # "device-removed" signal. Removal is completed as the connection is deleted, thus we + # set the result here to avoid hanging. + else: + interface_name: str = connection.get_interface_name() + if self._has_named_device(interface_name) or future_interface_removed.done(): + return + future_interface_removed.set_result(None) def _on_interface_removed(_nm_client, device): logger.debug( f"{device.get_iface()} was removed." ) + if future_interface_removed.done(): + return if device.get_iface() == connection.get_interface_name(): future_interface_removed.set_result(None) @@ -254,12 +265,28 @@ def _remove_connection_async(): return future_interface_removed + def _has_named_device(self, interface_name: str) -> bool: + """ + Returns whether NetworkManager currently has a device for the given + interface name. + + This must be called from the GLib main loop thread (e.g. from within a + NetworkManager async callback). + """ + for device in self._nm_client.get_devices(): + if device.get_iface() == interface_name: + break + else: + return False + return True + def get_active_connection(self, conn_id: str) -> Optional[NM.ActiveConnection]: """ Returns the specified active connection, if existing. :param conn_id: ID of the active connection. :return: the active connection if it was found. Otherwise, None. """ + def _get_active_connection(): active_connections = self._nm_client.get_active_connections() diff --git a/proton/vpn/backend/networkmanager/killswitch/wireguard/nmclient.py b/proton/vpn/backend/networkmanager/killswitch/wireguard/nmclient.py index 24ed342..ad338f3 100644 --- a/proton/vpn/backend/networkmanager/killswitch/wireguard/nmclient.py +++ b/proton/vpn/backend/networkmanager/killswitch/wireguard/nmclient.py @@ -28,6 +28,7 @@ from packaging.version import Version import gi + gi.require_version("NM", "1.0") from gi.repository import NM, GLib, Gio, GObject # pylint: disable=C0413 # noqa: E402 @@ -218,9 +219,9 @@ def get_physical_devices(self) -> List[NM.Device]: """Returns all the active ethernet/wifi devices.""" return [ device for device in self._nm_client.get_devices() if ( - device.get_device_type() in (NM.DeviceType.ETHERNET, NM.DeviceType.WIFI) - and device.get_state() is NM.DeviceState.ACTIVATED - and device.get_active_connection() # Maybe this is redundant. + device.get_device_type() in (NM.DeviceType.ETHERNET, NM.DeviceType.WIFI) + and device.get_state() is NM.DeviceState.ACTIVATED + and device.get_active_connection() # Maybe this is redundant. ) ] @@ -431,12 +432,23 @@ def _on_connection_removed(connection, result, _user_data): f"Error removing KS connection: {exc}" ).with_traceback(exc.__traceback__) ) + # If the connection to remove has no live device bound to it will not emit a + # "device-removed" signal. Removal is completed as the connection is deleted, thus we + # set the result here to avoid hanging. + else: + interface_name: str = connection.get_interface_name() + if self._has_named_device(interface_name) or future_interface_removed.done(): + return + future_interface_removed.set_result(None) def _on_interface_removed(_nm_client, device): logger.debug( f"{device.get_iface()} was removed." ) - if device.get_iface() == connection.get_interface_name(): + if ( + device.get_iface() == connection.get_interface_name() + and not future_interface_removed.done() + ): future_interface_removed.set_result(None) def _remove_connection_async(): @@ -457,12 +469,28 @@ def _remove_connection_async(): return future_interface_removed + def _has_named_device(self, interface_name: str) -> bool: + """ + Returns whether NetworkManager currently has a device for the given + interface name. + + This must be called from the GLib main loop thread (e.g. from within a + NetworkManager async callback). + """ + for device in self._nm_client.get_devices(): + if device.get_iface() == interface_name: + break + else: + return False + return True + def get_active_connection(self, conn_id: str) -> Optional[NM.ActiveConnection]: """ Returns the specified active connection, if existing. :param conn_id: ID of the active connection. :return: the active connection if it was found. Otherwise, None. """ + def _get_active_connection(): active_connections = self._nm_client.get_active_connections() diff --git a/tests/networkmanager/killswitch/default/test_nmclient.py b/tests/networkmanager/killswitch/default/test_nmclient.py new file mode 100644 index 0000000..dce8bbc --- /dev/null +++ b/tests/networkmanager/killswitch/default/test_nmclient.py @@ -0,0 +1,106 @@ +""" +Copyright (c) 2023 Proton AG + +This file is part of Proton VPN. + +Proton VPN is free software: you can redistribute it and/or modify +it under the terms of the GNU General Public License as published by +the Free Software Foundation, either version 3 of the License, or +(at your option) any later version. + +Proton VPN is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU General Public License for more details. + +You should have received a copy of the GNU General Public License +along with ProtonVPN. If not, see . +""" +from concurrent.futures import Future +from unittest.mock import Mock, patch + +import pytest + +from proton.vpn.backend.networkmanager.killswitch.default.nmclient import NMClient + +INTERFACE_NAME = "pvpnksintrf0" + + +def _run_synchronously(function, *args, **kwargs): + """Drop-in replacement for NMClient._run_on_glib_loop_thread that runs the + given function synchronously instead of scheduling it on the GLib loop.""" + future = Future() + future.set_running_or_notify_cancel() + try: + future.set_result(function(*args, **kwargs)) + except BaseException as exc: # pylint: disable=broad-except + future.set_exception(exc) + return future + + +@pytest.fixture +def nm_client(): + with patch.object(NMClient, "initialize_nm_client_singleton"), \ + patch.object(NMClient, "_run_on_glib_loop_thread", side_effect=_run_synchronously), \ + patch("proton.vpn.backend.networkmanager.killswitch.default.nmclient.GObject"): + client = NMClient() + client._nm_client = Mock() + yield client + + +def _build_connection(): + connection = Mock() + connection.get_interface_name.return_value = INTERFACE_NAME + connection.delete_finish.return_value = True + + # Simulate NetworkManager finishing the deletion synchronously by invoking + # the callback as soon as delete_async is called. + def delete_async(cancellable, callback, user_data): # noqa: ARG001 + callback(connection, Mock(), user_data) + + connection.delete_async.side_effect = delete_async + return connection + + +def test_remove_connection_resolves_immediately_when_there_is_no_interface(nm_client): + """When the connection being removed has no active interface, no + device-removed signal will ever be emitted, so the removal must be + considered complete as soon as the connection is deleted.""" + nm_client._nm_client.get_devices.return_value = [] # no devices at all + connection = _build_connection() + + future = nm_client.remove_connection_async(connection) + + assert future.done() + assert future.result() is None + + +def test_remove_connection_waits_for_device_removed_when_interface_is_active(nm_client): + """When the connection being removed has an active interface, the removal + is only complete once the corresponding device has actually been removed.""" + device = Mock() + device.get_iface.return_value = INTERFACE_NAME + nm_client._nm_client.get_devices.return_value = [device] + + # Capture the handler connected to the "device-removed" signal. + handlers = {} + + def connect(signal_name, handler): + handlers[signal_name] = handler + return 1 # handler id + + nm_client._nm_client.connect.side_effect = connect + + connection = _build_connection() + + future = nm_client.remove_connection_async(connection) + + # The connection was deleted, but the device is still present, so the + # future must not be resolved until the device is actually removed. + assert not future.done() + + # Simulate NetworkManager emitting the device-removed signal. + handlers["device-removed"](nm_client._nm_client, device) + + assert future.done() + assert future.result() is None diff --git a/tests/networkmanager/killswitch/wireguard/test_nmclient.py b/tests/networkmanager/killswitch/wireguard/test_nmclient.py new file mode 100644 index 0000000..17c9c68 --- /dev/null +++ b/tests/networkmanager/killswitch/wireguard/test_nmclient.py @@ -0,0 +1,106 @@ +""" +Copyright (c) 2023 Proton AG + +This file is part of Proton VPN. + +Proton VPN is free software: you can redistribute it and/or modify +it under the terms of the GNU General Public License as published by +the Free Software Foundation, either version 3 of the License, or +(at your option) any later version. + +Proton VPN is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU General Public License for more details. + +You should have received a copy of the GNU General Public License +along with ProtonVPN. If not, see . +""" +from concurrent.futures import Future +from unittest.mock import Mock, patch + +import pytest + +from proton.vpn.backend.networkmanager.killswitch.wireguard.nmclient import NMClient + +INTERFACE_NAME = "pvpnksintrf0" + + +def _run_synchronously(function, *args, **kwargs): + """Drop-in replacement for NMClient._run_on_glib_loop_thread that runs the + given function synchronously instead of scheduling it on the GLib loop.""" + future = Future() + future.set_running_or_notify_cancel() + try: + future.set_result(function(*args, **kwargs)) + except BaseException as exc: # pylint: disable=broad-except + future.set_exception(exc) + return future + + +@pytest.fixture +def nm_client(): + with patch.object(NMClient, "initialize_nm_client_singleton"), \ + patch.object(NMClient, "_run_on_glib_loop_thread", side_effect=_run_synchronously), \ + patch("proton.vpn.backend.networkmanager.killswitch.wireguard.nmclient.GObject"): + client = NMClient() + client._nm_client = Mock() + yield client + + +def _build_connection(): + connection = Mock() + connection.get_interface_name.return_value = INTERFACE_NAME + connection.delete_finish.return_value = True + + # Simulate NetworkManager finishing the deletion synchronously by invoking + # the callback as soon as delete_async is called. + def delete_async(cancellable, callback, user_data): # noqa: ARG001 + callback(connection, Mock(), user_data) + + connection.delete_async.side_effect = delete_async + return connection + + +def test_remove_connection_resolves_immediately_when_there_is_no_interface(nm_client): + """When the connection being removed has no active interface, no + device-removed signal will ever be emitted, so the removal must be + considered complete as soon as the connection is deleted.""" + nm_client._nm_client.get_devices.return_value = [] # no devices at all + connection = _build_connection() + + future = nm_client.remove_connection_async(connection) + + assert future.done() + assert future.result() is None + + +def test_remove_connection_waits_for_device_removed_when_interface_is_active(nm_client): + """When the connection being removed has an active interface, the removal + is only complete once the corresponding device has actually been removed.""" + device = Mock() + device.get_iface.return_value = INTERFACE_NAME + nm_client._nm_client.get_devices.return_value = [device] + + # Capture the handler connected to the "device-removed" signal. + handlers = {} + + def connect(signal_name, handler): + handlers[signal_name] = handler + return 1 # handler id + + nm_client._nm_client.connect.side_effect = connect + + connection = _build_connection() + + future = nm_client.remove_connection_async(connection) + + # The connection was deleted, but the device is still present, so the + # future must not be resolved until the device is actually removed. + assert not future.done() + + # Simulate NetworkManager emitting the device-removed signal. + handlers["device-removed"](nm_client._nm_client, device) + + assert future.done() + assert future.result() is None