self_elevation.cpp

int APIENTRY _tWinMain(HINSTANCE hInstance,
                     HINSTANCE hPrevInstance,
                     LPTSTR    lpCmdLine,
                     int       nCmdShow)
{
	UNREFERENCED_PARAMETER(hPrevInstance);
	UNREFERENCED_PARAMETER(lpCmdLine);

    return DialogBox(hInstance, MAKEINTRESOURCE(IDD_DLGMAIN), NULL, DlgProc);

    return 0;
}

BOOL IsProcessElevated()
{
    BOOL isElevated = FALSE;
    HANDLE hToken = NULL;
    DWORD errCode = 0;

    BOOL ret = OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken);
    errCode = GetLastError();
    ON_SCOPE_EXIT([&]()
    {
        if (hToken)
        {
            CloseHandle(hToken);
            hToken = NULL;
        }
    });

    if (!ret)
    {
        throw errCode;
    }

    TOKEN_ELEVATION elevationInfo;
    DWORD retSize = 0;
    if (!GetTokenInformation(hToken, TokenElevation, &elevationInfo,
                             sizeof(elevationInfo), &retSize))
    {
        errCode = GetLastError();
        throw errCode;
    }

    isElevated = elevationInfo.TokenIsElevated;

    return isElevated;
}

void OnClose(HWND hwnd)
{
    EndDialog(hwnd, FALSE);
}

BOOL OnInitDialog(HWND hwnd, HWND hwndFocus, LPARAM lParam)
{
    try
    {
        BOOL isElevated = IsProcessElevated();
        if (!isElevated)
        {
            HWND hBtn = GetDlgItem(hwnd, IDC_BTNELEVATION);
            Button_SetElevationRequiredState(hBtn, TRUE);
        }
    }
    catch (DWORD errCode)
    {
        MessageBox(hwnd, L"Retrieve Process Elevation Information Failed!",
                   L"Warning", MB_OK | MB_ICONEXCLAMATION);
        OutputDebugString(L"Error Code: ");
        OutputDebugString(std::to_wstring(static_cast<UINT64>(errCode)).c_str());
    }

    return TRUE;
}

void OnCommand(HWND hWnd, int id, HWND hwndCtl, UINT codeNotify)
{
    switch (id)
    {
        case IDC_BTNELEVATION:
            BOOL isElevated = IsProcessElevated();
            if (isElevated)
            {
                MessageBox(hWnd, L"Arealdy run at full administrator token",
                           L"warning", MB_OK | MB_ICONEXCLAMATION);
                return;
            } 
            else
            {
                wchar_t path[MAX_PATH];
                if (GetModuleFileName(NULL, path, MAX_PATH))
                {
                    SHELLEXECUTEINFO sei = {sizeof(SHELLEXECUTEINFO)};
                    sei.lpVerb = L"runas";
                    sei.lpFile = path;
                    sei.hwnd = hWnd;
                    sei.nShow = SW_NORMAL;

                    if (!ShellExecuteEx(&sei))
                    {
                        DWORD errCode = GetLastError();
                        if (errCode == ERROR_CANCELLED)
                        {
                            // cancelled elevation
                        }
                    }
                    else
                    {
                        EndDialog(hWnd, TRUE);
                    }
                }
            }

            break;
    }
}

INT_PTR CALLBACK DlgProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam)
{
    switch (msg)
    {
        HANDLE_MSG(hwnd, WM_INITDIALOG, OnInitDialog);
        HANDLE_MSG(hwnd, WM_COMMAND, OnCommand);
        HANDLE_MSG(hwnd, WM_CLOSE, OnClose);

        default:
            return FALSE;
    }

    return FALSE;
}