X-Git-Url: http://git.osdn.net/view?a=blobdiff_plain;f=fixscancodemap.cpp;h=249ba1a609fad93fbdab4fa4bae8e609a4788525;hb=HEAD;hp=4b2adb9fff8ccef013fa61a18652691fa490945b;hpb=f47d606f533c85d1c69af98699f80d1b8b1b03c3;p=yamy%2Fyamy.git

diff --git a/fixscancodemap.cpp b/fixscancodemap.cpp
old mode 100755
new mode 100644
index 4b2adb9..249ba1a
--- a/fixscancodemap.cpp
+++ b/fixscancodemap.cpp
@@ -1,15 +1,18 @@
 #include "fixscancodemap.h"
-#include "registry.h"
 #include "misc.h"
+#include "windowstool.h"
 #include <tchar.h>
 #include <tlhelp32.h>
+#include <process.h>
 
 #pragma runtime_checks( "", off )
-static DWORD invokeFunc(InjectInfo *info)
+static DWORD WINAPI invokeFunc(InjectInfo *info)
 {
 	BOOL ret;
 	HANDLE hToken;
 	HMODULE hAdvapi32;
+	DWORD result = 0;
+
 	FpImpersonateLoggedOnUser pImpersonateLoggedOnUser;
 	FpRevertToSelf pRevertToSelf;
 	FpOpenProcessToken pOpenProcessToken;
@@ -22,29 +25,44 @@ static DWORD invokeFunc(InjectInfo *info)
 
 	HANDLE hProcess = info->pOpenProcess(PROCESS_QUERY_INFORMATION, FALSE, info->pid_);
 	if (hProcess == NULL) {
-		return 1;
+		result = YAMY_ERROR_ON_OPEN_YAMY_PROCESS;
+		goto exit;
 	}
 
 	ret = pOpenProcessToken(hProcess, TOKEN_QUERY | TOKEN_DUPLICATE , &hToken);
 	if (ret == FALSE) {
-		return 2;
+		result = YAMY_ERROR_ON_OPEN_YAMY_TOKEN;
+		goto exit;
 	}
 
 	ret = pImpersonateLoggedOnUser(hToken);
 	if (ret == FALSE) {
-		return 3;
+		result = YAMY_ERROR_ON_IMPERSONATE;
+		goto exit;
 	}
 
-	info->pUpdate(0, 1);
+	if (info->isVistaOrLater_) {
+		info->pUpdate4(1);
+	} else {
+		info->pUpdate8(0, 1);
+	}
 
 	ret = pRevertToSelf();
 	if (ret == FALSE) {
-		return 4;
+		result = YAMY_ERROR_ON_REVERT_TO_SELF;
+		goto exit;
 	}
 
-	info->pCloseHandle(hToken);
-	info->pCloseHandle(hProcess);
-	return 0;
+exit:
+	if (hToken != NULL) {
+		info->pCloseHandle(hToken);
+	}
+
+	if (hProcess != NULL) {
+		info->pCloseHandle(hProcess);
+	}
+
+	return result;
 }
 static int afterFunc(int arg)
 {
@@ -61,7 +79,7 @@ const DWORD FixScancodeMap::s_fixEntry[] = {
 	0x003ae03a,
 	0x0029e029,
 	0x0070e070,
-	0x003b001e,
+	0x007be07b,
 };
 
 int FixScancodeMap::acquirePrivileges()
@@ -70,13 +88,13 @@ int FixScancodeMap::acquirePrivileges()
 	HANDLE hToken = NULL;
 
 	if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hToken)) {
-		ret = 5;
+		ret = YAMY_ERROR_ON_OPEN_CURRENT_PROCESS;
 		goto exit;
 	}
 
 	LUID luid;
 	if (!LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &luid)) {
-		ret = 6;
+		ret = YAMY_ERROR_ON_LOOKUP_PRIVILEGE;
 		goto exit;
 	}
 
@@ -86,7 +104,7 @@ int FixScancodeMap::acquirePrivileges()
 	tk_priv.Privileges[0].Luid = luid;
 
 	if (!AdjustTokenPrivileges(hToken, FALSE, &tk_priv, 0, NULL, NULL)) {
-		ret = 7;
+		ret = YAMY_ERROR_ON_ADJUST_PRIVILEGE;
 		goto exit;
 	}
 
@@ -117,7 +135,7 @@ DWORD FixScancodeMap::getWinLogonPid()
 
     BOOL bResult = Process32First(hSnap, &pe);
 	while (bResult){
-		if (!_tcscmp(pe.szExeFile, _T("winlogon.exe"))) {
+		if (!_tcsicmp(pe.szExeFile, _T("winlogon.exe"))) {
 			DWORD sessionId;
 
 			if (ProcessIdToSessionId(pe.th32ProcessID, &sessionId) != FALSE) {
@@ -135,96 +153,119 @@ DWORD FixScancodeMap::getWinLogonPid()
 }
 
 
+bool FixScancodeMap::clean(WlInfo wl)
+{
+	int ret = 0;
+
+	if (wl.m_hThread != NULL) {
+		DWORD result;
+
+		if (WaitForSingleObject(wl.m_hThread, 5000) == WAIT_TIMEOUT) {
+			return false;
+		}
+
+		GetExitCodeThread(wl.m_hThread, &result);
+		CloseHandle(wl.m_hThread);
+
+		if (wl.m_remoteMem != NULL && wl.m_hProcess != NULL) {
+			VirtualFreeEx(wl.m_hProcess, wl.m_remoteMem, 0, MEM_RELEASE);
+		}
+
+		if (wl.m_remoteInfo != NULL && wl.m_hProcess != NULL) {
+			VirtualFreeEx(wl.m_hProcess, wl.m_remoteInfo, 0, MEM_RELEASE);
+		}
+
+		if (wl.m_hProcess != NULL) {
+			CloseHandle(wl.m_hProcess);
+		}
+	}
+
+	return true;
+}
+
+
 int FixScancodeMap::injectThread(DWORD dwPID)
 {
 	int ret = 0;
 	DWORD err = 0;
 	BOOL wFlag;
+	WlInfo wi;
 
-	HANDLE hProcess = NULL;
-	LPVOID remoteMem = NULL;
-	LPVOID remoteInfo = NULL;
-	DWORD invokeFuncAddr = (DWORD)invokeFunc;
-	DWORD afterFuncAddr = (DWORD)afterFunc;
-	DWORD memSize =  afterFuncAddr - invokeFuncAddr;
+	wi.m_hProcess = NULL;
+	wi.m_remoteMem = NULL;
+	wi.m_remoteInfo = NULL;
+	wi.m_hThread = NULL;
 
-	if ((hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, dwPID)) == NULL) {
-		ret = 8;
+	ULONG_PTR invokeFuncAddr = (ULONG_PTR)invokeFunc;
+	ULONG_PTR afterFuncAddr = (ULONG_PTR)afterFunc;
+	SIZE_T memSize =  afterFuncAddr - invokeFuncAddr;
+
+	if ((wi.m_hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, dwPID)) == NULL) {
+		ret = YAMY_ERROR_ON_OPEN_WINLOGON_PROCESS;
 		goto exit;
 	}
 
-	remoteMem = VirtualAllocEx(hProcess, NULL, memSize, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
-	if (remoteMem == NULL) {
-		ret = 9;
+	wi.m_remoteMem = VirtualAllocEx(wi.m_hProcess, NULL, memSize, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
+	if (wi.m_remoteMem == NULL) {
+		ret = YAMY_ERROR_ON_VIRTUALALLOCEX;
 		err = GetLastError();
 		goto exit;
 	}
 
-	wFlag = WriteProcessMemory(hProcess, remoteMem, (char*)invokeFunc, memSize, (SIZE_T*)0);
+	wFlag = WriteProcessMemory(wi.m_hProcess, wi.m_remoteMem, (char*)invokeFunc, memSize, (SIZE_T*)0);
 	if (wFlag == FALSE) {
-		ret = 10;
+		ret = YAMY_ERROR_ON_WRITEPROCESSMEMORY;
 		goto exit;
 	}
 
-	remoteInfo = VirtualAllocEx(hProcess, NULL, sizeof(m_info), MEM_COMMIT, PAGE_READWRITE);
-	if (remoteInfo == NULL) {
-		ret = 11;
+	wi.m_remoteInfo = VirtualAllocEx(wi.m_hProcess, NULL, sizeof(m_info), MEM_COMMIT, PAGE_READWRITE);
+	if (wi.m_remoteInfo == NULL) {
+		ret = YAMY_ERROR_ON_VIRTUALALLOCEX;
 		err = GetLastError();
 		goto exit;
 	}
 
-	wFlag = WriteProcessMemory(hProcess, remoteInfo, (char*)&m_info, sizeof(m_info), (SIZE_T*)0);
+	wFlag = WriteProcessMemory(wi.m_hProcess, wi.m_remoteInfo, (char*)&m_info, sizeof(m_info), (SIZE_T*)0);
 	if (wFlag == FALSE) {
-		ret = 12;
+		ret = YAMY_ERROR_ON_WRITEPROCESSMEMORY;
 		goto exit;
 	}
 
-#if 0
-	TCHAR buf[1024];
-
-	_stprintf_s(buf, sizeof(buf)/sizeof(buf[0]),
-		_T("execute UpdatePerUserSystemParameters(), inject code to winlogon.exe?\r\n")
-		_T("invokeFunc=0x%p\r\n")
-		_T("afterFunc=0x%p\r\n")
-		_T("afterFunc - invokeFunc=%d\r\n")
-		_T("remoteMem=0x%p\r\n")
-		_T("remoteInfo=0x%p(size: %d)\r\n"),
-		invokeFunc, afterFunc, memSize, remoteMem, remoteInfo, sizeof(info));
-	if (MessageBox((HWND)NULL, buf, _T("upusp"), MB_OKCANCEL | MB_ICONSTOP) == IDCANCEL) {
-		(m_info.pUpdate)(0, 1);
+	wi.m_hThread = CreateRemoteThread(wi.m_hProcess, NULL, 0, 
+		(LPTHREAD_START_ROUTINE)wi.m_remoteMem, wi.m_remoteInfo, 0, NULL);
+	if (wi.m_hThread == NULL) {
+		ret = YAMY_ERROR_ON_CREATEREMOTETHREAD;
 		goto exit;
 	}
-#endif
 
-	HANDLE hThread = CreateRemoteThread(hProcess, NULL, 0, 
-		(LPTHREAD_START_ROUTINE)remoteMem, remoteInfo, 0, NULL);
-	if (hThread == NULL) {
-		ret = 13;
-		goto exit;
-	}
-
-	if (WaitForSingleObject(hThread, 5000) == WAIT_TIMEOUT) {
-		ret = 14;
-		goto exit;
+	if (WaitForSingleObject(wi.m_hThread, 5000) == WAIT_TIMEOUT) {
+		ret = YAMY_ERROR_TIMEOUT_INJECTION;
+		m_wlTrash.push_back(wi);
+		goto dirty_exit;
 	}
 	DWORD result = -1;
-	GetExitCodeThread(hThread, &result);
+	GetExitCodeThread(wi.m_hThread, &result);
 	ret = result;
-	CloseHandle(hThread);
+	CloseHandle(wi.m_hThread);
+	wi.m_hThread = NULL;
 
 exit:
-	if (remoteMem != NULL) {
-		VirtualFreeEx(hProcess, remoteMem, 0, MEM_RELEASE);
+	if (wi.m_remoteMem != NULL && wi.m_hProcess != NULL) {
+		VirtualFreeEx(wi.m_hProcess, wi.m_remoteMem, 0, MEM_RELEASE);
+		wi.m_remoteMem = NULL;
 	}
 
-	if (remoteInfo != NULL) {
-		VirtualFreeEx(hProcess, remoteInfo, 0, MEM_RELEASE);
+	if (wi.m_remoteInfo != NULL && wi.m_hProcess != NULL) {
+		VirtualFreeEx(wi.m_hProcess, wi.m_remoteInfo, 0, MEM_RELEASE);
+		wi.m_remoteInfo = NULL;
 	}
 
-	if (hProcess != NULL) {
-		CloseHandle(hProcess);
+	if (wi.m_hProcess != NULL) {
+		CloseHandle(wi.m_hProcess);
+		wi.m_hProcess = NULL;
 	}
 
+dirty_exit:
 	return ret;
 }
 
@@ -233,26 +274,27 @@ int FixScancodeMap::update()
 	MINIMIZEDMETRICS mm;
 	int result = 0;
 
-	result = acquirePrivileges();
-	if (result) {
+	if (m_errorOnConstruct) {
+		result = m_errorOnConstruct;
 		goto exit;
 	}
 
-	DWORD dwPID;
-	if ((dwPID = getWinLogonPid()) == 0) {
-		result = 15;
-		goto exit;
-	}
+	m_wlTrash.erase(remove_if(m_wlTrash.begin(), m_wlTrash.end(), FixScancodeMap::clean), m_wlTrash.end());
 
 	memset(&mm, 0, sizeof(mm));
 	mm.cbSize = sizeof(mm);
 	SystemParametersInfo(SPI_GETMINIMIZEDMETRICS, sizeof(mm), &mm, 0);
 
-	result = injectThread(dwPID);
-	if (result) {
-		goto exit;
+	result = injectThread(m_winlogonPid);
+	if (result == YAMY_ERROR_TIMEOUT_INJECTION) {
+		// retry once
+		result = injectThread(m_winlogonPid);
+		if (result == YAMY_SUCCESS) {
+			result = YAMY_ERROR_RETRY_INJECTION_SUCCESS;
+		}
 	}
 
+	mm.iArrange = ARW_HIDE;
 	SystemParametersInfo(SPI_SETMINIMIZEDMETRICS, sizeof(mm), &mm, 0);
 
 exit:
@@ -262,32 +304,29 @@ exit:
 int FixScancodeMap::fix()
 {
 	ScancodeMap *origMap, *fixMap;
-	Registry reg(HKEY_CURRENT_USER, _T("Keyboard Layout"));
-	// Windows7 RC not support Scancode Map on HKCU?
-	//Registry reg(HKEY_LOCAL_MACHINE, _T("SYSTEM\\CurrentControlSet\\Control\\Keyboard Layout"));
 	DWORD origSize, fixSize;
 	bool ret;
 	int result = 0;
 
 	// save original Scancode Map
-	ret = reg.read(_T("Scancode Map"), NULL, &origSize, NULL, 0);
+	ret = m_pReg->read(_T("Scancode Map"), NULL, &origSize, NULL, 0);
 	if (ret) {
 		origMap = reinterpret_cast<ScancodeMap*>(malloc(origSize));
 		if (origMap == NULL) {
-			result = 16;
+			result = YAMY_ERROR_NO_MEMORY;
 			goto exit;
 		}
 
-		ret = reg.read(_T("Scancode Map"), reinterpret_cast<BYTE*>(origMap), &origSize, NULL, 0);
+		ret = m_pReg->read(_T("Scancode Map"), reinterpret_cast<BYTE*>(origMap), &origSize, NULL, 0);
 		if (ret == false) {
-			result = 17;
+			result = YAMY_ERROR_ON_READ_SCANCODE_MAP;
 			goto exit;
 		}
 
 		fixSize = origSize;
 		fixMap = reinterpret_cast<ScancodeMap*>(malloc(origSize + s_fixEntryNum * sizeof(s_fixEntry[0])));
 		if (fixMap == NULL) {
-			result = 18;
+			result = YAMY_ERROR_NO_MEMORY;
 			goto exit;
 		}
 
@@ -299,7 +338,7 @@ int FixScancodeMap::fix()
 		fixSize = sizeof(ScancodeMap);
 		fixMap = reinterpret_cast<ScancodeMap*>(malloc(sizeof(ScancodeMap) + s_fixEntryNum * sizeof(s_fixEntry[0])));
 		if (fixMap == NULL) {
-			result = 19;
+			result = YAMY_ERROR_NO_MEMORY;
 			goto exit;
 		}
 
@@ -332,21 +371,21 @@ int FixScancodeMap::fix()
 		fixSize += 4;
 	}
 
-	ret = reg.write(_T("Scancode Map"), reinterpret_cast<BYTE*>(fixMap), fixSize);
+	ret = m_pReg->write(_T("Scancode Map"), reinterpret_cast<BYTE*>(fixMap), fixSize);
 	if (ret == false) {
-		result = 20;
+		result = YAMY_ERROR_ON_WRITE_SCANCODE_MAP;
 		goto exit;
 	}
 
 	result = update();
 
 	if (origMap) {
-		ret = reg.write(_T("Scancode Map"), reinterpret_cast<BYTE*>(origMap), origSize);
+		ret = m_pReg->write(_T("Scancode Map"), reinterpret_cast<BYTE*>(origMap), origSize);
 	} else {
-		ret = reg.remove(_T("Scancode Map"));
+		ret = m_pReg->remove(_T("Scancode Map"));
 	}
 	if (ret == false) {
-		result = 21;
+		result = YAMY_ERROR_ON_WRITE_SCANCODE_MAP;
 		goto exit;
 	}
 
@@ -365,7 +404,61 @@ int FixScancodeMap::restore()
 	return update();
 }
 
-FixScancodeMap::FixScancodeMap()
+int FixScancodeMap::escape(bool i_escape)
+{
+	if (i_escape) {
+		SetEvent(m_hFixEvent);
+	} else {
+		SetEvent(m_hRestoreEvent);
+	}
+	return 0;
+}
+
+unsigned int WINAPI FixScancodeMap::threadLoop(void *i_this)
+{
+	int err;
+	DWORD ret;
+	FixScancodeMap *This = reinterpret_cast<FixScancodeMap*>(i_this);
+	HANDLE handles[] = {This->m_hFixEvent, This->m_hRestoreEvent, This->m_hQuitEvent};
+	while ((ret = MsgWaitForMultipleObjects(NUMBER_OF(handles), &handles[0],
+		FALSE, INFINITE, QS_POSTMESSAGE)) != WAIT_FAILED) {
+		switch (ret) {
+		case WAIT_OBJECT_0:			// m_hFixEvent
+			ResetEvent(This->m_hFixEvent);
+			err = This->fix();
+			PostMessage(This->m_hwnd, This->m_messageOnFail, err, 1);
+			break;
+		case WAIT_OBJECT_0 + 1:		// m_hRestoreEvent
+			ResetEvent(This->m_hRestoreEvent);
+			err = This->restore();
+			PostMessage(This->m_hwnd, This->m_messageOnFail, err, 0);
+			break;
+		case WAIT_OBJECT_0 + 2:		// m_hQuiteEvent
+			ResetEvent(This->m_hQuitEvent);
+			// through below
+		default:
+			return 0;
+			break;
+		}
+	}
+	return 1;
+}
+
+int FixScancodeMap::init(HWND i_hwnd, UINT i_messageOnFail)
+{
+	m_hwnd = i_hwnd;
+	m_messageOnFail = i_messageOnFail;
+	return 0;
+}
+
+FixScancodeMap::FixScancodeMap() :
+	m_hwnd(NULL),
+	m_messageOnFail(WM_NULL),
+	m_errorOnConstruct(0),
+	m_winlogonPid(0),
+	m_regHKCU(HKEY_CURRENT_USER, _T("Keyboard Layout")),
+	m_regHKLM(HKEY_LOCAL_MACHINE, _T("SYSTEM\\CurrentControlSet\\Control\\Keyboard Layout")),
+	m_pReg(NULL)
 {
 	HMODULE hMod;
 
@@ -376,10 +469,20 @@ FixScancodeMap::FixScancodeMap()
 	memcpy(&m_info.revertToSelf_, "RevertToSelf", sizeof(m_info.revertToSelf_));
 	memcpy(&m_info.openProcessToken_, "OpenProcessToken", sizeof(m_info.openProcessToken_));
 
+	m_hFixEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
+	ASSERT(m_hFixEvent);
+	m_hRestoreEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
+	ASSERT(m_hRestoreEvent);
+	m_hQuitEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
+	ASSERT(m_hQuitEvent);
+
+	m_hThread = (HANDLE)_beginthreadex(NULL, 0, threadLoop, this, 0, &m_threadId);
+
 	hMod = GetModuleHandle(_T("user32.dll"));
 	if (hMod != NULL) {
-		m_info.pUpdate = (FpUpdatePerUserSystemParameters)GetProcAddress(hMod, "UpdatePerUserSystemParameters");
-		if (m_info.pUpdate == NULL) {
+		m_info.pUpdate4 = (FpUpdatePerUserSystemParameters4)GetProcAddress(hMod, "UpdatePerUserSystemParameters");
+		m_info.pUpdate8 = (FpUpdatePerUserSystemParameters8)m_info.pUpdate4;
+		if (m_info.pUpdate4 == NULL) {
 			return;
 		}
 	}
@@ -406,8 +509,37 @@ FixScancodeMap::FixScancodeMap()
 			return;
 		}
 	}
+
+	// Windows7 RC not support Scancode Map on HKCU?
+	if (checkWindowsVersion(6, 1) == FALSE) {
+		m_pReg = &m_regHKCU; // Vista or earlier
+	} else {
+		m_pReg = &m_regHKLM; // Windows7 or later
+	}
+
+	// prototype of UpdatePerUserSystemParameters() differ vista or earlier
+	if (checkWindowsVersion(6, 0) == FALSE) {
+		m_info.isVistaOrLater_ = 0; // before Vista
+	} else {
+		m_info.isVistaOrLater_ = 1; // Vista or later
+	}
+
+	m_errorOnConstruct = acquirePrivileges();
+	if (m_errorOnConstruct) {
+		goto exit;
+	}
+
+	if ((m_winlogonPid = getWinLogonPid()) == 0) {
+		m_errorOnConstruct = YAMY_ERROR_ON_GET_WINLOGON_PID;
+		goto exit;
+	}
+
+exit:
+	;
 }
 
 FixScancodeMap::~FixScancodeMap()
 {
+	SetEvent(m_hQuitEvent);
+	WaitForSingleObject(m_hThread, INFINITE);
 }