kern: refactor FindFreeArea region search logic per 20.0.0 changes

libraries/libmesosphere/include/mesosphere/kern_k_memory_block_manager.hpp

@@ -99,6 +99,7 @@ namespace ams::kern {
             Result Initialize(KProcessAddress st, KProcessAddress nd, KMemoryBlockSlabManager *slab_manager);
             void   Finalize(KMemoryBlockSlabManager *slab_manager);
 
+            static bool GetRegionForFindFreeArea(KProcessAddress *out_start, KProcessAddress *out_end, KProcessAddress region_start, size_t region_num_pages, size_t num_pages, size_t alignment, size_t offset, size_t guard_pages);
             KProcessAddress FindFreeArea(KProcessAddress region_start, size_t region_num_pages, size_t num_pages, size_t alignment, size_t offset, size_t guard_pages) const;
 
             void Update(KMemoryBlockManagerUpdateAllocator *allocator, KProcessAddress address, size_t num_pages, KMemoryState state, KMemoryPermission perm, KMemoryAttribute attr, KMemoryBlockDisableMergeAttribute set_disable_attr, KMemoryBlockDisableMergeAttribute clear_disable_attr);

libraries/libmesosphere/source/kern_k_memory_block_manager.cpp

@@ -118,29 +118,78 @@ namespace ams::kern {
         MESOSPHERE_ASSERT(m_memory_block_tree.empty());
     }
 
+    bool KMemoryBlockManager::GetRegionForFindFreeArea(KProcessAddress *out_start, KProcessAddress *out_end, KProcessAddress region_start, size_t region_num_pages, size_t num_pages, size_t alignment, size_t offset, size_t guard_pages) {
+        /* Check that there's room for the pages in the specified region. */
+        if (num_pages + 2 * guard_pages > region_num_pages) {
+            return false;
+        }
+
+        /* Determine the aligned start of the guarded region. */
+        const KProcessAddress guarded_start                     = region_start + guard_pages * PageSize;
+        const KProcessAddress aligned_guarded_start             = util::AlignDown(GetInteger(guarded_start), alignment);
+              KProcessAddress aligned_guarded_start_with_offset = aligned_guarded_start + offset;
+        if (guarded_start > aligned_guarded_start_with_offset) {
+            if (!util::CanAddWithoutOverflow<uintptr_t>(GetInteger(aligned_guarded_start), alignment)) {
+                return false;
+            }
+            aligned_guarded_start_with_offset += alignment;
+        }
+
+        /* Determine the aligned end of the guarded region. */
+        const KProcessAddress guarded_end                     = region_start + ((region_num_pages - (num_pages + guard_pages)) * PageSize);
+        const KProcessAddress aligned_guarded_end             = util::AlignDown(GetInteger(guarded_end), alignment);
+              KProcessAddress aligned_guarded_end_with_offset = aligned_guarded_end + offset;
+        if (aligned_guarded_end_with_offset > guarded_end) {
+            if (aligned_guarded_end < alignment) {
+                return false;
+            }
+            aligned_guarded_end_with_offset -= alignment;
+        }
+
+        /* Check that the extents are valid. */
+        if (aligned_guarded_end_with_offset < aligned_guarded_start_with_offset) {
+            return false;
+        }
+
+        /* Set the output extents. */
+        *out_start = aligned_guarded_start_with_offset;
+        *out_end   = aligned_guarded_end_with_offset;
+        return true;
+    }
+
     KProcessAddress KMemoryBlockManager::FindFreeArea(KProcessAddress region_start, size_t region_num_pages, size_t num_pages, size_t alignment, size_t offset, size_t guard_pages) const {
-        if (num_pages > 0) {
+        /* Determine the range to search in. */
-            const KProcessAddress region_end  = region_start + region_num_pages * PageSize;
+        KProcessAddress search_start = Null<KProcessAddress>;
-            const KProcessAddress region_last = region_end - 1;
+        KProcessAddress search_end   = Null<KProcessAddress>;
-            for (const_iterator it = this->FindIterator(region_start); it != m_memory_block_tree.cend(); it++) {
+        if (this->GetRegionForFindFreeArea(std::addressof(search_start), std::addressof(search_end), region_start, region_num_pages, num_pages, alignment, offset, guard_pages)) {
-                if (region_last < it->GetAddress()) {
+            /* Iterate over blocks in the search space, looking for a suitable one. */
+            for (const_iterator it = this->FindIterator(search_start); it != m_memory_block_tree.cend(); it++) {
+                /* If our block is past the end of our search space, we're done. */
+                if (search_end < it->GetAddress()) {
                     break;
                 }
+
+                /* We only want to consider free blocks. */
                 if (it->GetState() != KMemoryState_Free) {
                     continue;
                 }
 
-                KProcessAddress area = (it->GetAddress() <= GetInteger(region_start)) ? region_start : it->GetAddress();
+                /* Determine the candidate range. */
-                area += guard_pages * PageSize;
+                KProcessAddress candidate_start = Null<KProcessAddress>;
+                KProcessAddress candidate_end   = Null<KProcessAddress>;
+                if (!this->GetRegionForFindFreeArea(std::addressof(candidate_start), std::addressof(candidate_end), it->GetAddress(), it->GetNumPages(), num_pages, alignment, offset, guard_pages)) {
+                    continue;
+                }
 
-                const KProcessAddress offset_area = util::AlignDown(GetInteger(area), alignment) + offset;
+                /* Try the suggested candidate (coercing into the search region if needed). */
-                area = (area <= offset_area) ? offset_area : offset_area + alignment;
+                KProcessAddress candidate = candidate_start;
+                if (candidate < search_start) {
+                    candidate = search_start;
+                }
 
-                const KProcessAddress area_end = area + num_pages * PageSize + guard_pages * PageSize;
+                /* Check if the candidate is valid. */
-                const KProcessAddress area_last = area_end - 1;
+                if (candidate <= search_end && candidate <= candidate_end) {
-
+                    return candidate;
-                if (GetInteger(it->GetAddress()) <= GetInteger(area) && area < area_last && area_last <= region_last && GetInteger(area_last) <= GetInteger(it->GetLastAddress())) {
-                    return area;
                 }
             }
         }

libraries/libmesosphere/source/kern_k_page_table_base.cpp

@@ -1333,34 +1333,37 @@ namespace ams::kern {
     KProcessAddress KPageTableBase::FindFreeArea(KProcessAddress region_start, size_t region_num_pages, size_t num_pages, size_t alignment, size_t offset, size_t guard_pages) const {
         KProcessAddress address = Null<KProcessAddress>;
 
-        if (num_pages <= region_num_pages) {
+        KProcessAddress search_start = Null<KProcessAddress>;
+        KProcessAddress search_end   = Null<KProcessAddress>;
+        if (m_memory_block_manager.GetRegionForFindFreeArea(std::addressof(search_start), std::addressof(search_end), region_start, region_num_pages, num_pages, alignment, offset, guard_pages)) {
             if (this->IsAslrEnabled()) {
                 /* Try to directly find a free area up to 8 times. */
                 for (size_t i = 0; i < 8; i++) {
-                    const size_t random_offset = KSystemControl::GenerateRandomRange(0, (region_num_pages - num_pages - guard_pages) * PageSize / alignment) * alignment;
+                    const size_t random_offset = KSystemControl::GenerateRandomRange(0, (search_end - search_start) / alignment) * alignment;
-                    const KProcessAddress candidate = util::AlignDown(GetInteger(region_start + random_offset), alignment) + offset;
+                    const KProcessAddress candidate = search_start + random_offset;
 
                     KMemoryBlockManager::const_iterator it = m_memory_block_manager.FindIterator(candidate);
                     MESOSPHERE_ABORT_UNLESS(it != m_memory_block_manager.end());
 
                     if (it->GetState() != KMemoryState_Free) { continue; }
-                    if (!(region_start <= candidate)) { continue; }
                     if (!(it->GetAddress() + guard_pages * PageSize <= GetInteger(candidate))) { continue; }
                     if (!(candidate + (num_pages + guard_pages) * PageSize - 1 <= it->GetLastAddress())) { continue; }
-                    if (!(candidate + (num_pages + guard_pages) * PageSize - 1 <= region_start + region_num_pages * PageSize - 1)) { continue; }
 
                     address = candidate;
                     break;
                 }
+
                 /* Fall back to finding the first free area with a random offset. */
                 if (address == Null<KProcessAddress>) {
                     /* NOTE: Nintendo does not account for guard pages here. */
                     /* This may theoretically cause an offset to be chosen that cannot be mapped. */
                     /* We will account for guard pages. */
-                    const size_t offset_pages = KSystemControl::GenerateRandomRange(0, region_num_pages - num_pages - guard_pages);
+                    const size_t offset_blocks = KSystemControl::GenerateRandomRange(0, (search_end - search_start) / alignment);
-                    address = m_memory_block_manager.FindFreeArea(region_start + offset_pages * PageSize, region_num_pages - offset_pages, num_pages, alignment, offset, guard_pages);
+                    const auto   region_end    = region_start + region_num_pages * PageSize;
+                    address = m_memory_block_manager.FindFreeArea(search_start + offset_blocks * alignment, (region_end - (search_start + offset_blocks * alignment)) / PageSize, num_pages, alignment, offset, guard_pages);
                 }
             }
+
             /* Find the first free area. */
             if (address == Null<KProcessAddress>) {
                 address = m_memory_block_manager.FindFreeArea(region_start, region_num_pages, num_pages, alignment, offset, guard_pages);