NVMe: shutdown all controllers and fix IRQ handler

drivers/nvme/nvme.asm

@@ -76,7 +76,7 @@ local   AnythingLoadedSuccessfully db 0
 	test 	eax, eax
 	setne	[AnythingLoadedSuccessfully]
 	inc 	ebx
-	cmp 	ebx, dword [pcidevs_len]
+	cmp 	ebx, dword [num_pcidevs]
 	jne 	.loop
 	cmp 	[AnythingLoadedSuccessfully], 0
 	jz 	.err
@@ -536,7 +536,7 @@ proc nvme_readwrite stdcall, ns:dword, buf:dword, start_sector:qword, numsectors
 endp
 
 ; Detects NVMe devices on the PCI bus and stores them into 
-; [p_nvme_devices] and sets [pcidevs_len] to the appropriate
+; [p_nvme_devices] and sets [num_pcidevs] to the appropriate
 ; size based off how many NVMe devices there are.
 proc detect_nvme
 
@@ -571,13 +571,14 @@ proc detect_nvme
 	jnz 	.err
 
 @@:
-	cmp 	dword [pcidevs_len], TOTAL_PCIDEVS
+	cmp 	dword [num_pcidevs], TOTAL_PCIDEVS
 	jne 	@f
 	DEBUGF  DBG_INFO, "Can't add any more NVMe devices...\n"
 	jmp 	.exit_success
 
 @@:
-	inc 	dword [pcidevs_len]
+	inc 	dword [num_pcidevs]
+	add 	dword [num_pcidevs_sz], sizeof.pcidev
 	cmp     dword [p_nvme_devices], 0
 	jnz	@f ; was the pointer already allocated?
 	invoke  KernelAlloc, sizeof.pcidev * TOTAL_PCIDEVS
@@ -588,7 +589,7 @@ proc detect_nvme
 	DEBUGF  DBG_INFO, "nvme: Allocated memory for PCI devices at: 0x%x\n", eax
 
 @@:
-	mov 	ecx, dword [pcidevs_len]
+	mov 	ecx, dword [num_pcidevs]
 	dec 	ecx
 	mov 	edi, dword [p_nvme_devices]
 	mov 	edx, ecx
@@ -1281,19 +1282,24 @@ endp
 proc irq_handler
 
 	push 	ebx esi edi
-	mov 	edi, dword [p_nvme_devices]
-	mov 	esi, edi
-	sub 	esi, sizeof.pcidev
-	mov 	ebx, dword [pcidevs_len]
-	xor 	ecx, ecx
+	mov 	esi, dword [p_nvme_devices]
+	mov 	ebx, dword [num_pcidevs_sz]
+	add 	ebx, esi
 
 .check_who_raised_irq:
+	stdcall device_generated_interrupt, esi
+	test 	eax, eax
+	jnz 	@f
 	add 	esi, sizeof.pcidev
-	inc 	ecx
-	cmp 	ecx, ebx
-	; TODO: Apply solution given by @punk_joker of checking which device 
-	; generated an interrupt.
-	ja 	.not_our_irq
+	cmp 	esi, ebx
+	jbe 	.check_who_raised_irq
+
+	; Interrupt not handled by driver, return 0
+	pop 	edi esi ebx
+	xor 	eax, eax
+	ret
+
+@@:
 	mov 	edi, dword [esi + pcidev.io_addr]
 	mov 	dword [edi + NVME_MMIO.INTMS], 0x3
 	stdcall consume_cq_entries, esi, ADMIN_QUEUE
@@ -1307,12 +1313,28 @@ proc irq_handler
 	mov 	eax, 1
 	ret
 
-.not_our_irq:
-	; Interrupt not handled by driver, return 0
-	pop 	edi esi ebx
+endp
+
+proc device_generated_interrupt stdcall, pci:dword
+
+	mov 	edx, [pci]
+	mov 	edx, dword [edx + pcidev.queue_entries]
+	xor 	ecx, ecx
+
+@@:
+	mov 	ax, word [edx + ecx + NVM_QUEUE_ENTRY.head]
+	cmp 	ax, word [edx + ecx + NVM_QUEUE_ENTRY.tail]
+	jne 	@f
+	add 	ecx, sizeof.NVM_QUEUE_ENTRY
+	cmp 	ecx, LAST_QUEUE_ID * sizeof.NVM_QUEUE_ENTRY
+	jbe 	@b
 	xor 	eax, eax
 	ret
 
+@@:
+	mov 	eax, 1
+	ret
+
 endp
 
 ; Deletes the allocated I/O queues for all of the NVMe devices,
@@ -1362,13 +1384,7 @@ proc nvme_cleanup
 	cmp 	ebx, LAST_QUEUE_ID
 	jbe 	.get_queue
 	pop 	ebx
-	inc 	ebx
-	cmp 	ebx, dword [pcidevs_len]
-	jne 	.get_pcidev
 
-	; NOTE: This code has a bug! It only shuts down the last
-	; controller, not all of them. Move this inside the loop
-	; and check if the device is actually valid.
 	; Shutdown the controller
 	mov 	edi, dword [esi + pcidev.io_addr]
 	mov 	eax, dword [edi + NVME_MMIO.CC]
@@ -1382,6 +1398,10 @@ proc nvme_cleanup
 	test 	byte [edi + NVME_MMIO.CSTS], CSTS_SHST_SHUTDOWN_COMPLETE
 	jnz 	@b
 
+	inc 	ebx
+	cmp 	ebx, dword [num_pcidevs]
+	jne 	.get_pcidev
+
 .ret:
 	pop 	edi esi ebx
 	ret
@@ -1390,8 +1410,9 @@ endp
 
 ;all initialized data place here
 align 4
-	p_nvme_devices dd 0
-	pcidevs_len  dd 0
+	p_nvme_devices dd 0 ; Pointer to array of NVMe devices
+	num_pcidevs    dd 0 ; Number of NVMe devices
+	num_pcidevs_sz dd 0 ; Size in bytes
 	my_service   db "nvme",0  ;max 16 chars include zero
 	disk_functions:
 		dd 	 disk_functions.end - disk_functions