Browse Source

Add basic page allocator for Windows

ghorsington 3 years ago
parent
commit
deabaac389

+ 21 - 11
BepInEx.IL2CPP/Hook/Buffer/MemoryAllocator.cs

@@ -5,27 +5,30 @@ using MonoMod.Utils;
 namespace BepInEx.IL2CPP
 {
 	/// <summary>
-	///     A general purpose memory allocator for patching purposes.
-	///     Allows to allocate memory within the 2GB radius of a given address.
+	///     A general purpose page allocator for patching purposes.
+	///     Allows to allocate pages (4k memory chunks) within the 1GB radius of a given address.
 	/// </summary>
 	/// <remarks>Based on https://github.com/kubo/funchook</remarks>
-	internal abstract class MemoryAllocator
+	internal abstract class PageAllocator
 	{
 		/// <summary>
 		///     Common page size on Unix and Windows (4k).
 		/// </summary>
-		protected const int PAGE_SIZE = 0x1000;
+		public const int PAGE_SIZE = 0x1000;
 
 		/// <summary>
 		///     Allocation granularity on Windows (but can be reused in other implementations).
 		/// </summary>
 		protected const int ALLOCATION_UNIT = 0x100000;
 
-		private static MemoryAllocator instance;
-		public static MemoryAllocator Instance => instance ??= Init();
+		protected const int PAGES_PER_UNIT = ALLOCATION_UNIT / PAGE_SIZE;
 
-		public abstract IntPtr Allocate(IntPtr func);
-		public abstract void Free(IntPtr buffer);
+		private static PageAllocator instance;
+		public static PageAllocator Instance => instance ??= Init();
+
+		public abstract IntPtr Allocate(IntPtr hint);
+
+		public abstract void Free(IntPtr page);
 
 		[MethodImpl(MethodImplOptions.AggressiveInlining)]
 		protected static long RoundDown(long num, long unit)
@@ -39,12 +42,19 @@ namespace BepInEx.IL2CPP
 			return (num + unit - 1) & ~ (unit - 1);
 		}
 
-		private static MemoryAllocator Init()
+		[MethodImpl(MethodImplOptions.AggressiveInlining)]
+		protected static bool IsInRelJmpRange(IntPtr src, IntPtr dst)
+		{
+			long diff = dst.ToInt64() - src.ToInt64();
+			return int.MinValue <= diff && diff <= int.MaxValue;
+		}
+
+		private static PageAllocator Init()
 		{
 			if (PlatformHelper.Is(Platform.Windows))
-				return new WindowsMemoryAllocator();
+				return new WindowsPageAllocator();
 			if (PlatformHelper.Is(Platform.Unix))
-				return new UnixMemoryAllocator();
+				return new UnixPageAllocator();
 			throw new NotImplementedException();
 		}
 	}

+ 3 - 3
BepInEx.IL2CPP/Hook/Buffer/UnixMemoryAllocator.cs

@@ -5,14 +5,14 @@ namespace BepInEx.IL2CPP
 	/// <summary>
 	///     Based on https://github.com/kubo/funchook
 	/// </summary>
-	internal class UnixMemoryAllocator : MemoryAllocator
+	internal class UnixPageAllocator : PageAllocator
 	{
-		public override IntPtr Allocate(IntPtr func)
+		public override IntPtr Allocate(IntPtr hint)
 		{
 			throw new NotImplementedException();
 		}
 
-		public override void Free(IntPtr buffer)
+		public override void Free(IntPtr page)
 		{
 			throw new NotImplementedException();
 		}

+ 67 - 39
BepInEx.IL2CPP/Hook/Buffer/WindowsMemoryAllocator.cs

@@ -8,9 +8,9 @@ namespace BepInEx.IL2CPP
 	/// <summary>
 	///     Based on https://github.com/kubo/funchook
 	/// </summary>
-	internal class WindowsMemoryAllocator : MemoryAllocator
+	internal class WindowsPageAllocator : PageAllocator
 	{
-		private readonly LinkedList<IntPtr> allocatedChunks = new LinkedList<IntPtr>();
+		private readonly List<PageChunk> allocatedChunks = new List<PageChunk>();
 
 		/// <summary>
 		///     Allocates a single 64k chunk of memory near the given address
@@ -18,7 +18,7 @@ namespace BepInEx.IL2CPP
 		/// <param name="hint">Address near which to attempt allocate the chunk</param>
 		/// <returns>Allocated chunk</returns>
 		/// <exception cref="Win32Exception">Allocation failed</exception>
-		private IntPtr AllocateChunk(IntPtr hint)
+		private PageChunk AllocateChunk(IntPtr hint)
 		{
 			while (true)
 			{
@@ -51,18 +51,70 @@ namespace BepInEx.IL2CPP
 				throw new Win32Exception(error);
 			}
 
-			allocatedChunks.AddFirst(chunk);
-			return chunk;
+			var result = new PageChunk
+			{
+				BaseAddress = chunk
+			};
+			allocatedChunks.Add(result);
+			return result;
+		}
+
+		private IntPtr AllocPage(IntPtr hint)
+		{
+			foreach (var allocatedChunk in allocatedChunks)
+			{
+				// Small shortcut to speed up page lookup
+				if (allocatedChunk.UsedPages == PAGES_PER_UNIT)
+					continue;
+				for (var i = 0; i < allocatedChunk.Pages.Length; i++)
+				{
+					if (allocatedChunk.Pages[i])
+						continue;
+					var pageAddr = allocatedChunk.GetPage(i);
+					if (!IsInRelJmpRange(hint, pageAddr))
+						continue;
+					allocatedChunk.Pages[i] = true;
+					allocatedChunk.UsedPages++;
+					return pageAddr;
+				}
+			}
+
+			var chunk = AllocateChunk(hint);
+			chunk.Pages[0] = true;
+			chunk.UsedPages++;
+			return chunk.BaseAddress;
+		}
+
+		public override IntPtr Allocate(IntPtr hint)
+		{
+			var pageAddress = AllocPage(hint);
+			if (WinApi.VirtualAlloc(pageAddress, (UIntPtr)PAGE_SIZE, WinApi.AllocationType.MEM_COMMIT, WinApi.ProtectConstant.PAGE_READWRITE) == IntPtr.Zero)
+				throw new Win32Exception(Marshal.GetLastWin32Error());
+			return pageAddress;
 		}
 
-		public override IntPtr Allocate(IntPtr func)
+		public override void Free(IntPtr page)
 		{
-			throw new NotImplementedException();
+			foreach (var allocatedChunk in allocatedChunks)
+			{
+				long index = (page.ToInt64() - allocatedChunk.BaseAddress.ToInt64()) / PAGE_SIZE;
+				if (index < 0 || index > PAGES_PER_UNIT)
+					continue;
+				allocatedChunk.Pages[index] = false;
+				return;
+			}
 		}
 
-		public override void Free(IntPtr buffer)
+		private class PageChunk
 		{
-			throw new NotImplementedException();
+			public IntPtr BaseAddress;
+			public readonly bool[] Pages = new bool[PAGES_PER_UNIT];
+			public int UsedPages;
+
+			public IntPtr GetPage(int index)
+			{
+				return BaseAddress + index * PAGE_SIZE;
+			}
 		}
 
 		private static class WinApi
@@ -72,14 +124,8 @@ namespace BepInEx.IL2CPP
 			{
 				// ReSharper disable InconsistentNaming
 				MEM_COMMIT = 0x00001000,
-				MEM_RESERVE = 0x00002000,
-				MEM_RESET = 0x00080000,
-				MEM_RESET_UNDO = 0x1000000,
-				MEM_LARGE_PAGES = 0x20000000,
-				MEM_PHYSICAL = 0x00400000,
-				MEM_TOP_DOWN = 0x00100000,
-
-				MEM_WRITE_WATCH = 0x00200000
+
+				MEM_RESERVE = 0x00002000
 				// ReSharper restore InconsistentNaming
 			}
 
@@ -88,21 +134,14 @@ namespace BepInEx.IL2CPP
 			public enum FreeType : uint
 			{
 				// ReSharper disable InconsistentNaming
-				MEM_DECOMMIT = 0x00004000,
-				MEM_RELEASE = 0x00008000,
-				MEM_COALESCE_PLACEHOLDERS = 0x00000001,
-
-				MEM_PRESERVE_PLACEHOLDER = 0x00000002
+				MEM_RELEASE = 0x00008000
 				// ReSharper restore InconsistentNaming
 			}
 
 			public enum PageState : uint
 			{
 				// ReSharper disable InconsistentNaming
-				MEM_COMMIT = 0x1000,
-				MEM_FREE = 0x10000,
-
-				MEM_RESERVE = 0x2000
+				MEM_FREE = 0x10000
 				// ReSharper restore InconsistentNaming
 			}
 
@@ -110,20 +149,9 @@ namespace BepInEx.IL2CPP
 			public enum ProtectConstant : uint
 			{
 				// ReSharper disable InconsistentNaming
-				PAGE_EXECUTE = 0x10,
-				PAGE_EXECUTE_READ = 0x20,
-				PAGE_EXECUTE_READWRITE = 0x40,
-				PAGE_EXECUTE_WRITECOPY = 0x80,
 				PAGE_NOACCESS = 0x01,
-				PAGE_READONLY = 0x02,
-				PAGE_READWRITE = 0x04,
-				PAGE_WRITECOPY = 0x08,
-				PAGE_TARGETS_INVALID = 0x40000000,
-				PAGE_TARGETS_NO_UPDATE = 0x40000000,
-				PAGE_GUARD = 0x100,
-				PAGE_NOCACHE = 0x200,
-
-				PAGE_WRITECOMBINE = 0x400
+
+				PAGE_READWRITE = 0x04
 				// ReSharper restore InconsistentNaming
 			}