:: commit 82466973d309915bddd9e8b46b91e70c4fa408a1

diff --git a/stage23/drivers/disk.h b/stage23/drivers/disk.h
index 1357a315..4df24a71 100644
--- a/stage23/drivers/disk.h
+++ b/stage23/drivers/disk.h
@@ -6,6 +6,14 @@
 #include <stdbool.h>
 #include <lib/part.h>
 
+#if defined (uefi)
+
+#include <efi.h>
+
+bool disk_volume_from_efi_handle(struct volume *ret, EFI_HANDLE *efi_handle);
+
+#endif
+
 size_t disk_create_index(struct volume **ret);
 bool disk_read_sectors(struct volume *volume, void *buf, uint64_t block, size_t count);
 
diff --git a/stage23/drivers/disk.s2.c b/stage23/drivers/disk.s2.c
index ab3cc829..b67b7bf1 100644
--- a/stage23/drivers/disk.s2.c
+++ b/stage23/drivers/disk.s2.c
@@ -190,11 +190,68 @@ size_t disk_create_index(struct volume **ret) {
 
 #if defined (uefi)
 
+bool disk_volume_from_efi_handle(struct volume *ret, EFI_HANDLE *efi_handle) {
+    bool ok = false;
+
+    EFI_GUID disk_io_guid = DISK_IO_PROTOCOL;
+    EFI_GUID block_io_guid = BLOCK_IO_PROTOCOL;
+    EFI_DISK_IO *disk_io = NULL;
+    EFI_BLOCK_IO *block_io = NULL;
+
+    uefi_call_wrapper(gBS->HandleProtocol, 3, efi_handle, &disk_io_guid,
+                      &disk_io);
+    uefi_call_wrapper(gBS->HandleProtocol, 3, efi_handle, &block_io_guid,
+                      &block_io);
+
+    uint64_t signature = BUILD_ID;
+    uint64_t orig;
+
+    uefi_call_wrapper(disk_io->ReadDisk, 5, disk_io, block_io->Media->MediaId, 0,
+                      sizeof(uint64_t), &orig);
+    uefi_call_wrapper(disk_io->WriteDisk, 5, disk_io, block_io->Media->MediaId, 0,
+                      sizeof(uint64_t), &signature);
+
+    for (size_t i = 0; i < volume_index_i; i++) {
+        uint64_t compare;
+
+        EFI_DISK_IO *cur_disk_io = NULL;
+        EFI_BLOCK_IO *cur_block_io = NULL;
+
+        uefi_call_wrapper(gBS->HandleProtocol, 3, volume_index[i].drive,
+                          &disk_io_guid, &cur_disk_io);
+        uefi_call_wrapper(gBS->HandleProtocol, 3, volume_index[i].drive,
+                          &block_io_guid, &cur_block_io);
+
+        uefi_call_wrapper(cur_disk_io->ReadDisk, 5, cur_disk_io,
+                          cur_block_io->Media->MediaId,
+                          0 +
+                          volume_index[i].first_sect * volume_index[i].sector_size,
+                          sizeof(uint64_t), &compare);
+
+        if (compare == signature) {
+            *ret = volume_index[i];
+            ok = true;
+            break;
+        }
+    }
+
+    uefi_call_wrapper(disk_io->WriteDisk, 5, disk_io, block_io->Media->MediaId, 0,
+                      sizeof(uint64_t), &orig);
+
+    return ok;
+}
+
 bool disk_read_sectors(struct volume *volume, void *buf, uint64_t block, size_t count) {
     EFI_STATUS status;
 
-    status = uefi_call_wrapper(volume->drive->ReadBlocks, 5, volume->drive,
-                               volume->drive->Media->MediaId,
+    EFI_GUID block_io_guid = BLOCK_IO_PROTOCOL;
+    EFI_BLOCK_IO *block_io = NULL;
+
+    status = uefi_call_wrapper(gBS->HandleProtocol, 3, volume->drive,
+                               &block_io_guid, &block_io);
+
+    status = uefi_call_wrapper(block_io->ReadBlocks, 5, block_io,
+                               block_io->Media->MediaId,
                                block, count * volume->sector_size, buf);
 
     if (status != 0) {
@@ -226,20 +283,23 @@ size_t disk_create_index(struct volume **ret) {
     for (size_t i = 0; i < handles_size / sizeof(EFI_HANDLE); i++) {
         struct volume block = {0};
 
+        EFI_BLOCK_IO *block_io = NULL;
+
         status = uefi_call_wrapper(gBS->HandleProtocol, 3, handles[i],
-                                   &block_io_guid, &block.drive);
+                                   &block_io_guid, &block_io);
 
-        if (status != 0 || block.drive == NULL || block.drive->Media->LastBlock == 0)
+        if (status != 0 || block_io == NULL || block_io->Media->LastBlock == 0)
             continue;
 
-        if (block.drive->Media->LogicalPartition)
+        if (block_io->Media->LogicalPartition)
             continue;
 
         volume_count++;
 
-        block.sector_size = block.drive->Media->BlockSize;
+        block.drive = handles[i];
+        block.sector_size = block_io->Media->BlockSize;
         block.first_sect = 0;
-        block.sect_count = block.drive->Media->LastBlock + 1;
+        block.sect_count = block_io->Media->LastBlock + 1;
 
         for (int part = 0; ; part++) {
             struct volume p = {0};
@@ -257,7 +317,7 @@ size_t disk_create_index(struct volume **ret) {
     volume_index = ext_mem_alloc(sizeof(struct volume) * volume_count);
 
     for (size_t i = 0; i < handles_size / sizeof(EFI_HANDLE); i++) {
-        EFI_BLOCK_IO *drive;
+        EFI_BLOCK_IO *drive = NULL;
 
         status = uefi_call_wrapper(gBS->HandleProtocol, 3, handles[i],
                                    &block_io_guid, &drive);
@@ -270,7 +330,7 @@ size_t disk_create_index(struct volume **ret) {
 
         struct volume *block = &volume_index[volume_index_i++];
 
-        block->drive = drive;
+        block->drive = handles[i];
         block->partition = -1;
         block->sector_size = drive->Media->BlockSize;
         block->first_sect = 0;
diff --git a/stage23/entry.s3.c b/stage23/entry.s3.c
index 30e97c0e..46c48577 100644
--- a/stage23/entry.s3.c
+++ b/stage23/entry.s3.c
@@ -19,6 +19,7 @@
 #include <menu.h>
 #include <pxe/pxe.h>
 #include <pxe/tftp.h>
+#include <drivers/disk.h>
 
 void stage3_common(void);
 
@@ -33,7 +34,7 @@ EFI_STATUS efi_main(EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE *SystemTable) {
 
     print("Limine " LIMINE_VERSION "\n\n");
 
-    //volume_create_index();
+    volume_create_index();
 
     EFI_GUID loaded_img_prot_guid = EFI_LOADED_IMAGE_PROTOCOL_GUID;
     EFI_LOADED_IMAGE_PROTOCOL *loaded_image = NULL;
@@ -41,22 +42,18 @@ EFI_STATUS efi_main(EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE *SystemTable) {
     uefi_call_wrapper(gBS->HandleProtocol, 3, ImageHandle, &loaded_img_prot_guid,
                       &loaded_image);
 
-    EFI_GUID block_io_guid = BLOCK_IO_PROTOCOL;
-    EFI_BLOCK_IO *drive;
-
-    uefi_call_wrapper(gBS->HandleProtocol, 3, loaded_image->DeviceHandle,
-                      &block_io_guid, &drive);
-
-    struct volume boot_volume;
-    boot_volume.drive = drive;
-    boot_volume.partition = -1;
-    boot_volume.sector_size = drive->Media->BlockSize;
-    boot_volume.first_sect = 0;
-    boot_volume.sect_count = drive->Media->LastBlock + 1;
+    struct volume boot_volume = {0};
+    if (!disk_volume_from_efi_handle(&boot_volume, loaded_image->DeviceHandle)) {
+        panic("Can't determine boot disk");
+    }
 
-    if (!init_config_disk(&boot_volume)) {
-        print("Config file found and loaded.\n");
-    } else {
+    if (!volume_iterate_parts(boot_volume,
+        if (!init_config_disk(&_PART_)) {
+            print("Config file found and loaded.\n");
+            boot_partition = _PARTNUM_;
+            break;
+        }
+    )) {
         panic("Config file not found.");
     }
 
diff --git a/stage23/lib/part.h b/stage23/lib/part.h
index 6bebb57e..2abaa6bd 100644
--- a/stage23/lib/part.h
+++ b/stage23/lib/part.h
@@ -2,6 +2,7 @@
 #define __LIB__PART_H__
 
 #include <stdint.h>
+#include <stddef.h>
 #include <stdbool.h>
 #include <lib/guid.h>
 #if defined (uefi)
@@ -15,7 +16,7 @@
 #if defined (bios)
 typedef int drive_t;
 #elif defined (uefi)
-typedef EFI_BLOCK_IO *drive_t;
+typedef EFI_HANDLE drive_t;
 #endif
 
 struct volume {
@@ -33,6 +34,9 @@ struct volume {
     struct guid part_guid;
 };
 
+extern struct volume *volume_index;
+extern size_t volume_index_i;
+
 void volume_create_index(void);
 
 bool gpt_get_guid(struct guid *guid, struct volume *volume);
diff --git a/stage23/lib/part.s2.c b/stage23/lib/part.s2.c
index 61a4f8f4..9450fc9c 100644
--- a/stage23/lib/part.s2.c
+++ b/stage23/lib/part.s2.c
@@ -291,8 +291,8 @@ int part_get(struct volume *part, struct volume *volume, int partition) {
     return INVALID_TABLE;
 }
 
-static struct volume *volume_index = NULL;
-static size_t volume_index_i = 0;
+struct volume *volume_index = NULL;
+size_t volume_index_i = 0;
 
 void volume_create_index(void) {
     volume_index_i = disk_create_index(&volume_index);