314 lines
7.2 KiB
C++
314 lines
7.2 KiB
C++
#include "sgu_stdafx.h"
|
|
#include "sgu_mode_apply_patch.h"
|
|
#include "sgu_utils.h"
|
|
|
|
extern "C"
|
|
{
|
|
#include "bspatch.h"
|
|
}
|
|
|
|
#include "bzlib.h"
|
|
#include "zstd.h"
|
|
|
|
|
|
//--------------------------- Not Compress-----------------------------------
|
|
static int patchFileReader(const struct bspatch_stream* stream, void* buffer, int length)
|
|
{
|
|
FILE* fp = (FILE*)stream->opaque;
|
|
size_t n = fread(buffer, 1, (size_t)length, fp);
|
|
if (n != length)
|
|
return -1;
|
|
|
|
return 0;
|
|
}
|
|
|
|
//--------------------------- BZ2 -----------------------------------
|
|
static int bzPatchFileReader(const struct bspatch_stream* stream, void* buffer, int length)
|
|
{
|
|
BZFILE* bz2 = (BZFILE*)stream->opaque;
|
|
|
|
int bz2err;
|
|
int n = BZ2_bzRead(&bz2err, bz2, buffer, length);
|
|
if (n != length)
|
|
return -1;
|
|
|
|
return 0;
|
|
}
|
|
|
|
//--------------------------- ZSTD -----------------------------------
|
|
struct ZSTD_DecompressContext
|
|
{
|
|
ZSTD_DCtx* dctx;
|
|
|
|
FILE* fpInput;
|
|
|
|
void* inputBuf;
|
|
size_t inputBufSize;
|
|
size_t inputBufDataSize;
|
|
size_t inputBufDataReadedPos;
|
|
|
|
void* outputBuf;
|
|
size_t outputBufSize;
|
|
size_t outputBufDataSize;
|
|
size_t outputBufDataUploadedPos;
|
|
};
|
|
|
|
ZSTD_DecompressContext* createZSTDDecompressContext(FILE* fpInput)
|
|
{
|
|
ZSTD_DecompressContext* ctx = new ZSTD_DecompressContext();
|
|
ctx->dctx = ZSTD_createDCtx();
|
|
ctx->fpInput = fpInput;
|
|
|
|
ctx->inputBufSize = ZSTD_DStreamInSize();
|
|
ctx->inputBuf = malloc(ctx->inputBufSize);
|
|
ctx->inputBufDataSize = 0;
|
|
ctx->inputBufDataReadedPos = 0;
|
|
|
|
ctx->outputBufSize = ZSTD_DStreamOutSize();
|
|
ctx->outputBuf = malloc(ctx->outputBufSize);
|
|
ctx->outputBufDataSize = 0;
|
|
ctx->outputBufDataUploadedPos = 0;
|
|
|
|
return ctx;
|
|
}
|
|
|
|
void destoryZSTDDecompressContext(ZSTD_DecompressContext* ctx)
|
|
{
|
|
free(ctx->inputBuf);
|
|
free(ctx->outputBuf);
|
|
delete ctx;
|
|
}
|
|
|
|
static int zstdPatchFileReader(const struct bspatch_stream* stream, void* buffer, int length)
|
|
{
|
|
ZSTD_DecompressContext* ctx = (ZSTD_DecompressContext*)(stream->opaque);
|
|
if (length == 0)
|
|
{
|
|
return 0;
|
|
}
|
|
|
|
//total uploaded size
|
|
size_t uploadedSize = 0;
|
|
|
|
do {
|
|
//upload all decompressed data
|
|
if (ctx->outputBufDataSize > ctx->outputBufDataUploadedPos)
|
|
{
|
|
size_t remainDecompressedDataSize = ctx->outputBufDataSize - ctx->outputBufDataUploadedPos;
|
|
size_t remainNotUploadedSize = (size_t)length - uploadedSize;
|
|
|
|
size_t uploadSize = remainNotUploadedSize < remainDecompressedDataSize ? remainNotUploadedSize : remainDecompressedDataSize;
|
|
if (uploadSize > 0)
|
|
{
|
|
memcpy((uint8_t*)buffer + uploadedSize, ((uint8_t*)ctx->outputBuf) + ctx->outputBufDataUploadedPos, uploadSize);
|
|
|
|
uploadedSize += uploadSize;
|
|
ctx->outputBufDataUploadedPos += uploadSize;
|
|
}
|
|
}
|
|
|
|
if (uploadedSize >= length)
|
|
{
|
|
return 0;
|
|
}
|
|
|
|
//compressed data empty
|
|
if (ctx->inputBufDataSize == ctx->inputBufDataReadedPos)
|
|
{
|
|
//read file
|
|
ctx->inputBufDataSize = fread(ctx->inputBuf, 1, ctx->inputBufSize, ctx->fpInput);
|
|
ctx->inputBufDataReadedPos = 0;
|
|
|
|
if (ctx->inputBufDataSize == 0)
|
|
{
|
|
//TODO:
|
|
return -1;
|
|
}
|
|
}
|
|
|
|
//decompress data
|
|
ZSTD_inBuffer input = { ctx->inputBuf, ctx->inputBufDataSize, ctx->inputBufDataReadedPos };
|
|
ZSTD_outBuffer output = { ctx->outputBuf, ctx->outputBufSize, 0 };
|
|
|
|
const size_t ret = ZSTD_decompressStream(ctx->dctx, &output, &input);
|
|
if (ZSTD_isError(ret))
|
|
{
|
|
//TODO:
|
|
return -1;
|
|
}
|
|
ctx->inputBufDataReadedPos = input.pos;
|
|
|
|
ctx->outputBufDataUploadedPos = 0;
|
|
ctx->outputBufDataSize = output.pos;
|
|
} while (true);
|
|
|
|
return 0;
|
|
}
|
|
|
|
bool applyPatch(const char* inputFilename, const char* patchFilename, const char* outputFilename, int32_t compressMethod)
|
|
{
|
|
if (!inputFilename || inputFilename[0] == 0 || !patchFilename || patchFilename[0] == 0)
|
|
{
|
|
printf("Error: Input or Patch filename is not specified.\n");
|
|
return false;
|
|
}
|
|
if (!outputFilename || outputFilename[0] == 0)
|
|
{
|
|
printf("Error: Output filename is not specified.\n");
|
|
return false;
|
|
}
|
|
|
|
uint8_t* pInputBuffer = nullptr;
|
|
FILE* patchFile = nullptr;
|
|
FILE* outputFile = nullptr;
|
|
uint8_t* pOutputBuffer = nullptr;
|
|
|
|
bool result = false;
|
|
do {
|
|
// Open the input file
|
|
int64_t inputFileSize;
|
|
pInputBuffer = readFileToBuffer(inputFilename, inputFileSize);
|
|
if (pInputBuffer == nullptr)
|
|
{
|
|
printf("Error: Failed to read input file: %s\n", inputFilename);
|
|
break;
|
|
}
|
|
|
|
// Open the patch file
|
|
fopen_s(&patchFile, patchFilename, "rb");
|
|
if (!patchFile)
|
|
{
|
|
printf("Error: Could not open patch file %s\n", patchFilename);
|
|
break;
|
|
}
|
|
|
|
// Open the output file
|
|
fopen_s(&outputFile, outputFilename, "wb");
|
|
if (!outputFile)
|
|
{
|
|
printf("Error: Could not open output file %s\n", outputFilename);
|
|
break;
|
|
}
|
|
|
|
//read the patch file header and check for appropriate magic
|
|
uint8_t header[24];
|
|
if (fread(header, 1, 24, patchFile) != 24 || memcmp(header, "ENDSLEY/BSDIFF43", 16) != 0)
|
|
{
|
|
printf("Error: Corrupt patch file %s\n", patchFilename);
|
|
break;
|
|
}
|
|
|
|
// Read output file lengths from header
|
|
int64_t outputFileSize = *(int64_t*)(header + 16);
|
|
pOutputBuffer = new uint8_t[outputFileSize + 1];
|
|
if (pOutputBuffer == nullptr)
|
|
{
|
|
printf("Error: Failed to allocate memory[%lld] for output buffer\n", outputFileSize + 1);
|
|
break;
|
|
}
|
|
|
|
BZFILE* bzHandle = nullptr;
|
|
ZSTD_DecompressContext* zstdCtx = nullptr;
|
|
if (compressMethod == 1)
|
|
{
|
|
int bz2err;
|
|
bzHandle = BZ2_bzReadOpen(&bz2err, patchFile, 0, 0, NULL, 0);
|
|
if (bzHandle == nullptr)
|
|
{
|
|
printf("BZ2_bzReadOpen, bz2err=%d", bz2err);
|
|
break;
|
|
}
|
|
}
|
|
else if(compressMethod==2)
|
|
{
|
|
zstdCtx = createZSTDDecompressContext(patchFile);
|
|
if (zstdCtx == 0)
|
|
{
|
|
printf("create zstd context failed!");
|
|
break;
|
|
}
|
|
}
|
|
|
|
// Apply the patch
|
|
struct bspatch_stream stream;
|
|
if (compressMethod == 1)
|
|
{
|
|
stream.read = bzPatchFileReader;
|
|
stream.opaque = bzHandle;
|
|
}
|
|
else if (compressMethod == 2)
|
|
{
|
|
stream.read = zstdPatchFileReader;
|
|
stream.opaque = zstdCtx;
|
|
}
|
|
else
|
|
{
|
|
stream.read = patchFileReader;
|
|
stream.opaque = patchFile;
|
|
}
|
|
|
|
if (bspatch(pInputBuffer, inputFileSize, pOutputBuffer, outputFileSize, &stream) !=0 )
|
|
{
|
|
printf("Error: Failed to apply patch\n");
|
|
if (compressMethod == 1)
|
|
{
|
|
// Clean up the bzip2 reads
|
|
int bz2err;
|
|
BZ2_bzReadClose(&bz2err, bzHandle);
|
|
bzHandle = nullptr;
|
|
}
|
|
else if (compressMethod == 2)
|
|
{
|
|
destoryZSTDDecompressContext(zstdCtx);
|
|
zstdCtx = nullptr;
|
|
}
|
|
break;
|
|
}
|
|
|
|
if (compressMethod == 1)
|
|
{
|
|
// Clean up the bzip2 reads
|
|
int bz2err;
|
|
BZ2_bzReadClose(&bz2err, bzHandle);
|
|
bzHandle = nullptr;
|
|
}
|
|
else if (compressMethod == 2)
|
|
{
|
|
destoryZSTDDecompressContext(zstdCtx);
|
|
zstdCtx = nullptr;
|
|
}
|
|
|
|
// Write the output buffer to the output file
|
|
if (outputFileSize != fwrite(pOutputBuffer, 1, outputFileSize, outputFile))
|
|
{
|
|
printf("Error: Failed to write output file %s\n", outputFilename);
|
|
break;
|
|
}
|
|
|
|
result = true;
|
|
printf("Patch applied successfully. Output file: %s\n", outputFilename);
|
|
}while (false);
|
|
|
|
// Close all files and free buffers
|
|
freeFileBuffer(pInputBuffer);
|
|
pInputBuffer = nullptr;
|
|
|
|
if (pOutputBuffer != nullptr)
|
|
{
|
|
delete[] pOutputBuffer;
|
|
pOutputBuffer = nullptr;
|
|
}
|
|
|
|
if (patchFile != nullptr)
|
|
{
|
|
fclose(patchFile);
|
|
patchFile = nullptr;
|
|
}
|
|
if (outputFile != nullptr)
|
|
{
|
|
fclose(outputFile);
|
|
outputFile = nullptr;
|
|
}
|
|
|
|
return result;
|
|
} |