From f4ded13c9a9612c598a533769e69096bce79fbec Mon Sep 17 00:00:00 2001
From: Abdur-Rahman Mansoor <armansoor@missingno.dev>
Date: Thu, 6 Jun 2024 12:43:57 -0400
Subject: [PATCH] feat: implement `set_features` and `abort`

---
 drivers/nvme/nvme.asm | 62 +++++++++++++++++++++++++++++++++----------
 1 file changed, 48 insertions(+), 14 deletions(-)

diff --git a/drivers/nvme/nvme.asm b/drivers/nvme/nvme.asm
index 9dfa802..5102c4b 100644
--- a/drivers/nvme/nvme.asm
+++ b/drivers/nvme/nvme.asm
@@ -198,6 +198,43 @@ proc create_io_submission_queue stdcall, pci:dword, prp1:dword, qid:word, cqid:w
 
 endp
 
+; See page 95-96 of the NVMe 1.4 specification for reference
+proc abort stdcall, pci:dword, cid:word, sqid:word
+
+		push 	esi
+		mov 	esi, [pci]
+		stdcall memset, esi, 0, sizeof.SQ_ENTRY
+		stdcall set_cdw0, ADM_CMD_ABORT, 0 ; [TODO: Set CID to valid value]
+		mov 	dword [esi + SQ_ENTRY.cdw0], eax
+		mov 	eax, [cid]
+		shl 	eax, 16
+		or 		eax, [sqid]
+		mov 	dword [esi + SQ_ENTRY.cdw10], eax
+		stdcall sqytdbl_write, [pci], 0, 0
+		pop 	esi
+		ret
+
+endp
+
+; See page 205 of the NVMe 1.4 specification for reference
+proc set_features stdcall, pci:dword, dptr:dword, fid:byte
+		
+		push 	esi
+		mov 	esi, [pci]
+		mov 	esi, [esi + pcidev.sq_ptr]
+		stdcall memset, esi, 0, sizeof.SQ_ENTRY
+		stdcall set_cdw0, AMD_CMD_SET_FEATURES, 0 ; [TODO: Set CID to valid value]
+		mov 	eax, [dptr]
+		mov 	dword [esi + SQ_ENTRY.dptr], eax
+		mov 	al, [fid]
+		mov 	byte [esi + SQ_ENTRY.cdw10], al
+		; TODO: Set CDW10.SV and CDW14.UUID to valid value
+		stdcall sqytdbl_write, [pci], 0, 0
+		pop 	esi
+		ret
+
+endp
+
 ; See page 105 of the NVMe 1.4 specification for reference
 proc delete_io_completion_queue stdcall, pci:dword, qid:word
 
@@ -212,6 +249,7 @@ proc delete_io_completion_queue stdcall, pci:dword, qid:word
 		stdcall sqytdbl_write, [pci], [qid], 0 ; setting last param to 0 for now, change later
 		pop 	esi
 		ret
+
 endp
 
 ; See page 114-116 of the NVMe 1.4 specification for reference
@@ -442,29 +480,25 @@ proc nvme_init stdcall, pci:dword
 		mov 	dword [edi + NVME_MMIO.AQA], eax
 
 		; Configure Admin Submission/Completion Queue Base Address
-		invoke  AllocPage
+		invoke  KernelAlloc, 0x1000
 		test 	eax, eax
 		jz 		.exit_fail
-		;or 		eax, PG_SW or PG_NOCACHE
+		mov 	dword [esi + pcidev.sq_ptr], eax
+		DEBUGF  DBG_INFO, "(NVMe) ASQ: Virtual Address: 0x%x\n", eax
+		invoke  GetPhysAddr
+		DEBUGF  DBG_INFO, "(NVMe) ASQ: Physical Address: 0x%x\n", eax
 		mov 	dword [edi + NVME_MMIO.ASQ], eax
 		and 	dword [edi + NVME_MMIO.ASQ + 4], 0
-		invoke  AllocPage
+		invoke  KernelAlloc, 0x1000
 		test 	eax, eax
 		jz 		.exit_fail
-		;or 	 	eax, PG_SW or PG_NOCACHE
+		mov 	dword [esi + pcidev.cq_ptr], eax
+		DEBUGF  DBG_INFO, "(NVMe) ACQ: Virtual Address: 0x%x\n", eax
+		invoke  GetPhysAddr
+		DEBUGF  DBG_INFO, "(NVMe) ACQ: Physical Address: 0x%x\n", eax
 		mov 	dword [edi + NVME_MMIO.ACQ], eax
 		and 	dword [edi + NVME_MMIO.ACQ + 4], 0
 
-		invoke  KernelAlloc, 0x2000
-		test 	eax, eax
-		jz 	 	.exit_fail
-		mov 	ebx, eax
-		invoke  MapPage, ebx, dword [edi + NVME_MMIO.ASQ], PG_SW+PG_NOCACHE
-		mov 	dword [esi + pcidev.sq_ptr], eax
-		add 	ebx, 0x1000
-		invoke  MapPage, ebx, dword [edi + NVME_MMIO.ACQ], PG_SW+PG_NOCACHE
-		mov 	dword [esi + pcidev.cq_ptr], eax
-
 		; Attach interrupt handler
 		movzx 	eax, byte [esi + pcidev.iline]
 		DEBUGF  DBG_INFO, "(NVMe) Attaching interrupt handler to IRQ %u\n", eax