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