heap.h

Go to the documentation of this file.

 
#pragma once
 
#include <atomic>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <mutex>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>
 
#include <rex/platform.h>
#include <rex/memory/mapped_ptr.h>
#include <rex/memory/virtual_memory.h>
#include <rex/ppc/recompiled_code.h>
#include <rex/platform/lock_order.h>
#include <rex/types.h>
 
namespace rex::kernel {
 
inline constexpr u32 kPageSize4KiB = 4 * 1024;

inline constexpr u32 kPageSize64KiB = 64 * 1024;
 
inline constexpr u32 kPageSize16MiB = 16 * 1024 * 1024;
 
}  // namespace rex::kernel
 
namespace rex::kernel::detail {

template <typename T, bool IsOpaque = std::is_void_v<T> || std::is_function_v<T>>
struct PointeeAlignment {
  static constexpr u32 kValue = 1;
};

template <typename T>
struct PointeeAlignment<T, false> {
  static constexpr u32 kValue = static_cast<u32>(alignof(T));
};
 
}  // namespace rex::kernel::detail
 
namespace rex::kernel {
 
class GuestAllocator;
class GuestAddressSpace;
class WorkspaceHeap;

enum SystemHeapFlag : u8 {
  kSystemHeapVirtual = 1 << 0,
  kSystemHeapPhysical = 1 << 1,
 
  kSystemHeapDefault = kSystemHeapVirtual,
};

class ISystemHeap {
 public:
  virtual ~ISystemHeap() = default;

  virtual u32 SystemHeapAlloc(u32 size, u32 alignment = 0x20,
                              u32 system_heap_flags = kSystemHeapDefault) = 0;

  virtual void SystemHeapFree(u32 address, u32* out_region_size = nullptr) = 0;
 
  [[nodiscard]] virtual bool IsAddressMapped(u32 address) const = 0;

  virtual u32 AllocPhysical(u32 size, u32 page_size, u32 protect, u32 min_address, u32 max_address,
                            u32 alignment) = 0;
};
 
enum class HeapType : u8 {
  kGuestVirtual,
  kGuestXEX, 
  kGuestPhysical,
};
 
enum MemoryAllocationFlag : u8 {
  kMemoryAllocationReserve = 1 << 0,
  kMemoryAllocationCommit = 1 << 1,
  kMemoryAllocationTitle = 1 << 2,
};
 
enum MemoryProtectFlag : u8 {
  kMemoryProtectRead = 1 << 0,
  kMemoryProtectWrite = 1 << 1,
  kMemoryProtectNoCache = 1 << 2,
  kMemoryProtectWriteCombine = 1 << 3,
 
  kMemoryProtectNoAccess = 0,
};

struct HeapAllocation {
  u32 base_address; 
  u32 allocation_base;
  u32 allocation_protect;
  u32 allocation_size;
  u32 region_size;
  u32 state;
  u32 protect; 
};
 
union PageEntry {
  u64 qword; 
  struct {
    u32 base_address : 20;
    u32 region_page_count : 20;
    u32 allocation_protect : 4;
    u32 current_protect : 4;
    u32 state : 2;
    u32 title : 1;
    u32 reserved : 13;
  };
};
 
class REX_KERNEL_API BaseHeap {
 public:
  virtual ~BaseHeap();

  u32 GetHeapBase() const { return heap_base_; }
 
  u32 GetHeapSize() const { return heap_size_; }
 
  u32 GetPageSize() const { return page_size_; }
 
  HeapType GetHeapType() const { return heap_type_; }

  u32 GetHostAddressOffset() const { return host_address_offset_; }

  template <typename T = u8*>
  T TranslateRelative(size_t relative_address) const {
    return reinterpret_cast<T>(membase_ + heap_base_ + host_address_offset_ + relative_address);
  }

  virtual void Dispose();
 
  u32 GetTotalPageCount() const { return static_cast<u32>(page_table_.size()); }
 
  u32 GetUnreservedPageCount() const { return unreserved_page_count_; }
 
  u32 GetReservedPageCount() const { return GetTotalPageCount() - GetUnreservedPageCount(); }
 
  u32 GetCommittedPages() const { return committed_pages_.load(std::memory_order_relaxed); }

  void PlotProfile() const;

  virtual bool Alloc(u32 size, u32 alignment, u32 allocation_type, u32 protect, bool top_down,
                     u32* out_address);

  virtual bool AllocFixed(u32 base_address, u32 size, u32 alignment, u32 allocation_type,
                          u32 protect);

  virtual bool AllocRange(u32 low_address, u32 high_address, u32 size, u32 alignment,
                          u32 allocation_type, u32 protect, bool top_down, u32* out_address);

  virtual bool AllocSystemHeap(u32 size, u32 alignment, u32 allocation_type, u32 protect,
                               bool top_down, u32* out_address);

  virtual bool Decommit(u32 address, u32 size);

  virtual bool Release(u32 address, u32* out_region_size = nullptr);

  bool ReleaseRange(u32 address, u32 size, u32* out_address, u32* out_size);

  virtual bool Protect(u32 address, u32 size, u32 protect, u32* old_protect = nullptr);

  bool QueryRegionInfo(u32 base_address, HeapAllocation* out_info);

  bool QuerySize(u32 address, u32* out_size);

  bool QueryProtect(u32 address, u32* out_protect);

  rex::memory::PageAccess QueryRangeAccess(u32 low_address, u32 high_address);

  rex::memory::PageAccess GetHostPageAccess(u32 host_page_number) const;

  bool ReapplyHostPageAccess(void* host_address, bool is_write);
 
 protected:
  friend class PhysicalHeap;
 
  BaseHeap();

  rex::memory::PageAccess GetHostPageGuestAccess(u32 host_page_number) const;

  bool SyncHostPageAccess(u32 start_page_number, u32 end_page_number);

  virtual std::unique_lock<rex::platform::RankedMutex> AcquireHostPageReconcileLock() const;

  u32 GetSystemRangeBase() const;
 
  bool AllocFixedLocked(u32 base_address, u32 size, u32 alignment, u32 allocation_type,
                        u32 protect);
 
  bool AllocRangeLocked(u32 low_address, u32 high_address, u32 size, u32 alignment,
                        u32 allocation_type, u32 protect, bool top_down, u32* out_address);
 
  bool DecommitLocked(u32 address, u32 size);
 
  bool ReleaseLocked(u32 address, u32* out_region_size);
 
  bool ReleaseRangeLocked(u32 address, u32 size, u32* out_address, u32* out_size);
 
  bool ProtectLocked(u32 address, u32 size, u32 protect, u32* old_protect);

  bool CheckProtectRangeLocked(u32 address, u32 size, u32* out_start_page, u32* out_end_page);

  virtual bool IsHostPageWriteWatched([[maybe_unused]] u32 host_page_number) const { return false; }

  void OnPageCommitted(u32 page_number);

  void OnPageDecommitted(u32 page_number);

  void Initialize(kernel::GuestAddressSpace* memory, u8* membase, HeapType heap_type, u32 heap_base,
                  u32 heap_size, u32 page_size, u32 host_address_offset = 0);
 
  kernel::GuestAddressSpace*
      memory_; 
  u8* membase_;
  HeapType heap_type_;
  u32 heap_base_;
  u32 heap_size_;
  u32 page_size_;
  u32 page_size_shift_ = 0;
  u32 host_address_offset_;
  u32 unreserved_page_count_ = 0;
  std::atomic<u32> committed_pages_{0};
  const char* profile_pool_ = nullptr;
  const char* profile_plot_ = nullptr;
  std::recursive_mutex heap_mutex_; 
  std::vector<PageEntry> page_table_;
};
 
class REX_KERNEL_API VirtualHeap : public BaseHeap {
 public:
  VirtualHeap();
  ~VirtualHeap() override;

  void Initialize(kernel::GuestAddressSpace* memory, u8* membase, HeapType heap_type, u32 heap_base,
                  u32 heap_size, u32 page_size);
};

class REX_KERNEL_API PhysicalHeap : public BaseHeap {
 public:
  PhysicalHeap();
  ~PhysicalHeap() override;

  void Initialize(kernel::GuestAddressSpace* memory, u8* membase, HeapType heap_type, u32 heap_base,
                  u32 heap_size, u32 page_size, VirtualHeap* parent_heap);

  bool Alloc(u32 size, u32 alignment, u32 allocation_type, u32 protect, bool top_down,
             u32* out_address) override;
  bool AllocFixed(u32 base_address, u32 size, u32 alignment, u32 allocation_type,
                  u32 protect) override;
  bool AllocRange(u32 low_address, u32 high_address, u32 size, u32 alignment, u32 allocation_type,
                  u32 protect, bool top_down, u32* out_address) override;
  bool AllocSystemHeap(u32 size, u32 alignment, u32 allocation_type, u32 protect, bool top_down,
                       u32* out_address) override;
  bool Decommit(u32 address, u32 size) override;
  bool Release(u32 base_address, u32* out_region_size = nullptr) override;
  bool Protect(u32 address, u32 size, u32 protect, u32* old_protect = nullptr) override;

  void EnableAccessCallbacks(u32 physical_address, u32 length,
                             bool enable_invalidation_notifications, bool enable_data_providers);
  bool TriggerCallbacks(u32 virtual_address, u32 length, bool is_write, bool unwatch_exact_range,
                        bool unprotect = true);
 
  bool TriggerCallbacksLocked(u32 virtual_address, u32 length, bool is_write,
                              bool unwatch_exact_range, bool unprotect = true);

  u32 GetPhysicalAddress(u32 address) const;
 
 protected:
  bool IsHostPageWriteWatched(u32 host_page_number) const override;

  std::unique_lock<rex::platform::RankedMutex> AcquireHostPageReconcileLock() const override;
 
  VirtualHeap* parent_heap_; 
 
  u32 system_page_size_;
  u32 system_page_count_;

  struct SystemPageFlagsBlock {
    u64 notify_on_invalidation;
  };

  std::vector<SystemPageFlagsBlock> system_page_flags_;
};
 
enum class AccessViolationResult : u8 {
  kResumed,
  kMMIOMiss,
  kUnhandled,
};

class REX_KERNEL_API GuestAddressSpace : public ISystemHeap {
 public:
  GuestAddressSpace();
  ~GuestAddressSpace();

  bool Initialize();

  const std::filesystem::path& GetFileName() const { return file_name_; }
 
  u8* GetVirtualMembase() const { return virtual_membase_; }

  rex::memory::GuestMembase GetMembase() const {
    return rex::memory::GuestMembase(virtual_membase_);
  }

  template <typename T = u8*>
  T TranslateVirtual(u32 guest_address) const {
    assert(!LookupHeap(guest_address) ||
           LookupHeap(guest_address)->GetHostAddressOffset() ==
               rex::memory::detail::PhysicalHostOffset(guest_address));
    return rex::memory::GuestPtr<T>(virtual_membase_, guest_address);
  }

  template <typename T = u8*>
  T TryTranslateVirtual(u32 guest_address, u32 required = kMemoryProtectRead) const {
    static_assert(std::is_pointer_v<T>, "TryTranslateVirtual returns a pointer type");
    constexpr u32 kAlignment = detail::PointeeAlignment<std::remove_pointer_t<T>>::kValue;
    if (guest_address == 0 || guest_address % kAlignment != 0) {
      return nullptr;
    }
    const BaseHeap* heap = LookupHeap(guest_address);
    u32 protect = 0;
    if (heap == nullptr || !const_cast<BaseHeap*>(heap)->QueryProtect(guest_address, &protect) ||
        (protect & required) != required) {
      return nullptr;
    }
    return rex::memory::GuestPtr<T>(virtual_membase_, guest_address);
  }
 
  u8* GetPhysicalMembase() const { return physical_membase_; }

  template <typename T = u8*>
  T TranslatePhysical(u32 guest_address) const {
    return reinterpret_cast<T>(physical_membase_ + (guest_address & 0x1FFFFFFF));
  }

  u32 HostToGuestVirtual(const void* host_address) const;

  u32 GetPhysicalAddress(u32 address) const;
 
  void Zero(u32 address, u32 size) const;
 
  void Fill(u32 address, u32 size, u8 value) const;

  void Copy(u32 dest, u32 src, u32 size) const;

  AccessViolationResult HandleAccessViolation(void* host_address, bool is_write);

  using PhysicalMemoryInvalidationCallback = std::pair<u32, u32> (*)(void* context_ptr,
                                                                     u32 physical_address_start,
                                                                     u32 length, bool exact_range);
  void* RegisterPhysicalMemoryInvalidationCallback(PhysicalMemoryInvalidationCallback callback,
                                                   void* callback_context);
  void UnregisterPhysicalMemoryInvalidationCallback(void* callback_handle);

  void EnablePhysicalMemoryAccessCallbacks(u32 physical_address, u32 length,
                                           bool enable_invalidation_notifications,
                                           bool enable_data_providers);

  u32 SystemHeapAlloc(u32 size, u32 alignment = 0x20,
                      u32 system_heap_flags = kSystemHeapDefault) override;
 
  void SystemHeapFree(u32 address, u32* out_region_size = nullptr) override;
 
  u32 AllocPhysical(u32 size, u32 page_size, u32 protect, u32 min_address, u32 max_address,
                    u32 alignment) override;
 
  bool IsAddressMapped(u32 address) const override { return LookupHeap(address) != nullptr; }
 
  const BaseHeap* LookupHeap(u32 address) const;
 
  BaseHeap* LookupHeap(u32 address) {
    return const_cast<BaseHeap*>(
        const_cast<const kernel::GuestAddressSpace*>(this)->LookupHeap(address));
  }
 
  BaseHeap* LookupHeapByType(bool physical, u32 page_size);

  void GetHeapsPageStats(const BaseHeap* const* provided_heaps, size_t heaps_count,
                         u32& unreserved_pages, u32& reserved_pages, u32& used_pages,
                         u32& reserved_bytes);
 
  VirtualHeap* GetPhysicalHeap();

  void SetBudget(u32 megabytes);

  [[nodiscard]] u32 GetTitlePages() const;

  [[nodiscard]] u32 GetBudgetPages() const;

  void PlotProfile();

  bool InitializeFunctionTable(u32 code_base, u32 code_size, u32 image_base, u32 image_size);
  bool DestroyFunctionTable(u32 code_base);
  bool SetFunction(u32 guest_address, PPCFunc* host_function);

  [[nodiscard]] std::unique_lock<rex::platform::RankedMutex> LockWatches() {
    return std::unique_lock(watch_mutex_);
  }
 
 private:
#if REX_PLATFORM_MAC
  int MapViewsMac();
#endif
  int MapViews(u8* mapping_base);
  void UnmapViews();
  bool ReserveMMIOBlocks() const;
 
  bool AccessViolationCallback(void* host_address, bool is_write);

  bool ReserveBudgetPages(u32 pages);

  void ReleaseBudgetPages(u32 pages);
 
  std::filesystem::path file_name_;
  u32 system_page_size_ = 0;
  u32 system_allocation_granularity_ = 0;
  u8* virtual_membase_ = nullptr;
  u8* physical_membase_ = nullptr;
 
  struct FunctionTableEntry {
    u32 table_base;
    u32 code_base;
    u32 code_size;
    u32 thunk_reserve;
  };
  std::vector<FunctionTableEntry> function_tables_;
  mutable std::mutex function_tables_mutex_;
 
  rex::memory::FileMappingHandle mapping_ = rex::memory::kFileMappingHandleInvalid;
  u8* mapping_base_ = nullptr;
  union {
    struct {
      u8* v00000000;
      u8* v40000000;
      u8* v7f000000;
      u8* v80000000;
      u8* v90000000;
      u8* va0000000;
      u8* vc0000000;
      u8* ve0000000;
      u8* physical;
    };
    u8* all_views[9];
  } views_ = {{0, 0, 0, 0, 0, 0, 0, 0, 0}};
 
  struct {
    VirtualHeap v00000000;
    VirtualHeap v40000000;
    VirtualHeap v80000000;
    VirtualHeap v90000000;
 
    VirtualHeap physical;
    PhysicalHeap va0000000;
    PhysicalHeap vc0000000;
    PhysicalHeap ve0000000;
  } heaps_;
 
  friend class BaseHeap;
  friend class GuestAllocator;
  friend class WorkspaceHeap;
 
  friend class PhysicalHeap;
  rex::platform::RankedMutex watch_mutex_{rex::platform::LockRank::kLeaf};
  std::vector<std::unique_ptr<std::pair<PhysicalMemoryInvalidationCallback, void*>>>
      physical_memory_invalidation_callbacks_;
  std::atomic<u32> title_pages_{0};
  std::atomic<u32> budget_pages_{0};
  std::mutex allocators_mutex_;
  std::vector<GuestAllocator*> allocators_;      
  std::vector<const WorkspaceHeap*> workspaces_; 
};
 
}  // namespace rex::kernel