/*
 * Copyright (c) 2011 Nokia and others.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * which accompanies this distribution, and is available at
 * http://www.eclipse.org/legal/epl-v10.html
 *
 * Contributors:
 * Nokia - Initial API and implementation
 */

/*
 * WinHWPageProtMgr.cpp
 *
 *
 *  Created on: Aug 23, 2011
 *      Author: bkirk
 */

#include "WinHWPageProtMgr.h"

#include "BreakpointsService.h"
#include "WinProcess.h"

#include <winbase.h>
#include <winnt.h>

#define PAGE_INVALID 0

inline
DWORD WinHWPageProtMgr::sPageProtModeMap(unsigned char i) {
static const DWORD map[9]
  = { PAGE_INVALID, PAGE_NOACCESS, PAGE_READONLY, PAGE_NOACCESS,
	  PAGE_INVALID, PAGE_INVALID, PAGE_INVALID, PAGE_INVALID, PAGE_INVALID
    };
	return map[i];
}

WinHWPageProtMgr::WinHWPageProtMgr() {
}

bool WinHWPageProtMgr::IsAddressInProtectedPage(const ContextAddress& address) {
	ProtectedPages::iterator p;
	for (p = protectedPages.begin(); p != protectedPages.end(); ++p ) {
		ProtectedPage& page = p->second;
		if (address >= (ContextAddress)page.baseAddr
			&& address < (ContextAddress)page.baseAddr + page.size)
		{
			return true;
		}
	}
	return false;
}

TBreakpoint* WinHWPageProtMgr::FindWatchpointForAddress(
		const HANDLE& procHandle, const ContextAddress& address) {
	ProtectedPageWatchpoints::iterator w;
	for (w = watchpoints.begin(); w != watchpoints.end(); ++w) {
		TBreakpoint*& wp = w->second;
		if (address >= wp->address && address < wp->address + wp->size) {
			return wp;
		}
	}
	return NULL;
}

int WinHWPageProtMgr::SetPageProtectWatchpoint(TBreakpoint* wp) {
	if (!wp || !wp->process)
		return ERR_INV_CONTEXT;
	DWORD mode = sPageProtModeMap(wp->accessMode);
	if (mode == PAGE_INVALID)
		return ERR_PGPRO_INVALID_MODE;
	int err = ProtectPageForWatchpointAddress(wp->process->GetProcessHandle(),
								   wp->address, wp->size, mode);
	if (err != 0)
		return err;

	watchpoints.insert(AddressWatchpointPair(wp->address, wp));
	return 0;
}

int WinHWPageProtMgr::ClearPageProtectWatchpoint(TBreakpoint* wp) {
	if (!wp || !wp->process)
		return ERR_INV_CONTEXT;
	int err = UnProtectPageForWatchpoint(wp);
	if (err != 0)
		return err;

	watchpoints.erase(wp->address);
	return 0;
}


int WinHWPageProtMgr::ProtectPageForWatchpointAddress(const HANDLE& procHandle,
		const ContextAddress& address, const size_t size, const DWORD mode) {
	if (mode != PAGE_NOACCESS && mode != PAGE_READONLY)
		return ERR_PGPRO_INVALID_MODE;
	// use an already protected page if appropriate
	ProtectedPages::iterator iter;
	for (iter = protectedPages.begin(); iter != protectedPages.end(); ++iter ) {
		ProtectedPage& page = iter->second;
		if (address+size >= (ContextAddress)page.baseAddr
			&& address < (ContextAddress)page.baseAddr + page.size) {
			// if the incoming mode is stronger, it must be enforced.
			if (page.protectMode > mode) {
				DWORD priorDebugProtection;
				if (FAILED(::VirtualProtectEx(procHandle, page.baseAddr, page.size,
											  mode, &priorDebugProtection)))
					return ::GetLastError();
			}
			page.AddRef();
			return 0;
		}
	}

	MEMORY_BASIC_INFORMATION mbi;
	if (FAILED(::VirtualQueryEx(procHandle, (LPVOID)address, &mbi, sizeof(MEMORY_BASIC_INFORMATION))))
		return ::GetLastError();
	if (mbi.State != MEM_COMMIT)
		return ERR_PGPRO_UNSET;

	DWORD originalProtection;
	if (FAILED(::VirtualProtectEx(procHandle, mbi.BaseAddress, mbi.RegionSize,
								  mode, &originalProtection)))
		return ::GetLastError();

	protectedPages.insert(AddressPagePair((ContextAddress)mbi.BaseAddress,
			ProtectedPage(mbi.BaseAddress, mbi.RegionSize, originalProtection, mode)));

	return 0;
}

int WinHWPageProtMgr::DisableForPagesContainingAddress(const HANDLE& procH,
		const ContextAddress& addr, const DWORD size, const DWORD mode) {
	bool found = false;
	ProtectedPages::iterator p;
	for (p = protectedPages.begin(); p != protectedPages.end(); ++p ) {
		ProtectedPage& pg = p->second;
		const ContextAddress pgLo = (ContextAddress)pg.baseAddr;
		const ContextAddress pgHi = pgLo + pg.size;
		if (((addr + size) >= pgLo) && (addr < pgHi)) {
			if (!(mode == PAGE_READONLY && pg.protectMode == PAGE_READONLY)) {
				DWORD oldP;
				if (FAILED(::VirtualProtectEx(procH, pg.baseAddr, pg.size,
											  pg.origProtectMode, &oldP)))
					return ::GetLastError();
				found = true;
			}
			if (size == 0 || ((addr >= pgLo) && ((addr + size) < pgHi)))
				return 0;
		}
	}
	return found ? 0 : ERR_PGPRO_NO_ADDR;
}

int WinHWPageProtMgr::EnableForPagesContainingAddress(const HANDLE& procH,
		const ContextAddress& addr, const DWORD size, const DWORD mode) {
	bool found = false;
	ProtectedPages::iterator p;
	for (p = protectedPages.begin(); p != protectedPages.end(); ++p ) {
		ProtectedPage& pg = p->second;
		const ContextAddress pgLo = (ContextAddress)pg.baseAddr;
		const ContextAddress pgHi = pgLo + pg.size;
		if (((addr + size) >= pgLo) && (addr < pgHi)) {
			if (!(mode == PAGE_READONLY && pg.protectMode == PAGE_READONLY)) {
				DWORD oldP;
				if (FAILED(::VirtualProtectEx(procH, pg.baseAddr, pg.size,
											  pg.protectMode, &oldP)))
					return ::GetLastError();
				found = true;
			}
			if (size == 0 || ((addr >= pgLo) && ((addr + size) < pgHi)))
				return 0;
		}
	}
	return ERR_PGPRO_NO_ADDR;
}

int WinHWPageProtMgr::UnProtectPageForWatchpoint(const TBreakpoint* wp) {
	const HANDLE& procH = wp->process->GetProcessHandle();
	const ContextAddress& address = wp->address;
	const size_t size = wp->size;

	ProtectedPages::iterator p;
	for (p = protectedPages.begin(); p != protectedPages.end(); ++p ) {
		ProtectedPage& pg = p->second;
		if (address+size >= (ContextAddress)pg.baseAddr
			&& address < (ContextAddress)pg.baseAddr + pg.size)
		{
			int err = 0;
			DWORD priorDebugProtection;
			if (0 == pg.DecRef()) {
				if (FAILED(::VirtualProtectEx(procH, pg.baseAddr, pg.size,
											  pg.origProtectMode, &priorDebugProtection)))
					err = ::GetLastError();
				else
					protectedPages.erase(p);
			} else {
				err = ChangePageProtection(procH, wp, pg);
			}
			return err;
		}
	}
	return ERR_PGPRO_NO_ADDR;
}


/*
 * contract: call only if refCount on passed page is non-0 and if
 * watchpoint for page being unprotected has valid process
 * see if all remaining watchpoints in the passed page have a weaker
 * access setting than the watchpoint that's going away, and if so
 * change the protection.
 */
int WinHWPageProtMgr::ChangePageProtection(const HANDLE& procH,
		const TBreakpoint* wpGone, ProtectedPage& pg) {

	if (pg.refCount == 0)
		return ERR_PGPRO_UNCHANGED;

	if (!wpGone)
		return ERR_INV_CONTEXT;

    if (PAGE_READONLY == sPageProtModeMap(wpGone->accessMode))
    	return 0;

    DWORD otherProt = PAGE_INVALID;
	ProtectedPageWatchpoints::iterator w;
	for (w = watchpoints.begin(); w != watchpoints.end(); ++w) {
		TBreakpoint*& wp = w->second;
		if (wp == wpGone		// looking for other WPs; skip "this" one
			|| wp->process != wpGone->process
			|| wp->address + wp->size <= (ContextAddress)pg.baseAddr
			|| wp->address >= (ContextAddress)pg.baseAddr + pg.size)
			continue;

		DWORD wpProt = sPageProtModeMap(wp->accessMode);
		if (PAGE_INVALID == otherProt) {
			otherProt = wpProt;	// first "other-WP" in this page
		} else if (wpProt < otherProt) {
			otherProt = PAGE_NOACCESS;
			break;		// another WP in same page w/different access
		} else if (otherProt < wpProt) {
			break;		// another WP in same page w/different access
		}
		// getting here means all remaining WPs have the
		// the same protection as one another so far
	}

	// at this point, otherProt value possibilities are:
	// PAGE_NOACCESS:
	//		some other WP on this pg was already forcing NOACCESS
	// PAGE_READONLY:
	//		all other WPs on this pg were read-only; mode can be relaxed
	// PAGE_INVALID:
	//		well ... shouldn't happen w/refCount check at the top
	if (otherProt != PAGE_READONLY)
		return 0;

	DWORD dw;
	if (FAILED(::VirtualProtectEx(wpGone->process->GetProcessHandle(),
								  pg.baseAddr, pg.size, PAGE_READONLY, &dw)))
		return ::GetLastError();
	pg.protectMode = PAGE_READONLY;

    return 0;
}