Fix scm function storage

source/CCustomOpcodeSystem.cpp

@@ -618,10 +618,17 @@ namespace CLEO
     {
         auto cs = reinterpret_cast<CCustomScript*>(thread);
 
+        if (cs->GetScmFunction() == ScmFunction::Id_None)
+        {
+            SUSPEND("[%04X] used without preceding [0AB1]", opcode);
+        }
+
         ScmFunction* scmFunc = ScmFunction::Get(cs->GetScmFunction());
-        if (scmFunc == nullptr)
+        if (scmFunc == nullptr || scmFunc->caller != cs)
         {
-            SUSPEND("Invalid Cleo Call reference. [%04X] possibly used without preceding [0AB1]", opcode);
+            thread->ScmFunction = ScmFunction::Id_None; // clear invalid reference
+            thread->BaseIP      = nullptr;              // might be somebody's else, do not release during cleanup
+            SUSPEND("Cleo function call stack corruption detected");
         }
 
         // store return arguments
@@ -675,7 +682,6 @@ namespace CLEO
         // handle program flow
         scmFunc->Return(cs); // jump back to cleo_call, right after last input
                              // param. Return slot var args starts here
-        delete scmFunc;
 
         if (returnArgs)
         {
@@ -1069,7 +1075,7 @@ namespace CLEO
         static SCRIPT_VAR arguments[32];
         SCRIPT_VAR* locals       = thread->IsMission() ? missionLocals : thread->GetVarPtr();
         SCRIPT_VAR* localsEnd    = locals + 32;
-        SCRIPT_VAR* storedLocals = scmFunc->savedTls;
+        SCRIPT_VAR* storedLocals = &scmFunc->savedLocalVar[0];
 
         // collect arguments
         for (DWORD i = 0; i < nParams; i++)

source/CCustomScript.cpp

@@ -3,6 +3,7 @@
 #include "CleoBase.h"
 #include "crc32.h"
 #include "CScriptEngine.h"
+#include "ScmFunction.h"
 #include "ScriptUtils.h"
 
 using namespace CLEO;
@@ -173,9 +174,16 @@ CCustomScript::CCustomScript(const char* szFileName, bool bIsMiss, CRunningScrip
 
 CCustomScript::~CCustomScript()
 {
-    if (BaseIP) delete[] BaseIP;
+    // release entire call stack and retrieve original BaseIP
+    while (auto fun = ScmFunction::Get(this->ScmFunction))
+    {
+        fun->Return(this);
+    }
+
     CleoInstance.OpcodeSystem.scriptDeleteDelegate(this);
 
+    if (BaseIP && !m_parentScript) delete[] BaseIP;
+
     if (CleoInstance.ScriptEngine.LastScriptCreated == this) CleoInstance.ScriptEngine.LastScriptCreated = nullptr;
 }
 

source/ScmFunction.cpp

@@ -6,52 +6,80 @@
 namespace CLEO
 {
     ScmFunction* ScmFunction::store[Store_Size] = {0};
-    size_t ScmFunction::allocationPlace         = 0;
+    WORD ScmFunction::lastAllocIdx              = 0;
 
-    ScmFunction* ScmFunction::Get(unsigned short idx)
+    ScmFunction* ScmFunction::Get(WORD id)
     {
-        if (idx >= Store_Size) return nullptr;
+        if (id == Id_None) return nullptr;
+
+        auto idx = id - 1; // id to index
+
+        if (idx >= Store_Size || store[idx] == nullptr)
+        {
+            SHOW_ERROR("CLEO function with id %d not found in the storage!", id);
+            return nullptr;
+        }
 
         return store[idx];
     }
 
     void ScmFunction::Clear()
     {
-        for each (ScmFunction* scmFunc in store)
+        for (auto& scmFunc : store)
         {
-            if (scmFunc != nullptr) delete scmFunc;
+            if (scmFunc != nullptr)
+            {
+                ::operator delete(scmFunc); // delete without extra checks
+                scmFunc = nullptr;
+            }
         }
-        ScmFunction::allocationPlace = 0;
+        lastAllocIdx = 0;
     }
 
     void* ScmFunction::operator new(size_t size)
     {
-        size_t start_search = allocationPlace;
-        while (store[allocationPlace] != nullptr) // find first unused position in store
+        auto searchStart = lastAllocIdx;
+        while (store[lastAllocIdx] != nullptr) // find first unused position in store
         {
-            if (++allocationPlace >= Store_Size) allocationPlace = 0; // end of store reached
-            if (allocationPlace == start_search)
+            lastAllocIdx++;
+            if (lastAllocIdx >= Store_Size) lastAllocIdx = 0; // warp around
+
+            if (lastAllocIdx == searchStart)
             {
                 SHOW_ERROR("CLEO function storage stack overflow!");
                 throw std::bad_alloc(); // the store is filled up
             }
         }
-        ScmFunction* obj       = reinterpret_cast<ScmFunction*>(::operator new(size));
-        store[allocationPlace] = obj;
-        return obj;
+
+        auto func           = reinterpret_cast<ScmFunction*>(::operator new(size));
+        store[lastAllocIdx] = func;
+        return func;
     }
 
     void ScmFunction::operator delete(void* mem)
     {
-        store[reinterpret_cast<ScmFunction*>(mem)->thisScmFunctionId] = nullptr;
+        auto idx = reinterpret_cast<ScmFunction*>(mem)->thisScmFunctionId - 1; // id to index
+
+        if (idx >= Store_Size || store[idx] != mem)
+        {
+            SHOW_ERROR("Failed to delete corrupted CLEO function from storage!");
+            return; // keep allocated
+        }
+
+        store[idx] = nullptr;
         ::operator delete(mem);
     }
 
     ScmFunction::ScmFunction(CLEO::CRunningScript* thread, BYTE* callIP)
-        : prevScmFunctionId(reinterpret_cast<CCustomScript*>(thread)->GetScmFunction()), callIP(callIP)
     {
         auto cs = reinterpret_cast<CCustomScript*>(thread);
 
+        thisScmFunctionId = ScmFunction::lastAllocIdx + 1; // index to id
+        prevScmFunctionId = cs->GetScmFunction();
+
+        caller       = cs;
+        this->callIP = callIP;
+
         // create snapshot of current scope
         savedBaseIP   = cs->BaseIP;
         savedCodeSize = cs->IsCustom() ? cs->GetCodeSize() : -1;
@@ -60,7 +88,7 @@ namespace CLEO
         savedSP = cs->SP;
 
         auto scope = cs->IsMission() ? missionLocals : cs->LocalVar;
-        std::copy(scope, scope + _countof(CRunningScript::LocalVar), savedTls);
+        std::copy(scope, scope + _countof(CRunningScript::LocalVar), savedLocalVar.begin());
 
         savedCondResult = cs->bCondResult;
         savedLogicalOp  = cs->LogicalOp;
@@ -76,14 +104,7 @@ namespace CLEO
         cs->LogicalOp   = eLogicalOperation::NONE;
         cs->NotFlag     = false;
 
-        cs->SetScmFunction(thisScmFunctionId = (unsigned short)allocationPlace);
-    }
-
-    size_t ScmFunction::GetCallStackSize() const
-    {
-        if (prevScmFunctionId == 0) return 0;
-        auto parent = Get(prevScmFunctionId);
-        return parent ? parent->GetCallStackSize() + 1 : 1;
+        cs->SetScmFunction(thisScmFunctionId);
     }
 
     void ScmFunction::Return(CRunningScript* thread)
@@ -94,11 +115,11 @@ namespace CLEO
         if (cs->IsCustom()) cs->SetCodeSize(savedCodeSize);
 
         // restore parent scope's gosub call stack
-        std::copy(std::begin(savedStack), std::end(savedStack), std::begin(cs->Stack));
+        std::copy(savedStack.begin(), savedStack.end(), std::begin(cs->Stack));
         cs->SP = savedSP;
 
         // restore parent scope's local variables
-        std::copy(savedTls, savedTls + 32, cs->IsMission() ? missionLocals : cs->LocalVar);
+        std::copy(savedLocalVar.begin(), savedLocalVar.end(), cs->IsMission() ? missionLocals : cs->LocalVar);
 
         // process conditional result of just ended function in parent scope
         bool condResult = cs->bCondResult;
@@ -125,5 +146,13 @@ namespace CLEO
 
         cs->SetIp(retnAddress);
         cs->SetScmFunction(prevScmFunctionId);
+
+        delete this; // remove from storage
+    }
+
+    size_t ScmFunction::GetCallStackSize() const
+    {
+        auto parent = Get(prevScmFunctionId);
+        return parent ? 1 + parent->GetCallStackSize() : 0;
     }
 }; // namespace CLEO

source/ScmFunction.h

@@ -4,38 +4,44 @@
 
 namespace CLEO
 {
-    struct ScmFunction
+    class ScmFunction
     {
-        static const size_t Store_Size = 0x400;
-        static ScmFunction* store[Store_Size];
-        static size_t allocationPlace; // contains an index of last allocated object
-        static ScmFunction* Get(unsigned short idx);
-        static void Clear();
+      public:
+        static const WORD Id_None = 0;
+        static ScmFunction* Get(WORD id);
+        static void Clear(); // clear entire storage
+
+        void* operator new(size_t size);
+        void operator delete(void* mem);
+
+        ScmFunction(CRunningScript* thread, BYTE* callIP); // create
+        void Return(CRunningScript* thread);               // remove
+        size_t GetCallStackSize() const;
 
-        unsigned short prevScmFunctionId, thisScmFunctionId;
-        BYTE callArgCount = 0; // args provided to cleo_call
-        BYTE* callIP;          // address of 0AB1 for error messages
+        CRunningScript* caller = nullptr;
+        BYTE callArgCount      = 0;       // args provided to cleo_call
+        BYTE* callIP           = nullptr; // address of 0AB1 for error messages
 
         // saved nesting context state
-        void* savedBaseIP;
-        size_t savedCodeSize; // Custom Scripts only
-        BYTE* retnAddress;
-        BYTE* savedStack[8]; // gosub stack
-        WORD savedSP;
-        SCRIPT_VAR savedTls[32];
+        void* savedBaseIP                                                        = nullptr;
+        size_t savedCodeSize                                                     = 0; // Custom Scripts only
+        BYTE* retnAddress                                                        = nullptr;
+        std::array<BYTE*, _countof(CRunningScript::Stack)> savedStack            = {}; // gosub stack
+        WORD savedSP                                                             = 0;
+        std::array<SCRIPT_VAR, _countof(CRunningScript::LocalVar)> savedLocalVar = {};
         std::list<std::string> stringParams; // texts with this scope lifetime
-        bool savedCondResult;
-        eLogicalOperation savedLogicalOp;
-        bool savedNotFlag;
+        bool savedCondResult             = false;
+        eLogicalOperation savedLogicalOp = {};
+        bool savedNotFlag                = false;
         std::string savedScriptFileDir;  // modules switching
         std::string savedScriptFileName; // modules switching
 
-        void* operator new(size_t size);
-        void operator delete(void* mem);
-        ScmFunction(CRunningScript* thread, BYTE* callIP);
-
-        size_t GetCallStackSize() const;
+      private:
+        static const WORD Store_Size = 0x400;
+        static ScmFunction* store[Store_Size];
+        static WORD lastAllocIdx;
 
-        void Return(CRunningScript* thread);
+        WORD prevScmFunctionId = Id_None;
+        WORD thisScmFunctionId = Id_None;
     };
 } // namespace CLEO