From 7bba4781fbfb9690ec24ae5443e12492000b31ae Mon Sep 17 00:00:00 2001
From: Abdur-Rahman Mansoor <armansoor@missingno.dev>
Date: Tue, 25 Jun 2024 18:37:04 -0400
Subject: [PATCH] feat: implement disk registering code

---
 drivers/nvme/nvme.asm | 84 ++++++++++++++++++++++++++++++++++++++-----
 drivers/nvme/nvme.inc |  9 +++++
 2 files changed, 85 insertions(+), 8 deletions(-)

diff --git a/drivers/nvme/nvme.asm b/drivers/nvme/nvme.asm
index 085df68..0a57839 100644
--- a/drivers/nvme/nvme.asm
+++ b/drivers/nvme/nvme.asm
@@ -30,6 +30,12 @@ include "nvme.inc"
 include "macros.inc"
 include "lib.asm"
 
+struct DISKMEDIAINFO
+		flags 		dd ?
+		sectorsize  dd ?
+		capacity 	dq ?
+ends
+
 proc START c, reason:dword
 
       	cmp     [reason], DRV_ENTRY
@@ -58,6 +64,9 @@ proc START c, reason:dword
 		;inc 	ecx
 		;cmp 	ecx, dword [pcidevs_len]
 		;jne 	.loop
+		stdcall add_nvme_disk, [p_nvme_devices]
+		test 	eax, eax
+		jz 		.err
 		invoke  RegService, my_service, service_proc
 		ret
 
@@ -70,25 +79,68 @@ endp
 
 proc service_proc stdcall, ioctl:dword
 
-      	mov     ebx, [ioctl]
-      	mov     eax, [ebx+IOCTL.io_code]
+      	mov     esi, [ioctl]
+      	mov     eax, [esi + IOCTL.io_code]
       	cmp     eax, SRV_GETVERSION
-      	jne     @F
+      	jne     .ret
 
-      	mov     eax, [ebx+IOCTL.output]
-      	cmp     [ebx+IOCTL.out_size], 4
-      	jne     @F
+      	mov     eax, [esi + IOCTL.output]
+      	cmp     [esi + IOCTL.out_size], 4
+      	jne     .ret
       	mov     dword [eax], API_VERSION
       	xor     eax, eax
       	ret
 
-@@:
+.ret:
         or      eax, -1
         ret
 
 endp
 
-proc nvme_query_media stdcall
+proc add_nvme_disk stdcall, pci:dword
+
+		push 	esi
+		mov 	esi, [pci]
+		push 	0 ; null terminator
+		push 	dword [esi + pcidev.nsid]
+		push 	"n"
+		push 	dword [esi + pcidev.num]
+		push 	"nvme"
+		mov 	eax, esp
+		invoke  DiskAdd, disk_functions, eax, [esi + pcidev.nsinfo], 0
+		add 	esp, 20
+		test 	eax, eax
+		jz 		@f
+		DEBUGF  DBG_INFO, "nvme%u: Successfully registered disk\n", [esi + pcidev.num]
+		xor 	eax, eax
+		inc 	eax
+		pop 	esi
+		ret
+
+@@:
+		DEBUGF  DBG_INFO, "nvme%u: Failed to register disk\n", [esi + pcidev.num]
+		xor 	eax, eax
+		pop 	esi
+		ret
+
+endp
+
+proc nvme_query_media stdcall, userdata:dword, info:dword
+
+		push 	esi edi
+		mov 	esi, [userdata]
+		mov 	edi, [info]
+		mov 	dword [edi + DISKMEDIAINFO.flags], 0
+		mov 	eax, dword [esi + NSINFO.lbads]
+		mov 	dword [edi + DISKMEDIAINFO.sectorsize], eax
+		mov 	eax, dword [esi + NSINFO.capacity]
+		mov 	dword [edi + DISKMEDIAINFO.capacity], eax
+		mov 	eax, dword [esi + NSINFO.capacity + 4]
+		mov 	dword [edi + DISKMEDIAINFO.capacity + 4], eax
+		xor 	eax, eax
+		pop 	edi esi
+		ret
+
 endp
 
 proc set_cdw0 stdcall, pci:dword, y:dword, opcode:byte
@@ -733,14 +785,30 @@ proc nvme_init stdcall, pci:dword
 		mov 	edi, eax
 		invoke  GetPhysAddr
 		stdcall nvme_identify, [pci], [esi + pcidev.nsid], eax, CNS_IDNS
+		invoke  KernelAlloc, sizeof.NSINFO
+		test 	eax, eax
+		jz 		.exit_fail
+		mov 	ebx, eax
+		mov 	dword [esi + pcidev.nsinfo], eax
 		mov 	al, byte [edi + IDENTN.nsfeat]
+		mov 	byte [ebx + NSINFO.features], al
 		DEBUGF  DBG_INFO, "nvme%un%u: Namespace Features: 0x%x\n", [esi + pcidev.num], [esi + pcidev.nsid], al
+		mov 	dword [ebx + NSINFO.pci], esi
+		mov 	eax, dword [edi + IDENTN.nsze]
+		mov 	dword [ebx + NSINFO.size], eax
+		mov 	eax, dword [edi + IDENTN.nsze + 4]
+		mov 	dword [ebx + NSINFO.size + 4], eax
+		mov 	eax, dword [edi + IDENTN.ncap]
+		mov 	dword [ebx + NSINFO.capacity], eax 
+		mov 	eax, dword [edi + IDENTN.ncap + 4]
+		mov 	dword [ebx + NSINFO.capacity + 4], eax 
 		DEBUGF  DBG_INFO, "nvme%un%u: Namespace Size: %u + %u logical blocks\n", [esi + pcidev.num], [esi + pcidev.nsid], [edi + IDENTN.nsze], [edi + IDENTN.nsze + 4]
 		DEBUGF  DBG_INFO, "nvme%un%u: Namespace Capacity: %u + %u logical blocks\n", [esi + pcidev.num], [esi + pcidev.nsid], [edi + IDENTN.ncap], [edi + IDENTN.ncap + 4]
 		mov 	eax, dword [edi + IDENTN.lbaf0]
 		shr 	eax, 16 ; Get LBADS
 		and 	eax, 0xff
 		stdcall pow2, eax
+		mov 	dword [ebx + NSINFO.lbads], eax
 		DEBUGF  DBG_INFO, "nvme%un%u: Namespace LBA Data Size: %u\n", [esi + pcidev.num], [esi + pcidev.nsid], eax
 		invoke  KernelFree, edi
 		
diff --git a/drivers/nvme/nvme.inc b/drivers/nvme/nvme.inc
index f9759d0..b37c22a 100644
--- a/drivers/nvme/nvme.inc
+++ b/drivers/nvme/nvme.inc
@@ -303,6 +303,14 @@ struct CQ_ENTRY
 	status 			dw ?
 ends
 
+struct NSINFO
+	capacity 		dq ?
+	size 			dq ?
+	lbads 			dd ?
+	pci 			dd ?
+	features 		db ?
+ends
+
 struct pcidev 
 	bus 			db ?
 	devfn 			db ?
@@ -313,6 +321,7 @@ struct pcidev
 	queue_entries 	dd ?
 	version 		dd ?
 	nsid 			dd ?
+	nsinfo 			dd ?
 	nn 				dd ?
 	dstrd 			db ?
 					rb 3 ; align