[RecoverEmptyContainers] Fixed an exploit granting infinite container…

…s when drinking a Potion/brew and cancelling the ability/usage of it; Bumped version to v1.0.1

CHANGELOG.md

@@ -5,6 +5,9 @@
 * Initial release
 
 ## Recover Empty Containers
+### v1.0.1
+* Fixed an exploit where you could get infinite empty containers when cancelling the ability usage of the Potion/Brew
+
 ### v1.0.0
 * Initial release
 

RecoverEmptyContainers/Plugin.cs

@@ -10,7 +10,7 @@ namespace VMods.RecoverEmptyContainers
 	[BepInPlugin(PluginInfo.PLUGIN_GUID, PluginInfo.PLUGIN_NAME, PluginInfo.PLUGIN_VERSION)]
 	[BepInDependency("xyz.molenzwiebel.wetstone")]
 	[Reloadable]
-	public class Plugin : BasePlugin
+	public class Plugin : BasePlugin, IRunOnInitialized
 	{
 		#region Variables
 
@@ -28,20 +28,27 @@ public sealed override void Load()
 				return;
 			}
 			Utils.Initialize(Log, PluginInfo.PLUGIN_NAME);
+
 			RecoverEmptyContainersConfig.Initialize(Config);
 
 			_hooks = Harmony.CreateAndPatchAll(Assembly.GetExecutingAssembly());
 
 			Log.LogInfo($"Plugin {PluginInfo.PLUGIN_NAME} (v{PluginInfo.PLUGIN_VERSION}) is loaded!");
 		}
 
+		public void OnGameInitialized()
+		{
+			RecoverEmptyContainersSystem.Initialize();
+		}
+
 		public sealed override bool Unload()
 		{
 			if(VWorld.IsClient)
 			{
 				return true;
 			}
 			_hooks?.UnpatchSelf();
+			RecoverEmptyContainersSystem.Deinitialize();
 			Config.Clear();
 			Utils.Deinitialize();
 			return true;

RecoverEmptyContainers/RecoverEmptyContainers.csproj

@@ -4,7 +4,7 @@
     <AssemblyName>VMods.RecoverEmptyContainers</AssemblyName>
     <RootNamespace>VMods.RecoverEmptyContainers</RootNamespace>
     <Description>Allows a player to recover an empty container when consuming a potion/brew</Description>
-    <Version>1.0.0</Version>
+    <Version>1.0.1</Version>
     <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
     <LangVersion>latest</LangVersion>
     <GeneratePackageOnBuild>False</GeneratePackageOnBuild>
@@ -17,7 +17,9 @@
   </PropertyGroup>
 
   <ItemGroup>
+    <Compile Include="..\Shared\BuffSystemHook.cs" Link="Shared\BuffSystemHook.cs" />
     <Compile Include="..\Shared\Utils.cs" Link="Shared\Utils.cs" />
+    <Compile Include="..\Shared\VModCharacter.cs" Link="Shared\VModCharacter.cs" />
   </ItemGroup>
 
   <ItemGroup>

RecoverEmptyContainers/Systems/RecoverEmptyContainersSystem.cs

@@ -0,0 +1,149 @@
+using ProjectM;
+using ProjectM.Scripting;
+using System.Collections.Generic;
+using Unity.Entities;
+using VMods.Shared;
+using Wetstone.API;
+
+namespace VMods.RecoverEmptyContainers
+{
+	public static class RecoverEmptyContainersSystem
+	{
+		#region Variables
+
+		private static readonly Dictionary<PrefabGUID, (PrefabGUID itemGUID, int stackCount)> BuffToEmptyContainerMapping = new();
+
+		private static readonly Dictionary<PrefabGUID, PrefabGUID> RecipeItemToReturnedItemMapping = new()
+		{
+			[new PrefabGUID(-1322000172)] = new PrefabGUID(-810738866),// Water-filled Canteen -> Empty Canteen
+			[new PrefabGUID(-1382451936)] = new PrefabGUID(-437611596),// Water-filled Bottle -> Empty Glass Bottle
+		};
+
+		#endregion
+
+		#region Public Methods
+
+		public static void Initialize()
+		{
+			BuildBuffToEmptyContainerMapping();
+
+			BuffSystemHook.ProcessBuffEvent += OnProcessBuffEvent;
+		}
+
+		public static void Deinitialize()
+		{
+			BuffSystemHook.ProcessBuffEvent -= OnProcessBuffEvent;
+
+			BuffToEmptyContainerMapping.Clear();
+		}
+
+		#endregion
+
+		#region Private Methods
+
+		private static void OnProcessBuffEvent(Entity entity, PrefabGUID buffGUID)
+		{
+			var server = VWorld.Server;
+			var entityManager = server.EntityManager;
+
+			if(!entityManager.HasComponent<EntityCreator>(entity) || !entityManager.HasComponent<EntityOwner>(entity))
+			{
+				return;
+			}
+
+			var entityOwner = entityManager.GetComponentData<EntityOwner>(entity);
+			var entityCreator = entityManager.GetComponentData<EntityCreator>(entity);
+			var entityCreatorEntity = entityCreator.Creator._Entity;
+			if(!entityManager.HasComponent<AbilityOwner>(entityCreatorEntity))
+			{
+				return;
+			}
+
+			var abilityOwner = entityManager.GetComponentData<AbilityOwner>(entityCreatorEntity);
+			var abilityGroupEntity = abilityOwner.AbilityGroup._Entity;
+			if(!entityManager.HasComponent<PrefabGUID>(abilityGroupEntity))
+			{
+				return;
+			}
+
+			var abilityGroupPrefabGUID = entityManager.GetComponentData<PrefabGUID>(abilityGroupEntity);
+			if(!BuffToEmptyContainerMapping.TryGetValue(abilityGroupPrefabGUID, out var returnData))
+			{
+				return;
+			}
+
+			var gameDataSystem = server.GetExistingSystem<ServerScriptMapper>()._GameDataSystem;
+
+			Utils.TryGiveItem(entityManager, gameDataSystem.ItemHashLookupMap, entityOwner.Owner, returnData.itemGUID, returnData.stackCount, out _, out _, dropRemainder: true);
+		}
+
+		private static void BuildBuffToEmptyContainerMapping()
+		{
+			BuffToEmptyContainerMapping.Clear();
+
+			var server = VWorld.Server;
+			var entityManager = server.EntityManager;
+			var gameDataSystem = server.GetExistingSystem<ServerScriptMapper>()._GameDataSystem;
+			var prefabCollectionSystem = server.GetExistingSystem<PrefabCollectionSystem>();
+
+			var duplicateBuffs = new List<PrefabGUID>();
+
+			foreach(var recipeKvp in gameDataSystem.RecipeHashLookupMap)
+			{
+				var recipeEntity = recipeKvp.Value.Entity;
+				if(!entityManager.HasComponent<RecipeRequirementBuffer>(recipeEntity) || !entityManager.HasComponent<RecipeOutputBuffer>(recipeEntity))
+				{
+					continue;
+				}
+
+				// Find the returned item
+				PrefabGUID? returnItemGUID = null;
+				int returnItemStackCount = 0;
+				var requirementBuffer = entityManager.GetBuffer<RecipeRequirementBuffer>(recipeEntity);
+				foreach(var requirement in requirementBuffer)
+				{
+					if(RecipeItemToReturnedItemMapping.TryGetValue(requirement.Guid, out var prefabGUID))
+					{
+						returnItemGUID = prefabGUID;
+						returnItemStackCount = requirement.Stacks;
+						break;
+					}
+				}
+
+				// Check if we've found a returnable item
+				if(!returnItemGUID.HasValue)
+				{
+					continue;
+				}
+
+				// Find the buff that belongs to this item
+				var outputBuffer = entityManager.GetBuffer<RecipeOutputBuffer>(recipeEntity);
+				foreach(var output in outputBuffer)
+				{
+					var outputEntity = prefabCollectionSystem.PrefabLookupMap[output.Guid];
+					if(!entityManager.HasComponent<CastAbilityOnConsume>(outputEntity))
+					{
+						continue;
+					}
+
+					var castAbilityOnConsuime = entityManager.GetComponentData<CastAbilityOnConsume>(outputEntity);
+					var abilityGUID = castAbilityOnConsuime.AbilityGuid;
+
+#if DEBUG
+					Utils.Logger.LogMessage($"Matched Buff {prefabCollectionSystem.PrefabNameLookupMap[abilityGUID]} -> {prefabCollectionSystem.PrefabNameLookupMap[returnItemGUID.Value]}");
+#endif
+					if(BuffToEmptyContainerMapping.ContainsKey(abilityGUID))
+					{
+						Utils.Logger.LogWarning($"Found duplicate Buff {prefabCollectionSystem.PrefabNameLookupMap[abilityGUID]}. Remove it!");
+						duplicateBuffs.Add(abilityGUID);
+					}
+					BuffToEmptyContainerMapping[abilityGUID] = (returnItemGUID.Value, returnItemStackCount);
+				}
+			}
+
+			duplicateBuffs.ForEach(x => BuffToEmptyContainerMapping.Remove(x));
+		}
+
+		#endregion
+	}
+}

Shared/VModCharacter.cs

@@ -43,6 +43,18 @@ public VModCharacter(Entity userEntity, Entity charEntity, EntityManager? entity
 			};
 		}
 
+		public VModCharacter(Entity charEntity, EntityManager? entityManager = null)
+		{
+			entityManager ??= Utils.CurrentWorld.EntityManager;
+			Character = entityManager.Value.GetComponentData<PlayerCharacter>(charEntity);
+			User = entityManager.Value.GetComponentData<User>(Character.UserEntity._Entity);
+			FromCharacter = new FromCharacter()
+			{
+				User = Character.UserEntity._Entity,
+				Character = User.LocalCharacter._Entity,
+			};
+		}
+
 		#endregion
 
 		#region Public Methods

Thunderstone/RecoverEmptyContainers/manifest.json

@@ -1,7 +1,7 @@
 {
   "name": "VMods_Recover_Empty_Containers",
   "description": "A mod that allows players to recover empty containers when drinking potions or brews.",
-  "version_number": "1.0.0",
+  "version_number": "1.0.1",
   "dependencies": [
     "BepInEx-BepInExPack_V_Rising-1.0.0",
     "molenzwiebel-Wetstone-1.1.0"