#include "StdAfx.h"
#include "VBoxClient.h"
#include "../../GDBServerFoundation/HexHelpers.h"
#include <bzscore/path.h>

using namespace VBoxGDB;
using namespace BazisLib;

VBoxClient::VBoxClient( const Network::InternetAddress &addr)
	: m_Socket(addr)
	, m_b64Bit(false)
	, m_CPUCount(0)
	, m_BreakpointIDOfLastStopEvent(-1)
	, m_LastBlockReadAddr(-1LL)
	, m_TargetAddr(addr)
{
	m_BufferAllocationSize = 1024 * 1024;
	m_Socket.SetNoDelay(true);

	//m_Socket.StartLogging(L"VBoxDebugger.log");
}

VBoxGDB::VBoxClient::~VBoxClient()
{
	Disconnect();
}

bool VBoxGDB::VBoxClient::RecvTillCommandPrompt()
{
	const char debugPrompt[] = "VBoxDbg>";
	m_Buffer.SetLength(0);
	size_t done = m_Socket.RecvToMarker(m_Buffer.PreAllocate(m_BufferAllocationSize, false), m_BufferAllocationSize, debugPrompt, sizeof(debugPrompt) - 1);
	if (done == 0 || done == -1 || done < (sizeof(debugPrompt) - 1))
		return false;
	m_Buffer.SetLength(done - (sizeof(debugPrompt) - 1));
	return true;
}

bool VBoxGDB::VBoxClient::Initialize()
{
	if (!RecvTillCommandPrompt())
		return false;
	std::vector<DynamicStringA> result;

#ifdef USE_REG_WORKAROUND_DLL
	TCHAR tsz[MAX_PATH];
	GetModuleFileName(GetModuleHandle(0), tsz, __countof(tsz));

	DynamicStringA str = "loadplugin ";
	str += StringToANSIString(Path::GetDirectoryName(TempStrPointerWrapper(tsz))) + "/regfix64.dll";
	for (size_t i = 0; i < str.length(); i++)
		if (str[i] == '\\')
			str[i] = '/';
	result = ExecuteCommand(str.c_str());
	if (!IsSuccessfulReply(result))
	{
		ReportError("Cannot load regfix64.dll into VirtualBox. Debugging won't be possible.");
		return false;
	}
#endif

	result = ExecuteCommand("cpu ffff");
	for (size_t i = 0; i < result.size(); i++)
	{
		const char szPattern[] = "is out of range! Highest ID is ";
		size_t idx = result[i].find(szPattern);
		if (idx == -1)
			continue;
		m_CPUCount = atoi(result[i].c_str() + idx + sizeof(szPattern) - 1);
		break;
	}

	if (!m_CPUCount)
		m_CPUCount = 1;

	//TODO: detect x86/x64 mode

	return true;
}

std::vector<DynamicStringA> VBoxGDB::VBoxClient::ExecuteCommand( const char *pCommand )
{
	std::string str = pCommand;
	str += "\n";
	//VirtualBox debugger has a bug: it won't recognize the command if the "\n" suffix was sent separately
	m_Socket.Send(str.c_str(), str.length());

	std::vector<DynamicStringA> result;
	if (!RecvTillCommandPrompt())
			return result;

	for each(const TempStringA &str in m_Buffer.Split("\r\n"))
		result.push_back(str);

	return result;
}

void VBoxGDB::VBoxClient::ReportError( const char *pErrorDescription )
{
	printf("Error: %s\n", pErrorDescription);
}

void VBoxGDB::VBoxClient::RequestStopAsync()
{
	m_Socket.Send("stop\n", 5);
}

static bool FindNextRegVal(const DynamicStringA &str, size_t &initial, TempStringA &regName, TempStringA &regVal)
{
	size_t idx = str.find('=', initial);
	if (idx == -1 || !idx)
		return false;

	size_t regStart = 0, regEnd = idx;
	for (;;)
	{
		regStart = str.rfind(' ', regEnd - 1);
		if (regStart == -1)
			regStart = 0;
		else
			regStart++;

		if (regEnd != regStart || !regStart)
			break;

		regEnd = regStart - 1;
	}


	regName = str.substr(regStart, regEnd - regStart);

	size_t end = str.find_first_not_of("0x123456789ABCDEFabcdef", idx + 1);
	if (end == -1)
		end = str.length();

	regVal = str.substr(idx + 1, end - idx - 1);
	initial = end;
	return true;
}

std::map<DynamicStringA, ULONGLONG> VBoxGDB::VBoxClient::ReadRegisters()
{
	std::map<DynamicStringA, ULONGLONG> result;
	std::vector<DynamicStringA> lines = ExecuteCommand(m_b64Bit ? "rg64" : "rg32");
	for (size_t i = 0; i < lines.size(); i++)
	{
		size_t initial = 0;
		for( ;;)
		{
			TempStringA regName(ConstStringA(""));
			TempStringA regVal(ConstStringA(""));
			if (!FindNextRegVal(lines[i], initial, regName, regVal))
				break;

			result[regName] = GDBServerFoundation::HexHelpers::ParseHexString<ULONGLONG>(regVal);
		}
	}

	return result;
}

bool VBoxGDB::VBoxClient::WriteRegister( const char *pName, ULONGLONG value )
{
	DynamicStringA strCmd;
#ifdef USE_REG_WORKAROUND_DLL
	const char *pRegCmd = ".r";
#else
	const char *pRegCmd = "r";
#endif

	if (!m_b64Bit)
		strCmd.Format("%s %s 0x%08X", pRegCmd, pName, (unsigned)value);
	else
		strCmd.Format("%s %s 0x%016I64x", pRegCmd, pName, value);

	std::vector<DynamicStringA> result = ExecuteCommand(strCmd.c_str());
	return IsSuccessfulReply(result);
}

static bool StartsWith(const TempStringA &str, const char *pStr)
{
	size_t len = strlen(pStr);
	if (str.length() < len)
		return false;
	return !memcmp(str.GetConstBuffer(), pStr, len);
}

bool VBoxGDB::VBoxClient::WaitForStop()
{
	static const char szEvent[] = "dbgf event:";
	static const char szWarning[] = "warning: The VM is already halted...";
	static const char szBreakpoint[] = "Breakpoint ";

	for (;;)
	{
		if (!RecvTillCommandPrompt())
			return false;

		int idx = m_Buffer.find_first_not_of(" \r\n\t");
		if (idx == -1)
			continue;

		TempStringA &str = m_Buffer.substr(idx);
		if (StartsWith(str, szEvent))
		{
			TempStringA evt = str.substr(sizeof(szEvent));
			if (StartsWith(evt, szBreakpoint))
			{
				m_BreakpointIDOfLastStopEvent = atoi(evt.GetConstBuffer() + sizeof(szBreakpoint) - 1);
				AdjustPCAfterBreakpoint(m_BreakpointIDOfLastStopEvent);
			}
			else
				m_BreakpointIDOfLastStopEvent = -1;
			return true;
		}
		if (StartsWith(str, szWarning))
			return true;

		continue;
	}
}

void VBoxGDB::VBoxClient::Disconnect()
{
	m_Socket.Send("exit\n", 5);
	RecvTillCommandPrompt();
	m_Socket.Close();
}

bool VBoxGDB::VBoxClient::IsSuccessfulReply( const std::vector<DynamicStringA> &reply, bool report )
{
	for (size_t i = 0; i < reply.size(); i++)
	{
		if (reply[i].find("error:") != -1)
		{
			if (report)
				ReportError(reply[i].c_str());
			return false;
		}
	}
	return true;
}

bool VBoxGDB::VBoxClient::ReadMemoryBlockAligned( ULONGLONG address, BazisLib::BasicBuffer &buffer )
{
	if (address & 3)
	{
		ReportError("Error: trying to read from an unaligned address");
		return false;
	}

	DynamicStringA strCmd;
	if (!m_b64Bit)
		strCmd.Format("dd 0x%08X", (unsigned)address);
	else
		strCmd.Format("dd 0x%016I64x", address);

	std::vector<DynamicStringA> result = ExecuteCommand(strCmd.c_str());
	buffer.SetSize(0);
	ULONGLONG expectedAddr = address;
	for (size_t i = 0; i < result.size(); i++)
	{
		const DynamicStringA &line = result[i];
		int idx = line.find(':');
		if (idx == -1 || !idx)
			continue;
		size_t lineStart = line.find_first_not_of(" \t");
		if (lineStart == -1)
			continue;
		if (line[lineStart] != '%')
			continue;
		
		ULONGLONG thisAddr = GDBServerFoundation::HexHelpers::ParseHexString<ULONGLONG>(line.substr(lineStart + 1, idx - lineStart - 1));
		if (thisAddr != expectedAddr)
			break;

		size_t start = idx + 1;
		for (;;)
		{
			size_t numStart = line.find_first_not_of(' ', start);
			if (numStart == -1)
				break;
			size_t numEnd = line.find(' ', numStart);
			if (numEnd == -1)
				numEnd = line.length();

			unsigned num = GDBServerFoundation::HexHelpers::ParseHexString<unsigned>(line.substr(numStart, numEnd - numStart));
			C_ASSERT(sizeof(num) == 4);
			buffer.append(&num, 4, 4096);
			expectedAddr += 4;

			start = numEnd;
		}

	}
	return true;
}

bool VBoxGDB::VBoxClient::WriteMemory( ULONGLONG address, const void *pData, size_t length )
{
	m_LastBlockReadAddr = -1LL;
	const char *p = (const char *)pData;
	size_t remaining = length;
	DynamicStringA strVal, strCmd;
	while (remaining)
	{
		char mode;
		size_t todo = 0;

		if (remaining >= 8)
		{
			mode = 'q';
			strVal.Format("%016I64x", *((ULONGLONG *)p));
			todo = 8;
		}
		else if (remaining >= 4)
		{
			mode = 'd';
			strVal.Format("%08x", *((unsigned *)p));
			todo = 4;
		}
		else if (remaining >= 2)
		{
			mode = 'w';
			strVal.Format("%04x", *((unsigned short *)p));
			todo = 2;
		}
		else
		{
			mode = 'b';
			strVal.Format("%02x", *((unsigned char *)p));
			todo = 1;
		}

		if (m_b64Bit)
			strCmd.Format("e%c 0x%016I64x %s", mode, address, strVal.c_str());
		else
			strCmd.Format("e%c 0x%08x %s", mode, (unsigned)address, strVal.c_str());

		std::vector<DynamicStringA> result = ExecuteCommand(strCmd.c_str());
		if (!IsSuccessfulReply(result))
			return false;

		remaining -= todo;
		p += todo;
		address += todo;
	}
	return true;
}

bool VBoxGDB::VBoxClient::ReadMemory( ULONGLONG address, void *pData, size_t *pLength, bool allowCaching )
{
	size_t requested = *pLength, done = 0;

	char *p = (char *)pData;

	if (allowCaching && m_LastBlockReadAddr != -1LL && (address >= m_LastBlockReadAddr) && (address < (m_LastBlockReadAddr + m_LastBlockReadBuffer.GetSize())))
	{
		size_t offsetInBuffer = (size_t)(address - m_LastBlockReadAddr);
		size_t available = m_LastBlockReadBuffer.GetSize() - offsetInBuffer;

		size_t todo = min(requested, available);
		memcpy(p, m_LastBlockReadBuffer.GetData(offsetInBuffer), todo);
		done += todo;
		p += todo;
		address += todo;
	}

	unsigned off = address & 3;
	address &= ~3;

	while (done < requested)
	{
		if (!ReadMemoryBlockAligned(address, m_LastBlockReadBuffer))
			break;
		m_LastBlockReadAddr = address;
		if (!m_LastBlockReadBuffer.GetSize())
			break;
		ASSERT(!(m_LastBlockReadBuffer.GetSize() & 3));

		size_t todo = min(requested - done, m_LastBlockReadBuffer.GetSize() - off);
		memcpy(p, m_LastBlockReadBuffer.GetData(off), todo);
		p += todo;
		done += todo;
		address += todo;
		off = 0;
	}

	*pLength = done;
	return done != 0;
}

bool VBoxGDB::VBoxClient::StepAndWait()
{
	if (!IsSuccessfulReply(ExecuteCommand("t")))
		return false;
	return WaitForStop();
}

bool VBoxGDB::VBoxClient::ContinueAndWait()
{
	if (!IsSuccessfulReply(ExecuteCommand("g")))
		return false;
	return WaitForStop();
}

int VBoxGDB::VBoxClient::CreateBreakpoint( ULONGLONG addr )
{
	std::vector<DynamicStringA> result = ExecuteCommand(DynamicStringA::sFormat("bp 0x%I64x", addr).c_str());
	const char szText[] = "Set breakpoint ";
	for (size_t i = 0; i < result.size(); i++)
	{
		off_t idx = result[i].find(szText);
		if (idx == -1)
			continue;
		int bp = atoi(result[i].c_str() + idx + sizeof(szText) - 1);
		if (!bp)
			return -1;
		m_BreakpointMap[bp] = addr;
		return bp;
	}
	return -1;
}

bool VBoxGDB::VBoxClient::RemoveBreakpoint( int idx )
{
	std::vector<DynamicStringA> result = ExecuteCommand(DynamicStringA::sFormat("bc %x", idx).c_str());
	if (!IsSuccessfulReply(result))
		return false;
	BreakpointAddressMap::iterator it = m_BreakpointMap.find(idx);
	if (it != m_BreakpointMap.end())
	{
		m_BreakpointMap.erase(it);
	}
	return true;
}

void VBoxGDB::VBoxClient::AdjustPCAfterBreakpoint( int breakpointNumber )
{
	ULONGLONG val = GetCurrentPC();
	BreakpointAddressMap::iterator it = m_BreakpointMap.find(breakpointNumber);
	if (it == m_BreakpointMap.end())
		return;
	if (val == (it->second + 1))
	{
		//It looks like VirtualBox is using its own int3 handler to detect breakpoints.
		//This int3 handler determines the address of the breakpoint from the return address in the stack,
		//that actually points to the next instruction.

		//While this makes perfect sense, it makes debugging useless, as we don't want to skip the instruction
		//at the address where we have put the breakpoint.
		const char *pRegName = m_b64Bit ? "rip" : "eip";
		WriteRegister(pRegName, it->second);
	}
}

bool VBoxGDB::VBoxClient::ReadRegister( const char *pName, ULONGLONG *pValue )
{
	std::vector<DynamicStringA> result = ExecuteCommand(DynamicStringA::sFormat("r %s", pName).c_str());
	for (size_t i = 0; i < result.size(); i++)
	{
		TempStringA regName(ConstStringA(""));
		TempStringA regVal(ConstStringA(""));
		size_t initial = 0;
		if (!FindNextRegVal(result[i], initial, regName, regVal))
			continue;

		*pValue = GDBServerFoundation::HexHelpers::ParseHexString<ULONGLONG>(regVal);
		return true;
	}
	return false;

}

ULONGLONG VBoxGDB::VBoxClient::GetCurrentPC()
{
	const char *pRegName = m_b64Bit ? "rip" : "eip";
	ULONGLONG val;
	if (!ReadRegister(pRegName, &val))
		return 0;
	return val;
	
}
