From 76f94d74e1f028b23ec88c0899104aa97a0674da Mon Sep 17 00:00:00 2001
From: Abdur-Rahman Mansoor <armansoor@missingno.dev>
Date: Thu, 27 Jun 2024 19:19:06 -0400
Subject: [PATCH] nvme_init: check MQES

---
 drivers/nvme/nvme.asm | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/drivers/nvme/nvme.asm b/drivers/nvme/nvme.asm
index c5c3776..7693df8 100644
--- a/drivers/nvme/nvme.asm
+++ b/drivers/nvme/nvme.asm
@@ -449,7 +449,8 @@ proc nvme_write stdcall, pci:dword, qid:word, nsid:dword, dptr:dword, slba:qword
 		
 		; TODO: Use IDENTC.NOIOB to construct read/write commands that don't
 		; cross the I/O boundary to achieve optimal performance
-		; Also add DPTR/MPTR
+		; 
+		; Read AWUN/NAWUN
 		sub 	esp, sizeof.SQ_ENTRY
 		stdcall memset, esp, 0, sizeof.SQ_ENTRY
 		movzx 	ecx, [qid]
@@ -461,7 +462,7 @@ proc nvme_write stdcall, pci:dword, qid:word, nsid:dword, dptr:dword, slba:qword
 		mov 	dword [esp + SQ_ENTRY.nsid], eax
 		mov 	eax, dword [slba]
 		mov 	dword [esp + SQ_ENTRY.cdw10], eax
-		mov 	eax, dword [slba + 4]
+		mov 	eax, 0 ; SLBA + 32 (fix later)
 		mov 	dword [esp + SQ_ENTRY.cdw11], eax
 		movzx 	eax, [nlb]
 		mov 	word [esp + SQ_ENTRY.cdw12], ax
@@ -604,6 +605,11 @@ proc nvme_init stdcall, pci:dword
 			DEBUGF DBG_INFO, "(NVMe) CSTS: 0x%x\n", eax
 		end if
 
+		; check maximum queue entries supported
+		mov 	ax, word [edi + NVME_MMIO.CAP]
+		cmp 	ax, SQ_ENTRIES
+		jl 		.exit_fail
+		
 		; For some reason, bit 7 (No I/O command set supported) is also set to 1 despite bit 0 (NVM command set)
 	    ; being set to 1.. so I am not sure if bit 7 should be checked at all.. investigate later.
 		mov      eax, dword [edi + NVME_MMIO.CAP + 4]
@@ -763,11 +769,11 @@ proc nvme_init stdcall, pci:dword
 		mov 	eax, dword [esi + NVM_QUEUE_ENTRY.cq_ptr]
 		invoke 	GetPhysAddr
 		stdcall create_io_completion_queue, [pci], eax, 1, IEN_ON
-		DEBUGF  DBG_INFO, "nvme%u: Successfully created I/O completion queue 1", [edi + pcidev.num]
+		DEBUGF  DBG_INFO, "nvme%u: Successfully created I/O completion queue 1\n", [edi + pcidev.num]
 		mov 	eax, dword [esi + NVM_QUEUE_ENTRY.sq_ptr]
 		invoke 	GetPhysAddr
 		stdcall create_io_submission_queue, [pci], eax, 1, 1
-		DEBUGF  DBG_INFO, "nvme%u: Successfully created I/O submission queue 1", [edi + pcidev.num]
+		DEBUGF  DBG_INFO, "nvme%u: Successfully created I/O submission queue 1\n", [edi + pcidev.num]
 
 		if 1
 			stdcall determine_active_nsids, [pci]
@@ -816,8 +822,12 @@ proc nvme_init stdcall, pci:dword
 		DEBUGF  DBG_INFO, "nvme%un%u: Namespace LBA Data Size: %u\n", [esi + pcidev.num], [esi + pcidev.nsid], eax
 		invoke  KernelFree, edi
 		
-		mov 	eax, test_string
-		stdcall nvme_write, [pci], 1, [esi + pcidev.nsid], eax, 0, 1
+		;mov 	eax, test_string
+		invoke  KernelAlloc, 0x1000
+		test 	eax, eax
+		jz 		.exit_fail
+		invoke  GetPhysAddr
+		stdcall nvme_write, [pci], 1, [esi + pcidev.nsid], eax, 1, 1
 		DEBUGF  DBG_INFO, "nvme%u: Successfully initialized driver\n", [esi + pcidev.num]
       	xor     eax, eax
       	inc     eax