dnf-utils/sgutil/sgu_mode_apply_patch.cpp

#include "sgu_stdafx.h"
#include "sgu_mode_apply_patch.h"
#include "sgu_utils.h"

#include "bzlib.h"

extern "C"
{
#include "bspatch.h"
}


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;
}

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;
}

bool applyPatch(const char* inputFilename, const char* patchFilename, const char* outputFilename)
{
	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;
		}

		int bz2err;
		BZFILE* bzHandle = BZ2_bzReadOpen(&bz2err, patchFile, 0, 0, NULL, 0);
		if (bzHandle == nullptr)
		{
			printf("BZ2_bzReadOpen, bz2err=%d", bz2err);
			break;
		}

		// Apply the patch
		struct bspatch_stream stream;
		stream.read = bzPatchFileReader;
		stream.opaque = bzHandle;
		if (bspatch(pInputBuffer, inputFileSize, pOutputBuffer, outputFileSize, &stream))
		{
			printf("Error: Failed to apply patch\n");
			break;
		}

		// Clean up the bzip2 reads
		BZ2_bzReadClose(&bz2err, bzHandle);

		// 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;
}