path: root/win32/droplet/replace.cpp

/*
 * AUTHOR
 * N. Nielsen
 *
 * LICENSE
 * This software is in the public domain.
 *
 * The software is provided "as is", without warranty of any kind,
 * express or implied, including but not limited to the warranties
 * of merchantability, fitness for a particular purpose, and
 * noninfringement. In no event shall the author(s) be liable for any
 * claim, damages, or other liability, whether in an action of
 * contract, tort, or otherwise, arising from, out of, or in connection
 * with the software or the use or other dealings in the software.
 *
 * SUPPORT
 * Send bug reports to: <nielsen@memberwebs.com>
 */

#include "stdafx.h"
#include "Replace.h"
#include "common/errutil.h"
#include "common/rliberr.h"
#include "common/mystring.h"


// Error Macros: -------------------------------------------------
//  Display error messages and all that
//  Only valid in this context
//  Each of these jumps to the 'cleanup' label for wrapping
//  up any open objects.

#define RETURN(v) \
	{ ret = (v); goto cleanup; }
#define WINERR(f) \
	{ HRESULT hr = HRESULT_FROM_WIN32(::GetLastError()); RETURN(errorMessage(m_dlg.m_hWnd, hr, f)); }
#define WINERR_1(f, a) \
	{ HRESULT hr = HRESULT_FROM_WIN32(::GetLastError()); RETURN(errorMessage(m_dlg.m_hWnd, hr, f, a)); }
#define WINERR_E(e, f) \
	{ RETURN(errorMessage(m_dlg.m_hWnd, e, f)); }
#define REPERR(c, s) \
	{ RETURN(rlibError(m_dlg.m_hWnd, c, s)); }


// The extension appended to backup or temp files
LPCTSTR kBackupExt = _T(".x_r");

// (Con|De)struction: -------------------------------------------

Replace::Replace()
{
	m_outFile = NULL;
}

Replace::~Replace()
{
	ASSERT(m_outFile == NULL);
}


// initialize: ---------------------------------------------------
//  Loads the script from resources and sets up dialog
//  in another thread

HRESULT Replace::initialize()
{
	HRESULT ret = S_OK;

	{
		if(!m_droplet.load(GetModuleHandle(NULL)))
		{
			errorMessage(NULL, 0, _T("Couldn't access droplet rep script. The droplet is probably corrupted."));
			RETURN(E_FAIL);
		}

		m_dlg.setTitle(m_droplet.getTitle());

		r_context& ctx = m_droplet.getContext();

		int r = repInit(&ctx);
		if(r < 0)
			REPERR(r, &(ctx.script));

		// This starts the dialog in another thread
		// in order to keep it responsive
		m_dlg.startDialog();
	}

cleanup:
	return ret;
}


// terminate: ----------------------------------------------------
//  Undoes above

HRESULT Replace::terminate()
{
	m_dlg.stopDialog();
	r_context& ctx = m_droplet.getContext();
	repFree(&ctx);
	return S_OK;
}


// writeFile: ----------------------------------------------------
//  The callback called from within rlib when data output
//  is required. The argument passed through the r_stream
//  is a pointer to the Replace object.

int Replace::writeFile(r_stream* stream, byte* data, size_t len)
{
	Replace* replace = (Replace*)stream->arg;

	// The m_outFile member should be set (and only set)
	// whenever we're inside rlibRun
	ASSERT(replace->m_outFile != NULL);
	DWORD written = 0;

	// Double check that everything was written
	return WriteFile(replace->m_outFile, data, len, &written, NULL) && len == written;
}

// matchStatus: --------------------------------------------------
//  Callback whenever something is replaced in rlib
//  Updates the dialog.

int Replace::matchStatus(r_stream* stream, r_replace* repl)
{
	Replace* replace = (Replace*)stream->arg;
	ASSERT_PTR(replace);

	replace->m_dlg.onReplaced();
	return replace->m_dlg.isCancelled() ? 1 : 0;
}


// replaceFolder: ------------------------------------------------
//  Recursively processes a folder using replaceSingleFile

HRESULT Replace::replaceFolder(LPCTSTR folder)
{
	HANDLE hFindFile = NULL;
	WIN32_FIND_DATA findData;
	TCHAR old[260];			//  Save the current folder (as we'll change it)
	HRESULT ret = S_FALSE;

	{
		//  Backup the current directory
		if(!GetCurrentDirectory(260, old))
			WINERR("Couldn't get the current directory.");

		if(!SetCurrentDirectory(folder))
			WINERR_1("Couldn't change the current directory to %s", folder);

		if(hFindFile = FindFirstFile(_T("*.*"), &findData))
		{
			// And for every file
			do
			{
				HRESULT r = S_FALSE;
				if(findData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)
				{
					if(findData.cFileName[0] == '.' &&
					   (findData.cFileName[1] == '.' || findData.cFileName[1] == '\0'))
					   continue;

					r = replaceFolder(findData.cFileName);
				}
				else
				{
					r = replaceSingleFile(findData.cFileName);
				}

				if(FAILED(r))
					RETURN(r);
				if(r == S_OK)
					ret = r;
			}
			while(FindNextFile(hFindFile, &findData));
		}

		// Make sure we didn't miss anything
		if(::GetLastError() != 0 && ::GetLastError() != ERROR_NO_MORE_FILES)
			WINERR_1("Couldn't list the directory %s", folder);

		// Change directory back
		SetCurrentDirectory(old);
	}

cleanup:
	if(hFindFile)
		FindClose(hFindFile);

	return ret;
}


// replaceSingleFile: ----------------------------------------------
//  Process a file into a temp file or backup etc...
//  Calls replaceFilename for actual processing

HRESULT Replace::replaceSingleFile(LPCTSTR fileName)
{
	HRESULT ret = S_FALSE;
	string temp;
	USES_CONVERSION;

	r_context& ctx = m_droplet.getContext();

	{
		// Update the dialog display
		m_dlg.onNewFile(fileName);

		// This lets us use the filename from within the script
		ret = rlibSetVar(&(ctx.stream), "FILENAME", T2CA(fileName));
		if(ret < 0)
			REPERR(ret, &(ctx.script));

		// Make a temp/backup name
		temp.assign(fileName);
		temp.append(kBackupExt);

		// Now we do this 2 ways, if we're keeping backups then...
		if(m_droplet.keepBackups())
		{
			DeleteFile(temp);

			// Rename the original file to the backup
			if(!MoveFile(fileName, temp))
				WINERR_1("There was an error while creating the backup file: %s", temp.c_str());

			// Do a replacement operation back to the original
			// from the backup
			ret = replaceFilename(temp, fileName);

			// If there were no replacements
			if(ret != S_OK)
			{
				DeleteFile(fileName);

				// And rename the backup file back
				if(!MoveFile(temp, fileName))
					WINERR_1("There was an error while renaming the file: %s", fileName);
			}
		}

		// No backups, much faster (when there's no replacements)!
		else
		{
			// Do a replacement operation to a temp file
			ret = replaceFilename(fileName, temp);

			// If there were replacements
			if(ret == S_OK)
			{
				DeleteFile(fileName);

				// Rename temp to original
				if(!MoveFile(temp, fileName))
					WINERR_1("There was an error while renaming the file: %s", fileName);
			}

			// If no replacements
			else
			{
				// Remove temp file
				DeleteFile(temp);
			}
		}
	}

cleanup:

	return ret;
}

// replaceFilename: -------------------------------------------------
//  Process one file into another
//  Open the files and passes them to replaceFile

HRESULT Replace::replaceFilename(LPCTSTR fileIn, LPCTSTR fileOut)
{
	HANDLE in = INVALID_HANDLE_VALUE;
	HANDLE out = INVALID_HANDLE_VALUE;
	HRESULT ret = S_FALSE;

	{
		// Basic readonly open
		in = CreateFile(fileIn, GENERIC_READ, FILE_SHARE_READ, NULL,
						OPEN_EXISTING, 0, NULL);
		if(in == INVALID_HANDLE_VALUE)
			WINERR_1("There was an error opening the file: %s", fileIn);

		// We need to know whether we're NT or not below
		OSVERSIONINFO os;
		os.dwOSVersionInfoSize = sizeof(OSVERSIONINFO);
		GetVersionEx(&os);

		// This creates a new file for writing
		// If we're running on NT, then use exactly the same
		// attributes and permissions as original file
		out = CreateFile(fileOut, GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, 0,
						 os.dwPlatformId & VER_PLATFORM_WIN32_NT ? in : NULL);
		if(out == INVALID_HANDLE_VALUE)
			WINERR_1("There was an error opening the file: %s", fileIn);


		//  Do replacements!
		ret = replaceFile(in, out);

	}

cleanup:

	if(in != INVALID_HANDLE_VALUE)
		CloseHandle(in);
	if(out != INVALID_HANDLE_VALUE)
		CloseHandle(out);

	return ret;
}


// replaceFile: --------------------------------------------------------
//  Process one open file handle into another
//  This memory maps the input file and passes it to replaceBuffer

HRESULT Replace::replaceFile(HANDLE in, HANDLE out)
{
	HANDLE mapping = NULL;
	void* data = NULL;
	HRESULT ret = S_FALSE;

	{
		// Memory map the input file
		mapping = CreateFileMapping(in, NULL, PAGE_READONLY, 0, 0, NULL);
		if(!mapping)
			WINERR("There was an error reading the input file.");

		data = MapViewOfFile(mapping, FILE_MAP_READ, 0, 0, 0);
		if(!data)
			WINERR("There was an error reading the input file.");

		size_t size = GetFileSize(in, NULL);
		if(size == INVALID_FILE_SIZE)
			WINERR("There was an error reading the input file.");

		// And pass it on
		ret = replaceBuffer((byte*)data, size, out);
	}

cleanup:

	if(data)
		UnmapViewOfFile(data);
	if(mapping)
		CloseHandle(mapping);

	return ret;
}


// replaceBuffer: ----------------------------------------------------
//  Process a memory buffer into a file for output
//  Uses the rlib functions directly

HRESULT Replace::replaceBuffer(byte* data, size_t size, HANDLE out)
{
	ASSERT(m_outFile == NULL);

	size_t blockSize = m_droplet.getBuffSize() * 2;	//  Block size
	size_t batchSize;							//  Current batch size
	int r = R_IN;
	bool dirty = false;
	HRESULT ret = S_OK;

	{
		// If there's no blocksize then we process the
		// entire file at once
		if(blockSize == 0)
			blockSize = size;

		batchSize = blockSize;

		// Setup the rlib stream
		r_context& ctx = m_droplet.getContext();

		ctx.stream.nextIn = data;
		ctx.stream.availIn = blockSize;
		ctx.stream.fWrite = writeFile;
		ctx.stream.arg = this;
		ctx.stream.fMatch = matchStatus;

		// This is used by writeFile
		m_outFile = out;

		//  While we have more data to put in ...
		while(size > 0 || r == R_IN)
		{
			//  If rlib wants data then give it
			if(r == R_IN)
			{
 				batchSize = (blockSize > size) ? size : blockSize;

				ctx.stream.nextIn = data;
				ctx.stream.availIn = batchSize;

			}

			//  call rlib
			r = rlibRun(&(ctx.stream), &(ctx.script), (size - batchSize) <= 0);

			// Oops!
			if(r < 0)
				REPERR(r, &(ctx.script));

			if(ctx.stream.total > 0)
				dirty = true;

			// Increment last batch
			data += (batchSize - ctx.stream.availIn);
 			size -= (batchSize - ctx.stream.availIn);
		}

		// Clears and prepares for next file
		rlibClear(&ctx.stream);
	}

cleanup:

	m_outFile = NULL;

	if(FAILED(ret))
		return ret;

	return dirty ? S_OK : S_FALSE;
}