From 2f071876237ea3c48129323701a37d69f0d0f106 Mon Sep 17 00:00:00 2001
From: Abdur-Rahman Mansoor <armansoor@missingno.dev>
Date: Thu, 20 Jun 2024 15:55:32 -0400
Subject: [PATCH] fix fatal bugs in `write_admin_cmd` and dont handle IRQs if
 they dont belong to the NVMe device

---
 drivers/nvme/nvme.asm | 28 +++++++++++++++++++++-------
 drivers/nvme/nvme.inc |  2 +-
 2 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/drivers/nvme/nvme.asm b/drivers/nvme/nvme.asm
index c173eeb..94eab5a 100644
--- a/drivers/nvme/nvme.asm
+++ b/drivers/nvme/nvme.asm
@@ -730,21 +730,19 @@ proc write_admin_cmd stdcall, pci:dword, cmd:dword
 		lea 	edi, [edi + ecx]
 		stdcall memcpy, edi, esi, sizeof.SQ_ENTRY
 		
-		mov 	esi, dword [edi + pcidev.sq_ptr]
-		mov 	esi, dword [esi + pcidev.queue_entries]
-		mov 	edi, esi
-		mov 	ax, word [esi + NVM_QUEUE_ENTRY.tail]
+		mov 	edi, [pci]
+		mov 	esi, dword [edi + pcidev.io_addr]
+		mov 	edi, dword [edi + pcidev.queue_entries]
+		mov 	ax, word [edi + NVM_QUEUE_ENTRY.tail]
 		cmp 	ax, NVM_ASQS
 		jl		@f 		
 		xor 	ax, ax
 		
 @@:
-		mov 	esi, [pci]
-		mov 	esi, dword [esi + pcidev.io_addr]
 		inc 	ax
 		DEBUGF  DBG_INFO, "(NVMe) Writing to Admin Submission Queue: %u\n", ax
-		mov 	word [esi + 0x1000], ax
 		mov 	word [edi + NVM_QUEUE_ENTRY.tail], ax
+		mov 	word [esi + 0x1000], ax
 		pop 	edi esi
 		ret
 
@@ -789,11 +787,21 @@ proc irq_handler
 
 		push 	esi edi
 		mov 	esi, dword [p_nvme_devices]
+
+		; check if the NVMe device generated an interrupt
+		invoke 	PciRead16, dword [esi + pcidev.bus], dword [esi + pcidev.devfn], PCI_header00.status
+		test 	al, 1000b ; check interrupt status
+		jz 		.not_our_irq
+		
 		mov 	edi, esi
 		mov 	edi, dword [edi + pcidev.io_addr]
 		mov 	dword [edi + NVME_MMIO.INTMS], 0x1
 		mov 	esi, dword [esi + pcidev.queue_entries]
 		movzx 	ecx, word [esi + NVM_QUEUE_ENTRY.head]
+		mov 	dx, word [esi + NVM_QUEUE_ENTRY.tail]
+		DEBUGF  DBG_INFO, "IRQ (head): 0x%x, (tail): 0x%x\n", cx, dx
+		cmp 	cx, dx
+		je 		.end
 		mov 	edx, ecx
 		imul 	edx, sizeof.CQ_ENTRY
 		mov 	esi, dword [p_nvme_devices]
@@ -831,6 +839,12 @@ proc irq_handler
 		pop 	edi esi
 		ret
 
+.not_our_irq:
+		; Interrupt not handled by driver, return 0
+		xor 	eax, eax
+		pop 	edi esi
+		ret
+
 endp
 
 proc nvme_cleanup
diff --git a/drivers/nvme/nvme.inc b/drivers/nvme/nvme.inc
index afd9f43..25e3de0 100644
--- a/drivers/nvme/nvme.inc
+++ b/drivers/nvme/nvme.inc
@@ -18,7 +18,7 @@ VS140 = 0x00010400 ; (v1.4.0)
 
 NVM_CMDS = 64       ; Number of Commands
 NVM_MPS = 0         ; Memory Page Size (2 ^ (12 + MPS))
-NVM_ASQS = 4        ; Admin Submission Queue Size
+NVM_ASQS = 8        ; Admin Submission Queue Size
 NVM_ACQS = NVM_ASQS ; Admin Completion Queue Size
 
 ADMIN_QUEUE = 0 ; Admin Queue ID