BKTR rewrite part 4: add set substorage methods.

bktrReadIndirectStorage() is now finished. bktrReadSubStorage() needs to be completed.

include/core/bktr.h

@@ -25,8 +25,6 @@
 #ifndef __BKTR_H__
 #define __BKTR_H__
 
-#include "romfs.h"
-
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -48,6 +46,8 @@ extern "C" {
 
 #define BKTR_COMPRESSION_INVALID_PHYS_SIZE  UINT32_MAX
 
+#define BKTR_MAX_SUBSTORAGE_COUNT           2
+
 /// Used as the header for both BucketTreeOffsetNode and BucketTreeEntryNode.
 typedef struct {
     u32 index;  ///< BucketTreeOffsetNode / BucketTreeEntryNode index.
@@ -138,33 +138,58 @@ typedef struct {
 NXDT_ASSERT(BucketTreeTable, BKTR_NODE_SIZE);
 
 typedef enum {
-    BucketTreeStorageType_Indirect   = 0,
+    BucketTreeStorageType_Indirect   = 0,   ///< Uses two substorages: index 0 (points to the base NCA) and index 1 (AesCtrEx storage).
-    BucketTreeStorageType_AesCtrEx   = 1,
+                                            ///< All reads within storage index 0 use the calculated physical offsets for data decryption.
-    BucketTreeStorageType_Compressed = 2,
+    BucketTreeStorageType_AesCtrEx   = 1,   ///< Used as storage index 1 for BucketTreeStorageType_Indirect.
-    BucketTreeStorageType_Sparse     = 3,   ///< BucketTreeStorageType_Indirect, but uses virtual offsets for data decryption.
+    BucketTreeStorageType_Compressed = 2,   ///< Uses LZ4-compressed sections.
+    BucketTreeStorageType_Sparse     = 3,   ///< BucketTreeStorageType_Indirect with a twist. Storage index 0 points to the same NCA, and uses virtual offsets for data decryption.
+                                            ///< Zero-filled output is used for any reads within storage index 1.
     BucketTreeStorageType_Count      = 4    ///< Total values supported by this enum.
 } BucketTreeStorageType;
 
+typedef enum {
+    BucketTreeSubStorageType_Regular    = 0,    ///< Body storage with None, XTS or CTR crypto. Most common substorage type, used in all title types.
+                                                ///< May be used as substorage for all BucketTreeStorage types.
+    BucketTreeSubStorageType_Indirect   = 1,    ///< Indirect storage from patch NCAs. May be used as substorage for BucketTreeStorageType_Compressed only.
+    BucketTreeSubStorageType_AesCtrEx   = 2,    ///< AesCtrEx storage from patch NCAs. May be used as substorage for BucketTreeStorageType_Indirect only.
+    BucketTreeSubStorageType_Sparse     = 3,    ///< Sparse storage with CTR crypto, using virtual offsets as lower CTR IVs. Used in base applications only.
+                                                ///< May be used as substorage for BucketTreeStorageType_Compressed only.
+    BucketTreeSubStorageType_Compressed = 4,    ///< Compressed storage. If available, this is always the outmost storage type for any NCA. May be used by all title types.
+                                                ///< May be used as substorage for BucketTreeStorageType_Indirect only.
+    BucketTreeSubStorageType_Count      = 5     ///< Total values supported by this enum.
+} BucketTreeSubStorageType;
+
 typedef struct {
+    u8 index;                           ///< Substorage index.
     NcaFsSectionContext *nca_fs_ctx;    ///< NCA FS section context. Used to perform operations on the target NCA.
-    u8 storage_type;                    ///< BucketTreeStorageType.
+    u8 type;                 ///< BucketTreeSubStorageType.
-    BucketTreeTable *storage_table;     ///< Pointer to the dynamically allocated Bucket Tree Table for this storage.
+    void *bktr_ctx;                     ///< BucketTreeContext related to this storage. Only used if type > BucketTreeSubStorageType_Regular.
-    u64 node_size;                      ///< Node size for this type of Bucket Tree storage.
+} BucketTreeSubStorage;
-    u64 entry_size;                     ///< Size of each individual entry within BucketTreeEntryNode.
-    u32 offset_count;                   ///< Number of offsets available within each BucketTreeOffsetNode for this storage.
-    u32 entry_set_count;                ///< Number of BucketTreeEntryNode elements available in this storage.
-    u64 node_storage_size;              ///< Offset node segment size within 'storage_table'.
-    u64 entry_storage_size;             ///< Entry node segment size within 'storage_table'.
-    u64 start_offset;                   ///< Virtual storage start offset.
-    u64 end_offset;                     ///< Virtual storage end offset.
-    
-    // TODO: add sub storage
 
+typedef struct {
+    NcaFsSectionContext *nca_fs_ctx;                                ///< NCA FS section context. Used to perform operations on the target NCA.
+    u8 storage_type;                                                ///< BucketTreeStorageType.
+    BucketTreeTable *storage_table;                                 ///< Pointer to the dynamically allocated Bucket Tree Table for this storage.
+    u64 node_size;                                                  ///< Node size for this type of Bucket Tree storage.
+    u64 entry_size;                                                 ///< Size of each individual entry within BucketTreeEntryNode.
+    u32 offset_count;                                               ///< Number of offsets available within each BucketTreeOffsetNode for this storage.
+    u32 entry_set_count;                                            ///< Number of BucketTreeEntryNode elements available in this storage.
+    u64 node_storage_size;                                          ///< Offset node segment size within 'storage_table'.
+    u64 entry_storage_size;                                         ///< Entry node segment size within 'storage_table'.
+    u64 start_offset;                                               ///< Virtual storage start offset.
+    u64 end_offset;                                                 ///< Virtual storage end offset.
+    BucketTreeSubStorage substorages[BKTR_MAX_SUBSTORAGE_COUNT];    ///< Substorages required for this BucketTree storage. May be set after initializing this context.
 } BucketTreeContext;
 
-/// Initializes a Bucket Tree context using the provided NCA FS section context and a storage 
+/// Initializes a Bucket Tree context using the provided NCA FS section context and a storage type.
 bool bktrInitializeContext(BucketTreeContext *out, NcaFsSectionContext *nca_fs_ctx, u8 storage_type);
 
+/// Sets a BucketTreeSubStorageType_Regular substorage at index 0.
+bool bktrSetRegularSubStorage(BucketTreeContext *ctx, NcaFsSectionContext *nca_fs_ctx);
+
+/// Sets a substorage with type >= BucketTreeSubStorageType_Indirect and <= BucketTreeSubStorageType_Compressed at the provided index.
+bool bktrSetBucketTreeSubStorage(BucketTreeContext *parent_ctx, BucketTreeContext *child_ctx, u8 substorage_index);
+
 /// Reads data from a Bucket Tree storage using a previously initialized BucketTreeContext.
 bool bktrReadStorage(BucketTreeContext *ctx, void *out, u64 read_size, u64 offset);
 
@@ -193,104 +218,11 @@ NX_INLINE bool bktrIsBlockWithinStorageRange(BucketTreeContext *ctx, u64 size, u
     return (bktrIsValidContext(ctx) && size > 0 && ctx->start_offset <= offset && size <= (ctx->end_offset - offset));
 }
 
-
+NX_INLINE bool bktrIsValidSubstorage(BucketTreeSubStorage *substorage)
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-typedef struct {
-    RomFileSystemContext base_romfs_ctx;            ///< Base NCA RomFS context.
-    RomFileSystemContext patch_romfs_ctx;           ///< Update NCA RomFS context. Must be used with RomFS directory/file entry functions, because it holds the updated directory/file tables.
-    u64 offset;                                     ///< Patched RomFS image offset (relative to the start of the update NCA FS section).
-    u64 size;                                       ///< Patched RomFS image size.
-    u64 body_offset;                                ///< Patched RomFS image file data body offset (relative to the start of the RomFS).
-    BktrIndirectStorageBlock *indirect_block;       ///< BKTR Indirect Storage Block.
-    BktrAesCtrExStorageBlock *aes_ctr_ex_block;     ///< BKTR AesCtrEx Storage Block.
-    bool missing_base_romfs;                        ///< If true, only Patch RomFS data is used. Needed for games with base Program NCAs without a RomFS section (e.g. Fortnite, World of Tanks Blitz, etc.).
-    BktrIndirectStorageBlock *base_indirect_block;  ///< Base NCA Indirect Storage Block (sparse layer), if available.
-} BktrContext;
-
-/// Initializes a BKTR context.
-bool bktrInitializeContext(BktrContext *out, NcaFsSectionContext *base_nca_fs_ctx, NcaFsSectionContext *update_nca_fs_ctx);
-
-/// Reads raw filesystem data using a BKTR context.
-/// Input offset must be relative to the start of the patched RomFS image.
-bool bktrReadFileSystemData(BktrContext *ctx, void *out, u64 read_size, u64 offset);
-
-/// Reads data from a previously retrieved RomFileSystemFileEntry using a BKTR context.
-/// Input offset must be relative to the start of the RomFS file entry data.
-bool bktrReadFileEntryData(BktrContext *ctx, RomFileSystemFileEntry *file_entry, void *out, u64 read_size, u64 offset);
-
-/// Checks if a RomFS file entry is updated by the Patch RomFS.
-bool bktrIsFileEntryUpdated(BktrContext *ctx, RomFileSystemFileEntry *file_entry, bool *out);
-
-/// Miscellaneous functions.
-
-NX_INLINE void bktrFreeContext(BktrContext *ctx)
 {
-    if (!ctx) return;
+    return (substorage && substorage->index < BKTR_MAX_SUBSTORAGE_COUNT && substorage->nca_fs_ctx && substorage->type < BucketTreeSubStorageType_Count && \
-    romfsFreeContext(&(ctx->base_romfs_ctx));
+            ((substorage->type == BucketTreeSubStorageType_Regular && substorage->index == 0 && !substorage->bktr_ctx) || \
-    romfsFreeContext(&(ctx->patch_romfs_ctx));
+            (substorage->type > BucketTreeSubStorageType_Regular && substorage->bktr_ctx)));
-    if (ctx->indirect_block) free(ctx->indirect_block);
-    if (ctx->aes_ctr_ex_block) free(ctx->aes_ctr_ex_block);
-    if (ctx->base_indirect_block) free(ctx->base_indirect_block);
-    memset(ctx, 0, sizeof(BktrContext));
-}
-
-NX_INLINE RomFileSystemDirectoryEntry *bktrGetDirectoryEntryByOffset(BktrContext *ctx, u32 dir_entry_offset)
-{
-    return (ctx ? romfsGetDirectoryEntryByOffset(&(ctx->patch_romfs_ctx), dir_entry_offset) : NULL);
-}
-
-NX_INLINE RomFileSystemFileEntry *bktrGetFileEntryByOffset(BktrContext *ctx, u32 file_entry_offset)
-{
-    return (ctx ? romfsGetFileEntryByOffset(&(ctx->patch_romfs_ctx), file_entry_offset) : NULL);
-}
-
-NX_INLINE bool bktrGetTotalDataSize(BktrContext *ctx, u64 *out_size)
-{
-    return (ctx ? romfsGetTotalDataSize(&(ctx->patch_romfs_ctx), out_size) : false);
-}
-
-NX_INLINE bool bktrGetDirectoryDataSize(BktrContext *ctx, RomFileSystemDirectoryEntry *dir_entry, u64 *out_size)
-{
-    return (ctx ? romfsGetDirectoryDataSize(&(ctx->patch_romfs_ctx), dir_entry, out_size) : false);
-}
-
-NX_INLINE RomFileSystemDirectoryEntry *bktrGetDirectoryEntryByPath(BktrContext *ctx, const char *path)
-{
-    return (ctx ? romfsGetDirectoryEntryByPath(&(ctx->patch_romfs_ctx), path) : NULL);
-}
-
-NX_INLINE RomFileSystemFileEntry *bktrGetFileEntryByPath(BktrContext *ctx, const char *path)
-{
-    return (ctx ? romfsGetFileEntryByPath(&(ctx->patch_romfs_ctx), path) : NULL);
-}
-
-NX_INLINE bool bktrGeneratePathFromDirectoryEntry(BktrContext *ctx, RomFileSystemDirectoryEntry *dir_entry, char *out_path, size_t out_path_size, u8 illegal_char_replace_type)
-{
-    return (ctx ? romfsGeneratePathFromDirectoryEntry(&(ctx->patch_romfs_ctx), dir_entry, out_path, out_path_size, illegal_char_replace_type) : false);
-}
-
-NX_INLINE bool bktrGeneratePathFromFileEntry(BktrContext *ctx, RomFileSystemFileEntry *file_entry, char *out_path, size_t out_path_size, u8 illegal_char_replace_type)
-{
-    return (ctx ? romfsGeneratePathFromFileEntry(&(ctx->patch_romfs_ctx), file_entry, out_path, out_path_size, illegal_char_replace_type) : false);
 }
 
 #ifdef __cplusplus

source/core/bktr.c

@@ -159,6 +159,55 @@ bool bktrInitializeContext(BucketTreeContext *out, NcaFsSectionContext *nca_fs_c
     return success;
 }
 
+bool bktrSetRegularSubStorage(BucketTreeContext *ctx, NcaFsSectionContext *nca_fs_ctx)
+{
+    NcaContext *nca_ctx = NULL;
+    
+    if (!bktrIsValidContext(ctx) || !nca_fs_ctx || !nca_fs_ctx->enabled || nca_fs_ctx->section_type >= NcaFsSectionType_Invalid || \
+        !(nca_ctx = (NcaContext*)nca_fs_ctx->nca_ctx) || (nca_ctx->rights_id_available && !nca_ctx->titlekey_retrieved))
+    {
+        LOG_MSG("Invalid parameters!");
+        return false;
+    }
+    
+    /* Update the substorage. */
+    BucketTreeSubStorage *substorage = &(ctx->substorages[0]);
+    memset(substorage, 0, sizeof(BucketTreeSubStorage));
+    
+    substorage->index = 0;
+    substorage->nca_fs_ctx = nca_fs_ctx;
+    substorage->type = BucketTreeSubStorageType_Regular;
+    substorage->bktr_ctx = NULL;
+    
+    return true;
+}
+
+bool bktrSetBucketTreeSubStorage(BucketTreeContext *parent_ctx, BucketTreeContext *child_ctx, u8 substorage_index)
+{
+    if (!bktrIsValidContext(parent_ctx) || !bktrIsValidContext(child_ctx) || substorage_index >= BKTR_MAX_SUBSTORAGE_COUNT || \
+        (parent_ctx->storage_type != BucketTreeStorageType_Indirect && substorage_index != 0) || \
+        (parent_ctx->storage_type == BucketTreeStorageType_Indirect && child_ctx->storage_type != BucketTreeStorageType_Compressed && child_ctx->storage_type != BucketTreeStorageType_AesCtrEx) || \
+        (parent_ctx->storage_type == BucketTreeStorageType_Indirect && child_ctx->storage_type == BucketTreeStorageType_Compressed && substorage_index != 0) || \
+        (parent_ctx->storage_type == BucketTreeStorageType_Indirect && child_ctx->storage_type == BucketTreeStorageType_AesCtrEx && substorage_index != 1) || \
+        parent_ctx->storage_type == BucketTreeStorageType_AesCtrEx || parent_ctx->storage_type == BucketTreeStorageType_Sparse || \
+        (parent_ctx->storage_type == BucketTreeStorageType_Compressed && child_ctx->storage_type != BucketTreeStorageType_Indirect && child_ctx->storage_type != BucketTreeStorageType_Sparse))
+    {
+        LOG_MSG("Invalid parameters!");
+        return false;
+    }
+    
+    /* Update the substorage. */
+    BucketTreeSubStorage *substorage = &(parent_ctx->substorages[substorage_index]);
+    memset(substorage, 0, sizeof(BucketTreeSubStorage));
+    
+    substorage->index = substorage_index;
+    substorage->nca_fs_ctx = child_ctx->nca_fs_ctx;
+    substorage->type = (child_ctx->storage_type + 1);   /* Convert to BucketTreeSubStorageType value. */
+    substorage->bktr_ctx = child_ctx;
+    
+    return true;
+}
+
 bool bktrReadStorage(BucketTreeContext *ctx, void *out, u64 read_size, u64 offset)
 {
     if (!bktrIsBlockWithinStorageRange(ctx, read_size, offset) || !out)
@@ -204,29 +253,64 @@ end:
 
 
 
-static bool bktrReadIndirectStorage(BucketTreeVisitor *visitor, void *out, u64 read_size, u64 offset)
+
+
+static bool bktrReadSubStorage(BucketTreeSubStorage *substorage, void *out, u64 read_size, u64 offset)
 {
-    if (!bktrVisitorIsValid(visitor) || !out)
+    if (!bktrIsValidSubstorage(substorage) || !out || !read_size)
     {
         LOG_MSG("Invalid parameters!");
         return false;
     }
     
+    bool success = false;
+    
+    if (substorage->type == BucketTreeSubStorageType_Regular)
+    {
+        /* Perform a read on the target NCA. */
+        
+        
+        
+        // TODO: HANDLE SPARSE STORAGE VIRTUAL OFFSET DECRYPTION
+        // substorage->nca_fs_ctx->has_sparse_layer
+        
+        
+        success = ncaReadFsSection(substorage->nca_fs_ctx, out, read_size, offset);
+    } else {
+        /* Perform a read on the target BucketTree storage. */
+        success = bktrReadStorage(substorage->bktr_ctx, out, read_size, offset);
+    }
+    
+    if (!success) LOG_MSG("Failed to read 0x%lX-byte long chunk from offset 0x%lX!", read_size, offset);
+    
+    return success;
+}
+
+
+
+
+
+
+
+
+static bool bktrReadIndirectStorage(BucketTreeVisitor *visitor, void *out, u64 read_size, u64 offset)
+{
     BucketTreeContext *ctx = visitor->bktr_ctx;
     NcaFsSectionContext *nca_fs_ctx = ctx->nca_fs_ctx;
     bool is_sparse = (ctx->storage_type == BucketTreeStorageType_Sparse);
     
+    if (!out || !bktrIsValidSubstorage(&(ctx->substorages[0])) || (!is_sparse && !bktrIsValidSubstorage(&(ctx->substorages[1]))) || \
+        (!is_sparse && ctx->substorages[0].type != BucketTreeSubStorageType_Regular && ctx->substorages[0].type != BucketTreeStorageType_Compressed) || \
+        (is_sparse && ctx->substorages[0].type != BucketTreeSubStorageType_Regular) || (!is_sparse && ctx->substorages[1].type != BucketTreeSubStorageType_AesCtrEx))
+    {
+        LOG_MSG("Invalid parameters!");
+        return false;
+    }
+    
+    /* Validate Indirect Storage entry. */
     BucketTreeIndirectStorageEntry cur_entry = {0};
     memcpy(&cur_entry, visitor->entry, sizeof(BucketTreeIndirectStorageEntry));
     
-    
-    
-    // TODO: ADD BASE STORAGE CHECK IF NOT SPARSE
-    
-    
-    
-    
-    /* Validate Indirect Storage entry. */
     if (!bktrIsOffsetWithinStorageRange(ctx, cur_entry.virtual_offset) || cur_entry.virtual_offset > offset || cur_entry.storage_index > BucketTreeIndirectStorageIndex_Patch)
     {
         LOG_MSG("Invalid Indirect Storage entry! (0x%lX) (#1).", cur_entry.virtual_offset);
@@ -287,22 +371,16 @@ static bool bktrReadIndirectStorage(BucketTreeVisitor *visitor, void *out, u64 r
         if (cur_entry.storage_index == BucketTreeIndirectStorageIndex_Original)
         {
             /* Retrieve data from the original data storage. */
-            
+            /* This may either be a Regular/Compressed storage from the base NCA (Indirect) or a Regular storage from this very same NCA (Sparse). */
-            
+            success = bktrReadSubStorage(&(ctx->substorages[0]), out, read_size, data_offset);
-            
+            if (!success) LOG_MSG("Failed to read 0x%lX-byte long chunk from offset 0x%lX in original data storage!", read_size, data_offset);
-            // TODO: SET ORIGINAL DATA STORAGE
-            
-            
-            
         } else {
             if (!is_sparse)
             {
                 /* Retrieve data from the indirect data storage. */
-                
+                /* This must always be the AesCtrEx storage within this very same NCA (Indirect). */
-                
+                success = bktrReadSubStorage(&(ctx->substorages[1]), out, read_size, data_offset);
-                // TODO: SET INDIRECT DATA STORAGE
+                if (!success) LOG_MSG("Failed to read 0x%lX-byte long chunk from offset 0x%lX in AesCtrEx storage!", read_size, data_offset);
-                
-                
             } else {
                 /* Fill output buffer with zeroes (SparseStorage's ZeroStorage). */
                 memset(0, out, read_size);