mirrors/nstool
1
0
Fork 0
mirror of https://github.com/jakcron/nstool.git synced 2024-12-22 10:45:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

[nstool] Add HashTreeWrappedIFile

Browse source
This commit is contained in:
jakcron 2018-05-21 23:03:57 +08:00
parent c1f343d052
commit 18e4c175df
2 changed files with 257 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

210
programs/nstool/source/HashTreeWrappedIFile.cpp Normal file
View 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
View 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);
};