From 067d71f7eb591bc6ac77f11844360e9a02799c40 Mon Sep 17 00:00:00 2001
From: Petr Mikheev <ptmikheev@gmail.com>
Date: Sun, 23 Jan 2022 22:33:00 +0100
Subject: Fix heap use after free in components/lua/storage.cpp

---
 apps/openmw_test_suite/lua/test_storage.cpp |  2 +-
 components/lua/storage.cpp                  | 19 +++++++++++--------
 components/lua/storage.hpp                  |  8 ++++----
 3 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/apps/openmw_test_suite/lua/test_storage.cpp b/apps/openmw_test_suite/lua/test_storage.cpp
index fafba008a1..a1c6af6e0a 100644
--- a/apps/openmw_test_suite/lua/test_storage.cpp
+++ b/apps/openmw_test_suite/lua/test_storage.cpp
@@ -91,10 +91,10 @@ namespace
         mLua.safe_script("permanent:set('z', 4)");
 
         LuaUtil::LuaStorage storage2(mLua);
+        storage2.load(tmpFile);
         mLua["permanent"] = storage2.getMutableSection("permanent");
         mLua["temporary"] = storage2.getMutableSection("temporary");
 
-        storage2.load(tmpFile);
         EXPECT_EQ(get<int>(mLua, "permanent:get('x')"), 1);
         EXPECT_TRUE(get<bool>(mLua, "permanent:get('z') == nil"));
         EXPECT_TRUE(get<bool>(mLua, "temporary:get('y') == nil"));
diff --git a/components/lua/storage.cpp b/components/lua/storage.cpp
index def38fdd67..91a339cb03 100644
--- a/components/lua/storage.cpp
+++ b/components/lua/storage.cpp
@@ -121,7 +121,10 @@ namespace LuaUtil
         while (it != mData.end())
         {
             if (!it->second->mPermanent)
+            {
+                it->second->mValues.clear();
                 it = mData.erase(it);
+            }
             else
                 ++it;
         }
@@ -129,7 +132,7 @@ namespace LuaUtil
 
     void LuaStorage::load(const std::string& path)
     {
-        mData.clear();
+        assert(mData.empty());  // Shouldn't be used before loading
         try
         {
             Log(Debug::Info) << "Loading Lua storage \"" << path << "\" (" << std::filesystem::file_size(path) << " bytes)";
@@ -138,7 +141,7 @@ namespace LuaUtil
             sol::table data = deserialize(mLua, serializedData);
             for (const auto& [sectionName, sectionTable] : data)
             {
-                Section* section = getSection(sectionName.as<std::string_view>());
+                const std::shared_ptr<Section>& section = getSection(sectionName.as<std::string_view>());
                 for (const auto& [key, value] : sol::table(sectionTable))
                     section->set(key.as<std::string_view>(), value);
             }
@@ -164,26 +167,26 @@ namespace LuaUtil
         fout.close();
     }
 
-    LuaStorage::Section* LuaStorage::getSection(std::string_view sectionName)
+    const std::shared_ptr<LuaStorage::Section>& LuaStorage::getSection(std::string_view sectionName)
     {
         auto it = mData.find(sectionName);
         if (it != mData.end())
-            return it->second.get();
-        auto section = std::make_unique<Section>(this, std::string(sectionName));
+            return it->second;
+        auto section = std::make_shared<Section>(this, std::string(sectionName));
         sectionName = section->mSectionName;
         auto [newIt, _] = mData.emplace(sectionName, std::move(section));
-        return newIt->second.get();
+        return newIt->second;
     }
 
     sol::object LuaStorage::getReadOnlySection(std::string_view sectionName)
     {
-        Section* section = getSection(sectionName);
+        const std::shared_ptr<Section>& section = getSection(sectionName);
         return sol::make_object<SectionReadOnlyView>(mLua, SectionReadOnlyView{section, section->mChangeCounter});
     }
 
     sol::object LuaStorage::getMutableSection(std::string_view sectionName)
     {
-        Section* section = getSection(sectionName);
+        const std::shared_ptr<Section>& section = getSection(sectionName);
         return sol::make_object<SectionMutableView>(mLua, SectionMutableView{section, section->mChangeCounter});
     }
 
diff --git a/components/lua/storage.hpp b/components/lua/storage.hpp
index e847604aeb..c23e417f0f 100644
--- a/components/lua/storage.hpp
+++ b/components/lua/storage.hpp
@@ -60,19 +60,19 @@ namespace LuaUtil
         };
         struct SectionMutableView
         {
-            Section* mSection = nullptr;
+            std::shared_ptr<Section> mSection = nullptr;
             int64_t mLastCheck = 0;
         };
         struct SectionReadOnlyView
         {
-            Section* mSection = nullptr;
+            std::shared_ptr<Section> mSection = nullptr;
             int64_t mLastCheck = 0;
         };
 
-        Section* getSection(std::string_view sectionName);
+        const std::shared_ptr<Section>& getSection(std::string_view sectionName);
 
         lua_State* mLua;
-        std::map<std::string_view, std::unique_ptr<Section>> mData;
+        std::map<std::string_view, std::shared_ptr<Section>> mData;
         std::optional<ListenerFn> mListener;
     };
 
-- 
cgit v1.2.3