From 746ef2888d57539bd93c87e32f5383cea88c671f Mon Sep 17 00:00:00 2001
From: Viktor Dragomiretskyy <Viktor.Dragomiretskyy@gmail.com>
Date: Tue, 7 Jul 2020 14:07:19 +1200
Subject: [PATCH] Properly handle shutdown sequence in both prototype and dbus
 service.

---
 dbus-service.py |  6 +-----
 proto9x/init.py | 19 ++++++++++++++++++-
 proto9x/usb.py  |  8 ++++++++
 prototype.py    | 15 ---------------
 4 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/dbus-service.py b/dbus-service.py
index b36d466..0e743cc 100644
--- a/dbus-service.py
+++ b/dbus-service.py
@@ -159,10 +159,6 @@ usb.trace_enabled = True
 tls.trace_enabled = True
 
 svc = Device()
-try:
-    loop.run()
-finally:
-    reboot()
-    raise
+loop.run()
 
 print("Normal exit")
diff --git a/proto9x/init.py b/proto9x/init.py
index 44138bc..abcbb22 100644
--- a/proto9x/init.py
+++ b/proto9x/init.py
@@ -1,16 +1,33 @@
 
+import atexit
+
 from proto9x.usb import usb
 from proto9x.tls import tls
 from proto9x.sensor import sensor
+from proto9x.sensor import reboot
 from proto9x.flash import read_tls_flash
-from usb import core as usb_core
 
+def close():
+    if usb.dev is not None:
+        # Send the reboot command before closing the device.
+        # Without it the sensor seems to keep creating new TLS sessions and eventually runs out of memory.
+        # The reboot command may fail if we're shutting down because the device is already gone.
+        try:
+            reboot()
+        finally:
+            usb.close()
 
 def open():
     usb.open()
     usb.send_init()
+    # We must register atexit only after we opened usb device, 
+    # so that our handler is called before pyusb's one and we can still talk to the device
+    atexit.register(close)
 
     tls.parseTlsFlash(read_tls_flash())
     tls.open()
 
     sensor.open()
+
+
+
diff --git a/proto9x/usb.py b/proto9x/usb.py
index 5876f9b..47d6206 100644
--- a/proto9x/usb.py
+++ b/proto9x/usb.py
@@ -39,6 +39,14 @@ class Usb():
         self.thread.daemon = True
         self.thread.start()
 
+    def close(self):
+        if self.dev is not None:
+            try:
+                self.dev.reset()
+                self.dev = None
+            finally:
+                self.thread.join()
+
     def usb_dev(self):
         return self.dev
 
diff --git a/prototype.py b/prototype.py
index 9b787d8..8fa0e08 100644
--- a/prototype.py
+++ b/prototype.py
@@ -1,4 +1,3 @@
-
 from proto9x.tls import tls
 from proto9x.usb import usb
 from proto9x.db import db
@@ -80,17 +79,3 @@ def enroll(sid, finger):
 
     print('Created a finger record with dbid %d' % recid)
 
-# can't use atexit as it conflicts with atexit installed by libusb
-class Blah:
-    def __init__(self):
-        self.tls=tls
-
-    def __del__(self):
-        if usb.dev is not None:
-            print('Rebooting device...')
-            try:
-                reboot()
-            except:
-                pass
-
-blah=Blah()