CPP/Windows/MemoryLock.cpp

   1 // Common/MemoryLock.cpp\r
   2 \r
   3 #include "StdAfx.h"\r
   4 \r
   5 namespace NWindows {\r
   6 namespace NSecurity {\r
   7 \r
   8 #ifndef UNDER_CE\r
   9 \r
  10 #ifndef _UNICODE\r
  11 typedef BOOL (WINAPI * OpenProcessTokenP)(HANDLE ProcessHandle, DWORD DesiredAccess, PHANDLE TokenHandle);\r
  12 typedef BOOL (WINAPI * LookupPrivilegeValueP)(LPCTSTR lpSystemName, LPCTSTR lpName, PLUID  lpLuid);\r
  13 typedef BOOL (WINAPI * AdjustTokenPrivilegesP)(HANDLE TokenHandle, BOOL DisableAllPrivileges,\r
  14     PTOKEN_PRIVILEGES NewState, DWORD BufferLength, PTOKEN_PRIVILEGES PreviousState,PDWORD ReturnLength);\r
  15 #endif\r
  16 \r
  17 #ifdef _UNICODE\r
  18 bool EnableLockMemoryPrivilege(\r
  19 #else\r
  20 static bool EnableLockMemoryPrivilege2(HMODULE hModule,\r
  21 #endif\r
  22 bool enable)\r
  23 {\r
  24   #ifndef _UNICODE\r
  25   if (hModule == NULL)\r
  26     return false;\r
  27   OpenProcessTokenP openProcessToken = (OpenProcessTokenP)GetProcAddress(hModule, "OpenProcessToken");\r
  28   LookupPrivilegeValueP lookupPrivilegeValue = (LookupPrivilegeValueP)GetProcAddress(hModule, "LookupPrivilegeValueA" );\r
  29   AdjustTokenPrivilegesP adjustTokenPrivileges = (AdjustTokenPrivilegesP)GetProcAddress(hModule, "AdjustTokenPrivileges");\r
  30   if (openProcessToken == NULL || adjustTokenPrivileges == NULL || lookupPrivilegeValue == NULL)\r
  31     return false;\r
  32   #endif\r
  33 \r
  34   HANDLE token;\r
  35   if (!\r
  36     #ifdef _UNICODE\r
  37     ::OpenProcessToken\r
  38     #else\r
  39     openProcessToken\r
  40     #endif\r
  41     (::GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &token))\r
  42     return false;\r
  43   TOKEN_PRIVILEGES tp;\r
  44   bool res = false;\r
  45   if (\r
  46     #ifdef _UNICODE\r
  47     ::LookupPrivilegeValue\r
  48     #else\r
  49     lookupPrivilegeValue\r
  50     #endif\r
  51     (NULL, SE_LOCK_MEMORY_NAME, &(tp.Privileges[0].Luid)))\r
  52   {\r
  53     tp.PrivilegeCount = 1;\r
  54     tp.Privileges[0].Attributes = enable ? SE_PRIVILEGE_ENABLED: 0;\r
  55     if (\r
  56       #ifdef _UNICODE\r
  57       ::AdjustTokenPrivileges\r
  58       #else\r
  59       adjustTokenPrivileges\r
  60       #endif\r
  61       (token, FALSE, &tp, 0, NULL, NULL))\r
  62       res = (GetLastError() == ERROR_SUCCESS);\r
  63   }\r
  64   ::CloseHandle(token);\r
  65   return res;\r
  66 }\r
  67 \r
  68 #ifndef _UNICODE\r
  69 bool EnableLockMemoryPrivilege(bool enable)\r
  70 {\r
  71   HMODULE hModule = LoadLibrary(TEXT("Advapi32.dll"));\r
  72   if (hModule == NULL)\r
  73     return false;\r
  74   bool res = EnableLockMemoryPrivilege2(hModule, enable);\r
  75   ::FreeLibrary(hModule);\r
  76   return res;\r
  77 }\r
  78 #endif\r
  79 \r
  80 #endif\r
  81 \r
  82 }}\r