From 1cd88b1cd4976a16774696a9a11e77eba5728838 Mon Sep 17 00:00:00 2001 From: Pablo Curiel Date: Sun, 10 Jul 2022 06:41:18 +0200 Subject: [PATCH] bktr: handle compression in patches (part 3). Updated bktrIsBlockWithinIndirectStorageRange(). Got confirmation that everything is now working as it should. Big thanks to BigKahuna. --- code_templates/usb_romfs_dumper.c | 9 ++- include/core/bktr.h | 1 + source/core/bktr.c | 93 ++++++++++++++++++++++++++++++- source/core/romfs.c | 11 +++- 4 files changed, 109 insertions(+), 5 deletions(-) diff --git a/code_templates/usb_romfs_dumper.c b/code_templates/usb_romfs_dumper.c index 50b263f..cb9ecfc 100644 --- a/code_templates/usb_romfs_dumper.c +++ b/code_templates/usb_romfs_dumper.c @@ -125,7 +125,14 @@ static void read_thread_func(void *arg) } bool updated = false; - if (!romfsIsFileEntryUpdated(shared_data->romfs_ctx, file_entry, &updated) || !updated) + shared_data->read_error = !romfsIsFileEntryUpdated(shared_data->romfs_ctx, file_entry, &updated); + if (shared_data->read_error) + { + condvarWakeAll(&g_writeCondvar); + break; + } + + if (!updated) { shared_data->read_error = !romfsMoveToNextFileEntry(shared_data->romfs_ctx); if (shared_data->read_error) diff --git a/include/core/bktr.h b/include/core/bktr.h index 9dd9d3f..5c94479 100644 --- a/include/core/bktr.h +++ b/include/core/bktr.h @@ -203,6 +203,7 @@ bool bktrSetBucketTreeSubStorage(BucketTreeContext *parent_ctx, BucketTreeContex bool bktrReadStorage(BucketTreeContext *ctx, void *out, u64 read_size, u64 offset); /// Checks if the provided block extents are within the provided BucketTreeContext's Indirect Storage. +/// The storage type from the provided BucketTreeContext may only be BucketTreeStorageType_Indirect or BucketTreeStorageType_Compressed (with an underlying Indirect substorage). bool bktrIsBlockWithinIndirectStorageRange(BucketTreeContext *ctx, u64 offset, u64 size, bool *out); /// Helper inline functions. diff --git a/source/core/bktr.c b/source/core/bktr.c index 5bd8a89..99f7846 100644 --- a/source/core/bktr.c +++ b/source/core/bktr.c @@ -340,13 +340,13 @@ end: bool bktrIsBlockWithinIndirectStorageRange(BucketTreeContext *ctx, u64 offset, u64 size, bool *out) { - if (!bktrIsBlockWithinStorageRange(ctx, size, offset) || ctx->storage_type != BucketTreeStorageType_Indirect || !out) + if (!bktrIsBlockWithinStorageRange(ctx, size, offset) || (ctx->storage_type != BucketTreeStorageType_Indirect && ctx->storage_type != BucketTreeStorageType_Compressed) || \ + (ctx->storage_type == BucketTreeStorageType_Compressed && ctx->substorages[0].type != BucketTreeSubStorageType_Indirect) || !out) { LOG_MSG("Invalid parameters!"); return false; } - BucketTreeIndirectStorageEntry *start_entry = NULL, *end_entry = NULL; BucketTreeVisitor visitor = {0}; bool updated = false, success = false; @@ -357,6 +357,95 @@ bool bktrIsBlockWithinIndirectStorageRange(BucketTreeContext *ctx, u64 offset, u goto end; } + /* Check if we're dealing with a Compressed storage. */ + if (ctx->storage_type == BucketTreeStorageType_Compressed) + { + BucketTreeContext *indirect_storage = (BucketTreeContext*)ctx->substorages[0].bktr_ctx; + const u64 compressed_storage_base_offset = ctx->nca_fs_ctx->hash_region.size; + BucketTreeCompressedStorageEntry *start_entry = NULL, *end_entry = NULL; + + /* Validate start entry node. */ + start_entry = end_entry = (BucketTreeCompressedStorageEntry*)visitor.entry; + if (!bktrIsOffsetWithinStorageRange(ctx, (u64)start_entry->virtual_offset) || (u64)start_entry->virtual_offset > offset) + { + LOG_MSG("Invalid Compressed Storage entry! (0x%lX) (#1).", start_entry->virtual_offset); + goto end; + } + + /* Loop until we reach the upper bound of the requested block or find a match. */ + do { + u64 cur_entry_offset = 0; + + /* Check if we can move any further. */ + if (bktrVisitorCanMoveNext(&visitor)) + { + BucketTreeCompressedStorageEntry *tmp = end_entry; + + /* Retrieve next entry node. */ + if (!bktrVisitorMoveNext(&visitor)) + { + LOG_MSG("Failed to retrieve next Compressed Storage entry!"); + goto end; + } + + /* Validate next entry node. */ + end_entry = (BucketTreeCompressedStorageEntry*)visitor.entry; + if (!bktrIsOffsetWithinStorageRange(ctx, (u64)end_entry->virtual_offset) || (u64)end_entry->virtual_offset <= (u64)tmp->virtual_offset) + { + LOG_MSG("Invalid Indirect Storage entry! (0x%lX) (#2).", (u64)end_entry->virtual_offset); + goto end; + } + + /* Update current entry offset. */ + cur_entry_offset = (u64)end_entry->virtual_offset; + + /* Update start entry node. */ + start_entry = tmp; + } else { + /* Update current entry offset. */ + cur_entry_offset = ctx->end_offset; + + /* Update entry nodes. */ + start_entry = end_entry; + end_entry = NULL; + } + + /* Calculate indirect block extents. */ + u64 indirect_block_offset = compressed_storage_base_offset; + u64 indirect_block_size = (cur_entry_offset - (u64)start_entry->virtual_offset); + + if ((u64)start_entry->virtual_offset <= offset) + { + indirect_block_offset += ((offset - (u64)start_entry->virtual_offset) + (u64)start_entry->physical_offset); + indirect_block_size -= (offset - (u64)start_entry->virtual_offset); + } else { + indirect_block_offset += (u64)start_entry->physical_offset; + } + + if ((offset + size) <= cur_entry_offset) + { + indirect_block_size -= (cur_entry_offset - (offset + size)); + end_entry = NULL; /* Don't proceed any further, we have found our upper bound. */ + } + + /* Check if the current Compressed Storage entry node points to one or more Indirect Storage entry nodes with Patch storage index. */ + if (!bktrIsBlockWithinIndirectStorageRange(indirect_storage, indirect_block_offset, indirect_block_size, &updated)) + { + LOG_MSG("Failed to determine if 0x%lX-byte long Compressed storage block at offset 0x%lX is within Indirect Storage!", indirect_block_offset, indirect_block_size); + goto end; + } + } while(!updated && end_entry && (u64)end_entry->virtual_offset < (offset + size)); + + /* Update output values. */ + *out = updated; + success = true; + + goto end; + } + + /* Check the Indirect Storage. */ + BucketTreeIndirectStorageEntry *start_entry = NULL, *end_entry = NULL; + /* Validate start entry node. */ start_entry = end_entry = (BucketTreeIndirectStorageEntry*)visitor.entry; if (!bktrIsOffsetWithinStorageRange(ctx, start_entry->virtual_offset) || start_entry->virtual_offset > offset) diff --git a/source/core/romfs.c b/source/core/romfs.c index 1935c13..abfa271 100644 --- a/source/core/romfs.c +++ b/source/core/romfs.c @@ -235,6 +235,8 @@ bool romfsGetTotalDataSize(RomFileSystemContext *ctx, bool only_updated, u64 *ou /* Loop through all file entries. */ while(romfsCanMoveToNextFileEntry(ctx)) { + bool updated = false; + /* Get current file entry. */ if (!(file_entry = romfsGetCurrentFileEntry(ctx))) { @@ -243,8 +245,13 @@ bool romfsGetTotalDataSize(RomFileSystemContext *ctx, bool only_updated, u64 *ou } /* Update total data size, taking into account the only_updated flag. */ - bool updated = false; - if (!only_updated || (only_updated && romfsIsFileEntryUpdated(ctx, file_entry, &updated) && updated)) total_size += file_entry->size; + if (only_updated && !romfsIsFileEntryUpdated(ctx, file_entry, &updated)) + { + LOG_MSG("Failed to determine if file entry is updated or not! (0x%lX, 0x%lX).", ctx->cur_file_offset, ctx->file_table_size); + goto end; + } + + if (!only_updated || (only_updated && updated)) total_size += file_entry->size; /* Move to the next file entry. */ if (!romfsMoveToNextFileEntry(ctx))