Kernel: Rewrite the AHCI driver

There was a lot of stuff wrong with this driver and it was only getting
speeds of ~2.5 MB/s. Now 4K blocks get ~17 MB/s which is not good but
definitely better. I'm pretty sure the issue is qemu's emulation as 85%
of time is spent between sending the command and it to complete.
Large blocks are now supported too, qemu with block size 1M gets around
600 MB/s read speeds.

kernel/include/kernel/Storage/ATA/AHCI/Controller.h

@@ -22,6 +22,7 @@ namespace Kernel
 
 		virtual void handle_irq() override;
 
+		bool supports_64bit() const { return m_supports_64bit; }
 		uint32_t command_slot_count() const { return m_command_slot_count; }
 
 	private:
@@ -38,9 +39,8 @@ namespace Kernel
 
 		BAN::Array<AHCIDevice*, 32> m_devices;
 
+		bool m_supports_64bit { false };
 		uint32_t m_command_slot_count { 0 };
-
-		friend class ATAController;
 	};
 
 }

kernel/include/kernel/Storage/ATA/AHCI/Definitions.h

@@ -32,8 +32,6 @@
 namespace Kernel
 {
 
-	static constexpr uint32_t s_hba_prdt_count { 8 };
-
 	struct FISRegisterH2D
 	{
 		uint8_t  fis_type;			// FIS_TYPE_REGISTER_H2D
@@ -61,7 +59,8 @@ namespace Kernel
 		uint8_t  control;
 
 		uint8_t  __reserved1[4];
-	} __attribute__((packed));
+	};
+	static_assert(sizeof(FISRegisterH2D) == 20);
 
 	struct FISRegisterD2H
 	{
@@ -90,7 +89,8 @@ namespace Kernel
 		uint8_t  __reserved3[2];
 
 		uint8_t  __reserved4[4];
-	} __attribute__((packed));
+	};
+	static_assert(sizeof(FISRegisterD2H) == 20);
 
 	struct FISDataBI
 	{
@@ -102,7 +102,8 @@ namespace Kernel
 		uint8_t  __reserved1[2];
 
 		uint32_t data[0];			// Payload (1 - 2048 dwords)
-	} __attribute__((packed));
+	};
+	static_assert(sizeof(FISDataBI) == 4);
 
 	struct SetDeviceBitsD2H
 	{
@@ -117,7 +118,8 @@ namespace Kernel
 		uint8_t error;
 
 		uint32_t __reserved1;
-	} __attribute__((packed));
+	};
+	static_assert(sizeof(SetDeviceBitsD2H) == 8);
 
 	struct PIOSetupD2H
 	{
@@ -149,7 +151,8 @@ namespace Kernel
 
 		uint16_t tc;				// Transfer count
 		uint8_t  __reserved4[2];
-	} __attribute__((packed));
+	};
+	static_assert(sizeof(PIOSetupD2H) == 20);
 
 	struct DMASetupBI
 	{
@@ -163,8 +166,8 @@ namespace Kernel
 
 		uint8_t  __reserved1[2];
 
-		uint64_t dma_buffer_id;			// DMA Buffer Identifier. Used to Identify DMA buffer in host memory.
+		uint32_t dma_buffer_id_lo;		// DMA Buffer Identifier. Used to Identify DMA buffer in host memory.
-										// SATA Spec says host specific and not in Spec. Trying AHCI spec might work.
+		uint32_t dma_buffer_id_hi;		// SATA Spec says host specific and not in Spec. Trying AHCI spec might work.
 
 		uint32_t __reserved2;
 
@@ -173,7 +176,8 @@ namespace Kernel
 		uint32_t dma_transfer_count;	// Number of bytes to transfer. Bit 0 must be 0
 
 		uint32_t __reserved3;
-	} __attribute__((packed));
+	};
+	static_assert(sizeof(DMASetupBI) == 28);
 
 	struct HBAPortMemorySpace
 	{
@@ -196,7 +200,8 @@ namespace Kernel
 		uint32_t fbs;		// FIS-based switch control
 		uint32_t __reserved1[11];
 		uint32_t vendor[4];
-	} __attribute__((packed));
+	};
+	static_assert(sizeof(HBAPortMemorySpace) == 128);
 
 	struct HBAGeneralMemorySpace
 	{
@@ -217,7 +222,8 @@ namespace Kernel
 		uint8_t  vendor[0x100-0xA0];
 
 		HBAPortMemorySpace ports[0]; // 1 - 32 ports
-	} __attribute__((packed));
+	};
+	static_assert(sizeof(HBAGeneralMemorySpace) == 256);
 
 	struct ReceivedFIS
 	{
@@ -235,7 +241,8 @@ namespace Kernel
 		uint8_t				ufis[64];
 
 		uint8_t				__reserved[0x100-0xA0];
-	} __attribute__((packed));
+	};
+	static_assert(sizeof(ReceivedFIS) == 256);
 
 	struct HBACommandHeader
 	{
@@ -252,13 +259,14 @@ namespace Kernel
 
 		uint16_t prdtl;		// Physical region descriptor table length in entries
 
-		volatile uint32_t prdbc;	// Physical region descriptor byte count transferred
+		uint32_t prdbc;		// Physical region descriptor byte count transferred
 
 		uint32_t ctba;		// Command table descriptor base address
 		uint32_t ctbau;		// Command table descriptor base address upper 32 bits
 
 		uint32_t __reserved1[4];
-	} __attribute__((packed));
+	};
+	static_assert(sizeof(HBACommandHeader) == 32);
 
 	struct HBAPRDTEntry
 	{
@@ -269,15 +277,17 @@ namespace Kernel
 		uint32_t dbc : 22;		// Byte count, 4M max
 		uint32_t __reserved1 : 9;
 		uint32_t i : 1;			// Interrupt on completion
-	} __attribute__((packed));
+	};
+	static_assert(sizeof(HBAPRDTEntry) == 16);
 
 	struct HBACommandTable
 	{
 		uint8_t cfis[64];
 		uint8_t acmd[16];
 		uint8_t __reserved[48];
-		HBAPRDTEntry prdt_entry[s_hba_prdt_count];
+		HBAPRDTEntry prdt_entry[0];
-	} __attribute__((packed));
+	};
+	static_assert(sizeof(HBACommandTable) == 128);
 
 	enum class AHCIPortType
 	{

kernel/include/kernel/Storage/ATA/AHCI/Device.h

@@ -20,30 +20,36 @@ namespace Kernel
 			, m_port(port)
 		{ }
 		BAN::ErrorOr<void> initialize();
-		BAN::ErrorOr<void> allocate_buffers();
 		BAN::ErrorOr<void> rebase();
 		BAN::ErrorOr<void> read_identify_data();
 
+		paddr_t read_paddr(volatile uint32_t& lo, volatile uint32_t& hi) const;
+		void write_paddr(volatile uint32_t& lo, volatile uint32_t& hi, paddr_t paddr);
+
+		bool can_use_buffer_directly(BAN::ConstByteSpan buffer) const;
+
 		virtual BAN::ErrorOr<void> read_sectors_impl(uint64_t lba, uint64_t sector_count, BAN::ByteSpan) override;
 		virtual BAN::ErrorOr<void> write_sectors_impl(uint64_t lba, uint64_t sector_count, BAN::ConstByteSpan) override;
-		BAN::ErrorOr<void> send_command_and_block(uint64_t lba, uint64_t sector_count, paddr_t paddr, Command command);
+		BAN::ErrorOr<size_t> send_command_sync(uint64_t lba, BAN::ConstByteSpan buffer, Command command);
 
-		uint32_t find_free_command_slot();
+		BAN::ErrorOr<void> send_command_and_wait(uint32_t slot);
+		BAN::ErrorOr<uint32_t> find_free_command_slot();
 
 		void handle_irq();
 
 	private:
-		Mutex m_mutex;
+		static constexpr uint32_t m_max_hba_prdt_count { 64 };
 
+		BAN::Atomic<uint32_t> m_prev_is { 0 };
 		uint32_t m_free_slots { 0 };
 
 		BAN::RefPtr<AHCIController> m_controller;
 		volatile HBAPortMemorySpace* const m_port;
 
 		BAN::UniqPtr<DMARegion> m_dma_region;
-		// Intermediate read/write buffer
+
-		// TODO: can we read straight to user buffer?
+		Mutex m_temp_buffer_mutex;
-		BAN::UniqPtr<DMARegion> m_data_dma_region;
+		BAN::UniqPtr<DMARegion> m_temp_buffer;
 
 		SpinLock m_command_lock;
 		ThreadBlocker m_command_blocker;

kernel/kernel/Storage/ATA/AHCI/Controller.cpp

@@ -39,6 +39,7 @@ namespace Kernel
 		m_pci_device.enable_interrupt(0, *this);
 		abar_mem.ghc = abar_mem.ghc | SATA_GHC_INTERRUPT_ENABLE;
 
+		m_supports_64bit = !!(abar_mem.cap & (1u << 31));
 		m_command_slot_count = ((abar_mem.cap >> 8) & 0x1F) + 1;
 
 		uint32_t pi = abar_mem.pi;
@@ -85,28 +86,29 @@ namespace Kernel
 
 	void AHCIController::handle_irq()
 	{
-		auto& abar_mem = *(volatile HBAGeneralMemorySpace*)m_abar->vaddr();
+		auto& abar_mem = *reinterpret_cast<volatile HBAGeneralMemorySpace*>(m_abar->vaddr());
 
-		const uint32_t is = abar_mem.is;
+		while (uint32_t is = abar_mem.is)
+		{
 			abar_mem.is = is;
 
-		for (uint8_t i = 0; i < 32; i++)
+			while (is != 0)
 			{
-			if (is & (1 << i))
+				const size_t idx = __builtin_ctz(is);
-			{
+				if (auto& device = m_devices[idx])
-				if (m_devices[i])
+					device->handle_irq();
-					m_devices[i]->handle_irq();
 				else
-					dwarnln("ignoring interrupt to device {}", i);
+					dwarnln("ignoring interrupt for port {}", idx);
+				is &= ~(1u << idx);
 			}
 		}
 	}
 
 	BAN::Optional<AHCIPortType> AHCIController::check_port_type(volatile HBAPortMemorySpace& port)
 	{
-		uint32_t ssts = port.ssts;
+		const uint32_t ssts = port.ssts;
-		uint8_t ipm = (ssts >> 8) & 0x0F;
+		const uint8_t ipm = (ssts >> 8) & 0x0F;
-		uint8_t det = (ssts >> 0) & 0x0F;
+		const uint8_t det = (ssts >> 0) & 0x0F;
 
 		if (det != HBA_PORT_DET_PRESENT)
 			return {};

kernel/kernel/Storage/ATA/AHCI/Device.cpp

@@ -11,6 +11,13 @@ namespace Kernel
 
 	static constexpr uint64_t s_ata_timeout_ms = 1000;
 
+	static constexpr size_t align_up_to(size_t value, size_t alignment)
+	{
+		if (const size_t rem = value % alignment)
+			value += alignment - rem;
+		return value;
+	}
+
 	static void start_cmd(volatile HBAPortMemorySpace* port)
 	{
 		while (port->cmd & HBA_PxCMD_CR)
@@ -27,6 +34,27 @@ namespace Kernel
 			continue;
 	}
 
+	paddr_t AHCIDevice::read_paddr(volatile uint32_t& lo, volatile uint32_t& hi) const
+	{
+		if (!m_controller->supports_64bit())
+			return lo;
+		return (static_cast<uint64_t>(hi) << 32) | lo;
+	}
+
+	void AHCIDevice::write_paddr(volatile uint32_t& lo, volatile uint32_t& hi, paddr_t paddr)
+	{
+		if (!m_controller->supports_64bit())
+		{
+			ASSERT((paddr >> 32) == 0);
+			lo = paddr;
+		}
+		else
+		{
+			lo = paddr & 0xFFFFFFFF;
+			hi = paddr >> 32;
+		}
+	}
+
 	BAN::ErrorOr<BAN::RefPtr<AHCIDevice>> AHCIDevice::create(BAN::RefPtr<AHCIController> controller, volatile HBAPortMemorySpace* port)
 	{
 		auto* device_ptr = new AHCIDevice(controller, port);
@@ -37,9 +65,16 @@ namespace Kernel
 
 	BAN::ErrorOr<void> AHCIDevice::initialize()
 	{
-		TRY(allocate_buffers());
 		TRY(rebase());
 
+		m_temp_buffer = TRY(DMARegion::create(256 * 1024, PageTable::MemoryType::Normal));
+		memset(reinterpret_cast<void*>(m_temp_buffer->vaddr()), 0x00, m_temp_buffer->size());
+		if (!m_controller->supports_64bit() && m_temp_buffer->paddr() + m_temp_buffer->size() > 0x100000000)
+		{
+			dwarnln("cannot allocate 32 bit buffer and the controller does not support 64 bit");
+			return BAN::Error::from_errno(EFAULT);
+		}
+
 		if (const uint32_t command_slots = m_controller->command_slot_count(); command_slots < 32)
 			m_free_slots = (1u << command_slots) - 1;
 		else
@@ -49,50 +84,42 @@ namespace Kernel
 		m_port->ie = 0xFFFFFFFF;
 
 		TRY(read_identify_data());
-		TRY(detail::ATABaseDevice::initialize({ (const uint16_t*)m_data_dma_region->vaddr(), m_data_dma_region->size() / sizeof(uint16_t) }));
+		TRY(detail::ATABaseDevice::initialize({ reinterpret_cast<const uint16_t*>(m_temp_buffer->vaddr()), 256 }));
-
-		return {};
-	}
-
-	BAN::ErrorOr<void> AHCIDevice::allocate_buffers()
-	{
-		const uint32_t command_slot_count = m_controller->command_slot_count();
-		const size_t needed_bytes = (sizeof(HBACommandHeader) + sizeof(HBACommandTable)) * command_slot_count + sizeof(ReceivedFIS);
-
-		m_dma_region = TRY(DMARegion::create(needed_bytes));
-		memset((void*)m_dma_region->vaddr(), 0x00, m_dma_region->size());
-
-		m_data_dma_region = TRY(DMARegion::create(PAGE_SIZE, PageTable::Normal));
-		memset((void*)m_data_dma_region->vaddr(), 0x00, m_data_dma_region->size());
 
 		return {};
 	}
 
 	BAN::ErrorOr<void> AHCIDevice::rebase()
 	{
-		ASSERT(m_dma_region);
-
 		const uint32_t command_slot_count = m_controller->command_slot_count();
 
+		const size_t command_list_size = command_slot_count * sizeof(HBACommandHeader);
+
+		const size_t command_table_offset = align_up_to(command_list_size, 128);
+		const size_t command_table_entry_size = align_up_to(sizeof(HBACommandTable) + m_max_hba_prdt_count * sizeof(HBAPRDTEntry), 128);
+
+		const size_t fis_offset = align_up_to(command_table_offset + command_slot_count * command_table_entry_size, 256);
+
+		m_dma_region = TRY(DMARegion::create(fis_offset + sizeof(ReceivedFIS)));
+		if (!m_controller->supports_64bit() && m_dma_region->paddr() + m_dma_region->size() > 0x100000000)
+		{
+			dwarnln("cannot allocate 32 bit buffer and the controller does not support 64 bit");
+			return BAN::Error::from_errno(EFAULT);
+		}
+
+		memset(reinterpret_cast<void*>(m_dma_region->vaddr()), 0x00, m_dma_region->size());
+
 		stop_cmd(m_port);
 
-		const paddr_t fis_paddr = m_dma_region->paddr();
+		const paddr_t command_list_paddr = m_dma_region->paddr();
-		m_port->fb  = fis_paddr & 0xFFFFFFFF;
+		write_paddr(m_port->clb, m_port->clbu, command_list_paddr);
-		m_port->fbu = fis_paddr >> 32;
 
-		const paddr_t command_list_paddr = fis_paddr + sizeof(ReceivedFIS);
+		volatile auto* command_list = reinterpret_cast<volatile HBACommandHeader*>(m_dma_region->paddr_to_vaddr(command_list_paddr));
-		m_port->clb  = command_list_paddr & 0xFFFFFFFF;
+		const paddr_t command_table_base = command_list_paddr + command_slot_count * sizeof(HBACommandHeader);
-		m_port->clbu = command_list_paddr >> 32;
+		for (uint32_t slot = 0; slot < command_slot_count; slot++)
+			write_paddr(command_list[slot].ctba, command_list[slot].ctbau, command_table_base + slot * command_table_entry_size);
 
-		auto* command_headers = reinterpret_cast<volatile HBACommandHeader*>(m_dma_region->paddr_to_vaddr(command_list_paddr));
+		write_paddr(m_port->fb, m_port->fbu, m_dma_region->paddr() + fis_offset);
-		const paddr_t command_table_paddr = command_list_paddr + command_slot_count * sizeof(HBACommandHeader);
-		for (uint32_t i = 0; i < command_slot_count; i++)
-		{
-			const paddr_t command_table_entry_paddr = command_table_paddr + i * sizeof(HBACommandTable);
-			command_headers[i].prdtl = s_hba_prdt_count;
-			command_headers[i].ctba  = command_table_entry_paddr & 0xFFFFFFFF;
-			command_headers[i].ctbau = command_table_entry_paddr >> 32;
-		}
 
 		start_cmd(m_port);
 
@@ -101,39 +128,26 @@ namespace Kernel
 
 	BAN::ErrorOr<void> AHCIDevice::read_identify_data()
 	{
-		ASSERT(m_data_dma_region);
+		const auto slot = TRY(find_free_command_slot());
 
-		const auto slot = find_free_command_slot();
+		const vaddr_t command_header_vaddr = m_dma_region->paddr_to_vaddr(read_paddr(m_port->clb, m_port->clbu));
-
+		volatile auto& command_header = reinterpret_cast<volatile HBACommandHeader*>(command_header_vaddr)[slot];
-		const paddr_t command_header_paddr = (static_cast<paddr_t>(m_port->clbu) << 32) | m_port->clb;
-		volatile auto& command_header = reinterpret_cast<volatile HBACommandHeader*>(m_dma_region->paddr_to_vaddr(command_header_paddr))[slot];
 		command_header.cfl   = sizeof(FISRegisterH2D) / sizeof(uint32_t);
 		command_header.w     = 0;
 		command_header.prdtl = 1;
 
-		const paddr_t command_table_paddr = (static_cast<paddr_t>(command_header.ctbau) << 32) | command_header.ctba;
+		const vaddr_t command_table_vaddr = m_dma_region->paddr_to_vaddr(read_paddr(command_header.ctba, command_header.ctbau));
-		volatile auto& command_table = *reinterpret_cast<volatile HBACommandTable*>(m_dma_region->paddr_to_vaddr(command_table_paddr));
+		volatile auto& command_table = *reinterpret_cast<volatile HBACommandTable*>(command_table_vaddr);
-		command_table.prdt_entry[0].dba  = m_data_dma_region->paddr() & 0xFFFFFFFF;
+		write_paddr(command_table.prdt_entry[0].dba, command_table.prdt_entry[0].dbau, m_temp_buffer->paddr());
-		command_table.prdt_entry[0].dbau = m_data_dma_region->paddr() >> 32;
 		command_table.prdt_entry[0].dbc = 511;
 		command_table.prdt_entry[0].i   = 1;
 
-		volatile auto& command = *reinterpret_cast<volatile FISRegisterH2D*>(command_table.cfis);
+		volatile auto& command = *reinterpret_cast<volatile FISRegisterH2D*>(&command_table.cfis[0]);
 		command.fis_type = FIS_TYPE_REGISTER_H2D;
 		command.c        = 1;
 		command.command  = ATA_COMMAND_IDENTIFY;
 
-		SpinLockGuard _(m_command_lock);
+		TRY(send_command_and_wait(slot));
-
-		m_port->ci = 1u << slot;
-
-		while (m_port->ci & (1u << slot))
-		{
-			BlockableSpinLock block(m_command_lock);
-			m_command_blocker.block_indefinite(&block);
-		}
-
-		m_free_slots |= 1u << slot;
 
 		return {};
 	}
@@ -159,6 +173,7 @@ namespace Kernel
 	{
 		while (const uint32_t is = m_port->is)
 		{
+			m_prev_is |= is;
 			m_port->is = is;
 
 			SpinLockGuard _(m_command_lock);
@@ -172,18 +187,48 @@ namespace Kernel
 		}
 	}
 
+	bool AHCIDevice::can_use_buffer_directly(BAN::ConstByteSpan buffer) const
+	{
+		const vaddr_t buffer_vaddr = reinterpret_cast<vaddr_t>(buffer.data());
+		if (buffer_vaddr % 2)
+			return false;
+
+		if (m_controller->supports_64bit())
+			return true;
+
+		const vaddr_t buffer_base = buffer_vaddr & PAGE_ADDR_MASK;
+		for (size_t off = 0; off < buffer.size(); off++)
+			if (PageTable::kernel().physical_address_of(buffer_base + off) >= 0x100000000)
+				return false;
+
+		return true;
+	}
+
 	BAN::ErrorOr<void> AHCIDevice::read_sectors_impl(uint64_t lba, uint64_t sector_count, BAN::ByteSpan buffer)
 	{
 		ASSERT(buffer.size() >= sector_count * sector_size());
+		if (buffer.size() > sector_count * sector_size())
+			buffer = buffer.slice(0, sector_count * sector_size());
 
-		LockGuard _(m_mutex);
+		if (can_use_buffer_directly(buffer))
-
-		const size_t max_sectors = m_data_dma_region->size() / sector_size();
-		for (uint64_t sector_off = 0; sector_off < sector_count; sector_off += max_sectors)
 		{
-			const uint64_t to_read = BAN::Math::min<uint64_t>(sector_count - sector_off, max_sectors);
+			size_t sectors_done = 0;
-			TRY(send_command_and_block(lba + sector_off, to_read, m_data_dma_region->paddr(), Command::Read));
+			while (sectors_done < sector_count)
-			memcpy(buffer.data() + sector_off * sector_size(), reinterpret_cast<void*>(m_data_dma_region->vaddr()), to_read * sector_size());
+				sectors_done += TRY(send_command_sync(lba + sectors_done, buffer.slice(sectors_done * sector_size()), Command::Read));
+		}
+		else
+		{
+			LockGuard _(m_temp_buffer_mutex);
+
+			uint8_t* const temp_buffer = reinterpret_cast<uint8_t*>(m_temp_buffer->vaddr());
+			while (!buffer.empty())
+			{
+				const size_t max_bytes = BAN::Math::min(buffer.size(), m_temp_buffer->size());
+				const size_t sectors = TRY(send_command_sync(lba, { temp_buffer, max_bytes }, Command::Read));
+				memcpy(buffer.data(), temp_buffer, sectors * sector_size());
+				buffer = buffer.slice(sectors * sector_size());
+				lba += sectors;
+			}
 		}
 
 		return {};
@@ -192,31 +237,42 @@ namespace Kernel
 	BAN::ErrorOr<void> AHCIDevice::write_sectors_impl(uint64_t lba, uint64_t sector_count, BAN::ConstByteSpan buffer)
 	{
 		ASSERT(buffer.size() >= sector_count * sector_size());
+		if (buffer.size() > sector_count * sector_size())
+			buffer = buffer.slice(0, sector_count * sector_size());
 
-		LockGuard _(m_mutex);
+		if (can_use_buffer_directly(buffer))
-
-		const size_t max_sectors = m_data_dma_region->size() / sector_size();
-		for (uint64_t sector_off = 0; sector_off < sector_count; sector_off += max_sectors)
 		{
-			const uint64_t to_write = BAN::Math::min<uint64_t>(sector_count - sector_off, max_sectors);
+			size_t sectors_done = 0;
-			memcpy(reinterpret_cast<void*>(m_data_dma_region->vaddr()), buffer.data() + sector_off * sector_size(), to_write * sector_size());
+			while (sectors_done < sector_count)
-			TRY(send_command_and_block(lba + sector_off, to_write, m_data_dma_region->paddr(), Command::Write));
+				sectors_done += TRY(send_command_sync(lba + sectors_done, buffer.slice(sectors_done * sector_size()), Command::Write));
+		}
+		else
+		{
+			LockGuard _(m_temp_buffer_mutex);
+
+			uint8_t* const temp_buffer = reinterpret_cast<uint8_t*>(m_temp_buffer->vaddr());
+			while (!buffer.empty())
+			{
+				const size_t max_bytes = BAN::Math::min(buffer.size(), m_temp_buffer->size());
+				memcpy(temp_buffer, buffer.data(), max_bytes);
+				const size_t sectors = TRY(send_command_sync(lba, { temp_buffer, max_bytes }, Command::Write));
+				buffer = buffer.slice(sectors * sector_size());
+				lba += sectors;
+			}
 		}
 
 		return {};
 	}
 
-	BAN::ErrorOr<void> AHCIDevice::send_command_and_block(uint64_t lba, uint64_t sector_count, paddr_t paddr, Command command)
+	BAN::ErrorOr<size_t> AHCIDevice::send_command_sync(uint64_t lba, BAN::ConstByteSpan buffer, Command command)
 	{
 		ASSERT(m_dma_region);
-		ASSERT(0 < sector_count && sector_count <= 0xFFFF + 1);
 
-		const auto slot = find_free_command_slot();
+		const auto slot = TRY(find_free_command_slot());
 
-		const paddr_t command_header_paddr = (static_cast<paddr_t>(m_port->clbu) << 32) | m_port->clb;
+		const vaddr_t command_header_vaddr = m_dma_region->paddr_to_vaddr(read_paddr(m_port->clb, m_port->clbu));
-		volatile auto& command_header = reinterpret_cast<volatile HBACommandHeader*>(m_dma_region->paddr_to_vaddr(command_header_paddr))[slot];
+		volatile auto& command_header = reinterpret_cast<volatile HBACommandHeader*>(command_header_vaddr)[slot];
 		command_header.cfl = sizeof(FISRegisterH2D) / sizeof(uint32_t);
-		command_header.prdtl = 1;
 		switch (command)
 		{
 			case Command::Read:
@@ -229,19 +285,66 @@ namespace Kernel
 				ASSERT_NOT_REACHED();
 		}
 
-		const paddr_t command_table_paddr = (static_cast<paddr_t>(command_header.ctbau) << 32) | command_header.ctba;
+		const vaddr_t command_table_vaddr = m_dma_region->paddr_to_vaddr(read_paddr(command_header.ctba, command_header.ctbau));
-		volatile auto& command_table = *reinterpret_cast<HBACommandTable*>(m_dma_region->paddr_to_vaddr(command_table_paddr));
+		volatile auto& command_table = *reinterpret_cast<volatile HBACommandTable*>(command_table_vaddr);
 
-		command_table.prdt_entry[0].dba  = paddr & 0xFFFFFFFF;
+		size_t prdt_count = 0;
-		command_table.prdt_entry[0].dbau = paddr >> 32;
+		size_t total_bytes = 0;
-		command_table.prdt_entry[0].dbc  = sector_count * sector_size() - 1;
+		paddr_t extend_paddr = 0;
-		command_table.prdt_entry[0].i    = 1;
 
-		volatile auto& fis_command = *reinterpret_cast<volatile FISRegisterH2D*>(command_table.cfis);
+		while (!buffer.empty())
+		{
+			const auto to_paddr = [](vaddr_t vaddr) -> paddr_t {
+				return PageTable::kernel().physical_address_of(vaddr & PAGE_ADDR_MASK) + (vaddr % PAGE_SIZE);
+			};
+
+			const vaddr_t buffer_vaddr = reinterpret_cast<vaddr_t>(buffer.data());
+			const paddr_t buffer_paddr = to_paddr(buffer_vaddr);
+
+			const size_t bytes = BAN::Math::min(buffer.size(), PAGE_SIZE - buffer_vaddr % PAGE_SIZE);
+
+			bool can_extend = true;
+			if (prdt_count == 0)
+				can_extend = false;
+			else if (buffer_paddr != extend_paddr)
+				can_extend = false;
+			else if (command_table.prdt_entry[prdt_count - 1].dbc + bytes >= 0x400000)
+				can_extend = false;
+
+			if (can_extend)
+				command_table.prdt_entry[prdt_count - 1].dbc += bytes;
+			else
+			{
+				if (prdt_count >= m_max_hba_prdt_count)
+					break;
+				command_table.prdt_entry[prdt_count].dba  = buffer_paddr & 0xFFFFFFFF;
+				command_table.prdt_entry[prdt_count].dbau = buffer_paddr >> 32;
+				command_table.prdt_entry[prdt_count].dbc  = bytes - 1;
+				prdt_count++;
+			}
+
+			buffer = buffer.slice(bytes);
+			total_bytes += bytes;
+			extend_paddr = buffer_paddr + bytes;
+		}
+
+		if (const size_t rem = total_bytes % sector_size())
+		{
+			// TODO: this wont work with block sizes > PAGE_SIZE
+			ASSERT(rem < static_cast<size_t>(command_table.prdt_entry[prdt_count - 1].dbc + 1));
+			command_table.prdt_entry[prdt_count - 1].dbc -= rem;
+			total_bytes -= rem;
+		}
+
+		command_header.prdtl = prdt_count;
+
+		const size_t sector_count = total_bytes / sector_size();
+
+		volatile auto& fis_command = *reinterpret_cast<volatile FISRegisterH2D*>(&command_table.cfis[0]);
 		fis_command.fis_type = FIS_TYPE_REGISTER_H2D;
 		fis_command.c        = 1;
 
-		const bool needs_extended = lba >= (1 << 24) || sector_count > 0xFF;
+		const bool needs_extended = (lba + sector_count) > (1 << 24) || sector_count > 0xFF;
 		ASSERT (!needs_extended || (m_command_set & ATA_COMMANDSET_LBA48_SUPPORTED));
 
 		switch (command)
@@ -268,23 +371,53 @@ namespace Kernel
 		fis_command.count_lo = (sector_count >> 0) & 0xFF;
 		fis_command.count_hi = (sector_count >> 8) & 0xFF;
 
+		TRY(send_command_and_wait(slot));
+
+		return sector_count;
+	}
+
+	BAN::ErrorOr<void> AHCIDevice::send_command_and_wait(uint32_t slot)
+	{
+		const uint64_t timeout_ms = SystemTimer::get().ms_since_boot() + s_ata_timeout_ms;
+
 		SpinLockGuard _(m_command_lock);
 
-		m_port->ci = 1u << slot;
+		m_port->ci |= 1u << slot;
 
 		while (m_port->ci & (1u << slot))
 		{
+			if (SystemTimer::get().ms_since_boot() >= timeout_ms)
+			{
+				m_free_slots |= 1u << slot;
+				return BAN::Error::from_errno(ETIMEDOUT);
+			}
 			BlockableSpinLock block(m_command_lock);
-			m_command_blocker.block_indefinite(&block);
+			m_command_blocker.block_with_wake_time_ms(timeout_ms, &block);
 		}
 
 		m_free_slots |= 1u << slot;
 
+		constexpr uint32_t is_error =
+			(1 << 30) | // task file error
+			(1 << 29) | // host bus fatal error
+			(1 << 28) | // host bus data error
+			(1 << 27) | // interface fatal error
+			(1 << 26) | // interface non-fatal error
+			(1 << 24);  // overflow
+
+		if (const uint32_t is = (m_prev_is.exchange(0) | m_port->is); is & is_error)
+			return BAN::Error::from_errno(EFAULT);
+
+		if (m_port->tfd & (ATA_STATUS_ERR | ATA_STATUS_DF))
+			return BAN::Error::from_errno(EFAULT);
+
 		return {};
 	}
 
-	uint32_t AHCIDevice::find_free_command_slot()
+	BAN::ErrorOr<uint32_t> AHCIDevice::find_free_command_slot()
 	{
+		const uint64_t timeout_ms = SystemTimer::get().ms_since_boot() + s_ata_timeout_ms;
+
 		SpinLockGuard _(m_command_lock);
 
 		for (;;)
@@ -296,8 +429,11 @@ namespace Kernel
 				return slot;
 			}
 
+			if (SystemTimer::get().ms_since_boot() >= timeout_ms)
+				return BAN::Error::from_errno(ETIMEDOUT);
+
 			BlockableSpinLock block(m_command_lock);
-			m_command_blocker.block_indefinite(&block);
+			m_command_blocker.block_with_timeout_ms(timeout_ms, &block);
 		}
 	}