OSDN Git Service

update Scancode Map asynchronouslly by independent thread to avoid stall on Windows XP
[yamy/yamy.git] / fixscancodemap.cpp
1 #include "fixscancodemap.h"\r
2 #include "misc.h"\r
3 #include "windowstool.h"\r
4 #include <tchar.h>\r
5 #include <tlhelp32.h>\r
6 #include <process.h>\r
7 \r
8 #pragma runtime_checks( "", off )\r
9 static DWORD WINAPI invokeFunc(InjectInfo *info)\r
10 {\r
11         BOOL ret;\r
12         HANDLE hToken;\r
13         HMODULE hAdvapi32;\r
14         DWORD result = 0;\r
15 \r
16         FpImpersonateLoggedOnUser pImpersonateLoggedOnUser;\r
17         FpRevertToSelf pRevertToSelf;\r
18         FpOpenProcessToken pOpenProcessToken;\r
19 \r
20         hAdvapi32 = info->pGetModuleHandle(info->advapi32_);\r
21 \r
22         pImpersonateLoggedOnUser = (FpImpersonateLoggedOnUser)info->pGetProcAddress(hAdvapi32, info->impersonateLoggedOnUser_);\r
23         pRevertToSelf = (FpRevertToSelf)info->pGetProcAddress(hAdvapi32, info->revertToSelf_);\r
24         pOpenProcessToken = (FpOpenProcessToken)info->pGetProcAddress(hAdvapi32, info->openProcessToken_);\r
25 \r
26         HANDLE hProcess = info->pOpenProcess(PROCESS_QUERY_INFORMATION, FALSE, info->pid_);\r
27         if (hProcess == NULL) {\r
28                 result = 1;\r
29                 goto exit;\r
30         }\r
31 \r
32         ret = pOpenProcessToken(hProcess, TOKEN_QUERY | TOKEN_DUPLICATE , &hToken);\r
33         if (ret == FALSE) {\r
34                 result = 2;\r
35                 goto exit;\r
36         }\r
37 \r
38         ret = pImpersonateLoggedOnUser(hToken);\r
39         if (ret == FALSE) {\r
40                 result = 3;\r
41                 goto exit;\r
42         }\r
43 \r
44         if (info->isVistaOrLater_) {\r
45                 info->pUpdate4(1);\r
46         } else {\r
47                 info->pUpdate8(0, 1);\r
48         }\r
49 \r
50         ret = pRevertToSelf();\r
51         if (ret == FALSE) {\r
52                 result = 4;\r
53                 goto exit;\r
54         }\r
55 \r
56 exit:\r
57         if (hToken != NULL) {\r
58                 info->pCloseHandle(hToken);\r
59         }\r
60 \r
61         if (hProcess != NULL) {\r
62                 info->pCloseHandle(hProcess);\r
63         }\r
64 \r
65         return result;\r
66 }\r
67 static int afterFunc(int arg)\r
68 {\r
69         // dummy operation\r
70         // if this function empty, optimizer unify this with other empty functions.\r
71         // following code avoid it.\r
72         arg *= 710810; // non-sense operation\r
73         return arg;\r
74 }\r
75 #pragma runtime_checks( "", restore )\r
76 \r
77 const DWORD FixScancodeMap::s_fixEntryNum = 4;\r
78 const DWORD FixScancodeMap::s_fixEntry[] = {\r
79         0x003ae03a,\r
80         0x0029e029,\r
81         0x0070e070,\r
82         0x007be07b,\r
83 };\r
84 \r
85 int FixScancodeMap::acquirePrivileges()\r
86 {\r
87         int ret = 0;\r
88         HANDLE hToken = NULL;\r
89 \r
90         if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hToken)) {\r
91                 ret = 5;\r
92                 goto exit;\r
93         }\r
94 \r
95         LUID luid;\r
96         if (!LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &luid)) {\r
97                 ret = 6;\r
98                 goto exit;\r
99         }\r
100 \r
101         TOKEN_PRIVILEGES tk_priv;\r
102         tk_priv.PrivilegeCount = 1;\r
103         tk_priv.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;\r
104         tk_priv.Privileges[0].Luid = luid;\r
105 \r
106         if (!AdjustTokenPrivileges(hToken, FALSE, &tk_priv, 0, NULL, NULL)) {\r
107                 ret = 7;\r
108                 goto exit;\r
109         }\r
110 \r
111 exit:\r
112         if (hToken != NULL) {\r
113             CloseHandle(hToken);\r
114         }\r
115         return ret;\r
116 }\r
117 \r
118 \r
119 DWORD FixScancodeMap::getWinLogonPid()\r
120 {\r
121     DWORD pid = 0;\r
122         DWORD mySessionId = 0;\r
123 \r
124         if (ProcessIdToSessionId(GetCurrentProcessId(), &mySessionId) == FALSE) {\r
125                 return 0;\r
126         }\r
127 \r
128         HANDLE hSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);\r
129         if (hSnap == INVALID_HANDLE_VALUE) {\r
130                 return 0;\r
131         }\r
132 \r
133         PROCESSENTRY32 pe;\r
134         pe.dwSize = sizeof(pe);\r
135 \r
136     BOOL bResult = Process32First(hSnap, &pe);\r
137         while (bResult){\r
138                 if (!_tcsicmp(pe.szExeFile, _T("winlogon.exe"))) {\r
139                         DWORD sessionId;\r
140 \r
141                         if (ProcessIdToSessionId(pe.th32ProcessID, &sessionId) != FALSE) {\r
142                                 if (sessionId == mySessionId) {\r
143                                         pid = pe.th32ProcessID;\r
144                                         break;\r
145                                 }\r
146                         }\r
147                 }\r
148                 bResult = Process32Next(hSnap, &pe);\r
149         }\r
150 \r
151         CloseHandle(hSnap);\r
152         return pid;\r
153 }\r
154 \r
155 \r
156 bool FixScancodeMap::clean(WlInfo wl)\r
157 {\r
158         int ret = 0;\r
159 \r
160         if (wl.m_hThread != NULL) {\r
161                 DWORD result;\r
162 \r
163                 if (WaitForSingleObject(wl.m_hThread, 5000) == WAIT_TIMEOUT) {\r
164                         return false;\r
165                 }\r
166 \r
167                 GetExitCodeThread(wl.m_hThread, &result);\r
168                 CloseHandle(wl.m_hThread);\r
169 \r
170                 if (wl.m_remoteMem != NULL && wl.m_hProcess != NULL) {\r
171                         VirtualFreeEx(wl.m_hProcess, wl.m_remoteMem, 0, MEM_RELEASE);\r
172                 }\r
173 \r
174                 if (wl.m_remoteInfo != NULL && wl.m_hProcess != NULL) {\r
175                         VirtualFreeEx(wl.m_hProcess, wl.m_remoteInfo, 0, MEM_RELEASE);\r
176                 }\r
177 \r
178                 if (wl.m_hProcess != NULL) {\r
179                         CloseHandle(wl.m_hProcess);\r
180                 }\r
181         }\r
182 \r
183         return true;\r
184 }\r
185 \r
186 \r
187 int FixScancodeMap::injectThread(DWORD dwPID)\r
188 {\r
189         int ret = 0;\r
190         DWORD err = 0;\r
191         BOOL wFlag;\r
192         WlInfo wi;\r
193 \r
194         wi.m_hProcess = NULL;\r
195         wi.m_remoteMem = NULL;\r
196         wi.m_remoteInfo = NULL;\r
197         wi.m_hThread = NULL;\r
198 \r
199         ULONG_PTR invokeFuncAddr = (ULONG_PTR)invokeFunc;\r
200         ULONG_PTR afterFuncAddr = (ULONG_PTR)afterFunc;\r
201         SIZE_T memSize =  afterFuncAddr - invokeFuncAddr;\r
202 \r
203         if ((wi.m_hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, dwPID)) == NULL) {\r
204                 ret = 8;\r
205                 goto exit;\r
206         }\r
207 \r
208         wi.m_remoteMem = VirtualAllocEx(wi.m_hProcess, NULL, memSize, MEM_COMMIT, PAGE_EXECUTE_READWRITE);\r
209         if (wi.m_remoteMem == NULL) {\r
210                 ret = 9;\r
211                 err = GetLastError();\r
212                 goto exit;\r
213         }\r
214 \r
215         wFlag = WriteProcessMemory(wi.m_hProcess, wi.m_remoteMem, (char*)invokeFunc, memSize, (SIZE_T*)0);\r
216         if (wFlag == FALSE) {\r
217                 ret = 10;\r
218                 goto exit;\r
219         }\r
220 \r
221         wi.m_remoteInfo = VirtualAllocEx(wi.m_hProcess, NULL, sizeof(m_info), MEM_COMMIT, PAGE_READWRITE);\r
222         if (wi.m_remoteInfo == NULL) {\r
223                 ret = 11;\r
224                 err = GetLastError();\r
225                 goto exit;\r
226         }\r
227 \r
228         wFlag = WriteProcessMemory(wi.m_hProcess, wi.m_remoteInfo, (char*)&m_info, sizeof(m_info), (SIZE_T*)0);\r
229         if (wFlag == FALSE) {\r
230                 ret = 12;\r
231                 goto exit;\r
232         }\r
233 \r
234 #if 0\r
235         TCHAR buf[1024];\r
236 \r
237         _stprintf_s(buf, sizeof(buf)/sizeof(buf[0]),\r
238                 _T("execute UpdatePerUserSystemParameters(), inject code to winlogon.exe?\r\n")\r
239                 _T("invokeFunc=0x%p\r\n")\r
240                 _T("afterFunc=0x%p\r\n")\r
241                 _T("afterFunc - invokeFunc=%d\r\n")\r
242                 _T("remoteMem=0x%p\r\n")\r
243                 _T("remoteInfo=0x%p(size: %d)\r\n"),\r
244                 invokeFunc, afterFunc, memSize, m_remoteMem, m_remoteInfo, sizeof(m_info));\r
245         if (MessageBox((HWND)NULL, buf, _T("upusp"), MB_OKCANCEL | MB_ICONSTOP) == IDCANCEL) {\r
246                 (m_info.pUpdate)(0, 1);\r
247                 goto exit;\r
248         }\r
249 #endif\r
250 \r
251         wi.m_hThread = CreateRemoteThread(wi.m_hProcess, NULL, 0, \r
252                 (LPTHREAD_START_ROUTINE)wi.m_remoteMem, wi.m_remoteInfo, 0, NULL);\r
253         if (wi.m_hThread == NULL) {\r
254                 ret = 13;\r
255                 goto exit;\r
256         }\r
257 \r
258         if (WaitForSingleObject(wi.m_hThread, 5000) == WAIT_TIMEOUT) {\r
259                 ret = 14;\r
260                 m_wlTrash.push_back(wi);\r
261                 goto dirty_exit;\r
262         }\r
263         DWORD result = -1;\r
264         GetExitCodeThread(wi.m_hThread, &result);\r
265         ret = result;\r
266         CloseHandle(wi.m_hThread);\r
267         wi.m_hThread = NULL;\r
268 \r
269 exit:\r
270         if (wi.m_remoteMem != NULL && wi.m_hProcess != NULL) {\r
271                 VirtualFreeEx(wi.m_hProcess, wi.m_remoteMem, 0, MEM_RELEASE);\r
272                 wi.m_remoteMem = NULL;\r
273         }\r
274 \r
275         if (wi.m_remoteInfo != NULL && wi.m_hProcess != NULL) {\r
276                 VirtualFreeEx(wi.m_hProcess, wi.m_remoteInfo, 0, MEM_RELEASE);\r
277                 wi.m_remoteInfo = NULL;\r
278         }\r
279 \r
280         if (wi.m_hProcess != NULL) {\r
281                 CloseHandle(wi.m_hProcess);\r
282                 wi.m_hProcess = NULL;\r
283         }\r
284 \r
285 dirty_exit:\r
286         return ret;\r
287 }\r
288 \r
289 int FixScancodeMap::update()\r
290 {\r
291         MINIMIZEDMETRICS mm;\r
292         int result = 0;\r
293 \r
294         if (m_errorOnConstruct) {\r
295                 result = m_errorOnConstruct;\r
296                 goto exit;\r
297         }\r
298 \r
299         m_wlTrash.erase(remove_if(m_wlTrash.begin(), m_wlTrash.end(), FixScancodeMap::clean), m_wlTrash.end());\r
300 \r
301         memset(&mm, 0, sizeof(mm));\r
302         mm.cbSize = sizeof(mm);\r
303         SystemParametersInfo(SPI_GETMINIMIZEDMETRICS, sizeof(mm), &mm, 0);\r
304 \r
305         result = injectThread(m_winlogonPid);\r
306         if (result == 14) {\r
307                 // retry once\r
308                 result = injectThread(m_winlogonPid);\r
309                 if (result == 0) {\r
310                         result = 22;\r
311                 }\r
312         }\r
313 \r
314         mm.iArrange = ARW_HIDE;\r
315         SystemParametersInfo(SPI_SETMINIMIZEDMETRICS, sizeof(mm), &mm, 0);\r
316 \r
317 exit:\r
318         return result;\r
319 }\r
320 \r
321 int FixScancodeMap::fix()\r
322 {\r
323         ScancodeMap *origMap, *fixMap;\r
324         DWORD origSize, fixSize;\r
325         bool ret;\r
326         int result = 0;\r
327 \r
328         // save original Scancode Map\r
329         ret = m_pReg->read(_T("Scancode Map"), NULL, &origSize, NULL, 0);\r
330         if (ret) {\r
331                 origMap = reinterpret_cast<ScancodeMap*>(malloc(origSize));\r
332                 if (origMap == NULL) {\r
333                         result = 16;\r
334                         goto exit;\r
335                 }\r
336 \r
337                 ret = m_pReg->read(_T("Scancode Map"), reinterpret_cast<BYTE*>(origMap), &origSize, NULL, 0);\r
338                 if (ret == false) {\r
339                         result = 17;\r
340                         goto exit;\r
341                 }\r
342 \r
343                 fixSize = origSize;\r
344                 fixMap = reinterpret_cast<ScancodeMap*>(malloc(origSize + s_fixEntryNum * sizeof(s_fixEntry[0])));\r
345                 if (fixMap == NULL) {\r
346                         result = 18;\r
347                         goto exit;\r
348                 }\r
349 \r
350                 memcpy_s(fixMap, origSize + s_fixEntryNum, origMap, origSize);\r
351         } else {\r
352                 origSize = 0;\r
353                 origMap = NULL;\r
354 \r
355                 fixSize = sizeof(ScancodeMap);\r
356                 fixMap = reinterpret_cast<ScancodeMap*>(malloc(sizeof(ScancodeMap) + s_fixEntryNum * sizeof(s_fixEntry[0])));\r
357                 if (fixMap == NULL) {\r
358                         result = 19;\r
359                         goto exit;\r
360                 }\r
361 \r
362                 fixMap->header1 = 0;\r
363                 fixMap->header2 = 0;\r
364                 fixMap->count = 1;\r
365                 fixMap->entry[0] = 0;\r
366         }\r
367 \r
368         for (DWORD i = 0; i < s_fixEntryNum; i++) {\r
369                 bool skip = false;\r
370 \r
371                 if (origMap) {\r
372                         for (DWORD j = 0; j < origMap->count; j++) {\r
373                                 if (HIWORD(s_fixEntry[i]) == HIWORD(origMap->entry[j])) {\r
374                                         skip = true;\r
375                                 }\r
376                         }\r
377                 }\r
378 \r
379                 if (skip) {\r
380                         // s_fixEntry[i] found in original Scancode Map, so don't fix it\r
381                         continue;\r
382                 }\r
383 \r
384                 // add fix entry to fixMap\r
385                 fixMap->entry[fixMap->count - 1] = s_fixEntry[i];\r
386                 fixMap->entry[fixMap->count] = 0;\r
387                 fixMap->count++;\r
388                 fixSize += 4;\r
389         }\r
390 \r
391         ret = m_pReg->write(_T("Scancode Map"), reinterpret_cast<BYTE*>(fixMap), fixSize);\r
392         if (ret == false) {\r
393                 result = 20;\r
394                 goto exit;\r
395         }\r
396 \r
397         result = update();\r
398 \r
399         if (origMap) {\r
400                 ret = m_pReg->write(_T("Scancode Map"), reinterpret_cast<BYTE*>(origMap), origSize);\r
401         } else {\r
402                 ret = m_pReg->remove(_T("Scancode Map"));\r
403         }\r
404         if (ret == false) {\r
405                 result = 21;\r
406                 goto exit;\r
407         }\r
408 \r
409 exit:\r
410         free(origMap);\r
411         origMap = NULL;\r
412 \r
413         free(fixMap);\r
414         fixMap = NULL;\r
415 \r
416         return result;\r
417 }\r
418 \r
419 int FixScancodeMap::restore()\r
420 {\r
421         return update();\r
422 }\r
423 \r
424 int FixScancodeMap::escape(bool i_escape)\r
425 {\r
426         if (i_escape) {\r
427                 SetEvent(m_hFixEvent);\r
428         } else {\r
429                 SetEvent(m_hRestoreEvent);\r
430         }\r
431         return 0;\r
432 }\r
433 \r
434 unsigned int WINAPI FixScancodeMap::threadLoop(void *i_this)\r
435 {\r
436         int err;\r
437         DWORD ret;\r
438         FixScancodeMap *This = reinterpret_cast<FixScancodeMap*>(i_this);\r
439         HANDLE handles[] = {This->m_hFixEvent, This->m_hRestoreEvent, This->m_hQuitEvent};\r
440         while ((ret = MsgWaitForMultipleObjects(NUMBER_OF(handles), &handles[0],\r
441                 FALSE, INFINITE, QS_POSTMESSAGE)) != WAIT_FAILED) {\r
442                 switch (ret) {\r
443                 case WAIT_OBJECT_0:                     // m_hFixEvent\r
444                         ResetEvent(This->m_hFixEvent);\r
445                         err = This->fix();\r
446                         PostMessage(This->m_hwnd, This->m_messageOnFail, err, 1);\r
447                         break;\r
448                 case WAIT_OBJECT_0 + 1:         // m_hRestoreEvent\r
449                         ResetEvent(This->m_hRestoreEvent);\r
450                         err = This->restore();\r
451                         PostMessage(This->m_hwnd, This->m_messageOnFail, err, 0);\r
452                         break;\r
453                 case WAIT_OBJECT_0 + 2:         // m_hQuiteEvent\r
454                         ResetEvent(This->m_hQuitEvent);\r
455                         // through below\r
456                 default:\r
457                         return 0;\r
458                         break;\r
459                 }\r
460         }\r
461         return 1;\r
462 }\r
463 \r
464 int FixScancodeMap::init(HWND i_hwnd, UINT i_messageOnFail)\r
465 {\r
466         m_hwnd = i_hwnd;\r
467         m_messageOnFail = i_messageOnFail;\r
468         return 0;\r
469 }\r
470 \r
471 FixScancodeMap::FixScancodeMap() :\r
472         m_hwnd(NULL),\r
473         m_messageOnFail(WM_NULL),\r
474         m_errorOnConstruct(0),\r
475         m_winlogonPid(0),\r
476         m_regHKCU(HKEY_CURRENT_USER, _T("Keyboard Layout")),\r
477         m_regHKLM(HKEY_LOCAL_MACHINE, _T("SYSTEM\\CurrentControlSet\\Control\\Keyboard Layout")),\r
478         m_pReg(NULL)\r
479 {\r
480         HMODULE hMod;\r
481 \r
482         m_info.pid_ = GetCurrentProcessId();\r
483 \r
484         memcpy(&m_info.advapi32_, _T("advapi32.dll"), sizeof(m_info.advapi32_));\r
485         memcpy(&m_info.impersonateLoggedOnUser_, "ImpersonateLoggedOnUser", sizeof(m_info.impersonateLoggedOnUser_));\r
486         memcpy(&m_info.revertToSelf_, "RevertToSelf", sizeof(m_info.revertToSelf_));\r
487         memcpy(&m_info.openProcessToken_, "OpenProcessToken", sizeof(m_info.openProcessToken_));\r
488 \r
489         m_hFixEvent = CreateEvent(NULL, TRUE, FALSE, NULL);\r
490         ASSERT(m_hFixEvent);\r
491         m_hRestoreEvent = CreateEvent(NULL, TRUE, FALSE, NULL);\r
492         ASSERT(m_hRestoreEvent);\r
493         m_hQuitEvent = CreateEvent(NULL, TRUE, FALSE, NULL);\r
494         ASSERT(m_hQuitEvent);\r
495 \r
496         m_hThread = (HANDLE)_beginthreadex(NULL, 0, threadLoop, this, 0, &m_threadId);\r
497 \r
498         hMod = GetModuleHandle(_T("user32.dll"));\r
499         if (hMod != NULL) {\r
500                 m_info.pUpdate4 = (FpUpdatePerUserSystemParameters4)GetProcAddress(hMod, "UpdatePerUserSystemParameters");\r
501                 m_info.pUpdate8 = (FpUpdatePerUserSystemParameters8)m_info.pUpdate4;\r
502                 if (m_info.pUpdate4 == NULL) {\r
503                         return;\r
504                 }\r
505         }\r
506 \r
507         hMod = GetModuleHandle(_T("kernel32.dll"));\r
508         if (hMod != NULL) {\r
509                 m_info.pGetModuleHandle = (FpGetModuleHandleW)GetProcAddress(hMod, "GetModuleHandleW");\r
510                 if (m_info.pGetModuleHandle == NULL) {\r
511                         return;\r
512                 }\r
513 \r
514                 m_info.pGetProcAddress = (FpGetProcAddress)GetProcAddress(hMod, "GetProcAddress");\r
515                 if (m_info.pGetProcAddress == NULL) {\r
516                         return;\r
517                 }\r
518 \r
519                 m_info.pOpenProcess = (FpOpenProcess)GetProcAddress(hMod, "OpenProcess");\r
520                 if (m_info.pOpenProcess == NULL) {\r
521                         return;\r
522                 }\r
523 \r
524                 m_info.pCloseHandle = (FpCloseHandle)GetProcAddress(hMod, "CloseHandle");\r
525                 if (m_info.pCloseHandle == NULL) {\r
526                         return;\r
527                 }\r
528         }\r
529 \r
530         // Windows7 RC not support Scancode Map on HKCU?\r
531         if (checkWindowsVersion(6, 1) == FALSE) {\r
532                 m_pReg = &m_regHKCU; // Vista or earlier\r
533         } else {\r
534                 m_pReg = &m_regHKLM; // Windows7 or later\r
535         }\r
536 \r
537         // prototype of UpdatePerUserSystemParameters() differ vista or earlier\r
538         if (checkWindowsVersion(6, 0) == FALSE) {\r
539                 m_info.isVistaOrLater_ = 0; // before Vista\r
540         } else {\r
541                 m_info.isVistaOrLater_ = 1; // Vista or later\r
542         }\r
543 \r
544         m_errorOnConstruct = acquirePrivileges();\r
545         if (m_errorOnConstruct) {\r
546                 goto exit;\r
547         }\r
548 \r
549         if ((m_winlogonPid = getWinLogonPid()) == 0) {\r
550                 m_errorOnConstruct = 15;\r
551                 goto exit;\r
552         }\r
553 \r
554 exit:\r
555         ;\r
556 }\r
557 \r
558 FixScancodeMap::~FixScancodeMap()\r
559 {\r
560         SetEvent(m_hQuitEvent);\r
561         WaitForSingleObject(m_hThread, INFINITE);\r
562 }\r