1 files changed, 8 insertions, 6 deletions

diff --git a/apps/plugins/demystify.c b/apps/plugins/demystify.c
index 6a8f1eb0c1..468f1806ce 100644
--- a/apps/plugins/demystify.c
+++ b/apps/plugins/demystify.c
@@ -22,12 +22,11 @@
 ****************************************************************************/
 
 #include "plugin.h"
-#include "lib/pluginlib_exit.h"
 
 #ifdef HAVE_LCD_BITMAP
 #include "lib/pluginlib_actions.h"
 #include "lib/helper.h"
-
+PLUGIN_HEADER
 
 #define DEFAULT_WAIT_TIME 3
 #define DEFAULT_NB_POLYGONS 7
@@ -260,8 +259,10 @@ void polygons_draw(struct polygon_fifo * polygons, struct screen * display)
     }
 }
 
-void cleanup(void)
+void cleanup(void *parameter)
 {
+    (void)parameter;
+
     backlight_use_settings();
 #ifdef HAVE_REMOTE_LCD
     remote_backlight_use_settings();
@@ -393,6 +394,7 @@ int plugin_main(void)
         switch(action)
         {
             case DEMYSTIFY_QUIT:
+                cleanup(NULL);
                 return PLUGIN_OK;
 
             case DEMYSTIFY_ADD_POLYGON:
@@ -419,7 +421,9 @@ int plugin_main(void)
                 break;
 
             default:
-                exit_on_usb(action);
+                if (rb->default_event_handler_ex(action, cleanup, NULL)
+                    == SYS_USB_CONNECTED)
+                    return PLUGIN_USB_CONNECTED;
                 break;
         }
     }
@@ -432,8 +436,6 @@ enum plugin_status plugin_start(const void* parameter)
     int ret;
 
     (void)parameter;
-    atexit(cleanup);
-
 #if LCD_DEPTH > 1
     rb->lcd_set_backdrop(NULL);
 #endif