feat(papyrus-vm): add stack overflow protection (#2077)

misc/tests/test_dlc1chauruscocoonscript_stack_overflow.js

@@ -0,0 +1,23 @@
+const assert = require("node:assert");
+
+const main = async () => {
+  // dlc1chauruscocoonscript infinite recursion next tick after calling getFormEx
+  mp.callPapyrusFunction("global", "Game", "getFormEx", null, [33565538]);
+
+  await new Promise((resolve) => setTimeout(resolve, 500));
+};
+
+main().then(() => {
+  console.log("Test passed!");
+  process.exit(0);
+}).catch((err) => {
+  console.log("Test failed!")
+  console.error(err);
+  process.exit(1);
+});
+
+// Will output:
+// [error] ActivePexInstance::StartFunction - Stack overflow in script DLC1ChaurusCocoonScript, returning None
+// This is a sign that server isn't gonna crash
+
+// Triggers only after 2-nd start of the server (needs world to be initially here)

papyrus-vm/include/papyrus-vm/ActivePexInstance.h

@@ -15,6 +15,7 @@
 #include "VarValue.h"
 
 class VirtualMachine;
+class StackData;
 
 class ActivePexInstance
 {
@@ -39,7 +40,7 @@ class ActivePexInstance
 
   VarValue StartFunction(FunctionInfo& function,
                          std::vector<VarValue>& arguments,
-                         std::shared_ptr<StackIdHolder> stackIdHolder);
+                         std::shared_ptr<StackData> stackData);
 
   static uint8_t GetTypeByName(std::string typeRef);
   std::string GetActiveStateName() const;

papyrus-vm/include/papyrus-vm/VarValue.h

@@ -85,7 +85,7 @@ struct VarValue
   std::shared_ptr<std::string> stringHolder;
 
   int32_t GetMetaStackId() const;
-  void SetMetaStackIdHolder(std::shared_ptr<StackIdHolder> stackIdHolder);
+  void SetMetaStackIdHolder(const StackIdHolder& stackIdHolder);
   static VarValue AttachTestStackId(VarValue original = VarValue::None(),
                                     int32_t stackId = 108);
 

papyrus-vm/include/papyrus-vm/VirtualMachine.h

@@ -26,6 +26,27 @@ class StackIdHolder
   std::shared_ptr<MakeID> makeId;
 };
 
+class StackDepthHolder
+{
+public:
+  StackDepthHolder();
+
+  size_t GetStackDepth() const;
+  void IncreaseStackDepth();
+  void DecreaseStackDepth();
+
+  StackDepthHolder& operator=(const StackDepthHolder&) = delete;
+  StackDepthHolder(const StackDepthHolder&) = delete;
+
+private:
+  size_t depth = 0;
+};
+
+struct StackData
+{
+  StackIdHolder stackIdHolder;
+};
+
 struct VmExceptionInfo
 {
   std::string what;
@@ -37,7 +58,7 @@ class VirtualMachine
   friend class StackIdHolder;
 
 public:
-  using OnEnter = std::function<void(const StackIdHolder&)>;
+  using OnEnter = std::function<void(const StackData&)>;
   using ExceptionHandler = std::function<void(VmExceptionInfo)>;
   using MissingScriptHandler =
     std::function<std::optional<PexScript::Lazy>(std::string)>;
@@ -73,14 +94,14 @@ class VirtualMachine
 
   VarValue CallMethod(IGameObject* self, const char* methodName,
                       std::vector<VarValue>& arguments,
-                      std::shared_ptr<StackIdHolder> stackIdHolder = nullptr,
+                      std::shared_ptr<StackData> stackData = nullptr,
                       const std::vector<std::shared_ptr<ActivePexInstance>>*
                         activePexInstancesOverride = nullptr);
 
   VarValue CallStatic(const std::string& className,
                       const std::string& functionName,
                       std::vector<VarValue>& arguments,
-                      std::shared_ptr<StackIdHolder> stackIdHolder = nullptr);
+                      std::shared_ptr<StackData> stackData = nullptr);
 
   PexScript::Lazy GetPexByName(const std::string& name);
 

papyrus-vm/src/papyrus-vm-lib/ActivePexInstance.cpp

@@ -1,4 +1,5 @@
-#include "papyrus-vm/OpcodesImplementation.h"
+#include "ScopedTask.h"
+#include "papyrus-vm/OpcodesImplementation.h"
 #include "papyrus-vm/Utils.h"
 #include "papyrus-vm/VirtualMachine.h"
 #include <algorithm>
@@ -133,7 +134,7 @@ const std::string& ActivePexInstance::GetSourcePexName() const
 
 struct ActivePexInstance::ExecutionContext
 {
-  std::shared_ptr<StackIdHolder> stackIdHolder;
+  std::shared_ptr<StackData> stackData;
   std::vector<FunctionCode::Instruction> instructions;
   std::shared_ptr<std::vector<Local>> locals;
   bool needReturn = false;
@@ -314,7 +315,7 @@ void ActivePexInstance::ExecuteOpCode(ExecutionContext* ctx, uint8_t op,
         }
 
         auto res = parentVM->CallMethod(gameObject, (const char*)(*args[0]),
-                                        argsForCall, ctx->stackIdHolder,
+                                        argsForCall, ctx->stackData,
                                         &activePexInstancesForCallParent);
         if (EnsureCallResultIsSynchronous(res, ctx))
           *args[1] = res;
@@ -347,7 +348,7 @@ void ActivePexInstance::ExecuteOpCode(ExecutionContext* ctx, uint8_t op,
           auto nullableGameObject = static_cast<IGameObject*>(*object);
           auto res =
             parentVM->CallMethod(nullableGameObject, functionName.c_str(),
-                                 argsForCall, ctx->stackIdHolder);
+                                 argsForCall, ctx->stackData);
           spdlog::trace("callmethod object={} funcName={} result={}",
                         object->ToString(), functionName, res.ToString());
           if (EnsureCallResultIsSynchronous(res, ctx)) {
@@ -366,7 +367,7 @@ void ActivePexInstance::ExecuteOpCode(ExecutionContext* ctx, uint8_t op,
       const char* functionName = (const char*)(*args[1]);
       try {
         auto res = parentVM->CallStatic(className, functionName, argsForCall,
-                                        ctx->stackIdHolder);
+                                        ctx->stackData);
         if (EnsureCallResultIsSynchronous(res, ctx))
           *args[2] = res;
       } catch (std::exception& e) {
@@ -438,7 +439,7 @@ void ActivePexInstance::ExecuteOpCode(ExecutionContext* ctx, uint8_t op,
             // TODO: use of argsForCall looks incorrect. why use argsForCall
             // here? shoud be {} (empty args)
             *args[2] = inst->StartFunction(runProperty->readHandler,
-                                           argsForCall, ctx->stackIdHolder);
+                                           argsForCall, ctx->stackData);
             spdlog::trace("propget function called");
           } else {
             auto& instProps = inst->sourcePex.fn()->objectTable[0].properties;
@@ -524,7 +525,7 @@ void ActivePexInstance::ExecuteOpCode(ExecutionContext* ctx, uint8_t op,
             // TODO: use of argsForCall looks incorrect.
             // probably should only *args[2]
             inst->StartFunction(runProperty->writeHandler, argsForCall,
-                                ctx->stackIdHolder);
+                                ctx->stackData);
             spdlog::trace("propset function called");
           } else {
             auto& instProps = inst->sourcePex.fn()->objectTable[0].properties;
@@ -732,14 +733,35 @@ VarValue ActivePexInstance::ExecuteAll(
   return ctx.returnValue;
 }
 
-VarValue ActivePexInstance::StartFunction(
-  FunctionInfo& function, std::vector<VarValue>& arguments,
-  std::shared_ptr<StackIdHolder> stackIdHolder)
+VarValue ActivePexInstance::StartFunction(FunctionInfo& function,
+                                          std::vector<VarValue>& arguments,
+                                          std::shared_ptr<StackData> stackData)
 {
-  if (!stackIdHolder)
-    throw std::runtime_error("An empty stackIdHolder passed to StartFunction");
+  if (!stackData) {
+    throw std::runtime_error("An empty stackData passed to StartFunction");
+  }
+
+  thread_local StackDepthHolder g_stackDepthHolder;
+
+  g_stackDepthHolder.IncreaseStackDepth();
+
+  Viet::ScopedTask<StackDepthHolder> stackDepthDecreaseTask(
+    [](StackDepthHolder& stackDepthHolder) {
+      stackDepthHolder.DecreaseStackDepth();
+    },
+    g_stackDepthHolder);
+
+  constexpr size_t kMaxStackDepth = 128;
+
+  if (g_stackDepthHolder.GetStackDepth() >= kMaxStackDepth) {
+    spdlog::error("ActivePexInstance::StartFunction - Stack overflow in "
+                  "script {}, returning None",
+                  sourcePex.fn()->source);
+    return VarValue::None();
+  }
+
   auto locals = MakeLocals(function, arguments);
-  ExecutionContext ctx{ stackIdHolder, function.code.instructions, locals };
+  ExecutionContext ctx{ stackData, function.code.instructions, locals };
   return ExecuteAll(ctx);
 }
 

papyrus-vm/src/papyrus-vm-lib/VarValue.cpp

@@ -191,10 +191,9 @@ int32_t VarValue::GetMetaStackId() const
   return stackId;
 }
 
-void VarValue::SetMetaStackIdHolder(
-  std::shared_ptr<StackIdHolder> stackIdHolder_)
+void VarValue::SetMetaStackIdHolder(const StackIdHolder& stackIdHolder_)
 {
-  stackId = stackIdHolder_->GetStackId();
+  stackId = stackIdHolder_.GetStackId();
 }
 
 VarValue VarValue::AttachTestStackId(VarValue original, int32_t stackId)

papyrus-vm/src/papyrus-vm-lib/VirtualMachine.cpp

@@ -98,11 +98,13 @@ void VirtualMachine::SendEvent(std::shared_ptr<IGameObject> self,
     auto fn = scriptInstance->GetFunctionByName(
       eventName, scriptInstance->GetActiveStateName());
     if (fn.valid) {
-      auto stackIdHolder = std::make_shared<StackIdHolder>(*this);
-      if (enter)
-        enter(*stackIdHolder);
+      std::shared_ptr<StackData> stackData;
+      stackData.reset(new StackData{ StackIdHolder{ *this } });
+      if (enter) {
+        enter(*stackData);
+      }
       scriptInstance->StartFunction(
-        fn, const_cast<std::vector<VarValue>&>(arguments), stackIdHolder);
+        fn, const_cast<std::vector<VarValue>&>(arguments), stackData);
     }
   }
 }
@@ -115,8 +117,10 @@ void VirtualMachine::SendEvent(ActivePexInstance* instance,
   auto fn =
     instance->GetFunctionByName(eventName, instance->GetActiveStateName());
   if (fn.valid) {
+    std::shared_ptr<StackData> stackData;
+    stackData.reset(new StackData{ StackIdHolder{ *this } });
     instance->StartFunction(fn, const_cast<std::vector<VarValue>&>(arguments),
-                            std::make_shared<StackIdHolder>(*this));
+                            stackData);
   }
 }
 
@@ -140,15 +144,31 @@ int32_t StackIdHolder::GetStackId() const
   return stackId;
 }
 
+StackDepthHolder::StackDepthHolder() = default;
+
+size_t StackDepthHolder::GetStackDepth() const
+{
+  return depth;
+}
+
+void StackDepthHolder::IncreaseStackDepth()
+{
+  ++depth;
+}
+
+void StackDepthHolder::DecreaseStackDepth()
+{
+  --depth;
+}
+
 VarValue VirtualMachine::CallMethod(
   IGameObject* selfObj, const char* methodName,
-  std::vector<VarValue>& arguments,
-  std::shared_ptr<StackIdHolder> stackIdHolder,
+  std::vector<VarValue>& arguments, std::shared_ptr<StackData> stackData,
   const std::vector<std::shared_ptr<ActivePexInstance>>*
     activePexInstancesOverride)
 {
-  if (!stackIdHolder) {
-    stackIdHolder.reset(new StackIdHolder(*this));
+  if (!stackData) {
+    stackData.reset(new StackData{ StackIdHolder{ *this } });
   }
 
   if (!selfObj) {
@@ -175,8 +195,7 @@ VarValue VirtualMachine::CallMethod(
     }
 
     if (functionInfo.valid) {
-      return activeScript->StartFunction(functionInfo, arguments,
-                                         stackIdHolder);
+      return activeScript->StartFunction(functionInfo, arguments, stackData);
     }
   }
 
@@ -187,7 +206,7 @@ VarValue VirtualMachine::CallMethod(
   while (1) {
     if (auto f = nativeFunctions[ToLower(base)][ToLower(methodName)]) {
       auto self = VarValue(selfObj);
-      self.SetMetaStackIdHolder(stackIdHolder);
+      self.SetMetaStackIdHolder(stackData->stackIdHolder);
       return f(self, arguments);
     }
     auto it = allLoadedScripts.find(base);
@@ -204,13 +223,13 @@ VarValue VirtualMachine::CallMethod(
   throw std::runtime_error(e);
 }
 
-VarValue VirtualMachine::CallStatic(
-  const std::string& className, const std::string& functionName,
-  std::vector<VarValue>& arguments,
-  std::shared_ptr<StackIdHolder> stackIdHolder)
+VarValue VirtualMachine::CallStatic(const std::string& className,
+                                    const std::string& functionName,
+                                    std::vector<VarValue>& arguments,
+                                    std::shared_ptr<StackData> stackData)
 {
-  if (!stackIdHolder) {
-    stackIdHolder.reset(new StackIdHolder(*this));
+  if (!stackData) {
+    stackData.reset(new StackData{ StackIdHolder{ *this } });
   }
 
   VarValue result = VarValue::None();
@@ -223,7 +242,7 @@ VarValue VirtualMachine::CallStatic(
 
   if (f) {
     auto self = VarValue::None();
-    self.SetMetaStackIdHolder(stackIdHolder);
+    self.SetMetaStackIdHolder(stackData->stackIdHolder);
     return f(self, arguments);
   }
 
@@ -254,7 +273,7 @@ VarValue VirtualMachine::CallStatic(
                                std::string(functionName) + "'");
     }
 
-    result = instance->StartFunction(function, arguments, stackIdHolder);
+    result = instance->StartFunction(function, arguments, stackData);
   }
   if (!function.valid) {
     throw std::runtime_error("Function is not valid - '" +

skymp5-server/cpp/server_guest_lib/WorldState.cpp

@@ -766,9 +766,10 @@ void WorldState::SendPapyrusEvent(MpForm* form, const char* eventName,
     }
   }
 
-  VirtualMachine::OnEnter onEnter = [&](const StackIdHolder& holder) {
-    pImpl->policy->BeforeSendPapyrusEvent(form, eventName, arguments,
-                                          argumentsCount, holder.GetStackId());
+  VirtualMachine::OnEnter onEnter = [&](const StackData& stackData) {
+    pImpl->policy->BeforeSendPapyrusEvent(
+      form, eventName, arguments, argumentsCount,
+      stackData.stackIdHolder.GetStackId());
   };
   auto& vm = GetPapyrusVm();
   return vm.SendEvent(form->ToGameObject(), eventName, args, onEnter);