From e8ca05ca0add2096ff3b3d1edd11df278e663139 Mon Sep 17 00:00:00 2001
From: Ben Vanik <ben.vanik@gmail.com>
Date: Sun, 5 Jan 2014 11:19:02 -0800
Subject: [PATCH] MmQueryAllocationSize.

---
 src/alloy/memory.h                  |  2 ++
 src/xenia/cpu/xenon_memory.cc       | 36 +++++++++++++++++++++++++++++
 src/xenia/cpu/xenon_memory.h        |  2 ++
 src/xenia/kernel/xboxkrnl_memory.cc | 25 ++++++++++++++++++++
 4 files changed, 65 insertions(+)

diff --git a/src/alloy/memory.h b/src/alloy/memory.h
index 8ceb1ee63..9fa8c11fd 100644
--- a/src/alloy/memory.h
+++ b/src/alloy/memory.h
@@ -48,6 +48,8 @@ public:
       uint32_t alignment = 0x20) = 0;
   virtual int HeapFree(uint64_t address, size_t size) = 0;
 
+  virtual size_t QuerySize(uint64_t base_address) = 0;
+
   virtual int Protect(uint64_t address, size_t size, uint32_t access) = 0;
   virtual uint32_t QueryProtect(uint64_t address) = 0;
 
diff --git a/src/xenia/cpu/xenon_memory.cc b/src/xenia/cpu/xenon_memory.cc
index 909653f63..15678c744 100644
--- a/src/xenia/cpu/xenon_memory.cc
+++ b/src/xenia/cpu/xenon_memory.cc
@@ -99,6 +99,7 @@ public:
   uint64_t Alloc(uint64_t base_address, size_t size,
                  uint32_t flags, uint32_t alignment);
   uint64_t Free(uint64_t address, size_t size);
+  size_t QuerySize(uint64_t base_address);
 
   void Dump();
 
@@ -321,6 +322,26 @@ int XenonMemory::HeapFree(uint64_t address, size_t size) {
   }
 }
 
+size_t XenonMemory::QuerySize(uint64_t base_address) {
+  if (base_address >= XENON_MEMORY_VIRTUAL_HEAP_LOW &&
+      base_address < XENON_MEMORY_VIRTUAL_HEAP_HIGH) {
+    return virtual_heap_->QuerySize(base_address);
+  } else if (base_address >= XENON_MEMORY_PHYSICAL_HEAP_LOW &&
+             base_address < XENON_MEMORY_PHYSICAL_HEAP_HIGH) {
+    return physical_heap_->QuerySize(base_address);
+  } else {
+    // A placed address. Decommit.
+    uint8_t* p = Translate(base_address);
+    MEMORY_BASIC_INFORMATION mem_info;
+    if (VirtualQuery(p, &mem_info, sizeof(mem_info))) {
+      return mem_info.RegionSize;
+    } else {
+      // Error.
+      return 0;
+    }
+  }
+}
+
 int XenonMemory::Protect(uint64_t address, size_t size, uint32_t access) {
   uint8_t* p = Translate(address);
 
@@ -522,6 +543,21 @@ uint64_t XenonMemoryHeap::Free(uint64_t address, size_t size) {
   return (uint64_t)real_size;
 }
 
+size_t XenonMemoryHeap::QuerySize(uint64_t base_address) {
+  uint8_t* p = memory_->Translate(base_address);
+
+  // Heap allocated address.
+  size_t heap_guard_size = FLAGS_heap_guard_pages * 4096;
+  p -= heap_guard_size;
+  size_t real_size = mspace_usable_size(p);
+  real_size -= heap_guard_size * 2;
+  if (!real_size) {
+    return 0;
+  }
+
+  return real_size;
+}
+
 void XenonMemoryHeap::Dump() {
   XELOGI("XenonMemoryHeap::Dump - %s",
          is_physical_ ? "physical" : "virtual");
diff --git a/src/xenia/cpu/xenon_memory.h b/src/xenia/cpu/xenon_memory.h
index 1974b8124..96ba352fa 100644
--- a/src/xenia/cpu/xenon_memory.h
+++ b/src/xenia/cpu/xenon_memory.h
@@ -33,6 +33,8 @@ public:
       uint32_t alignment = 0x20);
   virtual int HeapFree(uint64_t address, size_t size);
 
+  virtual size_t QuerySize(uint64_t base_address);
+
   virtual int Protect(uint64_t address, size_t size, uint32_t access);
   virtual uint32_t QueryProtect(uint64_t address);
 
diff --git a/src/xenia/kernel/xboxkrnl_memory.cc b/src/xenia/kernel/xboxkrnl_memory.cc
index f08f3f65c..573658f13 100644
--- a/src/xenia/kernel/xboxkrnl_memory.cc
+++ b/src/xenia/kernel/xboxkrnl_memory.cc
@@ -330,6 +330,30 @@ SHIM_CALL MmQueryAddressProtect_shim(
 }
 
 
+uint32_t xeMmQueryAllocationSize(uint32_t base_address) {
+  KernelState* state = shared_kernel_state_;
+  XEASSERTNOTNULL(state);
+
+  size_t size = state->memory()->QuerySize(base_address);
+
+  return (uint32_t)size;
+}
+
+
+SHIM_CALL MmQueryAllocationSize_shim(
+    PPCContext* ppc_state, KernelState* state) {
+  uint32_t base_address = SHIM_GET_ARG_32(0);
+
+  XELOGD(
+      "MmQueryAllocationSize(%.8X)",
+      base_address);
+
+  uint32_t result = xeMmQueryAllocationSize(base_address);
+
+  SHIM_SET_RETURN(result);
+}
+
+
 SHIM_CALL MmQueryStatistics_shim(
     PPCContext* ppc_state, KernelState* state) {
   uint32_t stats_ptr = SHIM_GET_ARG_32(0);
@@ -434,6 +458,7 @@ void xe::kernel::xboxkrnl::RegisterMemoryExports(
   SHIM_SET_MAPPING("xboxkrnl.exe", MmAllocatePhysicalMemoryEx, state);
   SHIM_SET_MAPPING("xboxkrnl.exe", MmFreePhysicalMemory, state);
   SHIM_SET_MAPPING("xboxkrnl.exe", MmQueryAddressProtect, state);
+  SHIM_SET_MAPPING("xboxkrnl.exe", MmQueryAllocationSize, state);
   SHIM_SET_MAPPING("xboxkrnl.exe", MmQueryStatistics, state);
   SHIM_SET_MAPPING("xboxkrnl.exe", MmGetPhysicalAddress, state);
 }