From 7013ca673f97db342a8ecdf8cbf5b5ddb4dd4a38 Mon Sep 17 00:00:00 2001 From: Musa Mahmood Date: Sat, 6 Dec 2025 16:45:20 -0500 Subject: [PATCH] Fix major bugs in `Expandable_Arena` and `NTFS_MFT_read_raw` --- lib/Base/Expandable_Arena.cpp | 5 +- lib/Base/Hash_Table.h | 2 +- lib/OS/OS_Win32.cpp | 48 ++++++++++++++++--- lib/OS/OS_Win32.h | 5 +- lib/OS/OS_Win32_NTFS.cpp | 65 ++++++++++++++------------ src/Ex1.cpp | 86 +++++++++++++++++++++++------------ 6 files changed, 140 insertions(+), 71 deletions(-) diff --git a/lib/Base/Expandable_Arena.cpp b/lib/Base/Expandable_Arena.cpp index 92f1d1a..68bcde5 100644 --- a/lib/Base/Expandable_Arena.cpp +++ b/lib/Base/Expandable_Arena.cpp @@ -112,7 +112,7 @@ void arena_reset_to (ExpandableArena* arena_ex, Arena* last_arena, u8* starting_ void arena_reset (ExpandableArena* arena_ex, bool free_extra_pages) { if (!is_valid(arena_ex)) return; - // Free arenas in `next_arenas` + // Free expansion arenas in `next_arenas` for (s64 i = 0; i < arena_ex->next_arenas.count; i += 1) { release_arena(arena_ex->next_arenas[i], free_extra_pages); } @@ -131,8 +131,9 @@ void arena_reset (ExpandableArena* arena_ex, bool free_extra_pages) { } } +// #TODO: make an option to "FULL-DELETE" the expansion arenas as well. force_inline void arena_delete (ExpandableArena* arena_ex) { - array_free(arena_ex->next_arenas); arena_reset(arena_ex, true); + array_free(arena_ex->next_arenas); arena_delete((Arena*)arena_ex); } diff --git a/lib/Base/Hash_Table.h b/lib/Base/Hash_Table.h index 5a8d8b2..4fa372a 100644 --- a/lib/Base/Hash_Table.h +++ b/lib/Base/Hash_Table.h @@ -338,7 +338,7 @@ template U* table_find_or_add (Table* table, T ke return value; } - U new_value; + U new_value = {}; value = table_add(table, key, new_value); (*newly_added) = true; return value; diff --git a/lib/OS/OS_Win32.cpp b/lib/OS/OS_Win32.cpp index e9635d7..6ce4aa0 100644 --- a/lib/OS/OS_Win32.cpp +++ b/lib/OS/OS_Win32.cpp @@ -42,7 +42,7 @@ struct OS_System_Info { Array monitors; // Back with GPAllocator // #Drives - Table drives; // should we just store ptrs to OS_Drive? I think so.. + Table drives; // should we just store ptrs to OS_Drive? I think so.. // That way we can fetch the OS_Drive* and have the pointer be stable. Especially if it's b // backed by an arena. }; @@ -71,6 +71,40 @@ struct OS_State_Win32 { global OS_State_Win32 global_win32_state; internal b32 global_win32_is_quiet = 0; // No console output (`quiet` flag passed) + +Table* get_drive_table () { + return &global_win32_state.system_info.drives; +} + +ArrayView os_get_available_drives () { + // @Allocates: Recommended to set context allocator to temp(). + auto drive_table = get_drive_table(); + Array drives; + + // #hash_table_iterator : instead of writing this everywhere, just use this function + // to get an Array of drives. + for (s64 i = 0; i < drive_table->allocated; i += 1) { + Table_Entry* entry = &drive_table->entries[i]; // we should take ptr here if we want to modify? + + if (entry->hash > HASH_TABLE_FIRST_VALID_HASH) { + if (entry->value->label.data == nullptr) continue; // Some volumes may not be real and therefore have no label. + + array_add(drives, entry->value); + } + } + + return drives; +} + +OS_Drive* get_drive_from_label (string drive_label) { + // #TODO: Validate the input is in the correct format. We're looking for something like `C:\` + // not just the drive letter! + Table* drive_table = get_drive_table(); + OS_Drive** drive_ptr = table_find_pointer(drive_table, drive_label); + Assert(drive_ptr != nullptr); + return *drive_ptr; +} + internal LONG WINAPI Win32_Exception_Filter (EXCEPTION_POINTERS* exception_ptrs) { if (global_win32_is_quiet) { ExitProcess(1); } @@ -131,8 +165,7 @@ internal void Win32_Entry_Point (int argc, WCHAR **argv) { } // #cpuid { OS_System_Info* info = &global_win32_state.system_info; - // [ ] Extract input args - u32 length; + u32 length = 0; GetLogicalProcessorInformationEx(RelationProcessorCore, nullptr, (PDWORD)&length); u8* cpu_information_buffer = NewArray(length); GetLogicalProcessorInformationEx(RelationProcessorCore, // *sigh* @@ -144,7 +177,6 @@ internal void Win32_Entry_Point (int argc, WCHAR **argv) { u32 physical_cpu_count = 0; u32 max_performance = 0; u32 performance_core_count = 0; - // u32 efficient_core_count; while (offset < length) { SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX* cpu_information @@ -260,7 +292,7 @@ internal void thread_deinit (Thread* thread, bool zero_thread=false) { thread->os_thread.windows_thread = nullptr; // #TODO: before releasing arena, force-delete extra pages? - // array_reset(*thread->context->log_builder); + array_reset(*thread->context->log_builder); free_string_builder(thread->context->log_builder); release_arena(thread->context->error_arena); arena_delete(thread->context->temp); @@ -937,7 +969,7 @@ constexpr u64 Win32_Max_Path_Length = 260; bool Win32_Discover_Drives () { push_allocator(GPAllocator()); // Initialize drive_table if necessary. - Table* drive_table = &global_win32_state.system_info.drives; + Table* drive_table = get_drive_table(); if (!drive_table->allocated) { drive_table->allocator = GPAllocator(); // #TODO: #hash_table need a macro for string keys! @@ -977,7 +1009,9 @@ bool Win32_Discover_Drives () { current_index += (s32)(drive_label.count + 1); bool just_added = false; - Win32_Drive* drive = table_find_or_add(drive_table, drive_label, &just_added); + OS_Drive** drive_ptr = table_find_or_add(drive_table, drive_label, &just_added); + OS_Drive* drive = New(); + (*drive_ptr) = drive; // fill table slot with data. if (!just_added) { // delete old strings before updating // This is silly, but there's a small chance the volume has been renamed so... diff --git a/lib/OS/OS_Win32.h b/lib/OS/OS_Win32.h index 5af2c0b..b15939b 100644 --- a/lib/OS/OS_Win32.h +++ b/lib/OS/OS_Win32.h @@ -167,7 +167,6 @@ File_System Win32_filesystem_from_string (string s) { } struct Dense_FS; // #Temp forward declare! struct Win32_Drive { - // Arena* pool; string label; string volume_name; Win32_Drive_Type type; @@ -179,9 +178,11 @@ struct Win32_Drive { u32 file_system_flags; bool is_present; // Not sure if this should be here... - // f32 enumeration_time; // f64 last_seen_alive_timestamp; Dense_FS* data; + s64 bytes_accessed; + s64 file_count; + f32 time_to_enumerate; }; typedef Win32_Drive OS_Drive; diff --git a/lib/OS/OS_Win32_NTFS.cpp b/lib/OS/OS_Win32_NTFS.cpp index c4f3e78..57fe164 100644 --- a/lib/OS/OS_Win32_NTFS.cpp +++ b/lib/OS/OS_Win32_NTFS.cpp @@ -121,7 +121,7 @@ struct NTFS_MFT_Internal { #endif }; -NTFS_MFT_Internal* new_ntfs_drive () { // call with temp +NTFS_MFT_Internal* new_ntfs_mft_internal () { // call with temp NTFS_MFT_Internal* mft = New(true); mft->mft_file = ArrayView(NTFS_MFT_File_Record_Size); mft->mft_buffer = ArrayView(NTFS_MFT_File_Record_Size * NTFS_MFT_Files_Per_Buffer); // 64 MB @@ -142,14 +142,14 @@ bool NTFS_read_internal (NTFS_MFT_Internal* mft, void* buffer, u64 from, u64 cou } // #TODO: Release resources if we face an early return! -Dense_FS* NTFS_MFT_read_raw (string drive_path) { - Table* drive_table = &global_win32_state.system_info.drives; - bool just_added = false; - Win32_Drive* drive = table_find_or_add(drive_table, drive_path, &just_added); - Assert(just_added == false && drive != nullptr); - if (drive == nullptr) { - return nullptr; - } +// #TODO: Maybe this doesn't need to return a value? Return an Error* instead. +Dense_FS* NTFS_MFT_read_raw (OS_Drive* drive) { + auto start_time = GetUnixTimestamp(); + + Assert(drive != nullptr); + if (drive == nullptr) { return nullptr; } + + string drive_path = drive->label; Assert(context_allocator() != temp()); // pointless as we're releasing temp end-of-scope Allocator primary_allocator = context_allocator(); @@ -168,9 +168,9 @@ Dense_FS* NTFS_MFT_read_raw (string drive_path) { return nullptr; } - push_allocator(primary_allocator); - NTFS_MFT_Internal* mft = new_ntfs_drive(); + NTFS_MFT_Internal* mft = new_ntfs_mft_internal(); mft->handle = file_handle; + push_allocator(primary_allocator); bool success; NTFS_BootSector boot_sector; @@ -184,7 +184,7 @@ Dense_FS* NTFS_MFT_read_raw (string drive_path) { NTFS_FileRecordHeader* file_record_start = (NTFS_FileRecordHeader*)mft->mft_file.data; if (file_record_start->magic != 0x454C4946) { log_error("[NTFS_read_drive_raw] Magic number check failed! This drive is not NTFS or is corrupted!"); - return false; + return nullptr; } NTFS_AttributeHeader* attribute = (NTFS_AttributeHeader*)(mft->mft_file.data + file_record_start->firstAttributeOffset); @@ -203,7 +203,6 @@ Dense_FS* NTFS_MFT_read_raw (string drive_path) { attribute = (NTFS_AttributeHeader*) ((u8*) attribute + attribute->length); } // while (true) - Assert(data_attribute != nullptr); NTFS_RunHeader* dataRun = (NTFS_RunHeader*)((u8*)data_attribute + data_attribute->dataRunsOffset); @@ -213,17 +212,17 @@ Dense_FS* NTFS_MFT_read_raw (string drive_path) { while (((u8*)dataRun - (u8*)data_attribute) < data_attribute->length && dataRun->lengthFieldBytes) { u64 length = 0, offset = 0; - for (s64 i = 0; i < dataRun->lengthFieldBytes; i += 1) { + for (u8 i = 0; i < dataRun->lengthFieldBytes; i += 1) { length |= (u64)(((u8*)dataRun)[1 + i]) << (i * 8); } - for (s64 i = 0; i < dataRun->offsetFieldBytes; i += 1) { + for (u8 i = 0; i < dataRun->offsetFieldBytes; i += 1) { offset |= (u64)(((u8*)dataRun)[1 + dataRun->lengthFieldBytes + i]) << (i * 8); } if (offset & ((u64) 1 << (dataRun->offsetFieldBytes * 8 - 1))) { for (s64 i = dataRun->offsetFieldBytes; i < 8; i += 1) { - offset |= (u64)(0xFF << (i * 8)); + offset |= ((u64)0xFF << (u64)(i * 8)); } } @@ -251,11 +250,11 @@ Dense_FS* NTFS_MFT_read_raw (string drive_path) { // A file record may be blank or unused; just skip it. if (!fileRecord->inUse) continue; - NTFS_AttributeHeader* attribute = (NTFS_AttributeHeader*) ((u8*)fileRecord + fileRecord->firstAttributeOffset); + NTFS_AttributeHeader* attribute = (NTFS_AttributeHeader*)((u8*)fileRecord + fileRecord->firstAttributeOffset); Assert(fileRecord->magic == 0x454C4946); - if (file_record_start->magic != 0x454C4946) { + if (fileRecord->magic != 0x454C4946) { log_error("[NTFS_read_drive_raw] Magic number check failed! This drive is likely corrupted!"); - return false; + return nullptr; } // inner loop NTFS_File file = {}; @@ -268,7 +267,7 @@ Dense_FS* NTFS_MFT_read_raw (string drive_path) { file.record_id = fileRecord->recordNumber; file.name_count = fileNameAttribute->fileNameLength; file.name_data = (u16*)fileNameAttribute->fileName; - file.name_utf8 = wide_to_utf8(file.name_data, file.name_count); + // file.name_utf8 = wide_to_utf8(file.name_data, file.name_count); // @Allocates file.file_modtime = (u64)fileNameAttribute->modificationTime; file.is_directory = fileRecord->isDirectory; // We need to get size from the data attribute @@ -291,27 +290,33 @@ Dense_FS* NTFS_MFT_read_raw (string drive_path) { } // while: files_remaining } // while: outer loop - log("Found %lld files (bytes_accessed: %s)", mft->file_count, format_bytes(mft->bytes_accessed).data); - CloseHandle(file_handle); - return nullptr; + log("Found %lld files (bytes_accessed: %s)", mft->file_count, format_bytes(mft->bytes_accessed).data); + + drive->file_count = mft->file_count; + drive->bytes_accessed = mft->bytes_accessed; + drive->time_to_enumerate = (f32)(GetUnixTimestamp() - start_time); + + return drive->data; } struct NTFS_Enumeration_Task { Arena* pool; // small arena just for results - OS_Drive* drive; - string drive_path; // The drive path we want to enumerate + + ArrayView drives; + // Should be part of OS_Drive! }; s64 ntfs_enumeration_thread_proc (Thread* thread) { auto task = thread_task(NTFS_Enumeration_Task); - log("(Thread index: %lld) Task pointer: %p", thread->index, task); + log("[ntfs_enumeration_thread_proc] (Thread index: %lld) Task pointer: %p", thread->index, task); - // NTFS_MFT_read_raw|#TODO: + for_each(d, task->drives) { + auto result = NTFS_MFT_read_raw(task->drives[d]); + if (result == nullptr) return 1; + } - Sleep(100); // #temp - return 0; -} \ No newline at end of file +} diff --git a/src/Ex1.cpp b/src/Ex1.cpp index 08481a0..887234d 100644 --- a/src/Ex1.cpp +++ b/src/Ex1.cpp @@ -63,9 +63,8 @@ bool Ex1_check_key_combinations() { return false; } -// #define drive_table() &global_win32_state.system_info.drives void Ex1_Control_Panel () { using namespace ImGui; - Table* drive_table = &global_win32_state.system_info.drives; + Table* drive_table = get_drive_table(); Begin("Control Panel"); if (Button("Discover drives") || !table_is_valid(drive_table)) { @@ -73,32 +72,29 @@ void Ex1_Control_Panel () { using namespace ImGui; } Text("drive_table is valid: %d", table_is_valid(drive_table)); - // Basic Hash Table iterator: #hash_table_iterator - s32 drive_count = 0; - for (s64 i = 0; i < drive_table->allocated; i += 1) { - Table_Entry* entry = &drive_table->entries[i]; // we should take ptr here if we want to modify? - if (entry->hash > HASH_TABLE_FIRST_VALID_HASH) { - // #TODO #MOVE THIS + maybe don't check this every frame! - // entry->value.is_present = Win32_Drive_Exists(entry->value.label); - if (entry->value.label.data == nullptr) continue; - - Text(" > [%d] drive letter: %s (is_present: %d)", drive_count + 1, entry->value.label.data, entry->value.is_present); - drive_count += 1; + push_allocator(temp()); + ArrayView drives = os_get_available_drives(); + + for_each(i, drives) { + OS_Drive* drive = drives[i]; + + Text(" > [%d] drive letter: %s (is_present: %d)", drives.count + 1, drive->label.data, drive->is_present); + if (drive->time_to_enumerate != 0) { SameLine(); - push_allocator(temp()); - - if (Button(format_cstring("Read NTFS MFT Raw##%s", entry->value.label.data))) { - push_arena(thread_context()->arena); - // auto_release(thread_context()->arena); - auto dfs = NTFS_MFT_read_raw(entry->value.label); - if (dfs == nullptr) { - log("[NTFS_MFT_read_raw] operation failed"); - } + Text("Enumerated in %.2f seconds", drive->time_to_enumerate); + } + SameLine(); + if (Button(format_cstring("Read NTFS MFT Raw##%s", drive->label.data))) { + push_arena(thread_context()->arena); + // auto_release(thread_context()->arena); + auto dfs = NTFS_MFT_read_raw(drive); + if (dfs == nullptr) { + log("[NTFS_MFT_read_raw] operation failed"); } } } - if (drive_count > 0) { + if (drives.count > 0) { if (!ex1_ntfs.initialized) { Timed_Block_Print("Thread initialization (ntfs)"); push_allocator(GPAllocator()); ex1_ntfs.initialized = true; @@ -113,22 +109,54 @@ void Ex1_Control_Panel () { using namespace ImGui; } } - if (drive_count > 0 && ex1_ntfs.initialized && Button("Enumerate all NTFS drives")) { + if (drives.count > 0 && ex1_ntfs.initialized && Button("Enumerate all NTFS drives")) { // if drive count exceeds the number of threads, we need to group them so each thread // can enumerate multiple drives. // We need to distribute the drives across our available threads: - // see: #hash_table_iterator (what part of this can we macro out?): + Array> drive_split; + s32 thread_count = (s32)ex1_ntfs.threads.count; + drive_split.allocator = GPAllocator(); - if (drive_count > ex1_ntfs.threads.count) { - debug_break(); // #TODO: Fix this! + if (drives.count > thread_count) { + s64 drives_per_thread = (drives.count / thread_count); + s64 remainder = drives.count % thread_count; + s64 current_drive = 0; + push_allocator(GPAllocator()); + array_resize(drive_split, thread_count); + for_each(d, drive_split) { + if (d == drive_split.count) { + drive_split[d] = ArrayView(remainder); + } else { + drive_split[d] = ArrayView(drives_per_thread); + } + + for (s64 i = 0; i < drive_split[d].count; i += 1) { + drive_split[d][i] = drives[current_drive]; + current_drive += 1; + } + } + + debug_break(); // #TODO: Check that the work has been distributed correctly. + } else { // more threads than drives, or same amount + array_resize(drive_split, drives.count); + + for_each(d, drives) { + auto drive = drives[d]; + drive_split[d] = ArrayView(1); // Arrays of size one are sad :pensive: + + drive_split[d][0] = drive; + } } - for (s32 t = 0; t < drive_count; t += 1) { + + s64 active_thread_count = drive_split.count; + + for (s64 t = 0; t < active_thread_count; t += 1) { Thread* thread = &ex1_ntfs.threads[t]; Arena* thread_arena = next_arena(Arena_Reserve::Size_64K); push_arena(thread_arena); auto thread_data = New(); thread_data->pool = thread_arena; - // #TODO: fill out + thread_data->drives = drive_split[t]; thread_start(thread, thread_data); array_add(ex1_ntfs.threads_in_flight, thread); }