From 1f4b25a30929331da4c02bbcbfeb16418a6f0cfa Mon Sep 17 00:00:00 2001
From: Abdur-Rahman Mansoor <armansoor@missingno.dev>
Date: Thu, 27 Jun 2024 17:10:25 -0400
Subject: [PATCH] more fixes

---
 drivers/nvme/nvme.asm | 41 +++++++++++++++++++++++------------------
 1 file changed, 23 insertions(+), 18 deletions(-)

diff --git a/drivers/nvme/nvme.asm b/drivers/nvme/nvme.asm
index f6155de..ced720a 100644
--- a/drivers/nvme/nvme.asm
+++ b/drivers/nvme/nvme.asm
@@ -445,26 +445,26 @@ proc get_log_page stdcall, pci:dword, dptr:dword, lid:byte
 endp
 
 ; See page 269-271 of the NVMe 1.4 specification for reference
-proc nvme_write stdcall, pci:dword, qid:word, slba:qword, nlb:dword, dsm:byte
+proc nvme_write stdcall, pci:dword, qid:word, dptr:dword, slba:qword, nlb:word
 		
 		; 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
 		sub 	esp, sizeof.SQ_ENTRY
-		stdcall memset, esp, 0 sizeof.SQ_ENTRY
-		stdcall set_cdw0, [pci], [qid], NVM_CMD_WRITE
+		stdcall memset, esp, 0, sizeof.SQ_ENTRY
+		movzx 	ecx, [qid]
+		stdcall set_cdw0, [pci], ecx, NVM_CMD_WRITE
 		mov 	dword [esp + SQ_ENTRY.cdw0], eax ; CDW0
-
-		; Starting LBA (SLBA)
+		mov 	eax, [dptr]
+		mov 	dword [esp + SQ_ENTRY.dptr], eax
 		mov 	eax, dword [slba]
 		mov 	dword [esp + SQ_ENTRY.cdw10], eax
 		mov 	eax, dword [slba + 4]
 		mov 	dword [esp + SQ_ENTRY.cdw11], eax
-		mov 	ax, [nlb]
+		movzx 	eax, [nlb]
 		mov 	word [esp + SQ_ENTRY.cdw12], ax
-		mov 	al, [dsm]
-		mov 	byte [esp + SQ_ENTRY.cdw13], al
-		stdcall sqytdbl_write, [pci], [qid], esp
+		movzx 	ecx, [qid]
+		stdcall sqytdbl_write, [pci], ecx, esp
 		
 		add 	esp, sizeof.SQ_ENTRY
 		ret
@@ -819,6 +819,8 @@ 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, eax, 0, 1
 		DEBUGF  DBG_INFO, "nvme%u: Successfully initialized driver\n", [esi + pcidev.num]
       	xor     eax, eax
       	inc     eax
@@ -948,7 +950,9 @@ proc sqytdbl_write stdcall, pci:dword, y:word, cmd:dword
 		mov 	edi, [pci]
 		mov 	esi, dword [edi + pcidev.io_addr]
 		mov 	edi, dword [edi + pcidev.queue_entries]
-		movzx 	eax, word [edi + NVM_QUEUE_ENTRY.tail]
+		movzx 	ecx, [y]
+		imul 	ecx, sizeof.NVM_QUEUE_ENTRY
+		movzx 	eax, word [edi + ecx + NVM_QUEUE_ENTRY.tail]
 		cmp 	ax, NVM_ASQS
 		jl		@f 		
 		xor 	ax, ax
@@ -956,19 +960,19 @@ proc sqytdbl_write stdcall, pci:dword, y:word, cmd:dword
 @@:
 		mov 	esi, [pci]
 		inc 	ax
-		; 1000h + ((2y + 1) * (4 << CAP.DSTRD))
+		; 1000h + (2y * (4 << CAP.DSTRD))
 		movzx 	ebx, [y]
-		shl 	bl, 1
+		shl 	ebx, 1
 		mov 	edx, 4
 		mov 	cl, byte [esi + pcidev.dstrd]
-		shl 	dx, cl
-		imul 	dx, bx
-		add 	dx, 0x1000
-		;DEBUGF  DBG_INFO, "(NVMe) Writing to submission queue doorbell register 0x%x: %u\n", dx, ax
+		shl 	edx, cl
+		imul 	edx, ebx
+		add 	edx, 0x1000
 		mov 	esi, dword [esi + pcidev.io_addr]
 		mov 	word [esi + edx], ax
 		movzx 	ecx, [y]
-		mov 	word [edi + NVM_QUEUE_ENTRY.tail], ax
+		imul 	ecx, sizeof.NVM_QUEUE_ENTRY
+		mov 	word [edi + ecx + NVM_QUEUE_ENTRY.tail], ax
 		dec 	ax
 		stdcall nvme_cmd_wait, [pci], ecx, eax
 		pop 	edi esi ebx
@@ -1109,7 +1113,7 @@ proc irq_handler
 		pop 	ecx
 		inc 	ecx
 		cmp 	ecx, LAST_QUEUE_ID
-		jne 	@b
+		jng 	@b
 
 		; Interrupt handled by driver, return 1
 		mov 	dword [edi + NVME_MMIO.INTMC], 0x3
@@ -1163,6 +1167,7 @@ align 4
 			dd 		 0 ; use default cache size
 		.end:
 		if __DEBUG__
+		test_string  db "NVMe driver successfully wrote to disk! :D",0
 		include_debug_strings
 		end if