devices.cpp File Reference

#include "shellext.h"
#include "devices.h"
#include "resource.h"
#include "balance.h"
#include <stddef.h>
#include <uxtheme.h>
#include <setupapi.h>
#include <strsafe.h>
#include <mountmgr.h>
#include <algorithm>
#include "../btrfs.h"

Include dependency graph for devices.cpp:

Go to the source code of this file.

Macros
#define	ISOLATION_AWARE_ENABLED   1
#define	STRSAFE_NO_DEPRECATE

Functions
	DEFINE_GUID (GUID_DEVINTERFACE_HIDDEN_VOLUME, 0x7f108a28L, 0x9833, 0x4b3b, 0xb7, 0x80, 0x2c, 0x6b, 0x5f, 0xa5, 0xc0, 0x62)
static wstring	get_mountdev_name (const nt_handle &h)
static void	find_devices (HWND hwnd, const GUID *guid, const nt_handle &mountmgr, vector< device > &device_list)
static bool	sort_devices (device i, device j)
static INT_PTR CALLBACK	stub_DeviceAddDlgProc (HWND hwndDlg, UINT uMsg, WPARAM wParam, LPARAM lParam)
static INT_PTR CALLBACK	stub_DeviceResizeDlgProc (HWND hwndDlg, UINT uMsg, WPARAM wParam, LPARAM lParam)
void CALLBACK	AddDeviceW (HWND hwnd, HINSTANCE hinst, LPWSTR lpszCmdLine, int nCmdShow)
void CALLBACK	RemoveDeviceW (HWND hwnd, HINSTANCE hinst, LPWSTR lpszCmdLine, int nCmdShow)
void CALLBACK	ResizeDeviceW (HWND hwnd, HINSTANCE hinst, LPWSTR lpszCmdLine, int nCmdShow)

Macro Definition Documentation

ISOLATION_AWARE_ENABLED

#define ISOLATION_AWARE_ENABLED   1

Definition at line 18 of file devices.cpp.

STRSAFE_NO_DEPRECATE

#define STRSAFE_NO_DEPRECATE

Definition at line 19 of file devices.cpp.

Function Documentation

AddDeviceW()

void CALLBACK AddDeviceW	(	HWND	hwnd,
		HINSTANCE	hinst,
		LPWSTR	lpszCmdLine,
		int	nCmdShow
	)

Definition at line 878 of file devices.cpp.

                                                                                       {
    try {
        win_handle token;
        TOKEN_PRIVILEGES tp;
        LUID luid;

        set_dpi_aware();

        if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &token))
            throw last_error(GetLastError());

        if (!LookupPrivilegeValueW(nullptr, L"SeManageVolumePrivilege", &luid))
            throw last_error(GetLastError());

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

        if (!AdjustTokenPrivileges(token, false, &tp, sizeof(TOKEN_PRIVILEGES), nullptr, nullptr))
            throw last_error(GetLastError());

        BtrfsDeviceAdd bda(hinst, hwnd, lpszCmdLine);
        bda.ShowDialog();
    } catch (const exception& e) {
        error_message(hwnd, e.what());
    }
}

DEFINE_GUID()

DEFINE_GUID	(	GUID_DEVINTERFACE_HIDDEN_VOLUME	,
		0x7f108a28L	,
		0x9833	,
		0x4b3b	,
		0xb7	,
		0x80	,
		0x2c	,
		0x6b	,
		0x5f	,
		0xa5	,
		0xc0	,
		0x62
	)

find_devices()

static void find_devices ( HWND  hwnd, const GUID *  guid, const nt_handle &  mountmgr, vector< device > &  device_list  )

static

Definition at line 72 of file devices.cpp.

                                                                                                              {
    HDEVINFO h;

    static WCHAR dosdevices[] = L"\\DosDevices\\";

    h = SetupDiGetClassDevs(guid, nullptr, 0, DIGCF_DEVICEINTERFACE | DIGCF_PRESENT);

    if (h != INVALID_HANDLE_VALUE) {
        DWORD index = 0;
        SP_DEVICE_INTERFACE_DATA did;

        did.cbSize = sizeof(did);

        if (!SetupDiEnumDeviceInterfaces(h, nullptr, guid, index, &did))
            return;

        do {
            SP_DEVINFO_DATA dd;
            SP_DEVICE_INTERFACE_DETAIL_DATA_W* detail;
            DWORD size;

            dd.cbSize = sizeof(dd);

            SetupDiGetDeviceInterfaceDetailW(h, &did, nullptr, 0, &size, nullptr);

            detail = (SP_DEVICE_INTERFACE_DETAIL_DATA_W*)malloc(size);
            memset(detail, 0, size);

            detail->cbSize = sizeof(SP_DEVICE_INTERFACE_DETAIL_DATA_W);

            if (SetupDiGetDeviceInterfaceDetailW(h, &did, detail, size, &size, &dd)) {
                NTSTATUS Status;
                nt_handle file;
                device dev;
                STORAGE_DEVICE_NUMBER sdn;
                IO_STATUS_BLOCK iosb;
                UNICODE_STRING path;
                OBJECT_ATTRIBUTES attr;
                GET_LENGTH_INFORMATION gli;
                ULONG i;
                uint8_t sb[4096];

                path.Buffer = detail->DevicePath;
                path.Length = path.MaximumLength = (uint16_t)(wcslen(detail->DevicePath) * sizeof(WCHAR));

                if (path.Length > 4 * sizeof(WCHAR) && path.Buffer[0] == '\\' && path.Buffer[1] == '\\'  && path.Buffer[2] == '?'  && path.Buffer[3] == '\\')
                    path.Buffer[1] = '?';

                InitializeObjectAttributes(&attr, &path, 0, nullptr, nullptr);

                Status = NtOpenFile(&file, FILE_GENERIC_READ, &attr, &iosb, FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, FILE_SYNCHRONOUS_IO_ALERT);

                if (!NT_SUCCESS(Status)) {
                    free(detail);
                    index++;
                    continue;
                }

                dev.pnp_name = detail->DevicePath;

                Status = NtDeviceIoControlFile(file, nullptr, nullptr, nullptr, &iosb, IOCTL_DISK_GET_LENGTH_INFO, nullptr, 0, &gli, sizeof(GET_LENGTH_INFORMATION));
                if (!NT_SUCCESS(Status)) {
                    free(detail);
                    index++;
                    continue;
                }

                dev.size = gli.Length.QuadPart;

                Status = NtDeviceIoControlFile(file, nullptr, nullptr, nullptr, &iosb, IOCTL_STORAGE_GET_DEVICE_NUMBER, nullptr, 0, &sdn, sizeof(STORAGE_DEVICE_NUMBER));
                if (!NT_SUCCESS(Status)) {
                    dev.disk_num = 0xffffffff;
                    dev.part_num = 0xffffffff;
                } else {
                    dev.disk_num = sdn.DeviceNumber;
                    dev.part_num = sdn.PartitionNumber;
                }

                dev.friendly_name = L"";
                dev.drive = L"";
                dev.fstype = L"";
                dev.has_parts = false;
                dev.ignore = false;
                dev.multi_device = false;

                dev.is_disk = RtlCompareMemory(guid, &GUID_DEVINTERFACE_DISK, sizeof(GUID)) == sizeof(GUID);

                if (dev.is_disk) {
                    STORAGE_PROPERTY_QUERY spq;
                    STORAGE_DEVICE_DESCRIPTOR sdd, *sdd2;
                    ULONG dlisize;
                    DRIVE_LAYOUT_INFORMATION_EX* dli;

                    spq.PropertyId = StorageDeviceProperty;
                    spq.QueryType = PropertyStandardQuery;
                    spq.AdditionalParameters[0] = 0;

                    Status = NtDeviceIoControlFile(file, nullptr, nullptr, nullptr, &iosb, IOCTL_STORAGE_QUERY_PROPERTY,
                                                   &spq, sizeof(STORAGE_PROPERTY_QUERY), &sdd, sizeof(STORAGE_DEVICE_DESCRIPTOR));

                    if (NT_SUCCESS(Status) || Status == STATUS_BUFFER_OVERFLOW) {
                        sdd2 = (STORAGE_DEVICE_DESCRIPTOR*)malloc(sdd.Size);

                        Status = NtDeviceIoControlFile(file, nullptr, nullptr, nullptr, &iosb, IOCTL_STORAGE_QUERY_PROPERTY,
                                                       &spq, sizeof(STORAGE_PROPERTY_QUERY), sdd2, sdd.Size);
                        if (NT_SUCCESS(Status)) {
                            string desc2;

                            desc2 = "";

                            if (sdd2->VendorIdOffset != 0) {
                                desc2 += (char*)((uint8_t*)sdd2 + sdd2->VendorIdOffset);

                                while (desc2.length() > 0 && desc2[desc2.length() - 1] == ' ')
                                    desc2 = desc2.substr(0, desc2.length() - 1);
                            }

                            if (sdd2->ProductIdOffset != 0) {
                                if (sdd2->VendorIdOffset != 0 && desc2.length() != 0 && desc2[desc2.length() - 1] != ' ')
                                    desc2 += " ";

                                desc2 += (char*)((uint8_t*)sdd2 + sdd2->ProductIdOffset);

                                while (desc2.length() > 0 && desc2[desc2.length() - 1] == ' ')
                                    desc2 = desc2.substr(0, desc2.length() - 1);
                            }

                            if (sdd2->VendorIdOffset != 0 || sdd2->ProductIdOffset != 0) {
                                ULONG ss;

                                ss = MultiByteToWideChar(CP_OEMCP, MB_PRECOMPOSED, desc2.c_str(), -1, nullptr, 0);

                                if (ss > 0) {
                                    WCHAR* desc3 = (WCHAR*)malloc(ss * sizeof(WCHAR));

                                    if (MultiByteToWideChar(CP_OEMCP, MB_PRECOMPOSED, desc2.c_str(), -1, desc3, ss * sizeof(WCHAR)))
                                        dev.friendly_name = desc3;

                                    free(desc3);
                                }
                            }
                        }

                        free(sdd2);
                    }

                    dlisize = 0;
                    dli = nullptr;

                    do {
                        dlisize += 1024;

                        if (dli)
                            free(dli);

                        dli = (DRIVE_LAYOUT_INFORMATION_EX*)malloc(dlisize);

                        Status = NtDeviceIoControlFile(file, nullptr, nullptr, nullptr, &iosb, IOCTL_DISK_GET_DRIVE_LAYOUT_EX,
                                                       nullptr, 0, dli, dlisize);
                    } while (Status == STATUS_BUFFER_TOO_SMALL);

                    if (NT_SUCCESS(Status) && dli->PartitionCount > 0)
                        dev.has_parts = true;

                    free(dli);
                } else {
                    ULONG mmpsize;
                    MOUNTMGR_MOUNT_POINT* mmp;
                    MOUNTMGR_MOUNT_POINTS mmps;

                    mmpsize = sizeof(MOUNTMGR_MOUNT_POINT) + path.Length;

                    mmp = (MOUNTMGR_MOUNT_POINT*)malloc(mmpsize);

                    RtlZeroMemory(mmp, sizeof(MOUNTMGR_MOUNT_POINT));
                    mmp->DeviceNameOffset = sizeof(MOUNTMGR_MOUNT_POINT);
                    mmp->DeviceNameLength = path.Length;
                    RtlCopyMemory(&mmp[1], path.Buffer, path.Length);

                    Status = NtDeviceIoControlFile(mountmgr, nullptr, nullptr, nullptr, &iosb, IOCTL_MOUNTMGR_QUERY_POINTS,
                                                   mmp, mmpsize, &mmps, sizeof(MOUNTMGR_MOUNT_POINTS));
                    if (NT_SUCCESS(Status) || Status == STATUS_BUFFER_OVERFLOW) {
                        MOUNTMGR_MOUNT_POINTS* mmps2;

                        mmps2 = (MOUNTMGR_MOUNT_POINTS*)malloc(mmps.Size);

                        Status = NtDeviceIoControlFile(mountmgr, nullptr, nullptr, nullptr, &iosb, IOCTL_MOUNTMGR_QUERY_POINTS,
                                                    mmp, mmpsize, mmps2, mmps.Size);

                        if (NT_SUCCESS(Status)) {
                            ULONG i;

                            for (i = 0; i < mmps2->NumberOfMountPoints; i++) {
                                WCHAR* symlink = (WCHAR*)((uint8_t*)mmps2 + mmps2->MountPoints[i].SymbolicLinkNameOffset);

                                if (mmps2->MountPoints[i].SymbolicLinkNameLength == 0x1c &&
                                    RtlCompareMemory(symlink, dosdevices, wcslen(dosdevices) * sizeof(WCHAR)) == wcslen(dosdevices) * sizeof(WCHAR) &&
                                    symlink[13] == ':'
                                ) {
                                    WCHAR dr[3];

                                    dr[0] = symlink[12];
                                    dr[1] = ':';
                                    dr[2] = 0;

                                    dev.drive = dr;
                                    break;
                                }
                            }
                        }
                    }

                    free(mmp);
                }

                if (!dev.is_disk || !dev.has_parts) {
                    i = 0;
                    while (fs_ident[i].name) {
                        if (i == 0 || fs_ident[i].kboff != fs_ident[i-1].kboff) {
                            LARGE_INTEGER off;

                            off.QuadPart = fs_ident[i].kboff * 1024;
                            Status = NtReadFile(file, nullptr, nullptr, nullptr, &iosb, sb, sizeof(sb), &off, nullptr);
                        }

                        if (NT_SUCCESS(Status)) {
                            if (RtlCompareMemory(sb + fs_ident[i].sboff, fs_ident[i].magic, fs_ident[i].magiclen) == fs_ident[i].magiclen) {
                                dev.fstype = fs_ident[i].name;

                                if (dev.fstype == L"Btrfs") {
                                    superblock* bsb = (superblock*)sb;

                                    RtlCopyMemory(&dev.fs_uuid, &bsb->uuid, sizeof(BTRFS_UUID));
                                    RtlCopyMemory(&dev.dev_uuid, &bsb->dev_item.device_uuid, sizeof(BTRFS_UUID));
                                }

                                break;
                            }
                        }

                        i++;
                    }

                    if (dev.fstype == L"Btrfs" && RtlCompareMemory(guid, &GUID_DEVINTERFACE_DISK, sizeof(GUID)) != sizeof(GUID)) {
                        wstring name;
                        wstring pref = L"\\Device\\Btrfs{";

                        name = get_mountdev_name(file);

                        if (name.length() > pref.length() && RtlCompareMemory(name.c_str(), pref.c_str(), pref.length() * sizeof(WCHAR)) == pref.length() * sizeof(WCHAR))
                            dev.ignore = true;
                    }
                }

                device_list.push_back(dev);
            }

            free(detail);

            index++;
        } while (SetupDiEnumDeviceInterfaces(h, nullptr, guid, index, &did));

        SetupDiDestroyDeviceInfoList(h);
    } else
        throw last_error(GetLastError());
}

Referenced by BtrfsDeviceAdd::populate_device_tree().

get_mountdev_name()

static wstring get_mountdev_name ( const nt_handle &  h )

static

Definition at line 42 of file devices.cpp.

                                                      {
    NTSTATUS Status;
    IO_STATUS_BLOCK iosb;
    MOUNTDEV_NAME mdn, *mdn2;
    ULONG mdnsize;
    wstring name;

    Status = NtDeviceIoControlFile(h, nullptr, nullptr, nullptr, &iosb, IOCTL_MOUNTDEV_QUERY_DEVICE_NAME,
                                   nullptr, 0, &mdn, sizeof(MOUNTDEV_NAME));
    if (!NT_SUCCESS(Status) && Status != STATUS_BUFFER_OVERFLOW)
        return L"";

    mdnsize = offsetof(MOUNTDEV_NAME, Name[0]) + mdn.NameLength;

    mdn2 = (MOUNTDEV_NAME*)malloc(mdnsize);

    Status = NtDeviceIoControlFile(h, nullptr, nullptr, nullptr, &iosb, IOCTL_MOUNTDEV_QUERY_DEVICE_NAME,
                                   nullptr, 0, mdn2, mdnsize);
    if (!NT_SUCCESS(Status)) {
        free(mdn2);
        return L"";
    }

    name = wstring(mdn2->Name, mdn2->NameLength / sizeof(WCHAR));

    free(mdn2);

    return name;
}

Referenced by find_devices().

RemoveDeviceW()

void CALLBACK RemoveDeviceW	(	HWND	hwnd,
		HINSTANCE	hinst,
		LPWSTR	lpszCmdLine,
		int	nCmdShow
	)

Definition at line 906 of file devices.cpp.

                                                                                          {
    try {
        WCHAR *s, *vol, *dev;
        uint64_t devid;
        win_handle h, token;
        TOKEN_PRIVILEGES tp;
        LUID luid;
        NTSTATUS Status;
        IO_STATUS_BLOCK iosb;

        set_dpi_aware();

        if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &token))
            throw last_error(GetLastError());

        if (!LookupPrivilegeValueW(nullptr, L"SeManageVolumePrivilege", &luid))
            throw last_error(GetLastError());

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

        if (!AdjustTokenPrivileges(token, false, &tp, sizeof(TOKEN_PRIVILEGES), nullptr, nullptr))
            throw last_error(GetLastError());

        s = wcsstr(lpszCmdLine, L"|");
        if (!s)
            return;

        s[0] = 0;

        vol = lpszCmdLine;
        dev = &s[1];

        devid = _wtoi(dev);
        if (devid == 0)
            return;

        h = CreateFileW(vol, FILE_TRAVERSE | FILE_READ_ATTRIBUTES, FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, nullptr,
                        OPEN_EXISTING, FILE_FLAG_BACKUP_SEMANTICS | FILE_FLAG_OPEN_REPARSE_POINT, nullptr);

        if (h == INVALID_HANDLE_VALUE)
            throw last_error(GetLastError());

        Status = NtFsControlFile(h, nullptr, nullptr, nullptr, &iosb, FSCTL_BTRFS_REMOVE_DEVICE, &devid, sizeof(uint64_t), nullptr, 0);
        if (!NT_SUCCESS(Status)) {
            if (Status == STATUS_CANNOT_DELETE)
                throw string_error(IDS_CANNOT_REMOVE_RAID);
            else
                throw ntstatus_error(Status);

            return;
        }

        BtrfsBalance bb(vol, true);
        bb.ShowBalance(hwnd);
    } catch (const exception& e) {
        error_message(hwnd, e.what());
    }
}

ResizeDeviceW()

void CALLBACK ResizeDeviceW	(	HWND	hwnd,
		HINSTANCE	hinst,
		LPWSTR	lpszCmdLine,
		int	nCmdShow
	)

Definition at line 967 of file devices.cpp.

                                                                                          {
    try {
        WCHAR *s, *vol, *dev;
        uint64_t devid;
        win_handle token;
        TOKEN_PRIVILEGES tp;
        LUID luid;

        set_dpi_aware();

        if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &token))
            throw last_error(GetLastError());

        if (!LookupPrivilegeValueW(nullptr, L"SeManageVolumePrivilege", &luid))
            throw last_error(GetLastError());

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

        if (!AdjustTokenPrivileges(token, false, &tp, sizeof(TOKEN_PRIVILEGES), nullptr, nullptr))
            throw last_error(GetLastError());

        s = wcsstr(lpszCmdLine, L"|");
        if (!s)
            return;

        s[0] = 0;

        vol = lpszCmdLine;
        dev = &s[1];

        devid = _wtoi(dev);
        if (devid == 0)
            return;

        BtrfsDeviceResize bdr;
        bdr.ShowDialog(hwnd, vol, devid);
    } catch (const exception& e) {
        error_message(hwnd, e.what());
    }
}

sort_devices()

static bool sort_devices ( device  i, device  j  )

static

Definition at line 340 of file devices.cpp.

                                             {
    if (i.disk_num < j.disk_num)
        return true;

    if (i.disk_num == j.disk_num && i.part_num < j.part_num)
        return true;

    return false;
}

Referenced by BtrfsDeviceAdd::populate_device_tree().

stub_DeviceAddDlgProc()

static INT_PTR CALLBACK stub_DeviceAddDlgProc ( HWND  hwndDlg, UINT  uMsg, WPARAM  wParam, LPARAM  lParam  )

static

Definition at line 657 of file devices.cpp.

                                                                                                     {
    BtrfsDeviceAdd* bda;

    if (uMsg == WM_INITDIALOG) {
        SetWindowLongPtr(hwndDlg, GWLP_USERDATA, (LONG_PTR)lParam);
        bda = (BtrfsDeviceAdd*)lParam;
    } else {
        bda = (BtrfsDeviceAdd*)GetWindowLongPtr(hwndDlg, GWLP_USERDATA);
    }

    if (bda)
        return bda->DeviceAddDlgProc(hwndDlg, uMsg, wParam, lParam);
    else
        return false;
}

Referenced by BtrfsDeviceAdd::ShowDialog().

stub_DeviceResizeDlgProc()

static INT_PTR CALLBACK stub_DeviceResizeDlgProc ( HWND  hwndDlg, UINT  uMsg, WPARAM  wParam, LPARAM  lParam  )

static

Definition at line 852 of file devices.cpp.

                                                                                                        {
    BtrfsDeviceResize* bdr;

    if (uMsg == WM_INITDIALOG) {
        SetWindowLongPtr(hwndDlg, GWLP_USERDATA, (LONG_PTR)lParam);
        bdr = (BtrfsDeviceResize*)lParam;
    } else
        bdr = (BtrfsDeviceResize*)GetWindowLongPtr(hwndDlg, GWLP_USERDATA);

    if (bdr)
        return bdr->DeviceResizeDlgProc(hwndDlg, uMsg, wParam, lParam);
    else
        return false;
}

Referenced by BtrfsDeviceResize::ShowDialog().