Kernel/USB: Harden Descriptor memory allocation

The previous version of this was pretty bad and caused a lot of
odd behevaiour to occur. We now abstract a lot of the allocation
behind a `template`d pool class that handles all of the memory
allocation.

4 files changed, 129 insertions, 69 deletions

diff --git a/Kernel/Bus/USB/UHCI/UHCIController.cpp b/Kernel/Bus/USB/UHCI/UHCIController.cpp
index 796484765d..f396770434 100644
--- a/Kernel/Bus/USB/UHCI/UHCIController.cpp
+++ b/Kernel/Bus/USB/UHCI/UHCIController.cpp
@@ -16,8 +16,6 @@
 #include <Kernel/StdLib.h>
 #include <Kernel/Time/TimeManagement.h>
 
-static constexpr u8 MAXIMUM_NUMBER_OF_TDS = 128; // Upper pool limit. This consumes the second page we have allocated
-static constexpr u8 MAXIMUM_NUMBER_OF_QHS = 64;
 static constexpr u8 RETRY_COUNTER_RELOAD = 3;
 
 namespace Kernel::USB {
@@ -129,6 +127,7 @@ KResult UHCIController::reset()
 
     if (auto result = create_structures(); result.is_error())
         return result;
+
     setup_schedule();
 
     write_flbaseadd(m_framelist->physical_page(0)->paddr().get()); // Frame list (physical) address
@@ -144,21 +143,10 @@ KResult UHCIController::reset()
 
 UNMAP_AFTER_INIT KResult UHCIController::create_structures()
 {
-    // Let's allocate memory for both the QH and TD pools
-    // First the QH pool and all of the Interrupt QH's
-    auto maybe_qh_pool_vmobject = Memory::AnonymousVMObject::try_create_physically_contiguous_with_size(2 * PAGE_SIZE);
-    if (maybe_qh_pool_vmobject.is_error())
-        return maybe_qh_pool_vmobject.error();
-
-    m_qh_pool = MM.allocate_kernel_region_with_vmobject(maybe_qh_pool_vmobject.release_value(), 2 * PAGE_SIZE, "UHCI Queue Head Pool", Memory::Region::Access::Write);
-    memset(m_qh_pool->vaddr().as_ptr(), 0, 2 * PAGE_SIZE); // Zero out both pages
-
-    // Let's populate our free qh list (so we have some we can allocate later on)
-    m_free_qh_pool.resize(MAXIMUM_NUMBER_OF_TDS);
-    for (size_t i = 0; i < m_free_qh_pool.size(); i++) {
-        auto placement_addr = reinterpret_cast<void*>(m_qh_pool->vaddr().get() + (i * sizeof(QueueHead)));
-        auto paddr = static_cast<u32>(m_qh_pool->physical_page(0)->paddr().get() + (i * sizeof(QueueHead)));
-        m_free_qh_pool.at(i) = new (placement_addr) QueueHead(paddr);
+    m_queue_head_pool = UHCIDescriptorPool<QueueHead>::try_create("Queue Head Pool");
+    if (!m_queue_head_pool) {
+        dmesgln("UHCI: Failed to create Queue Head Pool!");
+        return ENOMEM;
     }
 
     // Create the Full Speed, Low Speed Control and Bulk Queue Heads
@@ -169,17 +157,27 @@ UNMAP_AFTER_INIT KResult UHCIController::create_structures()
     m_dummy_qh = allocate_queue_head();
 
     // Now the Transfer Descriptor pool
-    auto maybe_td_pool_vmobject = Memory::AnonymousVMObject::try_create_physically_contiguous_with_size(2 * PAGE_SIZE);
+    auto maybe_td_pool_vmobject = Memory::AnonymousVMObject::try_create_physically_contiguous_with_size(PAGE_SIZE);
     if (maybe_td_pool_vmobject.is_error())
         return maybe_td_pool_vmobject.error();
-    m_td_pool = MM.allocate_kernel_region_with_vmobject(maybe_td_pool_vmobject.release_value(), 2 * PAGE_SIZE, "UHCI Transfer Descriptor Pool", Memory::Region::Access::Write);
-    memset(m_td_pool->vaddr().as_ptr(), 0, 2 * PAGE_SIZE);
+
+    m_transfer_descriptor_pool = UHCIDescriptorPool<TransferDescriptor>::try_create("Transfer Descriptor Pool");
+    if (!m_transfer_descriptor_pool) {
+        dmesgln("UHCI: Failed to create Transfer Descriptor Pool!");
+        return ENOMEM;
+    }
+
+    m_isochronous_transfer_pool = MM.allocate_kernel_region_with_vmobject(*maybe_td_pool_vmobject.release_value(), PAGE_SIZE, "UHCI Isochronous Descriptor Pool", Memory::Region::Access::ReadWrite);
+    if (!m_isochronous_transfer_pool) {
+        dmesgln("UHCI: Failed to allocated Isochronous Descriptor Pool!");
+        return ENOMEM;
+    }
 
     // Set up the Isochronous Transfer Descriptor list
     m_iso_td_list.resize(UHCI_NUMBER_OF_ISOCHRONOUS_TDS);
     for (size_t i = 0; i < m_iso_td_list.size(); i++) {
-        auto placement_addr = reinterpret_cast<void*>(m_td_pool->vaddr().get() + (i * sizeof(Kernel::USB::TransferDescriptor)));
-        auto paddr = static_cast<u32>(m_td_pool->physical_page(0)->paddr().get() + (i * sizeof(Kernel::USB::TransferDescriptor)));
+        auto placement_addr = reinterpret_cast<void*>(m_isochronous_transfer_pool->vaddr().get() + (i * sizeof(Kernel::USB::TransferDescriptor)));
+        auto paddr = static_cast<u32>(m_isochronous_transfer_pool->physical_page(0)->paddr().get() + (i * sizeof(Kernel::USB::TransferDescriptor)));
 
         // Place a new Transfer Descriptor with a 1:1 in our region
         // The pointer returned by `new()` lines up exactly with the value
@@ -195,27 +193,10 @@ UNMAP_AFTER_INIT KResult UHCIController::create_structures()
             transfer_descriptor->print();
     }
 
-    m_free_td_pool.resize(MAXIMUM_NUMBER_OF_TDS);
-    for (size_t i = 0; i < m_free_td_pool.size(); i++) {
-        auto placement_addr = reinterpret_cast<void*>(m_td_pool->vaddr().offset(PAGE_SIZE).get() + (i * sizeof(Kernel::USB::TransferDescriptor)));
-        auto paddr = static_cast<u32>(m_td_pool->physical_page(1)->paddr().get() + (i * sizeof(Kernel::USB::TransferDescriptor)));
-
-        // Place a new Transfer Descriptor with a 1:1 in our region
-        // The pointer returned by `new()` lines up exactly with the value
-        // that we store in `paddr`, meaning our member functions directly
-        // access the raw descriptor (that we later send to the controller)
-        m_free_td_pool.at(i) = new (placement_addr) Kernel::USB::TransferDescriptor(paddr);
-
-        if constexpr (UHCI_VERBOSE_DEBUG) {
-            auto transfer_descriptor = m_free_td_pool.at(i);
-            transfer_descriptor->print();
-        }
-    }
-
     if constexpr (UHCI_DEBUG) {
         dbgln("UHCI: Pool information:");
-        dbgln("    qh_pool: {}, length: {}", PhysicalAddress(m_qh_pool->physical_page(0)->paddr()), m_qh_pool->range().size());
-        dbgln("    td_pool: {}, length: {}", PhysicalAddress(m_td_pool->physical_page(0)->paddr()), m_td_pool->range().size());
+        m_queue_head_pool->print_pool_information();
+        m_transfer_descriptor_pool->print_pool_information();
     }
 
     return KSuccess;
@@ -280,30 +261,14 @@ UNMAP_AFTER_INIT void UHCIController::setup_schedule()
     m_dummy_qh->print();
 }
 
-QueueHead* UHCIController::allocate_queue_head() const
+QueueHead* UHCIController::allocate_queue_head()
 {
-    for (QueueHead* queue_head : m_free_qh_pool) {
-        if (!queue_head->in_use()) {
-            queue_head->set_in_use(true);
-            dbgln_if(UHCI_DEBUG, "UHCI: Allocated a new Queue Head! Located @ {} ({})", VirtualAddress(queue_head), PhysicalAddress(queue_head->paddr()));
-            return queue_head;
-        }
-    }
-
-    return nullptr; // Huh!? We're outta queue heads!
+    return m_queue_head_pool->try_take_free_descriptor();
 }
 
-TransferDescriptor* UHCIController::allocate_transfer_descriptor() const
+TransferDescriptor* UHCIController::allocate_transfer_descriptor()
 {
-    for (TransferDescriptor* transfer_descriptor : m_free_td_pool) {
-        if (!transfer_descriptor->in_use()) {
-            transfer_descriptor->set_in_use(true);
-            dbgln_if(UHCI_DEBUG, "UHCI: Allocated a new Transfer Descriptor! Located @ {} ({})", VirtualAddress(transfer_descriptor), PhysicalAddress(transfer_descriptor->paddr()));
-            return transfer_descriptor;
-        }
-    }
-
-    return nullptr; // Huh?! We're outta TDs!!
+    return m_transfer_descriptor_pool->try_take_free_descriptor();
 }
 
 KResult UHCIController::stop()
@@ -421,8 +386,11 @@ void UHCIController::free_descriptor_chain(TransferDescriptor* first_descriptor)
     TransferDescriptor* descriptor = first_descriptor;
 
     while (descriptor) {
+        TransferDescriptor* next = descriptor->next_td();
+
         descriptor->free();
-        descriptor = descriptor->next_td();
+        m_transfer_descriptor_pool->release_to_pool(descriptor);
+        descriptor = next;
     }
 }
 
@@ -497,6 +465,7 @@ KResultOr<size_t> UHCIController::submit_control_transfer(Transfer& transfer)
 
     free_descriptor_chain(transfer_queue->get_first_td());
     transfer_queue->free();
+    m_queue_head_pool->release_to_pool(transfer_queue);
 
     return transfer_size;
 }
diff --git a/Kernel/Bus/USB/UHCI/UHCIController.h b/Kernel/Bus/USB/UHCI/UHCIController.h
index 102c4c1164..6e9e3559ce 100644
--- a/Kernel/Bus/USB/UHCI/UHCIController.h
+++ b/Kernel/Bus/USB/UHCI/UHCIController.h
@@ -11,6 +11,7 @@
 
 #include <AK/NonnullOwnPtr.h>
 #include <Kernel/Bus/PCI/Device.h>
+#include <Kernel/Bus/USB/UHCI/UHCIDescriptorPool.h>
 #include <Kernel/Bus/USB/UHCI/UHCIDescriptorTypes.h>
 #include <Kernel/Bus/USB/UHCI/UHCIRootHub.h>
 #include <Kernel/Bus/USB/USBController.h>
@@ -25,6 +26,9 @@ class UHCIController final
     : public USBController
     , public PCI::Device {
 
+    static constexpr u8 MAXIMUM_NUMBER_OF_TDS = 128; // Upper pool limit. This consumes the second page we have allocated
+    static constexpr u8 MAXIMUM_NUMBER_OF_QHS = 64;
+
 public:
     static constexpr u8 NUMBER_OF_ROOT_PORTS = 2;
     static KResultOr<NonnullRefPtr<UHCIController>> try_to_initialize(PCI::Address address);
@@ -77,8 +81,8 @@ private:
     KResult create_chain(Pipe& pipe, PacketID direction, Ptr32<u8>& buffer_address, size_t max_size, size_t transfer_size, TransferDescriptor** td_chain, TransferDescriptor** last_td);
     void free_descriptor_chain(TransferDescriptor* first_descriptor);
 
-    QueueHead* allocate_queue_head() const;
-    TransferDescriptor* allocate_transfer_descriptor() const;
+    QueueHead* allocate_queue_head();
+    TransferDescriptor* allocate_transfer_descriptor();
 
     void reset_port(u8);
 
@@ -86,9 +90,8 @@ private:
     IOAddress m_io_base;
 
     OwnPtr<UHCIRootHub> m_root_hub;
-
-    Vector<QueueHead*> m_free_qh_pool;
-    Vector<TransferDescriptor*> m_free_td_pool;
+    OwnPtr<UHCIDescriptorPool<QueueHead>> m_queue_head_pool;
+    OwnPtr<UHCIDescriptorPool<TransferDescriptor>> m_transfer_descriptor_pool;
     Vector<TransferDescriptor*> m_iso_td_list;
 
     QueueHead* m_interrupt_transfer_queue;
@@ -98,8 +101,7 @@ private:
     QueueHead* m_dummy_qh; // Needed for PIIX4 hack
 
     OwnPtr<Memory::Region> m_framelist;
-    OwnPtr<Memory::Region> m_qh_pool;
-    OwnPtr<Memory::Region> m_td_pool;
+    OwnPtr<Memory::Region> m_isochronous_transfer_pool;
 
     // Bitfield containing whether a given port should signal a change in reset or not.
     u8 m_port_reset_change_statuses { 0 };
diff --git a/Kernel/Bus/USB/UHCI/UHCIDescriptorPool.h b/Kernel/Bus/USB/UHCI/UHCIDescriptorPool.h
new file mode 100644
index 0000000000..d7358b40c1
--- /dev/null
+++ b/Kernel/Bus/USB/UHCI/UHCIDescriptorPool.h
@@ -0,0 +1,85 @@
+/*
+ * Copyright (c) 2021, Jesse Buhagiar <jooster669@gmail.com>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#pragma once
+
+#include <AK/NonnullOwnPtr.h>
+#include <AK/OwnPtr.h>
+#include <AK/Stack.h>
+#include <Kernel/Memory/MemoryManager.h>
+#include <Kernel/Memory/Region.h>
+#include <Kernel/StdLib.h>
+
+namespace Kernel::USB {
+
+// This pool is bound by PAGE_SIZE / sizeof(T). The underlying allocation for the pointers
+// is AK::Stack. As such, we never dynamically allocate any memory past the amount
+// that can fit in a single page.
+template<typename T>
+class UHCIDescriptorPool {
+
+    // Ensure that we can't get into a situation where we'll write past the page
+    // and blow up
+    static_assert(sizeof(T) <= PAGE_SIZE);
+
+public:
+    static OwnPtr<UHCIDescriptorPool<T>> try_create(const StringView name)
+    {
+        auto pool_memory_block = MM.allocate_kernel_region(PAGE_SIZE, "UHCI Descriptor Pool", Memory::Region::Access::ReadWrite);
+        if (!pool_memory_block)
+            return {};
+
+        return adopt_own_if_nonnull(new (nothrow) UHCIDescriptorPool(pool_memory_block.release_nonnull(), name));
+    }
+
+    UHCIDescriptorPool(NonnullOwnPtr<Memory::Region> pool_memory_block, const StringView& name)
+        : m_pool_name(name)
+        , m_pool_region(move(pool_memory_block))
+    {
+        // Go through the number of descriptors to create in the pool, and create a virtual/physical address mapping
+        for (size_t i = 0; i < PAGE_SIZE / sizeof(T); i++) {
+            auto placement_address = reinterpret_cast<void*>(m_pool_region->vaddr().get() + (i * sizeof(T)));
+            auto physical_address = static_cast<u32>(m_pool_region->physical_page(0)->paddr().get() + (i * sizeof(T)));
+            auto* object = new (placement_address) T(physical_address);
+            m_free_descriptor_stack.push(object); // Push the descriptor's pointer onto the free list
+        }
+    }
+
+public:
+    UHCIDescriptorPool() = delete;
+    ~UHCIDescriptorPool() = default;
+
+    [[nodiscard]] T* try_take_free_descriptor()
+    {
+        // We're out of descriptors!
+        if (m_free_descriptor_stack.is_empty())
+            return nullptr;
+
+        dbgln_if(UHCI_VERBOSE_DEBUG, "Got a free UHCI Descriptor @ {} from pool {}", m_free_descriptor_stack.top(), m_pool_name);
+        T* descriptor = m_free_descriptor_stack.top();
+        m_free_descriptor_stack.pop();
+
+        return descriptor;
+    }
+
+    void release_to_pool(T* ptr)
+    {
+        dbgln_if(UHCI_VERBOSE_DEBUG, "Returning descriptor @ {} to pool {}", ptr, m_pool_name);
+        if (!m_free_descriptor_stack.push(ptr))
+            dbgln("Failed to return descriptor to pool {}. Stack overflow!", m_pool_name);
+    }
+
+    void print_pool_information() const
+    {
+        dbgln("Pool {} allocated @ {}", m_pool_name, m_pool_region->physical_page(0)->paddr());
+    }
+
+private:
+    StringView m_pool_name;                                   // Name of this pool
+    NonnullOwnPtr<Memory::Region> m_pool_region;              // Memory region where descriptors actually reside
+    Stack<T*, PAGE_SIZE / sizeof(T)> m_free_descriptor_stack; // Stack of currently free descriptor pointers
+};
+}
diff --git a/Kernel/Bus/USB/UHCI/UHCIDescriptorTypes.h b/Kernel/Bus/USB/UHCI/UHCIDescriptorTypes.h
index 5525428a1a..9a9f8f0641 100644
--- a/Kernel/Bus/USB/UHCI/UHCIDescriptorTypes.h
+++ b/Kernel/Bus/USB/UHCI/UHCIDescriptorTypes.h
@@ -227,6 +227,8 @@ struct alignas(16) TransferDescriptor final {
         m_control_status = 0;
         m_token = 0;
         m_in_use = false;
+        m_next_td = nullptr;
+        m_prev_td = nullptr;
     }
 
 private:
@@ -346,6 +348,8 @@ struct alignas(16) QueueHead {
         m_element_link_ptr = 0;
         m_first_td = nullptr;
         m_transfer = nullptr;
+        m_next_qh = nullptr;
+        m_prev_qh = nullptr;
         m_in_use = false;
     }