Win32: Get large MMU pages for memory
Posted by Anonymous on Tue 16th Apr 2024 02:20
raw | new post

#define UNICODE
#define _UNICODE

#include <windows.h>
#include <ntsecapi.h>
#include <ntstatus.h>
#include <sddl.h>
#include <wchar.h>
#include <stdio.h>

void InitLsaString(PLSA_UNICODE_STRING LsaString, LPWSTR String)
{
    DWORD StringLength;

    if (String == NULL) {
        LsaString->Buffer = NULL;
        LsaString->Length = 0;
        LsaString->MaximumLength = 0;
        return;
    }

    StringLength = wcslen(String);
    LsaString->Buffer = String;
    LsaString->Length = (USHORT)StringLength * sizeof(WCHAR);
    LsaString->MaximumLength = (USHORT)(StringLength + 1) * sizeof(WCHAR);
}

NTSTATUS OpenPolicy(LPWSTR ServerName, DWORD DesiredAccess, PLSA_HANDLE PolicyHandle)
{
    LSA_OBJECT_ATTRIBUTES ObjectAttributes;
    LSA_UNICODE_STRING ServerString;
    PLSA_UNICODE_STRING Server = NULL;

    //
    // Always initialize the object attributes to all zeroes.
    //
    ZeroMemory(&ObjectAttributes, sizeof(ObjectAttributes));

    if (ServerName != NULL) {
        //
        // Make a LSA_UNICODE_STRING out of the LPWSTR passed in
        //
        InitLsaString(&ServerString, ServerName);
        Server = &ServerString;
    }

    //
    // Attempt to open the policy.
    //
    return LsaOpenPolicy(
        Server,
        &ObjectAttributes,
        DesiredAccess,
        PolicyHandle
    );
}

NTSTATUS SetPrivilegeOnAccount(LSA_HANDLE PolicyHandle, PSID AccountSid, LPWSTR PrivilegeName, BOOL bEnable)
{
    LSA_UNICODE_STRING PrivilegeString;

    //
    // Create a LSA_UNICODE_STRING for the privilege name.
    //
    InitLsaString(&PrivilegeString, PrivilegeName);

    //
    // grant or revoke the privilege, accordingly
    //
    if (bEnable) {
        return LsaAddAccountRights(
            PolicyHandle,       // open policy handle
            AccountSid,         // target SID
            &PrivilegeString,   // privileges
            1                   // privilege count
        );
    }
    else {
        return LsaRemoveAccountRights(
            PolicyHandle,       // open policy handle
            AccountSid,         // target SID
            FALSE,              // do not disable all rights
            &PrivilegeString,   // privileges
            1                   // privilege count
        );
    }
}

int main(int ac, char *av[])
{
    HANDLE hToken = NULL;

    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken))
    {
        printf("OpenProcessToken failed. GetLastError returned: %d\n", GetLastError());
        return -1;
    }

    DWORD dwBufferSize = 0;

    // Probe the buffer size reqired for PTOKEN_USER structure
    if (!GetTokenInformation(hToken, TokenUser, NULL, 0, &dwBufferSize) &&
        (GetLastError() != ERROR_INSUFFICIENT_BUFFER))
    {
        printf("GetTokenInformation failed. GetLastError returned: %d\n", GetLastError());

        // Cleanup
        CloseHandle(hToken);
        hToken = NULL;

        return -1;
    }

    PTOKEN_USER pTokenUser = (PTOKEN_USER) malloc(dwBufferSize);

    // Retrieve the token information in a TOKEN_USER structure
    if (!GetTokenInformation(
        hToken,
        TokenUser,
        pTokenUser,
        dwBufferSize,
        &dwBufferSize))
    {
        printf("GetTokenInformation failed. GetLastError returned: %d\n", GetLastError());

        // Cleanup
        CloseHandle(hToken);
        hToken = NULL;

        return -1;
    }

    // Print SID string
    LPWSTR strsid;
    ConvertSidToStringSid(pTokenUser->User.Sid, &strsid);
    printf("User SID: %S\n", strsid);

    // Cleanup
    CloseHandle(hToken);
    hToken = NULL;

    NTSTATUS status;
    LSA_HANDLE policyHandle;

    if (status = OpenPolicy(NULL, POLICY_CREATE_ACCOUNT | POLICY_LOOKUP_NAMES, &policyHandle))
    {
        printf("OpenPolicy %d\n", status);
    }
#if 0
    // Add new privelege to the account
    if (status = SetPrivilegeOnAccount(policyHandle, pTokenUser->User.Sid, SE_LOCK_MEMORY_NAME, TRUE))
    {
        printf("OpenPSetPrivilegeOnAccountolicy %d\n", status);
    }
#endif
    // Enable this priveledge for the current process
    hToken = NULL;
    TOKEN_PRIVILEGES tp;

    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, &hToken))
    {
        printf("OpenProcessToken #2 failed. GetLastError returned: %d\n", GetLastError());
        return -1;
    }

    tp.PrivilegeCount = 1;
    tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

    if (!LookupPrivilegeValue(NULL, SE_LOCK_MEMORY_NAME, &tp.Privileges[0].Luid))
    {
        printf("LookupPrivilegeValue failed. GetLastError returned: %d\n", GetLastError());
        return -1;
    }

    BOOL result = AdjustTokenPrivileges(hToken, FALSE, &tp, 0, (PTOKEN_PRIVILEGES)NULL, 0);
    DWORD error = GetLastError();

    if (!result || (error != ERROR_SUCCESS))
    {
        printf("AdjustTokenPrivileges failed. GetLastError returned: %d\n", error);
        return -1;
    }

    // Cleanup
    CloseHandle(hToken);
    hToken = NULL;

    SIZE_T pageSize = GetLargePageMinimum();
#define N_PAGES_TO_ALLOC 4
    // Finally allocate the memory
    char *largeBuffer = VirtualAlloc(NULL, pageSize * N_PAGES_TO_ALLOC, MEM_RESERVE | MEM_COMMIT | MEM_LARGE_PAGES, PAGE_READWRITE);
    if (largeBuffer)
    {
        printf("VirtualAlloc success, ptr = %p\n", largeBuffer);
    }
    else
    {
        printf("VirtualAlloc failed, error 0x%x\n", GetLastError());
    }
}