mirror of
https://github.com/jakcron/nstool
synced 2024-11-22 13:39:28 +00:00
[nstool] Add HashTreeWrappedIFile
This commit is contained in:
jakcron
parent
c1f343d052
commit
18e4c175df
2 changed files with 257 additions and 0 deletions
210
programs/nstool/source/HashTreeWrappedIFile.cpp
Normal file
210
programs/nstool/source/HashTreeWrappedIFile.cpp
Normal file
|
|
@ -0,0 +1,210 @@
|
|||
|
||||
|
||||
#include "HashTreeWrappedIFile.h"
|
||||
#include "CopiedIFile.h"
|
||||
#include "OffsetAdjustedIFile.h"
|
||||
|
||||
|
||||
HashTreeWrappedIFile::HashTreeWrappedIFile(fnd::IFile* file, const HashTreeMeta& hdr) :
|
||||
mOwnIFile(true),
|
||||
mFile(file),
|
||||
mData(nullptr),
|
||||
mDataHashLayer(),
|
||||
mAlignHashCalcToBlock(false)
|
||||
{
|
||||
initialiseDataLayer(hdr);
|
||||
}
|
||||
|
||||
HashTreeWrappedIFile::HashTreeWrappedIFile(fnd::IFile* file, bool ownIFile, const HashTreeMeta& hdr) :
|
||||
mOwnIFile(ownIFile),
|
||||
mFile(file),
|
||||
mData(nullptr),
|
||||
mDataHashLayer(),
|
||||
mAlignHashCalcToBlock(false)
|
||||
{
|
||||
initialiseDataLayer(hdr);
|
||||
}
|
||||
|
||||
HashTreeWrappedIFile::~HashTreeWrappedIFile()
|
||||
{
|
||||
if (mOwnIFile)
|
||||
{
|
||||
delete mFile;
|
||||
}
|
||||
delete mData;
|
||||
}
|
||||
|
||||
size_t HashTreeWrappedIFile::size()
|
||||
{
|
||||
return mData->size();
|
||||
}
|
||||
|
||||
void HashTreeWrappedIFile::seek(size_t offset)
|
||||
{
|
||||
mDataOffset = offset;
|
||||
}
|
||||
|
||||
void HashTreeWrappedIFile::read(byte_t* out, size_t len)
|
||||
{
|
||||
size_t offset_in_start_block = getOffsetInBlock(mDataOffset);
|
||||
size_t offset_in_end_block = getOffsetInBlock(offset_in_start_block + len);
|
||||
|
||||
size_t start_block = getOffsetBlock(mDataOffset);
|
||||
size_t block_num = align(offset_in_start_block + len, mDataBlockSize) / mDataBlockSize;
|
||||
|
||||
size_t scratch_block_capacity = mScratch.getSize() / mDataBlockSize;
|
||||
|
||||
size_t partial_last_block_num = block_num % scratch_block_capacity;
|
||||
bool has_partial_block_num = partial_last_block_num > 0;
|
||||
size_t read_iterations = (block_num / scratch_block_capacity) + has_partial_block_num;
|
||||
|
||||
|
||||
size_t block_read_len;
|
||||
size_t block_export_offset;
|
||||
size_t block_export_size;
|
||||
size_t block_export_pos = 0;
|
||||
for (size_t i = 0; i < read_iterations; i++)
|
||||
{
|
||||
// how many blocks to read from source file
|
||||
block_read_len = (i+1 == read_iterations && has_partial_block_num) ? partial_last_block_num : scratch_block_capacity;
|
||||
|
||||
// offset in this current read to copy from
|
||||
block_export_offset = (i == 0) ? offset_in_start_block : 0;
|
||||
|
||||
// size of current read to copy
|
||||
block_export_size = (block_read_len * mDataBlockSize) - block_export_offset;
|
||||
|
||||
// if last read, reduce the export size by one block less offset_in_end_block
|
||||
if (i+1 == read_iterations)
|
||||
{
|
||||
block_export_size -= (mDataBlockSize - offset_in_end_block);
|
||||
}
|
||||
|
||||
// read the blocks
|
||||
readData(start_block + (i * scratch_block_capacity), block_read_len);
|
||||
|
||||
// export the section of data that is relevant
|
||||
memcpy(out + block_export_pos, mScratch.getBytes() + block_export_offset, block_export_size);
|
||||
|
||||
// update export position
|
||||
block_export_pos += block_export_size;
|
||||
}
|
||||
|
||||
// update offset
|
||||
mDataOffset += len;
|
||||
}
|
||||
|
||||
void HashTreeWrappedIFile::read(byte_t* out, size_t offset, size_t len)
|
||||
{
|
||||
seek(offset);
|
||||
read(out, len);
|
||||
}
|
||||
|
||||
void HashTreeWrappedIFile::write(const byte_t* out, size_t len)
|
||||
{
|
||||
throw fnd::Exception(kModuleName, "write() not supported");
|
||||
}
|
||||
|
||||
void HashTreeWrappedIFile::write(const byte_t* out, size_t offset, size_t len)
|
||||
{
|
||||
throw fnd::Exception(kModuleName, "write() not supported");
|
||||
}
|
||||
|
||||
void HashTreeWrappedIFile::initialiseDataLayer(const HashTreeMeta& hdr)
|
||||
{
|
||||
crypto::sha::sSha256Hash hash;
|
||||
fnd::MemoryBlob cur, prev;
|
||||
|
||||
mAlignHashCalcToBlock = hdr.getAlignHashToBlock();
|
||||
|
||||
// copy master hash into prev
|
||||
prev.alloc(sizeof(crypto::sha::sSha256Hash) * hdr.getMasterHashList().getSize());
|
||||
for (size_t i = 0; i < hdr.getMasterHashList().getSize(); i++)
|
||||
{
|
||||
((crypto::sha::sSha256Hash*)prev.getBytes())[i] = hdr.getMasterHashList()[i];
|
||||
}
|
||||
|
||||
// check each hash layer
|
||||
for (size_t i = 0; i < hdr.getHashLayerInfo().getSize(); i++)
|
||||
{
|
||||
// get block size
|
||||
const HashTreeMeta::sLayer& layer = hdr.getHashLayerInfo()[i];
|
||||
|
||||
// allocate layer
|
||||
cur.alloc(align(layer.size, layer.block_size));
|
||||
|
||||
// read layer
|
||||
mFile->read(cur.getBytes(), layer.offset, layer.size);
|
||||
|
||||
// validate blocks
|
||||
size_t validate_size;
|
||||
for (size_t j = 0; j < cur.getSize() / layer.block_size; j++)
|
||||
{
|
||||
validate_size = mAlignHashCalcToBlock? layer.block_size : MIN(layer.size - (j * layer.block_size), layer.block_size);
|
||||
crypto::sha::Sha256(cur.getBytes() + (j * layer.block_size), validate_size, hash.bytes);
|
||||
if (hash.compare(prev.getBytes() + j * sizeof(crypto::sha::sSha256Hash)) == false)
|
||||
{
|
||||
mErrorSs << "Hash tree layer verification failed (layer: " << i << ", block: " << j << ")";
|
||||
throw fnd::Exception(kModuleName, mErrorSs.str());
|
||||
}
|
||||
}
|
||||
|
||||
// set prev to cur
|
||||
prev = cur;
|
||||
}
|
||||
|
||||
// save last layer as hash table for data layer
|
||||
crypto::sha::sSha256Hash* hash_list = (crypto::sha::sSha256Hash*)prev.getBytes();
|
||||
for (size_t i = 0; i < prev.getSize() / sizeof(crypto::sha::sSha256Hash); i++)
|
||||
{
|
||||
mDataHashLayer.addElement(hash_list[i]);
|
||||
}
|
||||
|
||||
// generate reader for data layer
|
||||
mData = new OffsetAdjustedIFile(mFile, false, hdr.getDataLayer().offset, hdr.getDataLayer().size);
|
||||
mDataOffset = 0;
|
||||
mDataBlockSize = hdr.getDataLayer().block_size;
|
||||
|
||||
// allocate scratchpad
|
||||
mScratch.alloc(mDataBlockSize * 0x10);
|
||||
}
|
||||
|
||||
void HashTreeWrappedIFile::readData(size_t block_offset, size_t block_num)
|
||||
{
|
||||
seek(block_offset * mDataBlockSize);
|
||||
crypto::sha::sSha256Hash hash;
|
||||
|
||||
// determine read size
|
||||
size_t read_len = 0;
|
||||
if ((block_offset + block_num) == getBlockNum(mData->size()))
|
||||
{
|
||||
read_len = (block_num-1) * mDataBlockSize + getRemanderBlockReadSize(mData->size());
|
||||
memset(mScratch.getBytes(), 0, block_num * mDataBlockSize);
|
||||
}
|
||||
else if ((block_offset + block_num) < getBlockNum(mData->size()))
|
||||
{
|
||||
read_len = block_num * mDataBlockSize;
|
||||
}
|
||||
else
|
||||
{
|
||||
throw fnd::Exception(kModuleName, "Out of bounds file read");
|
||||
}
|
||||
|
||||
// read
|
||||
mData->read(mScratch.getBytes(), block_offset * mDataBlockSize, read_len);
|
||||
|
||||
|
||||
|
||||
// validate blocks
|
||||
size_t validate_size;
|
||||
for (size_t i = 0; i < block_num; i++)
|
||||
{
|
||||
validate_size = mAlignHashCalcToBlock? mDataBlockSize : MIN(read_len - (i * mDataBlockSize), mDataBlockSize);
|
||||
crypto::sha::Sha256(mScratch.getBytes() + (i * mDataBlockSize), validate_size, hash.bytes);
|
||||
if (hash != mDataHashLayer[block_offset + i])
|
||||
{
|
||||
mErrorSs << "Hash tree layer verification failed (layer: data, block: " << (block_offset + i) << ", offset: 0x" << std::hex << ((block_offset + i) * mDataBlockSize) << ", size: 0x" << std::hex << validate_size <<")";
|
||||
throw fnd::Exception(kModuleName, mErrorSs.str());
|
||||
}
|
||||
}
|
||||
}
|
||||
47
programs/nstool/source/HashTreeWrappedIFile.h
Normal file
47
programs/nstool/source/HashTreeWrappedIFile.h
Normal file
|
|
@ -0,0 +1,47 @@
|
|||
#pragma once
|
||||
#include <sstream>
|
||||
#include <fnd/IFile.h>
|
||||
#include <fnd/MemoryBlob.h>
|
||||
#include <crypto/sha.h>
|
||||
#include "HashTreeMeta.h"
|
||||
|
||||
|
||||
class HashTreeWrappedIFile : public fnd::IFile
|
||||
{
|
||||
public:
|
||||
HashTreeWrappedIFile(fnd::IFile* file, const HashTreeMeta& hdr);
|
||||
HashTreeWrappedIFile(fnd::IFile* file, bool ownIFile, const HashTreeMeta& hdr);
|
||||
~HashTreeWrappedIFile();
|
||||
|
||||
size_t size();
|
||||
void seek(size_t offset);
|
||||
void read(byte_t* out, size_t len);
|
||||
void read(byte_t* out, size_t offset, size_t len);
|
||||
void write(const byte_t* out, size_t len);
|
||||
void write(const byte_t* out, size_t offset, size_t len);
|
||||
private:
|
||||
const std::string kModuleName = "HashTreeWrappedIFile";
|
||||
std::stringstream mErrorSs;
|
||||
|
||||
bool mOwnIFile;
|
||||
fnd::IFile* mFile;
|
||||
|
||||
|
||||
|
||||
// data file
|
||||
fnd::IFile* mData;
|
||||
size_t mDataOffset;
|
||||
size_t mDataBlockSize;
|
||||
fnd::List<crypto::sha::sSha256Hash> mDataHashLayer;
|
||||
bool mAlignHashCalcToBlock;
|
||||
|
||||
fnd::MemoryBlob mScratch;
|
||||
|
||||
inline size_t getOffsetBlock(size_t offset) const { return offset / mDataBlockSize; }
|
||||
inline size_t getOffsetInBlock(size_t offset) const { return offset % mDataBlockSize; }
|
||||
inline size_t getRemanderBlockReadSize(size_t total_size) const { return total_size % mDataBlockSize; }
|
||||
inline size_t getBlockNum(size_t total_size) const { return (total_size / mDataBlockSize) + (getRemanderBlockReadSize(total_size) > 0); }
|
||||
|
||||
void initialiseDataLayer(const HashTreeMeta& hdr);
|
||||
void readData(size_t block_offset, size_t block_num);
|
||||
};
|
||||
Loading…
Reference in a new issue