kern: implement KMemoryBlockManager::Update

libraries/libmesosphere/include/mesosphere/kern_k_memory_block.hpp

@@ -159,6 +159,8 @@ namespace ams::kern {
 
     enum KMemoryAttribute : u8 {
         KMemoryAttribute_None           = 0x00,
+        KMemoryAttribute_Mask           = 0x7F,
+        KMemoryAttribute_DontCareMask   = 0x80,
 
         KMemoryAttribute_Locked         = ams::svc::MemoryAttribute_Locked,
         KMemoryAttribute_IpcLocked      = ams::svc::MemoryAttribute_IpcLocked,
@@ -166,6 +168,54 @@ namespace ams::kern {
         KMemoryAttribute_Uncached       = ams::svc::MemoryAttribute_Uncached,
     };
 
+    static_assert((KMemoryAttribute_Mask & KMemoryAttribute_DontCareMask) == 0);
+    static_assert(static_cast<typename std::underlying_type<KMemoryAttribute>::type>(~(KMemoryAttribute_Mask | KMemoryAttribute_DontCareMask)) == 0);
+
+    struct KMemoryInfo {
+        uintptr_t address;
+        size_t size;
+        KMemoryState state;
+        KMemoryPermission perm;
+        KMemoryAttribute  attribute;
+        KMemoryPermission original_perm;
+        u16 ipc_lock_count;
+        u16 device_use_count;
+
+        constexpr ams::svc::MemoryInfo GetSvcMemoryInfo() const {
+            ams::svc::MemoryInfo svc_info = {};
+
+            svc_info.addr             = this->address;
+            svc_info.size             = this->size;
+            svc_info.state            = static_cast<ams::svc::MemoryState>(this->state & KMemoryState_Mask);
+            svc_info.attr             = static_cast<ams::svc::MemoryAttribute>(this->attribute & KMemoryAttribute_Mask);
+            svc_info.perm             = static_cast<ams::svc::MemoryPermission>(this->perm & KMemoryPermission_UserMask);
+            svc_info.ipc_refcount     = this->ipc_lock_count;
+            svc_info.device_refcount  = this->device_use_count;
+
+            return svc_info;
+        }
+
+        constexpr uintptr_t GetAddress() const {
+            return this->address;
+        }
+
+        constexpr size_t GetSize() const {
+            return this->size;
+        }
+
+        constexpr size_t GetNumPages() const {
+            return this->GetSize() / PageSize;
+        }
+
+        constexpr uintptr_t GetEndAddress() const {
+            return this->GetAddress() + this->GetSize();
+        }
+
+        constexpr uintptr_t GetLastAddress() const {
+            return this->GetEndAddress() - 1;
+        }
+    };
+
     class KMemoryBlock : public util::IntrusiveRedBlackTreeBaseNode<KMemoryBlock> {
         private:
             KProcessAddress address;
@@ -206,6 +256,21 @@ namespace ams::kern {
             constexpr KProcessAddress GetLastAddress() const {
                 return this->GetEndAddress() - 1;
             }
+
+            constexpr KMemoryInfo GetMemoryInfo() const {
+                KMemoryInfo info = {};
+
+                info.address          = GetInteger(this->GetAddress());
+                info.size             = this->GetSize();
+                info.state            = this->memory_state;
+                info.perm             = this->perm;
+                info.attribute        = this->attribute;
+                info.original_perm    = this->original_perm;
+                info.ipc_lock_count   = this->ipc_lock_count;
+                info.device_use_count = this->device_use_count;
+
+                return info;
+            }
         public:
             constexpr KMemoryBlock()
                 : address(), num_pages(), memory_state(KMemoryState_None), ipc_lock_count(), device_use_count(), perm(), original_perm(), attribute()
@@ -213,6 +278,12 @@ namespace ams::kern {
                 /* ... */
             }
 
+            constexpr KMemoryBlock(KProcessAddress addr, size_t np, KMemoryState ms, KMemoryPermission p, KMemoryAttribute attr)
+                : address(addr), num_pages(np), memory_state(ms), ipc_lock_count(0), device_use_count(0), perm(p), original_perm(KMemoryPermission_None), attribute(attr)
+            {
+                /* ... */
+            }
+
             constexpr void Initialize(KProcessAddress addr, size_t np, KMemoryState ms, KMemoryPermission p, KMemoryAttribute attr) {
                 MESOSPHERE_ASSERT_THIS();
                 this->address           = addr;
@@ -224,6 +295,67 @@ namespace ams::kern {
                 this->original_perm     = KMemoryPermission_None;
                 this->attribute         = attr;
             }
+
+            constexpr bool HasProperties(KMemoryState s, KMemoryPermission p, KMemoryAttribute a) const {
+                MESOSPHERE_ASSERT_THIS();
+                constexpr auto AttributeIgnoreMask = KMemoryAttribute_DontCareMask | KMemoryAttribute_IpcLocked | KMemoryAttribute_DeviceShared;
+                return this->memory_state == s && this->perm == p && (this->attribute | AttributeIgnoreMask) == (a | AttributeIgnoreMask);
+            }
+
+            constexpr bool HasSameProperties(const KMemoryBlock &rhs) const {
+                MESOSPHERE_ASSERT_THIS();
+
+                return this->memory_state     == rhs.memory_state   &&
+                       this->perm             == rhs.perm           &&
+                       this->original_perm    == rhs.original_perm  &&
+                       this->attribute        == rhs.attribute      &&
+                       this->ipc_lock_count   == rhs.ipc_lock_count &&
+                       this->device_use_count == rhs.device_use_count;
+            }
+
+            constexpr bool Contains(KProcessAddress addr) const {
+                MESOSPHERE_ASSERT_THIS();
+
+                return this->GetAddress() <= addr && addr <= this->GetEndAddress();
+            }
+
+            constexpr void Add(size_t np) {
+                MESOSPHERE_ASSERT_THIS();
+                MESOSPHERE_ASSERT(np > 0);
+                MESOSPHERE_ASSERT(this->GetAddress() + np * PageSize - 1 < this->GetEndAddress() + np * PageSize - 1);
+
+                this->num_pages += np;
+            }
+
+            constexpr void Update(KMemoryState s, KMemoryPermission p, KMemoryAttribute a) {
+                MESOSPHERE_ASSERT_THIS();
+                MESOSPHERE_ASSERT(this->original_perm == KMemoryPermission_None);
+                MESOSPHERE_ASSERT((this->attribute & KMemoryAttribute_IpcLocked) == 0);
+
+                this->memory_state = s;
+                this->perm         = p;
+                this->attribute    = static_cast<KMemoryAttribute>(a | (this->attribute & (KMemoryAttribute_IpcLocked | KMemoryAttribute_DeviceShared)));
+            }
+
+            constexpr void Split(KMemoryBlock *block, KProcessAddress addr) {
+                MESOSPHERE_ASSERT_THIS();
+                MESOSPHERE_ASSERT(this->GetAddress() < addr);
+                MESOSPHERE_ASSERT(this->Contains(addr));
+                MESOSPHERE_ASSERT(util::IsAligned(GetInteger(addr), PageSize));
+
+                block->address          = this->address;
+                block->num_pages        = (addr - this->GetAddress()) / PageSize;
+                block->memory_state     = this->memory_state;
+                block->ipc_lock_count   = this->ipc_lock_count;
+                block->device_use_count = this->device_use_count;
+                block->perm             = this->perm;
+                block->original_perm    = this->original_perm;
+                block->attribute        = this->attribute;
+
+                this->address = addr;
+                this->num_pages -= block->num_pages;
+            }
     };
+    static_assert(std::is_trivially_destructible<KMemoryBlock>::value);
 
 }

libraries/libmesosphere/include/mesosphere/kern_k_memory_block_manager.hpp

@@ -20,9 +20,63 @@
 
 namespace ams::kern {
 
+    class KMemoryBlockManagerUpdateAllocator {
+        public:
+            static constexpr size_t NumBlocks = 2;
+        private:
+            KMemoryBlock *blocks[NumBlocks];
+            size_t index;
+            KMemoryBlockSlabManager *slab_manager;
+            Result result;
+        public:
+            explicit KMemoryBlockManagerUpdateAllocator(KMemoryBlockSlabManager *sm) : blocks(), index(), slab_manager(sm), result(svc::ResultOutOfResource()) {
+                for (size_t i = 0; i < NumBlocks; i++) {
+                    this->blocks[i] = this->slab_manager->Allocate();
+                    if (this->blocks[i] == nullptr) {
+                        this->result = svc::ResultOutOfResource();
+                        return;
+                    }
+                }
+
+                this->result = ResultSuccess();
+            }
+
+            ~KMemoryBlockManagerUpdateAllocator() {
+                for (size_t i = 0; i < NumBlocks; i++) {
+                    if (this->blocks[i] != nullptr) {
+                        this->slab_manager->Free(this->blocks[i]);
+                    }
+                }
+            }
+
+            Result GetResult() const {
+                return this->result;
+            }
+
+            KMemoryBlock *Allocate() {
+                MESOSPHERE_ABORT_UNLESS(this->index < NumBlocks);
+                MESOSPHERE_ABORT_UNLESS(this->blocks[this->index] != nullptr);
+                KMemoryBlock *block = nullptr;
+                std::swap(block, this->blocks[this->index++]);
+                return block;
+            }
+
+            void Free(KMemoryBlock *block) {
+                MESOSPHERE_ABORT_UNLESS(this->index <= NumBlocks);
+                MESOSPHERE_ABORT_UNLESS(block != nullptr);
+                if (this->index == 0) {
+                    this->slab_manager->Free(block);
+                } else {
+                    this->blocks[--this->index] = block;
+                }
+            }
+    };
+
     class KMemoryBlockManager {
         public:
             using MemoryBlockTree = util::IntrusiveRedBlackTreeBaseTraits<KMemoryBlock>::TreeType<KMemoryBlock>;
+            using iterator = MemoryBlockTree::iterator;
+            using const_iterator = MemoryBlockTree::const_iterator;
         private:
             MemoryBlockTree memory_block_tree;
             KProcessAddress start;
@@ -33,18 +87,33 @@ namespace ams::kern {
             Result Initialize(KProcessAddress st, KProcessAddress nd, KMemoryBlockSlabManager *slab_manager);
             void   Finalize(KMemoryBlockSlabManager *slab_manager);
 
+            void Update(KMemoryBlockManagerUpdateAllocator *allocator, KProcessAddress address, size_t num_pages, KMemoryState state, KMemoryPermission perm, KMemoryAttribute attr);
+
+
+            iterator FindIterator(KProcessAddress address) const {
+                return this->memory_block_tree.find(KMemoryBlock(address, 1, KMemoryState_Free, KMemoryPermission_None, KMemoryAttribute_None));
+            }
+
+            const KMemoryBlock *FindBlock(KProcessAddress address) const {
+                if (const_iterator it = this->FindIterator(address); it != this->memory_block_tree.end()) {
+                    return std::addressof(*it);
+                }
+
+                return nullptr;
+            }
+
             /* Debug. */
             bool CheckState() const;
             void DumpBlocks() const;
     };
 
-    class KScopedMemoryBlockManagerVerifier {
+    class KScopedMemoryBlockManagerAuditor {
         private:
             KMemoryBlockManager *manager;
         public:
-            explicit ALWAYS_INLINE KScopedMemoryBlockManagerVerifier(KMemoryBlockManager *m) : manager(m) { MESOSPHERE_AUDIT(this->manager->CheckState()); }
+            explicit ALWAYS_INLINE KScopedMemoryBlockManagerAuditor(KMemoryBlockManager *m) : manager(m) { MESOSPHERE_AUDIT(this->manager->CheckState()); }
-            explicit ALWAYS_INLINE KScopedMemoryBlockManagerVerifier(KMemoryBlockManager &m) : KScopedMemoryBlockManagerVerifier(std::addressof(m)) { /* ... */ }
+            explicit ALWAYS_INLINE KScopedMemoryBlockManagerAuditor(KMemoryBlockManager &m) : KScopedMemoryBlockManagerAuditor(std::addressof(m)) { /* ... */ }
-            ALWAYS_INLINE ~KScopedMemoryBlockManagerVerifier() { MESOSPHERE_AUDIT(this->manager->CheckState()); }
+            ALWAYS_INLINE ~KScopedMemoryBlockManagerAuditor() { MESOSPHERE_AUDIT(this->manager->CheckState()); }
     };
 
 }

libraries/libmesosphere/include/mesosphere/kern_k_typed_address.hpp

@@ -85,6 +85,11 @@ namespace ams::kern {
                 return this->address >> shift;
             }
 
+            template<typename U>
+            constexpr ALWAYS_INLINE size_t operator/(U size) const { return this->address / size; }
+
+            /* constexpr ALWAYS_INLINE uintptr_t operator%(U align) const { return this->address % align; } */
+
             /* Comparison operators. */
             constexpr ALWAYS_INLINE bool operator==(KTypedAddress rhs) const {
                 return this->address == rhs.address;

libraries/libmesosphere/source/kern_k_memory_block_manager.cpp

@@ -47,6 +47,86 @@ namespace ams::kern {
         MESOSPHERE_ASSERT(this->memory_block_tree.empty());
     }
 
+    void KMemoryBlockManager::Update(KMemoryBlockManagerUpdateAllocator *allocator, KProcessAddress address, size_t num_pages, KMemoryState state, KMemoryPermission perm, KMemoryAttribute attr) {
+        /* Ensure for auditing that we never end up with an invalid tree. */
+        KScopedMemoryBlockManagerAuditor auditor(this);
+        MESOSPHERE_ASSERT(util::IsAligned(GetInteger(address), PageSize));
+        MESOSPHERE_ASSERT((attr & (KMemoryAttribute_IpcLocked | KMemoryAttribute_DeviceShared)) == 0);
+
+        KProcessAddress cur_address = address;
+        size_t remaining_pages = num_pages;
+        iterator it = this->FindIterator(address);
+
+        while (remaining_pages > 0) {
+            const size_t remaining_size = remaining_pages * PageSize;
+            KMemoryInfo cur_info = it->GetMemoryInfo();
+            if (it->HasProperties(state, perm, attr)) {
+                /* If we already have the right properties, just advance. */
+                if (cur_address + remaining_size < cur_info.GetEndAddress()) {
+                    remaining_pages = 0;
+                    cur_address += remaining_size;
+                } else {
+                    remaining_pages = (cur_address + remaining_size - cur_info.GetEndAddress()) / PageSize;
+                    cur_address = cur_info.GetEndAddress();
+                }
+            } else {
+                /* If we need to, create a new block before and insert it. */
+                if (cur_info.GetAddress() != 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();
+                }
+
+                /* If we need to, create a new block after and insert it. */
+                if (cur_info.GetSize() > remaining_size) {
+                    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();
+                }
+
+                /* Update block state. */
+                it->Update(state, perm, attr);
+                cur_address += cur_info.GetSize();
+                remaining_pages -= cur_info.GetNumPages();
+            }
+        }
+
+        /* Find the iterator now that we've updated. */
+        it = this->FindIterator(address);
+        if (address != this->start) {
+            it--;
+        }
+
+        /* Coalesce blocks that we can. */
+        while (true) {
+            iterator prev = it++;
+            if (it == this->memory_block_tree.end()) {
+                break;
+            }
+
+            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);
+                it = prev;
+            }
+
+            if (address + num_pages * PageSize < it->GetMemoryInfo().GetEndAddress()) {
+                break;
+            }
+        }
+    }
+
     /* Debug. */
     bool KMemoryBlockManager::CheckState() const {
         /* If we fail, we should dump blocks. */
@@ -56,12 +136,47 @@ namespace ams::kern {
         auto it   = this->memory_block_tree.cbegin();
         auto prev = it++;
         while (it != this->memory_block_tree.cend()) {
-            MESOSPHERE_TODO("Validate KMemoryBlock Tree");
+            const KMemoryInfo prev_info = prev->GetMemoryInfo();
+            const KMemoryInfo cur_info  = it->GetMemoryInfo();
+
+            /* Sequential blocks with same properties should be coalesced. */
+            if (prev->HasSameProperties(*it)) {
+                return false;
+            }
+
+            /* Sequential blocks should be sequential. */
+            if (prev_info.GetEndAddress() != cur_info.GetAddress()) {
+                return false;
+            }
+
+            /* If the block is ipc locked, it must have a count. */
+            if ((cur_info.attribute & KMemoryAttribute_IpcLocked) != 0 && cur_info.ipc_lock_count == 0) {
+                return false;
+            }
+
+            /* If the block is device shared, it must have a count. */
+            if ((cur_info.attribute & KMemoryAttribute_DeviceShared) != 0 && cur_info.device_use_count == 0) {
+                return false;
+            }
 
             /* Advance the iterator. */
             prev = it++;
         }
 
+        /* Our loop will miss checking the last block, potentially, so check it. */
+        if (prev != this->memory_block_tree.cend()) {
+            const KMemoryInfo prev_info = prev->GetMemoryInfo();
+            /* If the block is ipc locked, it must have a count. */
+            if ((prev_info.attribute & KMemoryAttribute_IpcLocked) != 0 && prev_info.ipc_lock_count == 0) {
+                return false;
+            }
+
+            /* If the block is device shared, it must have a count. */
+            if ((prev_info.attribute & KMemoryAttribute_DeviceShared) != 0 && prev_info.device_use_count == 0) {
+                return false;
+            }
+        }
+
         /* We're valid, so no need to print. */
         dump_guard.Cancel();
         return true;