Prevent plugins with active viewscreens from being unloaded

This requires plugins to pass plugin_self to Screen::show(), but
avoids the need to implement special checks in plugin_onstatechange
for the SC_BEGIN_UNLOAD event.

NEWS.rst

@@ -41,6 +41,8 @@ Internals
         ./dfhack +devel/print-args example
         "Dwarf Fortress.exe" +devel/print-args example
 
+- Prevented plugins with active viewscreens from being unloaded and causing a crash
+
 New Scripts
 -----------
 - `devel/save-version`: Displays DF version information about the current save

library/PluginManager.cpp

@@ -24,6 +24,7 @@ distribution.
 
 #include "modules/EventManager.h"
 #include "modules/Filesystem.h"
+#include "modules/Screen.h"
 #include "Internal.h"
 #include "Core.h"
 #include "MemAccess.h"
@@ -385,6 +386,12 @@ bool Plugin::unload(color_ostream &con)
     // if we are actually loaded
     if(state == PS_LOADED)
     {
+        if (Screen::hasActiveScreens(this))
+        {
+            con.printerr("Cannot unload plugin %s: has active viewscreens\n", name.c_str());
+            access->unlock();
+            return false;
+        }
         EventManager::unregisterAll(this);
         // notify the plugin about an attempt to shutdown
         if (plugin_onstatechange &&

library/include/modules/Screen.h

@@ -54,6 +54,7 @@ namespace df
 namespace DFHack
 {
     class Core;
+    class Plugin;
 
     typedef std::set<df::interface_key> interface_key_set;
 
@@ -208,9 +209,12 @@ namespace DFHack
         DFHACK_EXPORT bool findGraphicsTile(const std::string &page, int x, int y, int *ptile, int *pgs = NULL);
 
         // Push and remove viewscreens
-        DFHACK_EXPORT bool show(df::viewscreen *screen, df::viewscreen *before = NULL);
+        DFHACK_EXPORT bool show(df::viewscreen *screen, df::viewscreen *before = NULL, Plugin *p = NULL);
+        inline bool show(df::viewscreen *screen, Plugin *p)
+            { return show(screen, NULL, p); }
         DFHACK_EXPORT void dismiss(df::viewscreen *screen, bool to_first = false);
         DFHACK_EXPORT bool isDismissed(df::viewscreen *screen);
+        DFHACK_EXPORT bool hasActiveScreens(Plugin *p);
 
         /// Retrieve the string representation of the bound key.
         DFHACK_EXPORT std::string getKeyDisplay(df::interface_key key);

library/modules/Screen.cpp

@@ -282,7 +282,9 @@ bool Screen::findGraphicsTile(const std::string &pagename, int x, int y, int *pt
     return false;
 }
 
-bool Screen::show(df::viewscreen *screen, df::viewscreen *before)
+static std::map<df::viewscreen*, Plugin*> plugin_screens;
+
+bool Screen::show(df::viewscreen *screen, df::viewscreen *before, Plugin *plugin)
 {
     CHECK_NULL_POINTER(screen);
     CHECK_INVALID_ARGUMENT(!screen->parent && !screen->child);
@@ -306,6 +308,9 @@ bool Screen::show(df::viewscreen *screen, df::viewscreen *before)
     if (dfhack_viewscreen::is_instance(screen))
         static_cast<dfhack_viewscreen*>(screen)->onShow();
 
+    if (plugin)
+        plugin_screens[screen] = plugin;
+
     return true;
 }
 
@@ -313,6 +318,10 @@ void Screen::dismiss(df::viewscreen *screen, bool to_first)
 {
     CHECK_NULL_POINTER(screen);
 
+    auto it = plugin_screens.find(screen);
+    if (it != plugin_screens.end())
+        plugin_screens.erase(it);
+
     if (screen->breakdown_level != interface_breakdown_types::NONE)
         return;
 
@@ -332,6 +341,21 @@ bool Screen::isDismissed(df::viewscreen *screen)
     return screen->breakdown_level != interface_breakdown_types::NONE;
 }
 
+bool Screen::hasActiveScreens(Plugin *plugin)
+{
+    if (plugin_screens.empty())
+        return false;
+    df::viewscreen *screen = &gview->view;
+    while (screen)
+    {
+        auto it = plugin_screens.find(screen);
+        if (it != plugin_screens.end() && it->second == plugin)
+            return true;
+        screen = screen->child;
+    }
+    return false;
+}
+
 #ifdef _LINUX
 // Link to the libgraphics class directly:
 class DFHACK_EXPORT enabler_inputst {

plugins/autochop.cpp

@@ -600,7 +600,7 @@ struct autochop_hook : public df::viewscreen_dwarfmodest
         if (isInDesignationMenu() && input->count(interface_key::CUSTOM_C))
         {
             sendKey(interface_key::LEAVESCREEN);
-            Screen::show(new ViewscreenAutochop());
+            Screen::show(new ViewscreenAutochop(), plugin_self);
         }
         else
         {
@@ -643,7 +643,7 @@ command_result df_autochop (color_ostream &out, vector <string> & parameters)
             return CR_WRONG_USAGE;
     }
     if (Maps::IsValid())
-        Screen::show(new ViewscreenAutochop());
+        Screen::show(new ViewscreenAutochop(), plugin_self);
     return CR_OK;
 }
 

plugins/buildingplan.cpp

@@ -161,7 +161,7 @@ struct buildingplan_hook : public df::viewscreen_dwarfmodest
                 }
                 else if (input->count(interface_key::CUSTOM_SHIFT_M))
                 {
-                    Screen::show(new ViewscreenChooseMaterial(planner.getDefaultItemFilterForType(type)));
+                    Screen::show(new ViewscreenChooseMaterial(planner.getDefaultItemFilterForType(type)), plugin_self);
                 }
                 else if (input->count(interface_key::CUSTOM_SHIFT_Q))
                 {

plugins/command-prompt.cpp

@@ -318,7 +318,7 @@ command_result show_prompt(color_ostream &out, std::vector <std::string> & param
     std::string params;
     for(size_t i=0;i<parameters.size();i++)
         params+=parameters[i]+" ";
-    Screen::show(new viewscreen_commandpromptst(params));
+    Screen::show(new viewscreen_commandpromptst(params), plugin_self);
     return CR_OK;
 }
 bool hotkey_allow_all(df::viewscreen *top)
@@ -336,8 +336,6 @@ DFhackCExport command_result plugin_init ( color_ostream &out, std::vector <Plug
 
 DFhackCExport command_result plugin_onstatechange (color_ostream &out, state_change_event e)
 {
-    if (e == SC_BEGIN_UNLOAD && Gui::getCurFocus() == "dfhack/commandprompt")
-        return CR_FAILURE;
     return CR_OK;
 }
 

plugins/dwarfmonitor.cpp

@@ -1057,7 +1057,7 @@ public:
         {
             df::unit *selected_unit = (selected_column == 1) ? dwarf_activity_column.getFirstSelectedElem() : nullptr;
             Screen::dismiss(this);
-            Screen::show(new ViewscreenDwarfStats(selected_unit));
+            Screen::show(new ViewscreenDwarfStats(selected_unit), plugin_self);
         }
         else if  (input->count(interface_key::CUSTOM_SHIFT_Z))
         {
@@ -1665,7 +1665,7 @@ private:
 
 static void open_stats_srceen()
 {
-    Screen::show(new ViewscreenFortStats());
+    Screen::show(new ViewscreenFortStats(), plugin_self);
 }
 
 static void add_work_history(df::unit *unit, activity_type type)
@@ -1915,12 +1915,12 @@ static command_result dwarfmonitor_cmd(color_ostream &out, vector <string> & par
         else if (cmd == 's' || cmd == 'S')
         {
             if(Maps::IsValid())
-                Screen::show(new ViewscreenFortStats());
+                Screen::show(new ViewscreenFortStats(), plugin_self);
         }
         else if (cmd == 'p' || cmd == 'P')
         {
             if(Maps::IsValid())
-                Screen::show(new ViewscreenPreferences());
+                Screen::show(new ViewscreenPreferences(), plugin_self);
         }
         else if (cmd == 'r' || cmd == 'R')
         {

plugins/embark-tools.cpp

@@ -21,6 +21,9 @@ using namespace DFHack;
 using df::global::enabler;
 using df::global::gps;
 
+DFHACK_PLUGIN("embark-tools");
+DFHACK_PLUGIN_IS_ENABLED(is_enabled);
+
 #define FOR_ITER_TOOLS(iter) for(auto iter = tools.begin(); iter != tools.end(); iter++)
 
 void update_embark_sidebar (df::viewscreen_choose_start_sitest * screen)
@@ -684,7 +687,7 @@ struct choose_start_site_hook : df::viewscreen_choose_start_sitest
 
     void display_settings()
     {
-        Screen::show(new embark_tools_settings);
+        Screen::show(new embark_tools_settings, plugin_self);
     }
 
     inline bool is_valid_page()
@@ -734,9 +737,6 @@ struct choose_start_site_hook : df::viewscreen_choose_start_sitest
 IMPLEMENT_VMETHOD_INTERPOSE(choose_start_site_hook, feed);
 IMPLEMENT_VMETHOD_INTERPOSE(choose_start_site_hook, render);
 
-DFHACK_PLUGIN("embark-tools");
-DFHACK_PLUGIN_IS_ENABLED(is_enabled);
-
 command_result embark_tools_cmd (color_ostream &out, std::vector <std::string> & parameters);
 
 DFhackCExport command_result plugin_init (color_ostream &out, std::vector <PluginCommand> &commands)
@@ -783,14 +783,6 @@ DFhackCExport command_result plugin_enable (color_ostream &out, bool enable)
 
 DFhackCExport command_result plugin_onstatechange (color_ostream &out, state_change_event evt)
 {
-    if (evt == SC_BEGIN_UNLOAD)
-    {
-        if (Gui::getCurFocus() == "dfhack/embark-tools/options")
-        {
-            out.printerr("Settings screen active.\n");
-            return CR_FAILURE;
-        }
-    }
     return CR_OK;
 }
 

plugins/hotkeys.cpp

@@ -319,7 +319,7 @@ static command_result hotkeys_cmd(color_ostream &out, vector <string> & paramete
             if (Gui::getFocusString(top_screen) != "dfhack/viewscreen_hotkeys")
             {
                 find_active_keybindings(top_screen);
-                Screen::show(new ViewscreenHotkeys(top_screen));
+                Screen::show(new ViewscreenHotkeys(top_screen), plugin_self);
             }
         }
     }

plugins/manipulator.cpp

@@ -1773,14 +1773,14 @@ void viewscreen_unitlaborsst::feed(set<df::interface_key> *events)
 
     if (events->count(interface_key::CUSTOM_B))
     {
-        Screen::show(new viewscreen_unitbatchopst(units, true, &do_refresh_names));
+        Screen::show(new viewscreen_unitbatchopst(units, true, &do_refresh_names), plugin_self);
     }
 
     if (events->count(interface_key::CUSTOM_E))
     {
         vector<UnitInfo*> tmp;
         tmp.push_back(cur);
-        Screen::show(new viewscreen_unitbatchopst(tmp, false, &do_refresh_names));
+        Screen::show(new viewscreen_unitbatchopst(tmp, false, &do_refresh_names), plugin_self);
     }
 
     if (events->count(interface_key::CUSTOM_P))
@@ -1791,11 +1791,11 @@ void viewscreen_unitlaborsst::feed(set<df::interface_key> *events)
                 has_selected = true;
 
         if (has_selected) {
-            Screen::show(new viewscreen_unitprofessionset(units, true));
+            Screen::show(new viewscreen_unitprofessionset(units, true), plugin_self);
         } else {
             vector<UnitInfo*> tmp;
             tmp.push_back(cur);
-            Screen::show(new viewscreen_unitprofessionset(tmp, false));
+            Screen::show(new viewscreen_unitprofessionset(tmp, false), plugin_self);
         }
     }
 
@@ -2144,7 +2144,7 @@ struct unitlist_hook : df::viewscreen_unitlistst
         {
             if (units[page].size())
             {
-                Screen::show(new viewscreen_unitlaborsst(units[page], cursor_pos[page]));
+                Screen::show(new viewscreen_unitlaborsst(units[page], cursor_pos[page]), plugin_self);
                 return;
             }
         }

plugins/stocks.cpp

@@ -771,7 +771,7 @@ public:
         }
         else if (input->count(interface_key::HELP))
         {
-            Screen::show(new search_help);
+            Screen::show(new search_help, plugin_self);
         }
 
         bool key_processed = false;
@@ -1425,7 +1425,7 @@ struct stocks_hook : public df::viewscreen_storesst
         if (input->count(interface_key::CUSTOM_E))
         {
             Screen::dismiss(this);
-            Screen::show(new ViewscreenStocks());
+            Screen::show(new ViewscreenStocks(), plugin_self);
             return;
         }
         INTERPOSE_NEXT(feed)(input);
@@ -1457,7 +1457,7 @@ struct stocks_stockpile_hook : public df::viewscreen_dwarfmodest
 
         if (input->count(interface_key::CUSTOM_I))
         {
-            Screen::show(new ViewscreenStocks(sp));
+            Screen::show(new ViewscreenStocks(sp), plugin_self);
             return true;
         }
 
@@ -1531,7 +1531,7 @@ static command_result stocks_cmd(color_ostream &out, vector <string> & parameter
         }
         else if (toLower(parameters[0])[0] == 's')
         {
-            Screen::show(new ViewscreenStocks());
+            Screen::show(new ViewscreenStocks(), plugin_self);
             return CR_OK;
         }
     }
@@ -1557,10 +1557,6 @@ DFhackCExport command_result plugin_onstatechange(color_ostream &out, state_chan
     case SC_MAP_LOADED:
         ViewscreenStocks::reset();
         break;
-    case SC_BEGIN_UNLOAD:
-        if (Gui::getCurFocus().find("dfhack/stocks") == 0)
-            return CR_FAILURE;
-        break;
     default:
         break;
     }