mumble/overlay/HardHook_x86.cpp

/* Copyright (C) 2005-2011, Thorvald Natvig <thorvald@natvig.com>

   All rights reserved.

   Redistribution and use in source and binary forms, with or without
   modification, are permitted provided that the following conditions
   are met:

   - Redistributions of source code must retain the above copyright notice,
     this list of conditions and the following disclaimer.
   - Redistributions in binary form must reproduce the above copyright notice,
     this list of conditions and the following disclaimer in the documentation
     and/or other materials provided with the distribution.
   - Neither the name of the Mumble Developers nor the names of its
     contributors may be used to endorse or promote products derived from this
     software without specific prior written permission.

   THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
   ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
   LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
   A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE FOUNDATION OR
   CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
   EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
   PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
   PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
   LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
   NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
   SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include "HardHook.h"
#include "ods.h"

void *HardHook::pCode = NULL;
unsigned int HardHook::uiCode = 0;

const int HardHook::CODEREPLACESIZE = 6;
const int HardHook::CODEPROTECTSIZE = 16;

/**
 * @brief Constructs a new hook without actually injecting.
 */
HardHook::HardHook() : bTrampoline(false), call(0), baseptr(NULL) {
	for (int i = 0; i < CODEREPLACESIZE; ++i) {
		orig[i] = replace[i] = 0;
	}

//	assert(CODEREPLACESIZE == sizeof(orig) / sizeof(orig[0]));
//	assert(CODEREPLACESIZE == sizeof(replace) / sizeof(replace[0]));
}

/**
 * @brief Constructs a new hook by injecting given replacement function into func.
 * @see HardHook::setup
 * @param func Funktion to inject replacement into.
 * @param replacement Function to inject into func.
 */
HardHook::HardHook(voidFunc func, voidFunc replacement)
	: bTrampoline(false), call(0), baseptr(NULL) {
	for (int i = 0; i < CODEREPLACESIZE; ++i)
		orig[i] = replace[i] = 0;
	setup(func, replacement);
}

/**
 * @return Number of extra bytes.
 */
static unsigned int modrmbytes(unsigned char a, unsigned char b) {
	unsigned char lower = (a & 0x0f);
	if (a >= 0xc0) {
		return 0;
	} else if (a >= 0x80) {
		if ((lower == 4) || (lower == 12))
			return 5;
		else
			return 4;
	} else if (a >= 0x40) {
		if ((lower == 4) || (lower == 12))
			return 2;
		else
			return 1;

	} else {
		if ((lower == 4) || (lower == 12)) {
			if ((b & 0x07) == 0x05)
				return 5;
			else
				return 1;
		} else if ((lower == 5) || (lower == 13))
			return 4;
		return 0;
	}
}

/**
 * @brief Tries to construct a trampoline from original code.
 *
 * A trampoline is the replacement code that features the original code plus
 * a jump back to the original instructions that follow.
 * It is called to execute the original behavior. As it is a replacement for
 * the original, the original can then be overwritten.
 * The size of the trampoline is at least CODEREPLACESIZE. Thus, CODEREPLACESIZE
 * bytes of the original code can afterwards be overwritten (and the trampoline
 * called after those instructions for the original logic).
 * CODEREPLACESIZE has to be smaller than CODEPROTECTSIZE.
 *
 * As commands must not be destroyed they have to be disassembled to get their length.
 * All encountered commands will be part of the trampoline and stored in pCode (shared
 * for all trampolines).
 *
 * If code is encountered that can not be moved into the trampoline (conditionals etc.)
 * construction fails and NULL is returned. If enough commands can be saved the
 * trampoline is finalized by appending a jump back to the original code. The return value
 * in this case will be the address of the newly constructed trampoline.
 *
 * pCode + offset to trampoline:
 *     [SAVED CODE FROM ORIGINAL which is >= CODEREPLACESIZE bytes][JUMP BACK TO ORIGINAL CODE]
 *
 * @param porig Original code
 * @return Pointer to trampoline on success. NULL if trampoline construction failed.
 */
void *HardHook::cloneCode(void **porig) {

	if (! pCode || uiCode > 4000) {
		pCode = VirtualAlloc(NULL, 4096, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
		uiCode = 0;
	}
	// If we have no memory to clone to, return.
	if (! pCode) {
		return NULL;
	}

	unsigned char *o = (unsigned char *) *porig;

	DWORD origProtect;
	if (!VirtualProtect(o, CODEPROTECTSIZE, PAGE_EXECUTE_READ, &origProtect)) {
		fods("HardHook: CloneCode failed; failed to make original code read and executable");
		return NULL;
	}

	// Follow relative jumps to next instruction. On execution it doesn't make
	// a difference if we actually perform all the jumps or directly jump to the
	// end of the chain. Hence these jumps need not be part of the trampoline.
	while (*o == 0xe9) { // JMP
		unsigned char *tmp = o;
		int *iptr = reinterpret_cast<int *>(o+1);
		o += *iptr + 5;

		fods("HardHook: CloneCode: Skipping jump from %p to %p", *porig, o);
		*porig = o;

		// Assume jump took us out of our read enabled zone, get rights for the new one
		DWORD tempProtect;
		VirtualProtect(tmp, CODEPROTECTSIZE, origProtect, &tempProtect);
		if (!VirtualProtect(o, CODEPROTECTSIZE, PAGE_EXECUTE_READ, &origProtect)) {
			fods("HardHook: CloneCode failed; failed to make jump target code read and executable");
			return NULL;
		}
	}

	unsigned char *n = (unsigned char *) pCode;
	n += uiCode;
	unsigned int idx = 0;

	do {
		unsigned char opcode = o[idx];
		unsigned char a = o[idx+1];
		unsigned char b = o[idx+2];
		unsigned int extra = 0;

		switch (opcode) {
			case 0x50: // PUSH
			case 0x51:
			case 0x52:
			case 0x53:
			case 0x54:
			case 0x55:
			case 0x56:
			case 0x57:
			case 0x58: // POP
			case 0x59:
			case 0x5a:
			case 0x5b:
			case 0x5c:
			case 0x5d:
			case 0x5e:
			case 0x5f:
				break;
			case 0x6a: // PUSH immediate
				extra = 1;
				break;
			case 0x68: // PUSH immediate
				extra = 4;
				break;
			case 0x81: // CMP immediate
				extra = modrmbytes(a,b) + 5;
				break;
			case 0x83:	// CMP
				extra = modrmbytes(a,b) + 2;
				break;
			case 0x8b:	// MOV
				extra = modrmbytes(a,b) + 1;
				break;
			default: {
				int rmop = ((a>>3) & 7);
				if (opcode == 0xff && rmop == 6) { // PUSH memory
					extra = modrmbytes(a,b) + 1;
					break;
				}

				fods("HardHook: CloneCode failed; Unknown opcode %02x at %d: %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x",
						opcode, idx, o[0], o[1], o[2], o[3], o[4], o[5], o[6], o[7], o[8], o[9], o[10], o[11]);
				DWORD tempProtect;
				VirtualProtect(o, CODEPROTECTSIZE, origProtect, &tempProtect);
				return NULL;
				break;
			}
		}

		n[idx] = opcode;
		++idx;

		for (unsigned int i = 0; i < extra; ++i)
			n[idx+i] = o[idx+i];
		idx += extra;

	} while (idx < CODEREPLACESIZE);

	DWORD tempProtect;
	VirtualProtect(o, CODEPROTECTSIZE, origProtect, &tempProtect);

	// Add a relative jmp back to the original code, to after the copied code
	n[idx++] = 0xe9;
	int *iptr = reinterpret_cast<int *>(&n[idx]);
	const int JMP_OP_SIZE = 5;
	int offs = o - n - JMP_OP_SIZE;
	*iptr = offs;
	idx += 4;

	uiCode += idx;

	FlushInstructionCache(GetCurrentProcess(), n, idx);

	fods("HardHook: trampoline creation successful at %p", n);

	return n;
}

/**
 * @brief Makes sure the given replacement function is run once func is called.
 *
 * Tries to construct a trampoline for the given function (@see HardHook::cloneCode)
 * and then injects replacement function calling code into the first 6 bytes of the
 * original function (@see HardHook::inject).
 *
 * @param func Pointer to function to redirect.
 * @param replacement Pointer to code to redirect to.
 */
void HardHook::setup(voidFunc func, voidFunc replacement) {
	if (baseptr)
		return;

	fods("HardHook: Setup: Asked to replace %p with %p", func, replacement);

	unsigned char *fptr = reinterpret_cast<unsigned char *>(func);
	unsigned char *nptr = reinterpret_cast<unsigned char *>(replacement);

	call = (voidFunc) cloneCode((void **) &fptr);

	if (call) {
		bTrampoline = true;
	} else {
		// Could not create a trampoline. Use alternative method instead.
		// This alternative method is dependant on the replacement code
		// restoring before calling the original. Otherwise we get a jump recursion
		bTrampoline = false;
		call = func;
	}

	DWORD origProtect;
	if (VirtualProtect(fptr, CODEPROTECTSIZE, PAGE_EXECUTE_READ, &origProtect)) {
		replace[0] = 0x68; // PUSH immediate        1 Byte
		unsigned char **iptr = reinterpret_cast<unsigned char **>(&replace[1]);
		*iptr = nptr;      // (imm. value = nptr)   4 Byte
		replace[5] = 0xc3; // RETN                  1 Byte

		// Save original 6 bytes at start of original function
		for (int i = 0; i < CODEREPLACESIZE; ++i)
			orig[i] = fptr[i];

		baseptr = fptr;

		inject(true);

		DWORD tempProtect;
		VirtualProtect(fptr, CODEPROTECTSIZE, origProtect, &tempProtect);
	} else {
		fods("HardHook: setup failed; failed to make original code read and executable");
	}
}

void HardHook::setupInterface(IUnknown *unkn, LONG funcoffset, voidFunc replacement) {
	fods("HardHook: setupInterface: Replacing %p function #%ld", unkn, funcoffset);
	void **ptr = reinterpret_cast<void **>(unkn);
	ptr = reinterpret_cast<void **>(ptr[0]);
	setup(reinterpret_cast<voidFunc>(ptr[funcoffset]), replacement);
}

void HardHook::reset() {
	baseptr = 0;
	bTrampoline = false;
	call = NULL;
	for (int i = 0; i < CODEREPLACESIZE; ++i) {
		orig[i] = replace[i] = 0;
	}
}

/**
 * @brief Injects redirection code into the target function.
 *
 * Replaces the first 6 Bytes of the function indicated by baseptr
 * with the replacement code previously generated (usually a jump
 * to mumble code). If a trampoline is available this injection is not needed
 * as control flow was already permanently redirected by HardHook::setup .
 *
 * @param force Perform injection even when trampoline is available.
 */
void HardHook::inject(bool force) {
	if (! baseptr)
		return;
	if (! force && bTrampoline)
		return;

	DWORD origProtect;
	if (VirtualProtect(baseptr, CODEREPLACESIZE, PAGE_EXECUTE_READWRITE, &origProtect)) {
		for (int i = 0; i < CODEREPLACESIZE; ++i) {
			baseptr[i] = replace[i]; // Replace with jump to new code
		}

		DWORD tempProtect;
		VirtualProtect(baseptr, CODEREPLACESIZE, origProtect, &tempProtect);

		FlushInstructionCache(GetCurrentProcess(), baseptr, CODEREPLACESIZE);
	}

	// Verify that the injection was successful
	for (int i = 0; i < CODEREPLACESIZE; ++i) {
		if (baseptr[i] != replace[i]) {
			fods("HardHook: Injection failure noticed at byte %d", i);
		}
	}
}

/**
 * @brief Restores the original code in a target function.
 *
 * Restores the first 6 Bytes of the function indicated by baseptr
 * from previously stored original code in orig. If a trampoline is available this
 * restoration is not needed as trampoline will correctly restore control
 * flow.
 *
 * @param force If true injection will be reverted even when trampoline is available.
 */
void HardHook::restore(bool force) {

	if (! baseptr)
		return;
	if (! force && bTrampoline)
		return;

	DWORD origProtect;
	if (VirtualProtect(baseptr, CODEREPLACESIZE, PAGE_EXECUTE_READWRITE, &origProtect)) {
		for (int i = 0; i < CODEREPLACESIZE; ++i)
			baseptr[i] = orig[i];
		DWORD tempProtect;
		VirtualProtect(baseptr, CODEREPLACESIZE, origProtect, &tempProtect);

		FlushInstructionCache(GetCurrentProcess(), baseptr, CODEREPLACESIZE);
	}
}

void HardHook::print() {
	fods("HardHook: code replacement: %02x %02x %02x %02x %02x => %02x %02x %02x %02x %02x (currently effective: %02x %02x %02x %02x %02x)",
	     orig[0], orig[1], orig[2], orig[3], orig[4],
	     replace[0], replace[1], replace[2], replace[3], replace[4],
	     baseptr[0], baseptr[1], baseptr[2], baseptr[3], baseptr[4]);
}

/**
 * @brief Checks whether injected code is in good shape and injects if not yet injected.
 *
 * If injected code is not found injection is attempted unless 3rd party overwrote
 * original code at injection location.
 */
void HardHook::check() {
	if (memcmp(baseptr, replace, CODEREPLACESIZE) != 0) {
		// The instructions do not match our replacement instructions
		// If they match the original code, inject our hook.
		if (memcmp(baseptr, orig, CODEREPLACESIZE) == 0) {
			fods("HardHook: Reinjecting hook into function %p", baseptr);
			inject(true);
		} else {
			fods("HardHook: Function %p replaced by third party. Lost injected hook.");
		}
	}
}