Kernel: Clean up the AHCI code a bit

The AHCI code is not very good at OOM conditions, so this is a first
step towards OOM correctness. We should not allocate things inside C++
constructors because we can't catch OOM failures, so most allocation
code inside constructors is exported to a different function.

Also, don't use a HashMap for holding RefPtr of AHCIPort objects in
AHCIPortHandler because this structure is not very OOM-friendly. Instead
use a fixed Array of 32 RefPtrs, as at most we can have 32 AHCI ports
per AHCI controller.

5 files changed, 63 insertions, 44 deletions

diff --git a/Kernel/Storage/ATA/AHCIController.cpp b/Kernel/Storage/ATA/AHCIController.cpp
index 4d18840018..8ce6d1a8b9 100644
--- a/Kernel/Storage/ATA/AHCIController.cpp
+++ b/Kernel/Storage/ATA/AHCIController.cpp
@@ -18,7 +18,9 @@ namespace Kernel {
 
 NonnullRefPtr<AHCIController> AHCIController::initialize(PCI::DeviceIdentifier const& pci_device_identifier)
 {
-    return adopt_ref(*new AHCIController(pci_device_identifier));
+    auto controller = adopt_ref_if_nonnull(new (nothrow) AHCIController(pci_device_identifier)).release_nonnull();
+    controller->initialize_hba(pci_device_identifier);
+    return controller;
 }
 
 bool AHCIController::reset()
@@ -90,7 +92,6 @@ AHCIController::AHCIController(PCI::DeviceIdentifier const& pci_device_identifie
     , m_hba_region(default_hba_region())
     , m_capabilities(capabilities())
 {
-    initialize_hba(pci_device_identifier);
 }
 
 AHCI::HBADefinedCapabilities AHCIController::capabilities() const
@@ -153,8 +154,9 @@ void AHCIController::initialize_hba(PCI::DeviceIdentifier const& pci_device_iden
     PCI::enable_interrupt_line(pci_address());
     PCI::enable_bus_mastering(pci_address());
     enable_global_interrupts();
-    m_handlers.append(AHCIPortHandler::create(*this, pci_device_identifier.interrupt_line().value(),
-        AHCI::MaskedBitField((u32 volatile&)(hba().control_regs.pi))));
+
+    auto port_handler = AHCIPortHandler::create(*this, pci_device_identifier.interrupt_line().value(), AHCI::MaskedBitField((u32 volatile&)(hba().control_regs.pi))).release_value_but_fixme_should_propagate_errors();
+    m_handlers.append(move(port_handler));
 }
 
 void AHCIController::disable_global_interrupts() const
diff --git a/Kernel/Storage/ATA/AHCIPort.cpp b/Kernel/Storage/ATA/AHCIPort.cpp
index f72d307cac..7e4add1fd2 100644
--- a/Kernel/Storage/ATA/AHCIPort.cpp
+++ b/Kernel/Storage/ATA/AHCIPort.cpp
@@ -20,37 +20,46 @@
 
 namespace Kernel {
 
-NonnullRefPtr<AHCIPort> AHCIPort::create(AHCIPortHandler const& handler, volatile AHCI::PortRegisters& registers, u32 port_index)
+ErrorOr<NonnullRefPtr<AHCIPort>> AHCIPort::create(AHCIPortHandler const& handler, volatile AHCI::PortRegisters& registers, u32 port_index)
 {
-    return adopt_ref(*new AHCIPort(handler, registers, port_index));
+    auto port = TRY(adopt_nonnull_ref_or_enomem(new (nothrow) AHCIPort(handler, registers, port_index)));
+    TRY(port->allocate_resources_and_initialize_ports());
+    return port;
 }
 
-AHCIPort::AHCIPort(AHCIPortHandler const& handler, volatile AHCI::PortRegisters& registers, u32 port_index)
-    : m_port_index(port_index)
-    , m_port_registers(registers)
-    , m_parent_handler(handler)
-    , m_interrupt_status((u32 volatile&)m_port_registers.is)
-    , m_interrupt_enable((u32 volatile&)m_port_registers.ie)
+ErrorOr<void> AHCIPort::allocate_resources_and_initialize_ports()
 {
     if (is_interface_disabled()) {
         m_disabled_by_firmware = true;
-        return;
+        return {};
     }
 
-    m_fis_receive_page = MM.allocate_supervisor_physical_page().release_value_but_fixme_should_propagate_errors();
+    m_fis_receive_page = TRY(MM.allocate_supervisor_physical_page());
 
     for (size_t index = 0; index < 1; index++) {
-        m_dma_buffers.append(MM.allocate_supervisor_physical_page().release_value_but_fixme_should_propagate_errors());
+        auto dma_page = TRY(MM.allocate_supervisor_physical_page());
+        m_dma_buffers.append(move(dma_page));
     }
     for (size_t index = 0; index < 1; index++) {
-        m_command_table_pages.append(MM.allocate_supervisor_physical_page().release_value_but_fixme_should_propagate_errors());
+        auto command_table_page = TRY(MM.allocate_supervisor_physical_page());
+        m_command_table_pages.append(move(command_table_page));
     }
 
-    m_command_list_region = MM.allocate_dma_buffer_page("AHCI Port Command List", Memory::Region::Access::ReadWrite, m_command_list_page).release_value_but_fixme_should_propagate_errors();
+    m_command_list_region = TRY(MM.allocate_dma_buffer_page("AHCI Port Command List", Memory::Region::Access::ReadWrite, m_command_list_page));
 
     dbgln_if(AHCI_DEBUG, "AHCI Port {}: Command list page at {}", representative_port_index(), m_command_list_page->paddr());
     dbgln_if(AHCI_DEBUG, "AHCI Port {}: FIS receive page at {}", representative_port_index(), m_fis_receive_page->paddr());
     dbgln_if(AHCI_DEBUG, "AHCI Port {}: Command list region at {}", representative_port_index(), m_command_list_region->vaddr());
+    return {};
+}
+
+AHCIPort::AHCIPort(AHCIPortHandler const& handler, volatile AHCI::PortRegisters& registers, u32 port_index)
+    : m_port_index(port_index)
+    , m_port_registers(registers)
+    , m_parent_handler(handler)
+    , m_interrupt_status((u32 volatile&)m_port_registers.is)
+    , m_interrupt_enable((u32 volatile&)m_port_registers.ie)
+{
 }
 
 void AHCIPort::clear_sata_error_register() const
diff --git a/Kernel/Storage/ATA/AHCIPort.h b/Kernel/Storage/ATA/AHCIPort.h
index 6799a1d6ba..90ab7f0b68 100644
--- a/Kernel/Storage/ATA/AHCIPort.h
+++ b/Kernel/Storage/ATA/AHCIPort.h
@@ -37,7 +37,7 @@ class AHCIPort
     friend class AHCIController;
 
 public:
-    UNMAP_AFTER_INIT static NonnullRefPtr<AHCIPort> create(AHCIPortHandler const&, volatile AHCI::PortRegisters&, u32 port_index);
+    UNMAP_AFTER_INIT static ErrorOr<NonnullRefPtr<AHCIPort>> create(AHCIPortHandler const&, volatile AHCI::PortRegisters&, u32 port_index);
 
     u32 port_index() const { return m_port_index; }
     u32 representative_port_index() const { return port_index() + 1; }
@@ -51,6 +51,8 @@ public:
     void handle_interrupt();
 
 private:
+    ErrorOr<void> allocate_resources_and_initialize_ports();
+
     bool is_phy_enabled() const { return (m_port_registers.ssts & 0xf) == 3; }
     bool initialize();
 
diff --git a/Kernel/Storage/ATA/AHCIPortHandler.cpp b/Kernel/Storage/ATA/AHCIPortHandler.cpp
index 93ca767302..78ab3042df 100644
--- a/Kernel/Storage/ATA/AHCIPortHandler.cpp
+++ b/Kernel/Storage/ATA/AHCIPortHandler.cpp
@@ -9,57 +9,61 @@
 
 namespace Kernel {
 
-NonnullRefPtr<AHCIPortHandler> AHCIPortHandler::create(AHCIController& controller, u8 irq, AHCI::MaskedBitField taken_ports)
+ErrorOr<NonnullRefPtr<AHCIPortHandler>> AHCIPortHandler::create(AHCIController& controller, u8 irq, AHCI::MaskedBitField taken_ports)
 {
-    return adopt_ref(*new AHCIPortHandler(controller, irq, taken_ports));
+    auto port_handler = TRY(adopt_nonnull_ref_or_enomem(new (nothrow) AHCIPortHandler(controller, irq, taken_ports)));
+    // FIXME: Propagate errors from this method too.
+    port_handler->allocate_resources_and_initialize_ports();
+    return port_handler;
 }
 
-AHCIPortHandler::AHCIPortHandler(AHCIController& controller, u8 irq, AHCI::MaskedBitField taken_ports)
-    : IRQHandler(irq)
-    , m_parent_controller(controller)
-    , m_taken_ports(taken_ports)
-    , m_pending_ports_interrupts(create_pending_ports_interrupts_bitfield())
+void AHCIPortHandler::allocate_resources_and_initialize_ports()
 {
     // FIXME: Use the number of taken ports to determine how many pages we should allocate.
     for (size_t index = 0; index < (((size_t)AHCI::Limits::MaxPorts * 512) / PAGE_SIZE); index++) {
         m_identify_metadata_pages.append(MM.allocate_supervisor_physical_page().release_value_but_fixme_should_propagate_errors());
     }
 
-    dbgln_if(AHCI_DEBUG, "AHCI Port Handler: IRQ {}", irq);
-
     // Clear pending interrupts, if there are any!
     m_pending_ports_interrupts.set_all();
     enable_irq();
 
     if (kernel_command_line().ahci_reset_mode() == AHCIResetMode::Aggressive) {
-        for (auto index : taken_ports.to_vector()) {
-            auto port = AHCIPort::create(*this, static_cast<volatile AHCI::PortRegisters&>(controller.hba().port_regs[index]), index);
-            m_handled_ports.set(index, port);
+        for (auto index : m_taken_ports.to_vector()) {
+            auto port = AHCIPort::create(*this, static_cast<volatile AHCI::PortRegisters&>(m_parent_controller->hba().port_regs[index]), index).release_value_but_fixme_should_propagate_errors();
+            m_handled_ports[index] = port;
             port->reset();
         }
         return;
     }
-    for (auto index : taken_ports.to_vector()) {
-        auto port = AHCIPort::create(*this, static_cast<volatile AHCI::PortRegisters&>(controller.hba().port_regs[index]), index);
-        m_handled_ports.set(index, port);
+    for (auto index : m_taken_ports.to_vector()) {
+        auto port = AHCIPort::create(*this, static_cast<volatile AHCI::PortRegisters&>(m_parent_controller->hba().port_regs[index]), index).release_value_but_fixme_should_propagate_errors();
+        m_handled_ports[index] = port;
         port->initialize_without_reset();
     }
 }
 
+AHCIPortHandler::AHCIPortHandler(AHCIController& controller, u8 irq, AHCI::MaskedBitField taken_ports)
+    : IRQHandler(irq)
+    , m_parent_controller(controller)
+    , m_taken_ports(taken_ports)
+    , m_pending_ports_interrupts(create_pending_ports_interrupts_bitfield())
+{
+    dbgln_if(AHCI_DEBUG, "AHCI Port Handler: IRQ {}", irq);
+}
+
 void AHCIPortHandler::enumerate_ports(Function<void(AHCIPort const&)> callback) const
 {
-    for (auto& port : m_handled_ports) {
-        callback(*port.value);
+    for (auto port : m_handled_ports) {
+        if (port)
+            callback(*port);
     }
 }
 
 RefPtr<AHCIPort> AHCIPortHandler::port_at_index(u32 port_index) const
 {
     VERIFY(m_taken_ports.is_set_at(port_index));
-    auto it = m_handled_ports.find(port_index);
-    if (it == m_handled_ports.end())
-        return nullptr;
-    return (*it).value;
+    return m_handled_ports[port_index];
 }
 
 PhysicalAddress AHCIPortHandler::get_identify_metadata_physical_region(u32 port_index) const
@@ -86,10 +90,10 @@ bool AHCIPortHandler::handle_irq(RegisterState const&)
     if (m_pending_ports_interrupts.is_zeroed())
         return false;
     for (auto port_index : m_pending_ports_interrupts.to_vector()) {
-        auto port = m_handled_ports.get(port_index);
-        VERIFY(port.has_value());
+        auto port = m_handled_ports[port_index];
+        VERIFY(port);
         dbgln_if(AHCI_DEBUG, "AHCI Port Handler: Handling IRQ for port {}", port_index);
-        port.value()->handle_interrupt();
+        port->handle_interrupt();
         // We do this to clear the pending interrupt after we handled it.
         m_pending_ports_interrupts.set_at(port_index);
     }
diff --git a/Kernel/Storage/ATA/AHCIPortHandler.h b/Kernel/Storage/ATA/AHCIPortHandler.h
index 235b5b4162..30e69bf82b 100644
--- a/Kernel/Storage/ATA/AHCIPortHandler.h
+++ b/Kernel/Storage/ATA/AHCIPortHandler.h
@@ -31,7 +31,7 @@ class AHCIPortHandler final : public RefCounted<AHCIPortHandler>
     friend class AHCIController;
 
 public:
-    UNMAP_AFTER_INIT static NonnullRefPtr<AHCIPortHandler> create(AHCIController&, u8 irq, AHCI::MaskedBitField taken_ports);
+    UNMAP_AFTER_INIT static ErrorOr<NonnullRefPtr<AHCIPortHandler>> create(AHCIController&, u8 irq, AHCI::MaskedBitField taken_ports);
     virtual ~AHCIPortHandler() override;
 
     RefPtr<StorageDevice> device_at_port(size_t port_index) const;
@@ -46,6 +46,8 @@ public:
 private:
     UNMAP_AFTER_INIT AHCIPortHandler(AHCIController&, u8 irq, AHCI::MaskedBitField taken_ports);
 
+    void allocate_resources_and_initialize_ports();
+
     //^ IRQHandler
     virtual bool handle_irq(RegisterState const&) override;
 
@@ -63,7 +65,7 @@ private:
     RefPtr<AHCIPort> port_at_index(u32 port_index) const;
 
     // Data members
-    HashMap<u32, NonnullRefPtr<AHCIPort>> m_handled_ports;
+    Array<RefPtr<AHCIPort>, 32> m_handled_ports;
     NonnullRefPtr<AHCIController> m_parent_controller;
     NonnullRefPtrVector<Memory::PhysicalPage> m_identify_metadata_pages;
     AHCI::MaskedBitField m_taken_ports;