From 62626343245100fd8db71d1c5a43f7d740728c60 Mon Sep 17 00:00:00 2001
From: Abdur-Rahman Mansoor <armansoor@missingno.dev>
Date: Tue, 30 Jul 2024 13:48:50 -0400
Subject: [PATCH] feat: implement support for multiple NVMe drives

---
 drivers/nvme/nvme.asm | 143 +++++++++++++++++++++---------------------
 1 file changed, 72 insertions(+), 71 deletions(-)

diff --git a/drivers/nvme/nvme.asm b/drivers/nvme/nvme.asm
index 31b7c84..665f1be 100644
--- a/drivers/nvme/nvme.asm
+++ b/drivers/nvme/nvme.asm
@@ -62,28 +62,26 @@ proc START c, reason:dword
 	call    detect_nvme
 	test 	eax, eax
 	jz      .err
-	mov 	eax, dword [p_nvme_devices]
-	test 	eax, eax
-	jz 	.err
-	xor 	ecx, ecx
+	xor 	ebx, ebx
+	mov 	esi, dword [p_nvme_devices]
+	DEBUGF  DBG_INFO, "p_nvme_devices: 0x%x\n", [p_nvme_devices]
 
 .loop:
-	mov 	ebx, dword [p_nvme_devices]
-	stdcall device_is_compat, ebx
+	push 	ebx
+	lea 	esi, [esi + ebx * 4]
+	stdcall device_is_compat, esi
 	test 	eax, eax
 	jz 	@f
-	stdcall nvme_init, ebx
+	stdcall nvme_init, esi
 	test 	eax, eax
-	jz 	.err
-
-;@@:
-	;inc 	ecx
-	;cmp 	ecx, dword [pcidevs_len]
-	;jne 	.loop
-	stdcall add_nvme_disk, [p_nvme_devices]
-	test 	eax, eax
-	jz 	.err
+	jz 	@f
 
+@@:
+	pop 	ebx
+	inc 	ebx
+	cmp 	ebx, dword [pcidevs_len]
+	jne 	.loop
+	stdcall add_nvme_disk, dword [p_nvme_devices]
 	invoke  RegService, my_service, service_proc
 	ret
 
@@ -149,7 +147,7 @@ proc add_nvme_disk stdcall, pci:dword
 	ret
 
 @@:
-	DEBUGF  DBG_INFO, "nvme%u: Failed to register disk\n", [esi + pcidev.num]
+	DEBUGF  DBG_INFO, "nvme%un%u: Failed to register disk\n", [esi + pcidev.num], [esi + pcidev.nsid]
 	xor 	eax, eax
 	pop 	esi
 	ret
@@ -729,64 +727,58 @@ endp
 proc detect_nvme
 
       	invoke  GetPCIList
-	mov 	edx, eax
+	mov 	esi, eax
+	mov 	ebx, eax
 
 .check_dev:
-	mov 	ebx, dword [eax + PCIDEV.class]
-	and	ebx, 0x00ffff00 ; retrieve class/subclass code only
-	cmp 	ebx, 0x00010800 ; Mass Storage Controller - Non-Volatile Memory Controller
+	mov 	eax, dword [esi + PCIDEV.class]
+	and	eax, 0x00ffff00 ; retrieve class/subclass code only
+	cmp 	eax, 0x00010800 ; Mass Storage Controller - Non-Volatile Memory Controller
 	je	.found_dev
 
 .next_dev:
-	mov	eax, dword [eax + PCIDEV.fd]
-	cmp	eax, edx
+	mov	esi, dword [esi + PCIDEV.fd]
+	cmp	esi, ebx
 	jne 	.check_dev
-	jmp 	.exit_success
-
-.found_dev:
-	push    edx eax
-	;PDEBUGF DBG_INFO, "PCI(%u.%u.%u): Detected NVMe device...\n", byte [eax + PCIDEV.bus], byte [eax + PCIDEV.devfn]
-	cmp 	dword [pcidevs_len], TOTAL_PCIDEVS
-	jne 	@f
-	pop 	eax edx
-	jmp 	.exit_success
-
-@@:
-	inc 	dword [pcidevs_len]
-	mov     ebx, dword [p_nvme_devices]
-	test    ebx, ebx
-	jnz	@f
-	invoke  KernelAlloc, sizeof.pcidev
-	test 	eax, eax
-	jz 	.err_no_mem
-	mov     dword [p_nvme_devices], eax 
-	;DEBUGF  DBG_INFO, "(NVMe) Allocated pcidev struct at 0x%x\n", [p_nvme_devices]
-
-@@:
-	mov 	ecx, dword [pcidevs_len]
-	dec 	ecx
-	pop	eax
-	mov 	ebx, dword [p_nvme_devices]
-
-	movzx 	edx, byte [eax + PCIDEV.bus]
-	mov 	byte [ebx + pcidev.bus], dl
-	movzx 	edx, byte [eax + PCIDEV.devfn]
-	mov 	byte [ebx + pcidev.devfn], dl
-	mov 	dword [ebx + pcidev.num], ecx
-
-	pop 	edx
-	jmp	.next_dev
-
-.err_no_mem:
-	pop     eax edx
-	xor 	eax, eax
-	ret
 
 .exit_success:
 	xor 	eax, eax
 	inc 	eax
 	ret
 
+.found_dev:
+	cmp 	dword [pcidevs_len], TOTAL_PCIDEVS
+	jne 	@f
+	DEBUGF  DBG_INFO, "Can't add any more NVMe devices...\n"
+	jmp 	.exit_success
+
+@@:
+	inc 	dword [pcidevs_len]
+	cmp     dword [p_nvme_devices], 0
+	jnz	@f ; was the pointer already allocated?
+	invoke  KernelAlloc, sizeof.pcidev * TOTAL_PCIDEVS
+	test 	eax, eax
+	jz 	.err_no_mem
+	mov     dword [p_nvme_devices], eax 
+
+@@:
+	mov 	ecx, dword [pcidevs_len]
+	dec 	ecx
+	mov 	edi, dword [p_nvme_devices]
+	lea 	edi, [edi + ecx * 4]
+
+	movzx 	eax, byte [esi + PCIDEV.bus]
+	mov 	byte [edi + pcidev.bus], al
+	movzx 	eax, byte [esi + PCIDEV.devfn]
+	mov 	byte [edi + pcidev.devfn], al
+	mov 	dword [edi + pcidev.num], ecx
+
+	jmp	.next_dev
+
+.err_no_mem:
+	xor 	eax, eax
+	ret
+
 endp
 
 proc device_is_compat stdcall, pci:dword
@@ -1416,28 +1408,37 @@ endp
 
 proc irq_handler
 
-	push 	esi edi
-	mov 	esi, dword [p_nvme_devices]
+	push 	ebx esi edi
+	mov 	edi, dword [p_nvme_devices]
+	mov 	esi, edi
+	sub 	esi, sizeof.pcidev
+	mov 	ebx, dword [pcidevs_len]
+	xor 	ecx, ecx
 
-	; check if a NVMe device generated an interrupt
+.check_who_raised_irq:
+	add 	esi, sizeof.pcidev
+	inc 	ecx
+	cmp 	ecx, ebx
+	ja 	.not_our_irq
 	invoke 	PciRead16, dword [esi + pcidev.bus], dword [esi + pcidev.devfn], PCI_header00.status
 	test 	al, 1000b ; check interrupt status
-	jz 	.not_our_irq
+	jz 	.check_who_raised_irq
 	
+	mov 	esi, edi
 	mov 	edi, dword [esi + pcidev.io_addr]
 	mov 	dword [edi + NVME_MMIO.INTMS], 0x3
 
 	mov 	eax, dword [spinlock]
 	test 	eax, eax
 	jz 	@f ; not locked, so it must be an I/O command
-	stdcall consume_cq_entries, [p_nvme_devices], 0
+	stdcall consume_cq_entries, esi, 0
 
 @@:
-	stdcall consume_cq_entries, [p_nvme_devices], 1
+	stdcall consume_cq_entries, esi, 1
 
 	; Interrupt handled by driver, return 1
 	mov 	dword [edi + NVME_MMIO.INTMC], 0x3
-	pop 	edi esi
+	pop 	edi esi ebx
 	xor 	eax, eax
 	xchg 	eax, dword [spinlock] ; unlock spinlock
 	mov 	eax, 1
@@ -1445,7 +1446,7 @@ proc irq_handler
 
 .not_our_irq:
 	; Interrupt not handled by driver, return 0
-	pop 	edi esi
+	pop 	edi esi ebx
 	xor 	eax, eax
 	ret