Skip to content

Commit

Permalink
refactor some xmemory stuff
Browse files Browse the repository at this point in the history
  • Loading branch information
Jnnshschl committed Jan 7, 2024
1 parent 3d8eeda commit b381acb
Show file tree
Hide file tree
Showing 3 changed files with 78 additions and 65 deletions.
18 changes: 5 additions & 13 deletions AmeisenBotX.Memory/Structs/AllocationPool.cs
Original file line number Diff line number Diff line change
Expand Up @@ -4,21 +4,13 @@

namespace AmeisenBotX.Memory.Structs
{
public readonly struct AllocationPool
public readonly struct AllocationPool(nint address, int size)
{
public AllocationPool(nint address, int size)
{
Address = address;
Size = size;

Allocations = [];
}

public nint Address { get; }
public nint Address { get; } = address;

public SortedList<int, int> Allocations { get; }
public SortedList<int, int> Allocations { get; } = [];

public int Size { get; }
public int Size { get; } = size;

/// <summary>
/// Free a memory block in the pool
Expand Down Expand Up @@ -80,7 +72,7 @@ private bool GetNextFreeBlock(int size, out int offset)
KeyValuePair<int, int> allocation = Allocations.ElementAt(i);
int allocationEnd = allocation.Key + allocation.Value;

// when there is a next element, used it as the limiter, if not use the
// when there is a next element, use it as the limiter, if not use the
// whole remaining space
int memoryLeft = i + 1 < Allocations.Count
? Allocations.ElementAt(i + 1).Key - allocationEnd
Expand Down
52 changes: 38 additions & 14 deletions AmeisenBotX.Memory/Win32/Win32Imports.cs
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,9 @@ namespace AmeisenBotX.Memory.Win32
{
public static unsafe partial class Win32Imports
{
public const int GWL_EXSTYLE = -0x14;
public const int GWL_STYLE = -16;
public const int STARTF_USESHOWWINDOW = 1;
public const int SW_SHOWMINNOACTIVE = 7;
public const int SW_SHOWNOACTIVATE = 4;
public const int SWP_NOACTIVATE = 0x10;
public const int SWP_NOZORDER = 0x4;

Expand Down Expand Up @@ -138,6 +136,15 @@ public enum WindowStyle : uint
[return: MarshalAs(UnmanagedType.Bool)]
public static partial bool GetWindowRect(nint windowHandle, ref Rect rectangle);

[LibraryImport("ntdll", SetLastError = true)]
public static partial int NtOpenProcess
(
out nint ProcessHandle,
int DesiredAccess,
ref ObjectAttributes ObjectAttributes,
ref ClientId ClientId
);

[LibraryImport("kernel32", EntryPoint = "CreateProcessA")]
[return: MarshalAs(UnmanagedType.Bool)]
internal static partial bool CreateProcess
Expand Down Expand Up @@ -203,10 +210,10 @@ out MemoryProtectionFlag oldProtection
[LibraryImport("ntdll")]
internal static partial int NtReadVirtualMemory
(
nint processHandle,
nint baseAddress,
void* buffer,
int size,
nint processHandle,
nint baseAddress,
void* buffer,
int size,
out nint numberOfBytesRead
);

Expand All @@ -219,10 +226,10 @@ out nint numberOfBytesRead
[LibraryImport("ntdll")]
internal static partial int NtWriteVirtualMemory
(
nint processHandle,
nint baseAddress,
void* buffer,
int size,
nint processHandle,
nint baseAddress,
void* buffer,
int size,
out nint numberOfBytesWritten
);

Expand All @@ -241,13 +248,30 @@ internal static partial nint SetWindowPos
(
nint windowHandle,
nint windowHandleInsertAfter,
int x,
int y,
int cx,
int cy,
int x,
int y,
int cx,
int cy,
int wFlags
);

[StructLayout(LayoutKind.Sequential)]
public struct ClientId
{
public nint UniqueProcess;
public nint UniqueThread;
}

[StructLayout(LayoutKind.Sequential)]
public struct ObjectAttributes
{
public int Length;
public nint RootDirectory;
public nint ObjectName;
public uint Attributes;
public nint SecurityDescriptor;
public nint SecurityQualityOfService;
}
[StructLayout(LayoutKind.Sequential)]
public struct ProcessInformation
{
Expand Down
73 changes: 35 additions & 38 deletions AmeisenBotX.Memory/XMemory.cs
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ public unsafe class XMemory : IMemoryApi
private const int FASM_PASSES = 100;

// initial memory pool size
private const int INITIAL_POOL_SIZE = 0; // 16384;
private const int INITIAL_POOL_SIZE = 16384;

// lock needs to be static as FASM isn't thread safe
private static readonly object fasmLock = new();
Expand Down Expand Up @@ -55,12 +55,9 @@ public ulong RpmCallCount
{
get
{
unchecked
{
ulong val = rpmCalls;
rpmCalls = 0;
return val;
}
ulong val = rpmCalls;
rpmCalls = 0;
return val;
}
}

Expand All @@ -69,17 +66,16 @@ public ulong WpmCallCount
{
get
{
unchecked
{
ulong val = wpmCalls;
wpmCalls = 0;
return val;
}
ulong val = wpmCalls;
wpmCalls = 0;
return val;
}
}

private List<AllocationPool> AllocationPools { get; set; }

private bool Initialized { get; set; }

///<inheritdoc cref="IMemoryApi.AllocateMemory"/>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public bool AllocateMemory(uint size, out nint address)
Expand All @@ -102,9 +98,9 @@ public bool AllocateMemory(uint size, out nint address)
// we need a new pool
nint newPoolSize = Math.Max((int)size, INITIAL_POOL_SIZE);
nint newPoolAddress = nint.Zero;
NtAllocateVirtualMemory(ProcessHandle, ref newPoolAddress, 0, ref newPoolSize, AllocationType.Commit, MemoryProtectionFlag.ExecuteReadWrite);
int result = NtAllocateVirtualMemory(ProcessHandle, ref newPoolAddress, 0, ref newPoolSize, AllocationType.Commit, MemoryProtectionFlag.ExecuteReadWrite);

if (newPoolAddress != nint.Zero)
if (result == 0 && newPoolAddress != nint.Zero)
{
AllocationPool pool = new(newPoolAddress, newPoolSize.ToInt32());
AllocationPools.Add(pool);
Expand All @@ -127,10 +123,8 @@ public bool AllocateMemory(uint size, out nint address)
public void Dispose()
{
GC.SuppressFinalize(this);

NtClose(MainThreadHandle);
NtClose(ProcessHandle);

FreeAllMemory();
}

Expand Down Expand Up @@ -186,11 +180,18 @@ public bool FreeMemory(nint address)
{
AmeisenLogger.I.Log("XMemory", $"Freed {size} bytes in Pool[{i}] at: 0x{address:X}");

// if (AllocationPools[i].Allocations.Count == 0 &&
// VirtualFreeEx(ProcessHandle, AllocationPools[i].Address, 0,
// AllocationType.Release)) { AmeisenLogger.I.Log("XMemory", $"Freed
// Pool[{i}] with {AllocationPools[i].Size} bytes at: 0x{address:X}");
// AllocationPools.RemoveAt(i); }
// pool freeing is not needed at the moment, disabling it to reduce memory new allocations
if (false && AllocationPools[i].Allocations.Count == 0)
{
nint addr = AllocationPools[i].Address;
nint s = 0;

if (NtFreeVirtualMemory(ProcessHandle, ref addr, ref s, AllocationType.Release) == 0)
{
AmeisenLogger.I.Log("XMemory", $"Freed Pool[{i}] with {AllocationPools[i].Size} bytes at: 0x{addr:X}");
AllocationPools.RemoveAt(i);
}
}

return true;
}
Expand Down Expand Up @@ -228,9 +229,6 @@ public Rect GetWindowPosition()

return rect;
}

private bool Initialized { get; set; }

///<inheritdoc cref="IMemoryApi.Init"/>
public virtual bool Init(Process process, nint processHandle, nint mainThreadHandle)
{
Expand Down Expand Up @@ -317,19 +315,6 @@ public bool InjectAssembly(IEnumerable<string> asm, nint address, bool patchMemP
}
}

///<inheritdoc cref="IMemoryApi.ProtectMemory"/>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public bool ProtectMemory(nint address, uint size, MemoryProtectionFlag memoryProtection, out MemoryProtectionFlag oldMemoryProtection)
{
#if DEBUG
if (!Initialized) { throw new InvalidOperationException("call Init() before you do anything with this class"); }
if (address == nint.Zero) { throw new ArgumentOutOfRangeException(nameof(address), "address must be > 0"); }
if (size <= 0) { throw new ArgumentOutOfRangeException(nameof(size), "size must be > 0"); }
#endif
nint s = new(size);
return NtProtectVirtualMemory(ProcessHandle, ref address, ref s, memoryProtection, out oldMemoryProtection) == 0;
}

///<inheritdoc cref="IMemoryApi.PatchMemory"/>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public void PatchMemory<T>(nint address, T data) where T : unmanaged
Expand All @@ -347,6 +332,18 @@ public void PatchMemory<T>(nint address, T data) where T : unmanaged
}
}

///<inheritdoc cref="IMemoryApi.ProtectMemory"/>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public bool ProtectMemory(nint address, uint size, MemoryProtectionFlag memoryProtection, out MemoryProtectionFlag oldMemoryProtection)
{
#if DEBUG
if (!Initialized) { throw new InvalidOperationException("call Init() before you do anything with this class"); }
if (address == nint.Zero) { throw new ArgumentOutOfRangeException(nameof(address), "address must be > 0"); }
if (size <= 0) { throw new ArgumentOutOfRangeException(nameof(size), "size must be > 0"); }
#endif
nint s = new(size);
return NtProtectVirtualMemory(ProcessHandle, ref address, ref s, memoryProtection, out oldMemoryProtection) == 0;
}
///<inheritdoc cref="IMemoryApi.Read"/>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public bool Read<T>(nint address, out T value) where T : unmanaged
Expand Down

0 comments on commit b381acb

Please sign in to comment.