From 87afd80486069bc23691ec40beaee692948a285c Mon Sep 17 00:00:00 2001
From: Abdur-Rahman Mansoor <armansoor@missingno.dev>
Date: Mon, 8 Jul 2024 14:49:43 -0400
Subject: [PATCH] fix writing to PRP entries

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

diff --git a/drivers/nvme/nvme.asm b/drivers/nvme/nvme.asm
index 12653a5..2f8f61b 100644
--- a/drivers/nvme/nvme.asm
+++ b/drivers/nvme/nvme.asm
@@ -611,6 +611,59 @@ proc alloc_dptr stdcall, ns:dword, prps_ptr:dword, start_sector:qword, numsector
 	
 endp
 
+proc write_prp_buf stdcall, ns:dword, prp:dword, dst:dword, numsectors_ptr:dword
+
+	push 	esi edi ebx
+	mov 	esi, [prp]
+	test 	esi, esi
+	jz 	.err
+	invoke  KernelAlloc, PAGE_SIZE
+	test 	eax, eax
+	jz 	.err
+	mov 	esi, eax
+	mov 	edi, [prp]
+	invoke  MapPage, edi, esi, PG_SW+PG_NOCACHE
+	mov 	esi, edi
+	mov 	edi, [dst]
+	; determine if the number of sectors crosses a page boundary
+	mov 	eax, [ns]
+	mov 	ebx, [numsectors_ptr]
+	mov 	ecx, dword [eax + NSINFO.lbads]
+	mov 	eax, dword [eax + NSINFO.pg_sectors]
+	cmp 	eax, dword [ebx]
+	jae 	@f
+	mov 	eax, dword [ebx]
+	imul 	ecx, eax
+	shr 	ecx, 2
+	sub 	dword [ebx], eax
+
+@@:
+	mov 	ecx, PAGE_SIZE / 4
+	sub 	dword [ebx], eax
+
+.copy:
+	rep 	movsd
+	sub 	esi, ecx
+	invoke  KernelFree, esi
+	invoke  FreePage, [prp]
+	pop 	ebx edi esi
+	xor 	eax, eax
+	inc 	eax
+	ret
+
+.err:
+	mov 	esi, [prp]
+	test 	esi, esi
+	jz 	@f
+	invoke  FreePage, esi
+
+@@:
+	pop 	ebx edi esi
+	xor 	eax, eax
+	ret
+
+endp
+
 ; note that prp_list is required to be a physical address
 proc write_prp_list_buf stdcall, prp_list:dword, dst:dword, nprps:dword
 
@@ -627,7 +680,7 @@ proc write_prp_list_buf stdcall, prp_list:dword, dst:dword, nprps:dword
 	mov 	ecx, [nprps]
 	cmp 	ecx, PAGE_SIZE
 	jae 	@f
-	shl 	ecx, 12 + NVM_MPS
+	shl 	ecx, (12 + NVM_MPS) - 2
 	rep 	stosd
 	pop 	edi
 	invoke  KernelFree, edi
@@ -635,17 +688,19 @@ proc write_prp_list_buf stdcall, prp_list:dword, dst:dword, nprps:dword
 	xor 	eax, eax
 	inc 	eax
 	ret
-	
+
 ; more than or equal to 4096 PRP entries
 @@:
-	mov 	ecx, PAGE_SIZE * PAGE_SIZE - 12
+	; WIP: Work on this later
+	if 0
+	mov 	ecx, PAGE_SIZE * PAGE_SIZE / 4 - 3
 	rep 	stosd
-	sub 	ecx, PAGE_SIZE
 	stdcall write_prp_list_buf, [prp_list], edi, ecx
 	mov 	esi, eax
 	pop 	edi
 	invoke  KernelFree, edi
 	mov 	eax, esi
+	end if
 
 .err:
 	pop 	edi esi
@@ -662,11 +717,14 @@ nvme_write:
 	mov 	edx, NVM_CMD_WRITE
 
 proc nvme_readwrite stdcall, ns:dword, dst:dword, start_sector:qword, numsectors_ptr:dword
+
 	push 	ebx esi edi
-	sub 	esp, 12
+	sub 	esp, 16
 	mov 	ebx, [numsectors_ptr]
 	mov 	ebx, [ebx]
-	mov 	dword [esp + 8], edx
+	mov 	dword [esp + 8], edx ; command type (read or write)
+	mov 	dword [esp + 12], ebx ; save original numsectors value
+
 	; Note that [esp] will contain the value of PRP1 and [esp + 4] will
 	; contain the value of PRP2 (after this call, if it completes successfully)
 	stdcall alloc_dptr, [ns], esp, dword [start_sector], dword [start_sector + 4], ebx
@@ -685,38 +743,44 @@ proc nvme_readwrite stdcall, ns:dword, dst:dword, start_sector:qword, numsectors
 	                        dword [esp + 8]
 
 	; assume command completes successfully for now
+	; Write PRP1
+	stdcall write_prp_buf, [ns], dword [esp], [dst], [numsectors_ptr]
+	test 	eax, eax
+	jz 	.end
 
-	; only 1-2 pages are used, which makes our life easier
-	mov 	ecx, dword [esi + NSINFO.lbads]
-	imul 	ecx, ebx
-	mov 	edi, [dst]
-	mov 	esi, dword [esp]
-	rep 	movsd
-	invoke  KernelFree, dword [esp]
-	mov 	esi, dword [esp + 4]
+	; check if PRP2 should be a PRP list or a regular entry
 	mov 	ecx, dword [esi + NSINFO.pg_sectors]
+	mov 	eax, PAGE_SIZE
+	xor 	edx, edx
+	div 	ecx
+	mov 	ecx, eax
+	mov 	esi, dword [esp + 4]
 	cmp 	ebx, ecx
 	ja 	.is_prp_list
-	xor 	eax, eax
-	test 	esi, esi
-	jz 	.end
-	rep 	movsd
-	invoke  KernelFree, dword [esp + 4]
-	xor 	eax, eax
+
+	; PRP2 is a PRP entry
+	mov 	edi, [dst]
+	mov 	eax, dword [esp + 12]
+	mov 	edx, [ns]
+	imul 	eax, dword [edx + NSINFO.lbads]
+	mov 	ebx, [numsectors_ptr]
+	sub 	eax, dword [ebx]
+	add 	edi, eax
+	stdcall write_prp_buf, [ns], dword [esp + 4], edi, [numsectors_ptr]
 	jmp 	.end
 
 .is_prp_list:
-	xor 	edx, edx
+	mov 	esi, [ns]
 	mov 	eax, ebx
+	xor 	edx, edx
 	mov 	ecx, dword [esi + NSINFO.lbads]
 	div 	ecx
 	mov 	ebx, eax
 	stdcall write_prp_list_buf, dword [esp + 4], [dst], ebx
-	test 	eax, eax
 	stdcall free_prp_list, dword [esp + 4], ebx, FALSE
 
 .end:
-	add 	esp, 12
+	add 	esp, 16
 	pop 	edi esi ebx
 	ret