myvtm/Framework/libtoolkit/modCheck.cpp

#include "modCheck.h"
#include <set>
#include <string>
#include <list>
#include <winpr/synch.h>
#include <winpr/crt.h>
#include <winpr/thread.h>

using namespace std;

class modCheckLock
{
public:
	modCheckLock() {
		static bool isInit = false;
		if (!isInit)
		{
			InitializeCriticalSection(&cs_);
			isInit = true;
		}
		EnterCriticalSection(&cs_);
	}
	~modCheckLock() { 
		LeaveCriticalSection(&cs_);
	}
private:
	modCheckLock(const modCheckLock&);
	modCheckLock& operator=(const modCheckLock&);
private:
	static CRITICAL_SECTION cs_;
};

CRITICAL_SECTION modCheckLock::cs_;

void InitResource(toolkitResource &m_resource)
{
	ZeroMemory(&m_resource, sizeof(toolkitResource));
	m_resource.g_mgr = 0;
	m_resource.g_initialized = 0;
	m_resource.hprintf_index = 0;
	ZeroMemory(m_resource.hprintf_buffer, sizeof(m_resource.hprintf_buffer));
	ZeroMemory(m_resource.szCrashID, sizeof(m_resource.szCrashID));
}

class modInfo
{
public:
	modInfo(string name)
		:m_name(name)
	{
		InitResource(m_resource);
	}
	bool operator == (const string name){
		return name == this->m_name;
	}
	void AddModThread(DWORD threadId){
		m_threadArr.insert(threadId);
	}
	bool findModThread(DWORD threadId){
		return m_threadArr.end() != m_threadArr.find(threadId);
	}
	toolkitResource *getResource()
	{
		return &m_resource;
	}
	string m_name;
	friend BOOL toolkit_setAssign(void *assign);
	friend BOOL toolkit_setThreadGroupByAssign(void *assign);

	set<DWORD> m_threadArr;
private:
	toolkitResource m_resource;
	void *m_assign;
};

list<modInfo*> g_modInfoArr;

modInfo* findModInfo(string name)
{
	for (list<modInfo*>::iterator i = g_modInfoArr.begin(); i != g_modInfoArr.end(); i++)
	{
		if (**i == name)
			return *i;
	}
	return NULL;
}

modInfo *findModInfoByThreadId()
{
	DWORD threadId = GetCurrentThreadId();
	for (list<modInfo*>::iterator i = g_modInfoArr.begin(); i != g_modInfoArr.end(); i++)
	{
		if ((*i)->findModThread(threadId))
			return *i;
	}
	return NULL;
}

TOOLKIT_API BOOL toolkit_checkThreadInName(DWORD threadId, const char *name)
{
	if (NULL == name)
		return false;

	modCheckLock curLock;
	for (list<modInfo*>::iterator i = g_modInfoArr.begin(); i != g_modInfoArr.end(); i++)
	{
		//TODO: should separate it or it would iterate all mod.
		if (string(name) == (*i)->m_name && (*i)->findModThread(threadId))
			return true;
	}
	return false;
}

TOOLKIT_API BOOL toolkit_CreateModuleInfo(const char *name)
{
	if (NULL == name)
		return false;

	if (NULL != findModInfo(name))
		toolkit_DestoryModuleInfo(name);

	modCheckLock curLock;
	modInfo *curModInfo = new modInfo(name);

	g_modInfoArr.push_back(curModInfo);

	return true;
}


TOOLKIT_API BOOL toolkit_DestoryModuleInfo(const char *name)
{
	if (NULL == name)
		return false;

	modCheckLock curLock;
	for (list<modInfo*>::iterator i = g_modInfoArr.begin(); i != g_modInfoArr.end(); i++)
	{
		if (**i == name)
		{
			delete *i;
			g_modInfoArr.erase(i);
			return true;
		}
	}
	return false;
}


TOOLKIT_API BOOL toolkit_SetthreadGroup(DWORD threadId, const char *name)
{
	if (NULL == name)
		return false;

	modCheckLock curLock;
	modInfo *curModInfo = findModInfo(name);
	if (NULL == curModInfo)
		return false;

	curModInfo->AddModThread(threadId);
	return true;
}


TOOLKIT_API BOOL toolkit_TerminateThreadExcept(DWORD exceptThreadId, const char *name)
{
	if (NULL == name)
		return false;

	modCheckLock curLock;
	modInfo *curModInfo = findModInfo(name);
	if (NULL == curModInfo)
		return false;

	for (set<DWORD>::iterator i = curModInfo->m_threadArr.begin(); i != curModInfo->m_threadArr.end(); i++)
	{
		if (*i != exceptThreadId)
		{
			HANDLE currThread = OpenThread(PROCESS_ALL_ACCESS, FALSE, *i);
			if (currThread != NULL)	//
			{
				TerminateThread(currThread, -1);
			}
		}
	}
	return true;
}

toolkitResource* toolkit_getResource()
{
	static toolkitResource *spShellResource = NULL;
	modCheckLock curLock;
	modInfo *curMod = findModInfoByThreadId();
	if (NULL != curMod)
		return curMod->getResource();
	else if ((NULL != spShellResource))
		return spShellResource;
	else
	{
		spShellResource = new toolkitResource();
		InitResource(*spShellResource);
		return spShellResource;
	}
	return NULL;
}

BOOL toolkit_setAssign(void *assign)
{
	if (NULL == assign)
		return FALSE;
	modCheckLock curLock;
	modInfo *curMod = findModInfoByThreadId();
	if (NULL != curMod)
	{
		curMod->m_assign = assign;
		return TRUE;
	}
	return FALSE;
}

BOOL toolkit_setThreadGroupByAssign(void *assign)
{
	if (NULL == assign)
		return FALSE;

	modCheckLock curLock;
	for (list<modInfo*>::iterator i = g_modInfoArr.begin(); i != g_modInfoArr.end(); i++)
	{
		if ((*i)->m_assign == assign)
		{
			(*i)->AddModThread(GetCurrentThreadId());
			return TRUE;
		}
	}
	return FALSE;
}