From 09cd81752d1d53de4cc1c63f950dc717acc20ad8 Mon Sep 17 00:00:00 2001
From: Adeeb Shihadeh <adeebshihadeh@gmail.com>
Date: Thu, 24 Aug 2023 19:15:37 -0700
Subject: [PATCH] python lib: close USB context properly (#1606)

* python lib: close USB context properly

* fix jungle

* rm that

---------

Co-authored-by: Comma Device <device@comma.ai>
---
 board/jungle/__init__.py |  2 +-
 python/__init__.py       | 12 +++++++-----
 python/dfu.py            | 27 ++++++++++++++++++++-------
 3 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/board/jungle/__init__.py b/board/jungle/__init__.py
index 1c90fa00c..84be11932 100644
--- a/board/jungle/__init__.py
+++ b/board/jungle/__init__.py
@@ -50,7 +50,7 @@ class PandaJungle(Panda):
 
   @classmethod
   def spi_connect(cls, serial, ignore_version=False):
-    return None, None, None, None
+    return None, None, None, None, None
 
   def flash(self, fn=None, code=None, reconnect=True):
     if not fn:
diff --git a/python/__init__.py b/python/__init__.py
index aa9be43a7..56900d26b 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -280,6 +280,8 @@ class Panda:
     if self._handle_open:
       self._handle.close()
       self._handle_open = False
+      if self._context is not None:
+        self._context.close()
 
   def connect(self, claim=True, wait=False):
     self.close()
@@ -287,9 +289,9 @@ class Panda:
     self._handle = None
     while self._handle is None:
       # try USB first, then SPI
-      self._handle, serial, self.bootstub, bcd = self.usb_connect(self._connect_serial, claim=claim)
+      self._context, self._handle, serial, self.bootstub, bcd = self.usb_connect(self._connect_serial, claim=claim)
       if self._handle is None:
-        self._handle, serial, self.bootstub, bcd = self.spi_connect(self._connect_serial)
+        self._context, self._handle, serial, self.bootstub, bcd = self.spi_connect(self._connect_serial)
       if not wait:
         break
 
@@ -365,7 +367,7 @@ class Panda:
         err = f"panda protocol mismatch: expected {handle.PROTOCOL_VERSION}, got {spi_version}. reflash panda"
         raise PandaProtocolMismatch(err)
 
-    return handle, spi_serial, bootstub, None
+    return None, handle, spi_serial, bootstub, None
 
   @classmethod
   def usb_connect(cls, serial, claim=True):
@@ -407,7 +409,7 @@ class Panda:
     else:
       context.close()
 
-    return usb_handle, usb_serial, bootstub, bcd
+    return context, usb_handle, usb_serial, bootstub, bcd
 
   @classmethod
   def list(cls): # noqa: A003
@@ -436,7 +438,7 @@ class Panda:
 
   @classmethod
   def spi_list(cls):
-    _, serial, _, _ = cls.spi_connect(None, ignore_version=True)
+    _, _, serial, _, _ = cls.spi_connect(None, ignore_version=True)
     if serial is not None:
       return [serial, ]
     return []
diff --git a/python/dfu.py b/python/dfu.py
index 39303d528..fc196e4a8 100644
--- a/python/dfu.py
+++ b/python/dfu.py
@@ -14,9 +14,9 @@ class PandaDFU:
   def __init__(self, dfu_serial: Optional[str]):
     # try USB, then SPI
     handle: Optional[BaseSTBootloaderHandle]
-    handle = PandaDFU.usb_connect(dfu_serial)
+    self._context, handle = PandaDFU.usb_connect(dfu_serial)
     if handle is None:
-      handle = PandaDFU.spi_connect(dfu_serial)
+      self._context, handle = PandaDFU.spi_connect(dfu_serial)
 
     if handle is None:
       raise Exception(f"failed to open DFU device {dfu_serial}")
@@ -24,8 +24,21 @@ class PandaDFU:
     self._handle: BaseSTBootloaderHandle = handle
     self._mcu_type: McuType = self._handle.get_mcu_type()
 
+  def __enter__(self):
+    return self
+
+  def __exit__(self, *args):
+    self.close()
+
+  def close(self):
+    if self._handle is not None:
+      self._handle.close()
+      self._handle = None
+      if self._context is not None:
+        self._context.close()
+
   @staticmethod
-  def usb_connect(dfu_serial: Optional[str]) -> Optional[STBootloaderUSBHandle]:
+  def usb_connect(dfu_serial: Optional[str]):
     handle = None
     context = usb1.USBContext()
     context.open()
@@ -40,10 +53,10 @@ class PandaDFU:
           handle = STBootloaderUSBHandle(device, device.open())
           break
 
-    return handle
+    return context, handle
 
   @staticmethod
-  def spi_connect(dfu_serial: Optional[str]) -> Optional[STBootloaderSPIHandle]:
+  def spi_connect(dfu_serial: Optional[str]):
     handle = None
     this_dfu_serial = None
 
@@ -56,7 +69,7 @@ class PandaDFU:
     if dfu_serial is not None and dfu_serial != this_dfu_serial:
       handle = None
 
-    return handle
+    return None, handle
 
   @staticmethod
   def list() -> List[str]: # noqa: A003
@@ -82,7 +95,7 @@ class PandaDFU:
   @staticmethod
   def spi_list() -> List[str]:
     try:
-      h = PandaDFU.spi_connect(None)
+      _, h = PandaDFU.spi_connect(None)
       if h is not None:
         dfu_serial = PandaDFU.st_serial_to_dfu_serial(h.get_uid(), h.get_mcu_type())
         return [dfu_serial, ]