From 46d79387e83d48d32773da9a1fea3b595f7aef86 Mon Sep 17 00:00:00 2001
From: Michael Scire <SciresM@gmail.com>
Date: Thu, 16 Apr 2020 17:58:51 -0700
Subject: [PATCH] mesosphere: implement KMemoryBlockManager::UpdateLock

---
 .../kern_k_memory_block_manager.hpp           |  1 +
 .../source/kern_k_memory_block_manager.cpp    | 81 +++++++++++++++++++
 2 files changed, 82 insertions(+)

diff --git a/libraries/libmesosphere/include/mesosphere/kern_k_memory_block_manager.hpp b/libraries/libmesosphere/include/mesosphere/kern_k_memory_block_manager.hpp
index ca24b2b99..19e992611 100644
--- a/libraries/libmesosphere/include/mesosphere/kern_k_memory_block_manager.hpp
+++ b/libraries/libmesosphere/include/mesosphere/kern_k_memory_block_manager.hpp
@@ -94,6 +94,7 @@ namespace ams::kern {
             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);
+            void UpdateLock(KMemoryBlockManagerUpdateAllocator *allocator, KProcessAddress address, size_t num_pages, void (KMemoryBlock::*lock_func)(KMemoryPermission new_perm), KMemoryPermission perm);
 
             iterator FindIterator(KProcessAddress address) const {
                 return this->memory_block_tree.find(KMemoryBlock(address, 1, KMemoryState_Free, KMemoryPermission_None, KMemoryAttribute_None));
diff --git a/libraries/libmesosphere/source/kern_k_memory_block_manager.cpp b/libraries/libmesosphere/source/kern_k_memory_block_manager.cpp
index 016247ba9..1f818e94d 100644
--- a/libraries/libmesosphere/source/kern_k_memory_block_manager.cpp
+++ b/libraries/libmesosphere/source/kern_k_memory_block_manager.cpp
@@ -197,6 +197,7 @@ namespace ams::kern {
                 cur_address += cur_info.GetSize();
                 remaining_pages -= cur_info.GetNumPages();
             }
+            it++;
         }
 
         /* Find the iterator now that we've updated. */
@@ -227,6 +228,86 @@ namespace ams::kern {
         }
     }
 
+    void KMemoryBlockManager::UpdateLock(KMemoryBlockManagerUpdateAllocator *allocator, KProcessAddress address, size_t num_pages, void (KMemoryBlock::*lock_func)(KMemoryPermission new_perm), KMemoryPermission perm) {
+        /* Ensure for auditing that we never end up with an invalid tree. */
+        KScopedMemoryBlockManagerAuditor auditor(this);
+        MESOSPHERE_ASSERT(util::IsAligned(GetInteger(address), PageSize));
+
+        KProcessAddress cur_address = address;
+        size_t remaining_pages = num_pages;
+        iterator it   = this->FindIterator(address);
+        iterator prev = it, next = it;
+        bool check_coalesce_prev = false, check_coalesce_next = false;
+
+        while (remaining_pages > 0) {
+            const size_t remaining_size = remaining_pages * PageSize;
+            KMemoryInfo cur_info = it->GetMemoryInfo();
+
+            /* If we need to, create a new block before and insert it. */
+            if (cur_info.address != GetInteger(cur_address)) {
+                KMemoryBlock *new_block = allocator->Allocate();
+
+                it->Split(new_block, cur_address);
+                it = this->memory_block_tree.insert(*new_block);
+                it++;
+
+                cur_info = it->GetMemoryInfo();
+                cur_address = cur_info.GetAddress();
+            } else if (cur_address == address && cur_address != this->start_address) {
+                /* If there's a previous, we should check for coalescing. */
+                check_coalesce_prev = true;
+                prev--;
+            } else if (cur_info.GetSize() > remaining_size) {
+                /* If we need to, create a new block after and insert it. */
+                KMemoryBlock *new_block = allocator->Allocate();
+
+                it->Split(new_block, cur_address + remaining_size);
+                it = this->memory_block_tree.insert(*new_block);
+
+                cur_info = it->GetMemoryInfo();
+            } else if (cur_info.GetSize() == remaining_size) {
+                /* Otherwise if we can map precisely, we may need to check for coalescing against next block. */
+                next = it;
+                ++next;
+                if (next != this->memory_block_tree.end()) {
+                    check_coalesce_next = true;
+                }
+            }
+
+            /* Call the locked update function. */
+            std::addressof(*it)->*lock_func(perm);
+            cur_address += cur_info.GetSize();
+            remaining_pages -= cur_info.GetNumPages();
+            it++;
+        }
+
+        /* If we should try to coalesce prev, do so. */
+        if (check_coalesce_prev) {
+            it = prev;
+            it++;
+            if (prev->HasSameProperties(*it)) {
+                KMemoryBlock *block = std::addressof(*it);
+                const size_t pages = it->GetNumPages();
+                this->memory_block_tree.erase(it);
+                allocator->Free(block);
+                prev->Add(pages);
+            }
+        }
+
+        /* If we should try to coalesce next, do so. */
+        if (check_coalesce_next) {
+            it = next;
+            it--;
+            if (it->HasSameProperties(*next)) {
+                KMemoryBlock *block = std::addressof(*next);
+                const size_t pages = next->GetNumPages();
+                this->memory_block_tree.erase(next);
+                allocator->Free(block);
+                it->Add(pages);
+            }
+        }
+    }
+
     /* Debug. */
     bool KMemoryBlockManager::CheckState() const {
         /* If we fail, we should dump blocks. */