From 70db081abfc1887766193852035d750490b4e2ae Mon Sep 17 00:00:00 2001
From: Simon Holesch <simon@holesch.de>
Date: Fri, 26 Jan 2024 16:01:50 +0100
Subject: [PATCH] USB/IP: Restore the default USB driver on exit

This solves the problem when you stop the exporter and want to access
the board locally, then it wasn't possible, because the USB device was
still bound to the usbip-host driver.

Fixed by loading the default USB driver on exit.
---
 not_my_board/_usbip.py | 33 ++++++++++++++++++++++++++++-----
 tests/test_usbip.py    |  6 ++++++
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/not_my_board/_usbip.py b/not_my_board/_usbip.py
index 8b7ccff..75e1f10 100644
--- a/not_my_board/_usbip.py
+++ b/not_my_board/_usbip.py
@@ -30,6 +30,7 @@ def __init__(self, devices):
     async def __aenter__(self):
         async with contextlib.AsyncExitStack() as stack:
             for _, device in self._devices.items():
+                stack.callback(device.restore_default_usb_driver)
                 await stack.enter_async_context(
                     util.background_task(_refresh_task(device))
                 )
@@ -176,6 +177,24 @@ async def _bind_usbip_host_driver(self):
         (usbip_host_driver / "match_busid").write_text(f"add {self._busid}")
         (usbip_host_driver / "bind").write_text(self._busid)
 
+    def restore_default_usb_driver(self):
+        driver_path = self._sysfs_path / "driver"
+        if driver_path.exists():
+            driver_name = driver_path.resolve().name
+            if driver_name == "usbip-host":
+                logger.info(
+                    'Unbinding USB device %s from driver "%s"', self._busid, driver_name
+                )
+                (driver_path / "unbind").write_text(self._busid)
+                self._bind_default_usb_driver()
+        elif self._sysfs_path.exists():
+            self._bind_default_usb_driver()
+
+    def _bind_default_usb_driver(self):
+        logger.info("Binding USB device %s to default driver", self._busid)
+        probe_path = pathlib.Path("/sys/bus/usb/drivers_probe")
+        probe_path.write_text(self._busid)
+
     @property
     def busid(self):
         return self._busid.encode("utf-8")
@@ -467,12 +486,16 @@ async def _refresh_task(device):
 
     tmp_path = pipe_path.with_name(pipe_path.name + ".new")
     os.mkfifo(tmp_path)
-    tmp_path.replace(pipe_path)
 
-    async with _open_read_pipe(pipe_path, "r+b", buffering=0) as pipe:
-        while True:
-            await pipe.read(4096)
-            device.refresh()
+    try:
+        tmp_path.replace(pipe_path)
+
+        async with _open_read_pipe(pipe_path, "r+b", buffering=0) as pipe:
+            while True:
+                await pipe.read(4096)
+                device.refresh()
+    finally:
+        tmp_path.unlink()
 
 
 @contextlib.asynccontextmanager
diff --git a/tests/test_usbip.py b/tests/test_usbip.py
index d041b90..84645cd 100644
--- a/tests/test_usbip.py
+++ b/tests/test_usbip.py
@@ -172,6 +172,12 @@ async def test_usb_forwarding(vms):
                 await vms.client.ssh("not-my-board detach qemu-usb")
                 await vms.client.ssh("! test -e /sys/bus/usb/devices/2-1")
 
+    # When the exporter is killed, then it should clean up and restore the
+    # default USB driver.
+    result = await vms.exporter.ssh("readlink /sys/bus/usb/devices/2-1/driver")
+    driver_name = pathlib.Path(result.stdout).name
+    assert driver_name == "usb"
+
 
 ShResult = collections.namedtuple("ShResult", ["stdout", "stderr", "returncode"])