/*
 * Copyright 2009, Ingo Weinhold, ingo_weinhold@gmx.de.
 * Distributed under the terms of the MIT License.
 */

#include "BreakpointManager.h"

#include <algorithm>

#include <AutoDeleter.h>

#include <kernel.h>
#include <util/AutoLock.h>
#include <vm/vm.h>


//#define TRACE_BREAKPOINT_MANAGER
#ifdef TRACE_BREAKPOINT_MANAGER
#	define TRACE(x...) dprintf(x)
#else
#	define TRACE(x...) do {} while (false)
#endif


// soft limit for the number of breakpoints
const int32 kMaxBreakpointCount = 10240;


BreakpointManager::InstalledBreakpoint::InstalledBreakpoint(addr_t address)
	:
	breakpoint(NULL),
	address(address)
{
}


// #pragma mark -


BreakpointManager::BreakpointManager()
	:
	fBreakpointCount(0),
	fWatchpointCount(0)
{
	rw_lock_init(&fLock, "breakpoint manager");
}


BreakpointManager::~BreakpointManager()
{
	WriteLocker locker(fLock);

	// delete the installed breakpoint objects
	BreakpointTree::Iterator it = fBreakpoints.GetIterator();
	while (InstalledBreakpoint* installedBreakpoint = it.Next()) {
		it.Remove();

		// delete underlying software breakpoint
		if (installedBreakpoint->breakpoint->software)
			delete installedBreakpoint->breakpoint;

		delete installedBreakpoint;
	}

	// delete the watchpoints
	while (InstalledWatchpoint* watchpoint = fWatchpoints.RemoveHead())
		delete watchpoint;

	// delete the hardware breakpoint objects
	while (Breakpoint* breakpoint = fHardwareBreakpoints.RemoveHead())
		delete breakpoint;

	rw_lock_destroy(&fLock);
}


status_t
BreakpointManager::Init()
{
	// create objects for the hardware breakpoints
	for (int32 i = 0; i < DEBUG_MAX_BREAKPOINTS; i++) {
		Breakpoint* breakpoint = new(std::nothrow) Breakpoint;
		if (breakpoint == NULL)
			return B_NO_MEMORY;

		breakpoint->address = 0;
		breakpoint->installedBreakpoint = NULL;
		breakpoint->used = false;
		breakpoint->software = false;

		fHardwareBreakpoints.Add(breakpoint);
	}

	return B_OK;
}


status_t
BreakpointManager::InstallBreakpoint(void* _address)
{
	const addr_t address = (addr_t)_address;

	WriteLocker locker(fLock);

	if (fBreakpointCount >= kMaxBreakpointCount)
		return B_BUSY;

	// check whether there's already a breakpoint at the address
	InstalledBreakpoint* installed = fBreakpoints.Lookup(address);
	if (installed != NULL)
		return B_BAD_VALUE;

	// create the breakpoint object
	installed = new(std::nothrow) InstalledBreakpoint(address);
	if (installed == NULL)
		return B_NO_MEMORY;
	ObjectDeleter<InstalledBreakpoint> installedDeleter(installed);

	// If we still have enough hardware breakpoints left, install a hardware
	// breakpoint.
	Breakpoint* breakpoint = _GetUnusedHardwareBreakpoint(false);
	if (breakpoint != NULL) {
		status_t error = _InstallHardwareBreakpoint(breakpoint, address);
		if (error != B_OK)
			return error;

		breakpoint->installedBreakpoint = installed;
		installed->breakpoint = breakpoint;
	} else {
		// install a software breakpoint
		status_t error = _InstallSoftwareBreakpoint(installed, address);
		if (error != B_OK)
			return error;
	}

	fBreakpoints.Insert(installed);
	installedDeleter.Detach();
	fBreakpointCount++;

	return B_OK;
}


status_t
BreakpointManager::UninstallBreakpoint(void* _address)
{
	const addr_t address = (addr_t)_address;

	WriteLocker locker(fLock);

	InstalledBreakpoint* installed = fBreakpoints.Lookup(address);
	if (installed == NULL)
		return B_BAD_VALUE;

	if (installed->breakpoint->software)
		_UninstallSoftwareBreakpoint(installed->breakpoint);
	else
		_UninstallHardwareBreakpoint(installed->breakpoint);

	fBreakpoints.Remove(installed);
	delete installed;
	fBreakpointCount--;

	return B_OK;
}


status_t
BreakpointManager::InstallWatchpoint(void* _address, uint32 type, int32 length)
{
	const addr_t address = (addr_t)_address;

	WriteLocker locker(fLock);

	InstalledWatchpoint* watchpoint = _FindWatchpoint(address);
	if (watchpoint != NULL)
		return B_BAD_VALUE;

#if DEBUG_SHARED_BREAK_AND_WATCHPOINTS
	// We need at least one hardware breakpoint for our breakpoint management.
	if (fWatchpointCount + 1 >= DEBUG_MAX_WATCHPOINTS)
		return B_BUSY;
#else
	if (fWatchpointCount >= DEBUG_MAX_WATCHPOINTS)
		return B_BUSY;
#endif

	watchpoint = new(std::nothrow) InstalledWatchpoint;
	if (watchpoint == NULL)
		return B_NO_MEMORY;
	ObjectDeleter<InstalledWatchpoint> watchpointDeleter(watchpoint);

	status_t error = _InstallWatchpoint(watchpoint, address, type, length);
	if (error != B_OK)
		return error;

	fWatchpoints.Add(watchpointDeleter.Detach());
	fWatchpointCount++;
	return B_OK;
}


status_t
BreakpointManager::UninstallWatchpoint(void* address)
{
	WriteLocker locker(fLock);

	InstalledWatchpoint* watchpoint = _FindWatchpoint((addr_t)address);
	if (watchpoint == NULL)
		return B_BAD_VALUE;

	ObjectDeleter<InstalledWatchpoint> deleter(watchpoint);
	fWatchpoints.Remove(watchpoint);
	fWatchpointCount--;

	return _UninstallWatchpoint(watchpoint);
}


void
BreakpointManager::RemoveAllBreakpoints()
{
	WriteLocker locker(fLock);

	// remove the breakpoints
	BreakpointTree::Iterator it = fBreakpoints.GetIterator();
	while (InstalledBreakpoint* installedBreakpoint = it.Next()) {
		it.Remove();

		// uninstall underlying hard/software breakpoint
		if (installedBreakpoint->breakpoint->software)
			_UninstallSoftwareBreakpoint(installedBreakpoint->breakpoint);
		else
			_UninstallHardwareBreakpoint(installedBreakpoint->breakpoint);

		delete installedBreakpoint;
	}

	// remove the watchpoints
	while (InstalledWatchpoint* watchpoint = fWatchpoints.RemoveHead()) {
		_UninstallWatchpoint(watchpoint);
		delete watchpoint;
	}
}


/*!	\brief Returns whether the given address can be accessed in principle.
	No check whether there's an actually accessible area is performed, though.
*/
/*static*/ bool
BreakpointManager::CanAccessAddress(const void* _address, bool write)
{
	const addr_t address = (addr_t)_address;

	// user addresses are always fine
	if (IS_USER_ADDRESS(address))
		return true;

	return false;
}


/*!	\brief Reads data from user memory.

	Tries to read \a size bytes of data from user memory address \a address
	into the supplied buffer \a buffer. If only a part could be read the
	function won't fail. The number of bytes actually read is return through
	\a bytesRead.

	\param address The user memory address from which to read.
	\param buffer The buffer into which to write.
	\param size The number of bytes to read.
	\param bytesRead Will be set to the number of bytes actually read.
	\return \c B_OK, if reading went fine. Then \a bytesRead will be set to
			the amount of data actually read. An error indicates that nothing
			has been read.
*/
status_t
BreakpointManager::ReadMemory(const void* _address, void* buffer, size_t size,
	size_t& bytesRead)
{
	const addr_t address = (addr_t)_address;

	ReadLocker locker(fLock);

	status_t error = _ReadMemory(address, buffer, size, bytesRead);
	if (error != B_OK)
		return error;

	// If we have software breakpoints installed, fix the buffer not to contain
	// any of them.

	// address of the first possibly intersecting software breakpoint
	const addr_t startAddress
		= std::max(address, (addr_t)DEBUG_SOFTWARE_BREAKPOINT_SIZE - 1)
			- (DEBUG_SOFTWARE_BREAKPOINT_SIZE - 1);

	BreakpointTree::Iterator it = fBreakpoints.GetIterator(startAddress, true,
		true);
	while (InstalledBreakpoint* installed = it.Next()) {
		Breakpoint* breakpoint = installed->breakpoint;
		if (breakpoint->address >= address + size)
			break;

		if (breakpoint->software) {
			// Software breakpoint intersects -- replace the read data with
			// the data saved in the breakpoint object.
			addr_t minAddress = std::max(breakpoint->address, address);
			size_t toCopy = std::min(address + size,
					breakpoint->address + DEBUG_SOFTWARE_BREAKPOINT_SIZE)
				- minAddress;
			memcpy((uint8*)buffer + (minAddress - address),
				breakpoint->softwareData + (minAddress - breakpoint->address),
				toCopy);
		}
	}

	return B_OK;
}


status_t
BreakpointManager::WriteMemory(void* _address, const void* _buffer, size_t size,
	size_t& bytesWritten)
{
	bytesWritten = 0;

	if (size == 0)
		return B_OK;

	addr_t address = (addr_t)_address;
	const uint8* buffer = (uint8*)_buffer;

	WriteLocker locker(fLock);

	// We don't want to overwrite software breakpoints, so things are a bit more
	// complicated. We iterate through the intersecting software breakpoints,
	// writing the memory between them normally, but skipping the breakpoints
	// itself. We write into their softwareData instead.

	// Get the first breakpoint -- if it starts before the address, we'll
	// handle it separately to make things in the main loop simpler.
	const addr_t startAddress
		= std::max(address, (addr_t)DEBUG_SOFTWARE_BREAKPOINT_SIZE - 1)
			- (DEBUG_SOFTWARE_BREAKPOINT_SIZE - 1);

	BreakpointTree::Iterator it = fBreakpoints.GetIterator(startAddress, true,
		true);
	InstalledBreakpoint* installed = it.Next();
	while (installed != NULL) {
		Breakpoint* breakpoint = installed->breakpoint;
		if (breakpoint->address >= address)
			break;

		if (breakpoint->software) {
			// We've got a breakpoint that is partially intersecting with the
			// beginning of the address range to write.
			size_t toCopy = std::min(address + size,
					breakpoint->address + DEBUG_SOFTWARE_BREAKPOINT_SIZE)
				- address;
			memcpy(breakpoint->softwareData + (address - breakpoint->address),
				buffer, toCopy);

			address += toCopy;
			size -= toCopy;
			bytesWritten += toCopy;
			buffer += toCopy;
		}

		installed = it.Next();
	}

	// loop through the breakpoints intersecting with the range
	while (installed != NULL) {
		Breakpoint* breakpoint = installed->breakpoint;
		if (breakpoint->address >= address + size)
			break;

		if (breakpoint->software) {
			// write the data up to the breakpoint (if any)
			size_t toCopy = breakpoint->address - address;
			if (toCopy > 0) {
				size_t chunkWritten;
				status_t error = _WriteMemory(address, buffer, toCopy,
					chunkWritten);
				if (error != B_OK)
					return bytesWritten > 0 ? B_OK : error;

				address += chunkWritten;
				size -= chunkWritten;
				bytesWritten += chunkWritten;
				buffer += chunkWritten;

				if (chunkWritten < toCopy)
					return B_OK;
			}

			// write to the breakpoint data
			toCopy = std::min(size, (size_t)DEBUG_SOFTWARE_BREAKPOINT_SIZE);
			memcpy(breakpoint->softwareData, buffer, toCopy);

			address += toCopy;
			size -= toCopy;
			bytesWritten += toCopy;
			buffer += toCopy;
		}

		installed = it.Next();
	}

	// write remaining data
	if (size > 0) {
		size_t chunkWritten;
		status_t error = _WriteMemory(address, buffer, size, chunkWritten);
		if (error != B_OK)
			return bytesWritten > 0 ? B_OK : error;

		bytesWritten += chunkWritten;
	}

	return B_OK;
}


void
BreakpointManager::PrepareToContinue(void* _address)
{
	const addr_t address = (addr_t)_address;

	WriteLocker locker(fLock);

	// Check whether there's a software breakpoint at the continuation address.
	InstalledBreakpoint* installed = fBreakpoints.Lookup(address);
	if (installed == NULL || !installed->breakpoint->software)
		return;

	// We need to replace the software breakpoint by a hardware one, or
	// we can't continue the thread.
	Breakpoint* breakpoint = _GetUnusedHardwareBreakpoint(true);
	if (breakpoint == NULL) {
		dprintf("Failed to allocate a hardware breakpoint.\n");
		return;
	}

	status_t error = _InstallHardwareBreakpoint(breakpoint, address);
	if (error != B_OK)
		return;

	_UninstallSoftwareBreakpoint(installed->breakpoint);

	breakpoint->installedBreakpoint = installed;
	installed->breakpoint = breakpoint;
}


BreakpointManager::Breakpoint*
BreakpointManager::_GetUnusedHardwareBreakpoint(bool force)
{
	// try to find a free one first
	for (BreakpointList::Iterator it = fHardwareBreakpoints.GetIterator();
			Breakpoint* breakpoint = it.Next();) {
		if (!breakpoint->used)
			return breakpoint;
	}

	if (!force)
		return NULL;

	// replace one by a software breakpoint
	for (BreakpointList::Iterator it = fHardwareBreakpoints.GetIterator();
			Breakpoint* breakpoint = it.Next();) {
		if (breakpoint->installedBreakpoint == NULL)
			continue;

		status_t error = _InstallSoftwareBreakpoint(
			breakpoint->installedBreakpoint, breakpoint->address);
		if (error != B_OK)
			continue;

		if (_UninstallHardwareBreakpoint(breakpoint) == B_OK)
			return breakpoint;
	}

	return NULL;
}


status_t
BreakpointManager::_InstallSoftwareBreakpoint(InstalledBreakpoint* installed,
	addr_t address)
{
	Breakpoint* breakpoint = new(std::nothrow) Breakpoint;
	if (breakpoint == NULL)
		return B_NO_MEMORY;
	ObjectDeleter<Breakpoint> breakpointDeleter(breakpoint);

	breakpoint->address = address;
	breakpoint->installedBreakpoint = installed;
	breakpoint->used = true;
	breakpoint->software = true;

	// save the memory where the software breakpoint shall be installed
	size_t bytesTransferred;
	status_t error = _ReadMemory(address, breakpoint->softwareData,
		DEBUG_SOFTWARE_BREAKPOINT_SIZE, bytesTransferred);
	if (error != B_OK)
		return error;
	if (bytesTransferred != DEBUG_SOFTWARE_BREAKPOINT_SIZE)
		return B_BAD_ADDRESS;

	// write the breakpoint code
	error = _WriteMemory(address, DEBUG_SOFTWARE_BREAKPOINT,
		DEBUG_SOFTWARE_BREAKPOINT_SIZE, bytesTransferred);
	if (error != B_OK)
		return error;

	if (bytesTransferred < DEBUG_SOFTWARE_BREAKPOINT_SIZE) {
		// breakpoint written partially only -- undo the written part
		if (bytesTransferred > 0) {
			size_t dummy;
			_WriteMemory(address, breakpoint->softwareData, bytesTransferred,
				dummy);
		}
		return B_BAD_ADDRESS;
	}

	installed->breakpoint = breakpoint;
	breakpointDeleter.Detach();

	TRACE("installed software breakpoint at %#lx\n", address);

	return B_OK;
}


status_t
BreakpointManager::_UninstallSoftwareBreakpoint(Breakpoint* breakpoint)
{
	size_t bytesWritten;
	_WriteMemory(breakpoint->address, breakpoint->softwareData,
		DEBUG_SOFTWARE_BREAKPOINT_SIZE, bytesWritten);

	TRACE("uninstalled software breakpoint at %#lx\n", breakpoint->address);

	delete breakpoint;
	return B_OK;
}


status_t
BreakpointManager::_InstallHardwareBreakpoint(Breakpoint* breakpoint,
	addr_t address)
{
	status_t error = arch_set_breakpoint((void*)address);
	if (error != B_OK)
		return error;

	// move to the tail of the list
	fHardwareBreakpoints.Remove(breakpoint);
	fHardwareBreakpoints.Add(breakpoint);

	TRACE("installed hardware breakpoint at %#lx\n", address);

	breakpoint->address = address;
	breakpoint->used = true;
	return B_OK;
}


status_t
BreakpointManager::_UninstallHardwareBreakpoint(Breakpoint* breakpoint)
{
	status_t error = arch_clear_breakpoint((void*)breakpoint->address);
	if (error != B_OK)
		return error;

	TRACE("uninstalled hardware breakpoint at %#lx\n", breakpoint->address);

	breakpoint->used = false;
	breakpoint->installedBreakpoint = NULL;
	return B_OK;
}


BreakpointManager::InstalledWatchpoint*
BreakpointManager::_FindWatchpoint(addr_t address) const
{
	for (InstalledWatchpointList::ConstIterator it = fWatchpoints.GetIterator();
		InstalledWatchpoint* watchpoint = it.Next();) {
		if (address == watchpoint->address)
			return watchpoint;
	}

	return NULL;
}


status_t
BreakpointManager::_InstallWatchpoint(InstalledWatchpoint* watchpoint,
	addr_t address, uint32 type, int32 length)
{
#if DEBUG_SHARED_BREAK_AND_WATCHPOINTS
	// We need a hardware breakpoint.
	watchpoint->breakpoint = _GetUnusedHardwareBreakpoint(true);
	if (watchpoint->breakpoint == NULL) {
		dprintf("Failed to allocate a hardware breakpoint for watchpoint.\n");
		return B_BUSY;
	}
#endif

	status_t error = arch_set_watchpoint((void*)address, type, length);
	if (error != B_OK)
		return error;

	watchpoint->address = address;

#if DEBUG_SHARED_BREAK_AND_WATCHPOINTS
	watchpoint->breakpoint->used = true;
#endif

	return B_OK;
}


status_t
BreakpointManager::_UninstallWatchpoint(InstalledWatchpoint* watchpoint)
{
#if DEBUG_SHARED_BREAK_AND_WATCHPOINTS
	watchpoint->breakpoint->used = false;
#endif

	return arch_clear_watchpoint((void*)watchpoint->address);
}


status_t
BreakpointManager::_ReadMemory(const addr_t _address, void* _buffer,
	size_t size, size_t& bytesRead)
{
	const uint8* address = (const uint8*)_address;
	uint8* buffer = (uint8*)_buffer;

	// check the parameters
	if (!CanAccessAddress(address, false))
		return B_BAD_ADDRESS;
	if (size <= 0)
		return B_BAD_VALUE;

	// If the region to be read crosses page boundaries, we split it up into
	// smaller chunks.
	status_t error = B_OK;
	bytesRead = 0;
	while (size > 0) {
		// check whether we're still in user address space
		if (!CanAccessAddress(address, false)) {
			error = B_BAD_ADDRESS;
			break;
		}

		// don't cross page boundaries in a single read
		int32 toRead = size;
		int32 maxRead = B_PAGE_SIZE - (addr_t)address % B_PAGE_SIZE;
		if (toRead > maxRead)
			toRead = maxRead;

		error = user_memcpy(buffer, address, toRead);
		if (error != B_OK)
			break;

		bytesRead += toRead;
		address += toRead;
		buffer += toRead;
		size -= toRead;
	}

	// If reading fails, we only fail, if we haven't read anything yet.
	if (error != B_OK) {
		if (bytesRead > 0)
			return B_OK;
		return error;
	}

	return B_OK;
}


status_t
BreakpointManager::_WriteMemory(addr_t _address, const void* _buffer,
	size_t size, size_t& bytesWritten)
{
	uint8* address = (uint8*)_address;
	const uint8* buffer = (const uint8*)_buffer;

	// check the parameters
	if (!CanAccessAddress(address, true))
		return B_BAD_ADDRESS;
	if (size <= 0)
		return B_BAD_VALUE;

	// If the region to be written crosses area boundaries, we split it up into
	// smaller chunks.
	status_t error = B_OK;
	bytesWritten = 0;
	while (size > 0) {
		// check whether we're still in user address space
		if (!CanAccessAddress(address, true)) {
			error = B_BAD_ADDRESS;
			break;
		}

		// get the area for the address (we need to use _user_area_for(), since
		// we're looking for a user area)
		area_id area = _user_area_for(address);
		if (area < 0) {
			TRACE("BreakpointManager::_WriteMemory(): area not found for "
				"address: %p: %lx\n", address, area);
			error = area;
			break;
		}

		area_info areaInfo;
		status_t error = get_area_info(area, &areaInfo);
		if (error != B_OK) {
			TRACE("BreakpointManager::_WriteMemory(): failed to get info for "
				"area %ld: %lx\n", area, error);
			error = B_BAD_ADDRESS;
			break;
		}

		// restrict this round of writing to the found area
		int32 toWrite = size;
		int32 maxWrite = (uint8*)areaInfo.address + areaInfo.size - address;
		if (toWrite > maxWrite)
			toWrite = maxWrite;

		// if the area is read-only, we temporarily need to make it writable
		bool protectionChanged = false;
		if (!(areaInfo.protection & (B_WRITE_AREA | B_KERNEL_WRITE_AREA))) {
			error = set_area_protection(area,
				areaInfo.protection | B_WRITE_AREA);
			if (error != B_OK) {
				TRACE("BreakpointManager::_WriteMemory(): failed to set new "
					"protection for area %ld: %lx\n", area, error);
				break;
			}
			protectionChanged = true;
		}

		// copy the memory
		error = user_memcpy(address, buffer, toWrite);

		// reset the area protection
		if (protectionChanged)
			set_area_protection(area, areaInfo.protection);

		if (error != B_OK) {
			TRACE("BreakpointManager::_WriteMemory(): user_memcpy() failed: "
				"%lx\n", error);
			break;
		}

		bytesWritten += toWrite;
		address += toWrite;
		buffer += toWrite;
		size -= toWrite;
	}

	// If writing fails, we only fail, if we haven't written anything yet.
	if (error != B_OK) {
		if (bytesWritten > 0)
			return B_OK;
		return error;
	}

	return B_OK;
}