From 1639a132a9cc3461b3af47d49593ca298871a8f1 Mon Sep 17 00:00:00 2001
From: Deukhoofd <Deukhoofd@gmail.com>
Date: Sun, 15 May 2022 10:08:15 +0200
Subject: [PATCH] Fixes for use after free

---
 src/Battling/Library/CaptureLibrary.cpp                 | 8 +++++---
 src/Battling/Pokemon/Pokemon.cpp                        | 2 +-
 src/ScriptResolving/AngelScript/AngelScriptResolver.cpp | 2 +-
 src/ScriptResolving/AngelScript/AngelScriptResolver.hpp | 2 +-
 tests/ScriptTests/Angelscript/ItemUseScriptTests.cpp    | 4 ++--
 5 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/src/Battling/Library/CaptureLibrary.cpp b/src/Battling/Library/CaptureLibrary.cpp
index 3d6889c..b1aca06 100644
--- a/src/Battling/Library/CaptureLibrary.cpp
+++ b/src/Battling/Library/CaptureLibrary.cpp
@@ -11,9 +11,11 @@ namespace PkmnLib::Battling {
         auto rate = pokemon->GetSpecies()->GetCaptureRate();
 
         u8 bonusBall = 1;
-        auto* itemScript = dynamic_cast<PkmnItemUseScript*>(
-            pokemon->GetLibrary()->GetScriptResolver()->LoadItemScript(catchItem).GetValue());
-        itemScript->ModifyPokeballCatchBonus(pokemon, &bonusBall);
+        auto itemScript = pokemon->GetLibrary()->GetScriptResolver()->LoadItemScript(catchItem).As<PkmnItemUseScript>();
+        if (!itemScript.HasValue()) {
+            return {};
+        }
+        itemScript.GetValue()->ModifyPokeballCatchBonus(pokemon, &bonusBall);
 
         u8 bonusStatus = 1;
         PKMN_HOOK(ModifyCaptureRateBonus, pokemon, pokemon, catchItem, &bonusStatus);
diff --git a/src/Battling/Pokemon/Pokemon.cpp b/src/Battling/Pokemon/Pokemon.cpp
index c0634c4..7c1034d 100644
--- a/src/Battling/Pokemon/Pokemon.cpp
+++ b/src/Battling/Pokemon/Pokemon.cpp
@@ -83,7 +83,7 @@ void PkmnLib::Battling::Pokemon::AttemptCapture(PkmnLib::Library::Item* catchIte
     Ensure(_battleData.Side.HasValue());
     Ensure(!IsFainted());
     Ensure(IsUsable());
-    Ensure(!GetBattleSide().GetValue()->IsSlotUnfillabe(this)) auto captureLibrary =
+    auto captureLibrary =
         GetLibrary().ForceAs<const BattleLibrary>()->GetCaptureLibrary();
     auto result = captureLibrary->TryCatch(this, catchItem, _battleData.Battle.GetValue()->GetRandom());
     _battleData.Battle.GetValue()->TriggerEventListener<CaptureAttemptEvent>(this, result);
diff --git a/src/ScriptResolving/AngelScript/AngelScriptResolver.cpp b/src/ScriptResolving/AngelScript/AngelScriptResolver.cpp
index 3beff00..b210d99 100644
--- a/src/ScriptResolving/AngelScript/AngelScriptResolver.cpp
+++ b/src/ScriptResolving/AngelScript/AngelScriptResolver.cpp
@@ -193,7 +193,7 @@ AngelScriptResolver::LoadScript(const ArbUt::OptionalBorrowedPtr<void>& owner, S
     return new AngelScriptScript(owner, ownerType, this, t.value(), obj, _contextPool);
 }
 
-ArbUt::OptionalUniquePtr<CreatureLib::Battling::ItemUseScript>
+ArbUt::OptionalBorrowedPtr<CreatureLib::Battling::ItemUseScript>
 AngelScriptResolver::LoadItemScript(const CreatureLib::Library::Item* item) {
     auto v = this->_itemUseScripts.TryGet(item);
     if (v.has_value()) {
diff --git a/src/ScriptResolving/AngelScript/AngelScriptResolver.hpp b/src/ScriptResolving/AngelScript/AngelScriptResolver.hpp
index 5c2758c..008a799 100644
--- a/src/ScriptResolving/AngelScript/AngelScriptResolver.hpp
+++ b/src/ScriptResolving/AngelScript/AngelScriptResolver.hpp
@@ -69,7 +69,7 @@ public:
     ArbUt::OptionalUniquePtr<CreatureLib::Battling::BattleScript>
     LoadScript(const ArbUt::OptionalBorrowedPtr<void>& owner, ScriptCategory category,
                const ArbUt::StringView& scriptName) override;
-    ArbUt::OptionalUniquePtr<CreatureLib::Battling::ItemUseScript>
+    ArbUt::OptionalBorrowedPtr<CreatureLib::Battling::ItemUseScript>
     LoadItemScript(const CreatureLib::Library::Item* item) override;
 
     ArbUt::OptionalBorrowedPtr<const PkmnLib::Battling::EvolutionScript>
diff --git a/tests/ScriptTests/Angelscript/ItemUseScriptTests.cpp b/tests/ScriptTests/Angelscript/ItemUseScriptTests.cpp
index 42723dc..5d98a2d 100644
--- a/tests/ScriptTests/Angelscript/ItemUseScriptTests.cpp
+++ b/tests/ScriptTests/Angelscript/ItemUseScriptTests.cpp
@@ -45,9 +45,9 @@ static AngelScriptItemUseScript* GetScript(PkmnLib::Battling::BattleLibrary* mai
                                            new CreatureLib::Library::SecondaryEffect(100, name, {}), nullptr, {});
 
     auto s = lib->LoadItemScript(&item);
-    auto script = dynamic_cast<AngelScriptItemUseScript*>(s.TakeOwnership());
+    auto script = s.As<AngelScriptItemUseScript>();
     REQUIRE(script != nullptr);
-    return script;
+    return script.GetValue();
 }
 
 TEST_CASE("Invoke isItemUsable item use script function on empty class") {