Skip to content

Commit

Permalink
thcrap_tasofro: Cleanup and optimize TasofroFile
Browse files Browse the repository at this point in the history
  • Loading branch information
zero318 committed Jun 14, 2024
1 parent 274900f commit 8b17bf9
Show file tree
Hide file tree
Showing 5 changed files with 147 additions and 132 deletions.
27 changes: 17 additions & 10 deletions thcrap_tasofro/src/crypt.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,9 @@

ICrypt* ICrypt::instance = nullptr;

#define ASSUME_OVERALLOCATED_BUFFERS 1
#define USE_UNROLLED_CRYPT_LOOPS 1

DWORD CryptTh135::cryptBlock(BYTE *TH_RESTRICT Data, DWORD FileSize, const DWORD *TH_RESTRICT Key)
{
if (FileSize) {
Expand All @@ -39,7 +42,7 @@ DWORD CryptTh135::cryptBlock(BYTE *TH_RESTRICT Data, DWORD FileSize, const DWORD

void CryptTh135::uncryptBlock(BYTE *TH_RESTRICT Data, DWORD FileSize, const DWORD *TH_RESTRICT Key)
{
this->cryptBlock(Data, FileSize, Key);
this->cryptBlock(Data, AlignUpToMultipleOf2(FileSize, 16), Key);
}

// Normalized Hash
Expand All @@ -66,8 +69,6 @@ DWORD CryptTh135::SpecialFNVHash(const char *begin, const char *end, DWORD initH
void CryptTh135::convertKey(DWORD*)
{}

#define USE_UNROLLED_CRYPT_LOOPS 1

static constexpr int8_t sse_mask_array[] = {
-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
Expand Down Expand Up @@ -153,15 +154,15 @@ DWORD cryptBlockInternal(BYTE *TH_RESTRICT Data, DWORD FileSize, const DWORD *TH
TH_UNREACHABLE;
case 3:
temp = cur_data[2];
cur_data[2] = Aux >> 16;
cur_data[2] = (BYTE)(Aux >> 16);
Aux ^= (temp ^ key_byte_ptr[2]) << 16;
case 2:
temp = cur_data[1];
cur_data[1] = Aux >> 8;
cur_data[1] = (BYTE)(Aux >> 8);
Aux ^= (temp ^ key_byte_ptr[1]) << 8;
case 1:
temp = cur_data[0];
cur_data[0] = Aux;
cur_data[0] = (BYTE)Aux;
Aux ^= temp ^ key_byte_ptr[0];
}
}
Expand Down Expand Up @@ -201,21 +202,26 @@ void CryptTh145::uncryptBlock(BYTE *TH_RESTRICT Data, DWORD FileSize, const DWOR
{

if (FileSize) {
#if !ASSUME_OVERALLOCATED_BUFFERS
DWORD aux = Key[0];
size_t i = 0;
if (FileSize >= 16) {
#else

FileSize = AlignUpToMultipleOf2(FileSize, 16);
#endif
__m128 key_wide = _mm_loadu_ps((float*)Key);
__m128 aux_wide = key_wide;
BYTE *TH_RESTRICT data_end = Data + FileSize - 16;
do {
__m128 temp = _mm_loadu_ps((float*)Data);
aux_wide = _mm_shuffle_ps(_mm_movelh_ps(aux_wide, temp), temp, 0x98);
__m128 xor_data = key_wide;
_mm_storeu_ps((float*)Data, _mm_xor_ps(_mm_xor_ps(xor_data, temp), aux_wide));
_mm_storeu_ps((float*)Data, _mm_xor_ps(_mm_xor_ps(key_wide, temp), _mm_shuffle_ps(_mm_movelh_ps(aux_wide, temp), temp, 0x98)));
*(__m128i*)&aux_wide = _mm_srli_si128(*(__m128i*)&temp, 12);
} while ((Data += 16) <= data_end);

#if !ASSUME_OVERALLOCATED_BUFFERS
aux = _mm_cvtsi128_si32(*(__m128i*)&aux_wide);
}

#if USE_UNROLLED_CRYPT_LOOPS
const BYTE *TH_RESTRICT key_byte_ptr = (const BYTE *TH_RESTRICT)Key;
if (size_t dwords_width_remaining = FileSize & 0b1100) {
Expand Down Expand Up @@ -265,6 +271,7 @@ void CryptTh145::uncryptBlock(BYTE *TH_RESTRICT Data, DWORD FileSize, const DWOR
aux |= temp << 24;
} while (--remaining);
}
#endif
#endif
}
}
Expand Down
170 changes: 88 additions & 82 deletions thcrap_tasofro/src/tasofro_file.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,12 @@ void TasofroFile::tls_set(TasofroFile *file)
}

TasofroFile::TasofroFile()
: init_done(false)
{
memset(this, 0, sizeof(file_rep_t));
// Without this the compiler insists
// on generating individual stores
// to dodge around padding bytes, even
// when default initializing.
memset(this, 0, sizeof(TasofroFile));
}

TasofroFile::~TasofroFile()
Expand Down Expand Up @@ -65,29 +68,6 @@ void TasofroFile::clear()
this->init_done = false;
}

bool TasofroFile::need_orig_file() const
{
// If dat_dump is enabled, we *always* need the original buffer
if unexpected(runconfig_dat_dump_get()) {
return true;
}

// We have a replacement file, no need to read the original one (even if we have
// hook callbacks, they will be called on the replacement file)
if (this->rep_buffer) {
return false;
}

// We have a hook callback. It will need the original file.
if (this->hooks) {
return true;
}

// Nothing to do, we just let the game read its original file
// with its own code.
return false;
}

bool TasofroFile::need_replace() const
{
// If dat_dump is enabled, we always want to run the patching code
Expand Down Expand Up @@ -115,7 +95,9 @@ void TasofroFile::read_from_ReadFile(HANDLE hFile, DWORD fileOffset, DWORD size)
{
DWORD nbOfBytesRead;

SAFE_FREE(this->game_buffer);
if (this->game_buffer) {
free(this->game_buffer);
}
this->game_buffer = malloc(size);
this->pre_json_size = size;

Expand All @@ -124,69 +106,94 @@ void TasofroFile::read_from_ReadFile(HANDLE hFile, DWORD fileOffset, DWORD size)
SetFilePointer(hFile, oldOffset, NULL, FILE_BEGIN);
}

void TasofroFile::init_buffer()
{
const char *dat_dump = runconfig_dat_dump_get();
if (dat_dump) {
DumpDatFile(dat_dump, this->name, this->game_buffer, this->pre_json_size);
}

if (!this->hooks) {
// No patches to apply
SAFE_FREE(this->game_buffer);
return ;
}

if (!this->rep_buffer) {
this->rep_buffer = malloc(POST_JSON_SIZE(this));
memcpy(this->rep_buffer, this->game_buffer, this->pre_json_size);
}
// We filled rep_buffer, we won't need this one anymore
SAFE_FREE(this->game_buffer);

// Patch the game
if (!patchhooks_run(this->hooks, this->rep_buffer, POST_JSON_SIZE(this), this->pre_json_size, this->name, this->patch)) {
// Remove the hooks so that next calls won't try to patch the file.
// This is only an optimization, avoiding to re-read the file, but this
// optimization is required as long as we don't detect open/close calls
// and try to re-open a file on every read call.
SAFE_FREE(this->hooks);
}
}

void TasofroFile::replace_ReadFile_init(ReadFileStack *stack,
std::function<void (TasofroFile *fr, BYTE *buffer, DWORD size)>& decrypt,
std::function<void (TasofroFile *fr, BYTE *buffer, DWORD size)>& crypt)
crypt_func_t decrypt,
crypt_func_t crypt)
{
// In order to get the correct offset, we wait until the 1st read
// of a file by the game, then remember this offset for future calls.
this->offset = SetFilePointer(stack->hFile, 0, nullptr, FILE_CURRENT);
this->offset = SetFilePointer(stack->hFile, 0, NULL, FILE_CURRENT);

// Local variables used to get MSVC to store
// things in stack locals and registers
const char* dat_dump = runconfig_dat_dump_get();
void* rep_buffer = this->rep_buffer;
size_t pre_json_size = this->pre_json_size;

// Return true if we need to read the original file, false otherwise.
auto need_orig_file = [=]() {
// If dat_dump is enabled, we *always* need the original buffer
if unexpected(dat_dump) {
return true;
}
// We have a replacement file, no need to read the original one (even if we have
// hook callbacks, they will be called on the replacement file)
if (rep_buffer) {
return false;
}
// replace_ReadFile won't be run if the other conditions are false,
// so have a hook callback. It will need the original file.
return true;
};

if (need_orig_file()) {
// This buffer is slightly over-allocated
// to simplify the decrypt code
void* game_buffer = malloc(AlignUpToMultipleOf2(pre_json_size + this->patch_size, 16));

if (this->need_orig_file()) {
this->read_from_ReadFile(stack->hFile, this->offset, this->pre_json_size);
decrypt(this, (BYTE*)this->game_buffer, this->pre_json_size);
// This variable is static just to avoid
// allocating stack space
static DWORD nbOfBytesRead;
ReadFile(stack->hFile, game_buffer, pre_json_size, &nbOfBytesRead, NULL);
// Rewind the file pointer
SetFilePointer(stack->hFile, this->offset, NULL, FILE_BEGIN);

decrypt(this, (BYTE*)game_buffer, pre_json_size);

if unexpected(dat_dump) {
DumpDatFile(dat_dump, this->name, game_buffer, pre_json_size);
}
// If there are hooks and no replacement file
// then just pass the original file to the
// hook functions.
if (!rep_buffer && this->hooks) {
rep_buffer = this->rep_buffer = game_buffer;
// MSVC is too dumb to jump over the duplicate
// condition checks without this goto
goto run_hooks;
} else {
// No patches to apply or rep_buffer is already
// filled, we won't need this one anymore
free(game_buffer);
}
}
this->init_buffer();
if (this->rep_buffer) {
crypt(this, (BYTE*)this->rep_buffer, POST_JSON_SIZE(this));
if (rep_buffer) {
if (this->hooks) {
run_hooks:
// Patch the game
patchhooks_run(this->hooks, rep_buffer, pre_json_size + this->patch_size, pre_json_size, this->name, this->patch);
// Free the hook array now since they
// won't be run again anyway
free(this->hooks);
this->hooks = NULL;
}

crypt(this, (BYTE*)rep_buffer, pre_json_size + this->patch_size);
}
this->init_done = true;
}

int TasofroFile::replace_ReadFile_write(x86_reg_t *regs, ReadFileStack *stack)
{
DWORD offset = SetFilePointer(stack->hFile, 0, NULL, FILE_CURRENT) - this->offset;
DWORD offset_ = SetFilePointer(stack->hFile, 0, NULL, FILE_CURRENT);
DWORD size;
if (offset <= POST_JSON_SIZE(this)) {
size = MIN(POST_JSON_SIZE(this) - offset, stack->nNumberOfBytesToRead);
}
else {
size = 0;
DWORD size = 0;
ptrdiff_t remaining_size = POST_JSON_SIZE(this) - offset;
if (remaining_size >= 0) {
size = stack->nNumberOfBytesToRead;
size = __min((DWORD)remaining_size, size);
memcpy(stack->lpBuffer, (BYTE*)this->rep_buffer + offset, size);
SetFilePointer(stack->hFile, size, NULL, FILE_CURRENT);
}

memcpy(stack->lpBuffer, (BYTE*)this->rep_buffer + offset, size);
SetFilePointer(stack->hFile, size, NULL, FILE_CURRENT);
*stack->lpNumberOfBytesRead = size;

regs->eax = 1;
Expand All @@ -195,12 +202,12 @@ int TasofroFile::replace_ReadFile_write(x86_reg_t *regs, ReadFileStack *stack)
}

int TasofroFile::replace_ReadFile(x86_reg_t *regs,
std::function<void (TasofroFile *fr, BYTE *buffer, DWORD size)> decrypt,
std::function<void (TasofroFile *fr, BYTE *buffer, DWORD size)> crypt)
crypt_func_t decrypt,
crypt_func_t crypt)
{
ReadFileStack *stack = (ReadFileStack*)(regs->esp + sizeof(void*));

if (stack->lpOverlapped) {
if unexpected(stack->lpOverlapped) {
// Overlapped operations are not supported.
// We'd better leave that file alone rather than ignoring that.
return 1;
Expand All @@ -210,9 +217,8 @@ int TasofroFile::replace_ReadFile(x86_reg_t *regs,
this->replace_ReadFile_init(stack, decrypt, crypt);
}

int ret = 1;
if (this->rep_buffer) {
ret = this->replace_ReadFile_write(regs, stack);
if unexpected(this->rep_buffer) {
return this->replace_ReadFile_write(regs, stack);
}
return ret;
return 1;
}
20 changes: 9 additions & 11 deletions thcrap_tasofro/src/tasofro_file.h
Original file line number Diff line number Diff line change
Expand Up @@ -24,16 +24,17 @@ struct ReadFileStack
LPOVERLAPPED lpOverlapped;
};

struct TasofroFile;

func_ptr_typedef(void, , crypt_func_t)(TasofroFile *fr, BYTE *buffer, size_t size);

struct TasofroFile : public file_rep_t
{
// Offset of this file in the archive
size_t offset;
// True if the file have already been patched
// and re-encrypted
bool init_done;
// Mutex to prevent concurrent access
// to the file by several threads
std::mutex mutex;

static TasofroFile* tls_get();
static void tls_set(TasofroFile *file);
Expand All @@ -55,18 +56,15 @@ struct TasofroFile : public file_rep_t
// and patching the file if needed, or let the game do its own ReadFile Call.
// You must use the return value of this function as the return value for
// your breakpoint.
[[nodiscard]] int replace_ReadFile(x86_reg_t *regs,
std::function<void (TasofroFile *fr, BYTE *buffer, DWORD size)> decrypt,
std::function<void (TasofroFile *fr, BYTE *buffer, DWORD size)> crypt);
TH_NODISCARD int replace_ReadFile(x86_reg_t *regs,
crypt_func_t decrypt,
crypt_func_t crypt);

private:
void read_from_ReadFile(HANDLE hFile, DWORD fileOffset, DWORD size);
void init_buffer();
// Return true if we need to read the original file, false otherwise.
bool need_orig_file() const;

void replace_ReadFile_init(ReadFileStack *stack,
std::function<void (TasofroFile *fr, BYTE *buffer, DWORD size)>& decrypt,
std::function<void (TasofroFile *fr, BYTE *buffer, DWORD size)>& crypt);
crypt_func_t decrypt,
crypt_func_t crypt);
int replace_ReadFile_write(x86_reg_t *regs, ReadFileStack *stack);
};
Loading

0 comments on commit 8b17bf9

Please sign in to comment.