dll/shellext/shellbtrfs/devices.cpp

   1 /* Copyright (c) Mark Harmstone 2016-17
   2  *
   3  * This file is part of WinBtrfs.
   4  *
   5  * WinBtrfs is free software: you can redistribute it and/or modify
   6  * it under the terms of the GNU Lesser General Public Licence as published by
   7  * the Free Software Foundation, either version 3 of the Licence, or
   8  * (at your option) any later version.
   9  *
  10  * WinBtrfs is distributed in the hope that it will be useful,
  11  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  12  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13  * GNU Lesser General Public Licence for more details.
  14  *
  15  * You should have received a copy of the GNU Lesser General Public Licence
  16  * along with WinBtrfs.  If not, see <http://www.gnu.org/licenses/>. */
  17
  18 #define ISOLATION_AWARE_ENABLED 1
  19 #define STRSAFE_NO_DEPRECATE
  20
  21 #include "shellext.h"
  22 #include "devices.h"
  23 #include "resource.h"
  24 #include "balance.h"
  25 #include <stddef.h>
  26 #include <uxtheme.h>
  27 #include <setupapi.h>
  28 #include <strsafe.h>
  29 #include <mountmgr.h>
  30 #include <algorithm>
  31 #ifndef __REACTOS__
  32 #include "../btrfs.h"
  33 #else
  34 #include <ntddstor.h>
  35 #include <ndk/rtlfuncs.h>
  36 #include <ndk/obfuncs.h>
  37 #include "btrfs.h"
  38 #endif
  39
  40 DEFINE_GUID(GUID_DEVINTERFACE_HIDDEN_VOLUME, 0x7f108a28L, 0x9833, 0x4b3b, 0xb7, 0x80, 0x2c, 0x6b, 0x5f, 0xa5, 0xc0, 0x62);
  41
  42 static std::wstring get_mountdev_name(HANDLE h) {
  43     NTSTATUS Status;
  44     IO_STATUS_BLOCK iosb;
  45     MOUNTDEV_NAME mdn, *mdn2;
  46     ULONG mdnsize;
  47     std::wstring name;
  48
  49     Status = NtDeviceIoControlFile(h, NULL, NULL, NULL, &iosb, IOCTL_MOUNTDEV_QUERY_DEVICE_NAME,
  50                                    NULL, 0, &mdn, sizeof(MOUNTDEV_NAME));
  51     if (!NT_SUCCESS(Status) && Status != STATUS_BUFFER_OVERFLOW)
  52         return L"";
  53
  54     mdnsize = offsetof(MOUNTDEV_NAME, Name[0]) + mdn.NameLength;
  55
  56     mdn2 = (MOUNTDEV_NAME*)malloc(mdnsize);
  57
  58     Status = NtDeviceIoControlFile(h, NULL, NULL, NULL, &iosb, IOCTL_MOUNTDEV_QUERY_DEVICE_NAME,
  59                                    NULL, 0, mdn2, mdnsize);
  60     if (!NT_SUCCESS(Status)) {
  61         free(mdn2);
  62         return L"";
  63     }
  64
  65     name = std::wstring(mdn2->Name, mdn2->NameLength / sizeof(WCHAR));
  66
  67     free(mdn2);
  68
  69     return name;
  70 }
  71
  72 static void find_devices(HWND hwnd, const GUID* guid, HANDLE mountmgr, std::vector<device>* device_list) {
  73     HDEVINFO h;
  74
  75     static WCHAR dosdevices[] = L"\\DosDevices\\";
  76
  77     h = SetupDiGetClassDevs(guid, NULL, 0, DIGCF_DEVICEINTERFACE | DIGCF_PRESENT);
  78
  79     if (h != INVALID_HANDLE_VALUE) {
  80         DWORD index = 0;
  81         SP_DEVICE_INTERFACE_DATA did;
  82
  83         did.cbSize = sizeof(did);
  84
  85         if (!SetupDiEnumDeviceInterfaces(h, NULL, guid, index, &did))
  86             return;
  87
  88         do {
  89             SP_DEVINFO_DATA dd;
  90             SP_DEVICE_INTERFACE_DETAIL_DATA_W* detail;
  91             DWORD size;
  92
  93             dd.cbSize = sizeof(dd);
  94
  95             SetupDiGetDeviceInterfaceDetailW(h, &did, NULL, 0, &size, NULL);
  96
  97             detail = (SP_DEVICE_INTERFACE_DETAIL_DATA_W*)malloc(size);
  98             memset(detail, 0, size);
  99
 100             detail->cbSize = sizeof(SP_DEVICE_INTERFACE_DETAIL_DATA_W);
 101
 102             if (SetupDiGetDeviceInterfaceDetailW(h, &did, detail, size, &size, &dd)) {
 103                 NTSTATUS Status;
 104                 HANDLE file;
 105                 device dev;
 106                 STORAGE_DEVICE_NUMBER sdn;
 107                 IO_STATUS_BLOCK iosb;
 108                 UNICODE_STRING path;
 109                 OBJECT_ATTRIBUTES attr;
 110                 GET_LENGTH_INFORMATION gli;
 111                 ULONG i;
 112                 UINT8 sb[4096];
 113
 114                 path.Buffer = detail->DevicePath;
 115                 path.Length = path.MaximumLength = wcslen(detail->DevicePath) * sizeof(WCHAR);
 116
 117                 if (path.Length > 4 * sizeof(WCHAR) && path.Buffer[0] == '\\' && path.Buffer[1] == '\\'  && path.Buffer[2] == '?'  && path.Buffer[3] == '\\')
 118                     path.Buffer[1] = '?';
 119
 120                 InitializeObjectAttributes(&attr, &path, 0, NULL, NULL);
 121
 122                 Status = NtOpenFile(&file, FILE_GENERIC_READ, &attr, &iosb, FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, FILE_SYNCHRONOUS_IO_ALERT);
 123
 124                 if (!NT_SUCCESS(Status))
 125                     goto nextitem2;
 126
 127                 dev.pnp_name = detail->DevicePath;
 128
 129                 Status = NtDeviceIoControlFile(file, NULL, NULL, NULL, &iosb, IOCTL_DISK_GET_LENGTH_INFO, NULL, 0, &gli, sizeof(GET_LENGTH_INFORMATION));
 130                 if (!NT_SUCCESS(Status))
 131                     goto nextitem;
 132
 133                 dev.size = gli.Length.QuadPart;
 134
 135                 Status = NtDeviceIoControlFile(file, NULL, NULL, NULL, &iosb, IOCTL_STORAGE_GET_DEVICE_NUMBER, NULL, 0, &sdn, sizeof(STORAGE_DEVICE_NUMBER));
 136                 if (!NT_SUCCESS(Status)) {
 137                     dev.disk_num = 0xffffffff;
 138                     dev.part_num = 0xffffffff;
 139                 } else {
 140                     dev.disk_num = sdn.DeviceNumber;
 141                     dev.part_num = sdn.PartitionNumber;
 142                 }
 143
 144                 dev.friendly_name = L"";
 145                 dev.drive = L"";
 146                 dev.fstype = L"";
 147                 dev.has_parts = FALSE;
 148                 dev.ignore = FALSE;
 149                 dev.multi_device = FALSE;
 150
 151                 dev.is_disk = RtlCompareMemory(guid, &GUID_DEVINTERFACE_DISK, sizeof(GUID)) == sizeof(GUID);
 152
 153                 if (dev.is_disk) {
 154                     STORAGE_PROPERTY_QUERY spq;
 155                     STORAGE_DEVICE_DESCRIPTOR sdd, *sdd2;
 156                     ULONG dlisize;
 157                     DRIVE_LAYOUT_INFORMATION_EX* dli;
 158
 159                     spq.PropertyId = StorageDeviceProperty;
 160                     spq.QueryType = PropertyStandardQuery;
 161                     spq.AdditionalParameters[0] = 0;
 162
 163                     Status = NtDeviceIoControlFile(file, NULL, NULL, NULL, &iosb, IOCTL_STORAGE_QUERY_PROPERTY,
 164                                                    &spq, sizeof(STORAGE_PROPERTY_QUERY), &sdd, sizeof(STORAGE_DEVICE_DESCRIPTOR));
 165
 166                     if (NT_SUCCESS(Status) || Status == STATUS_BUFFER_OVERFLOW) {
 167                         sdd2 = (STORAGE_DEVICE_DESCRIPTOR*)malloc(sdd.Size);
 168
 169                         Status = NtDeviceIoControlFile(file, NULL, NULL, NULL, &iosb, IOCTL_STORAGE_QUERY_PROPERTY,
 170                                                     &spq, sizeof(STORAGE_PROPERTY_QUERY), sdd2, sdd.Size);
 171                         if (NT_SUCCESS(Status)) {
 172                             std::string desc2;
 173
 174                             desc2 = "";
 175
 176                             if (sdd2->VendorIdOffset != 0) {
 177                                 desc2 += (char*)((UINT8*)sdd2 + sdd2->VendorIdOffset);
 178
 179                                 while (desc2.length() > 0 && desc2[desc2.length() - 1] == ' ')
 180                                     desc2 = desc2.substr(0, desc2.length() - 1);
 181                             }
 182
 183                             if (sdd2->ProductIdOffset != 0) {
 184                                 if (sdd2->VendorIdOffset != 0 && desc2.length() != 0 && desc2[desc2.length() - 1] != ' ')
 185                                     desc2 += " ";
 186
 187                                 desc2 += (char*)((UINT8*)sdd2 + sdd2->ProductIdOffset);
 188
 189                                 while (desc2.length() > 0 && desc2[desc2.length() - 1] == ' ')
 190                                     desc2 = desc2.substr(0, desc2.length() - 1);
 191                             }
 192
 193                             if (sdd2->VendorIdOffset != 0 || sdd2->ProductIdOffset != 0) {
 194                                 ULONG ss;
 195
 196                                 ss = MultiByteToWideChar(CP_OEMCP, MB_PRECOMPOSED, desc2.c_str(), -1, NULL, 0);
 197
 198                                 if (ss > 0) {
 199                                     WCHAR* desc3 = (WCHAR*)malloc(ss * sizeof(WCHAR));
 200
 201                                     if (MultiByteToWideChar(CP_OEMCP, MB_PRECOMPOSED, desc2.c_str(), -1, desc3, ss * sizeof(WCHAR)))
 202                                         dev.friendly_name = desc3;
 203
 204                                     free(desc3);
 205                                 }
 206                             }
 207                         }
 208
 209                         free(sdd2);
 210                     }
 211
 212                     dlisize = 0;
 213                     dli = NULL;
 214
 215                     do {
 216                         dlisize += 1024;
 217
 218                         if (dli)
 219                             free(dli);
 220
 221                         dli = (DRIVE_LAYOUT_INFORMATION_EX*)malloc(dlisize);
 222
 223                         Status = NtDeviceIoControlFile(file, NULL, NULL, NULL, &iosb, IOCTL_DISK_GET_DRIVE_LAYOUT_EX,
 224                                                        NULL, 0, dli, dlisize);
 225                     } while (Status == STATUS_BUFFER_TOO_SMALL);
 226
 227                     if (NT_SUCCESS(Status) && dli->PartitionCount > 0)
 228                         dev.has_parts = TRUE;
 229
 230                     free(dli);
 231                 } else {
 232                     ULONG mmpsize;
 233                     MOUNTMGR_MOUNT_POINT* mmp;
 234                     MOUNTMGR_MOUNT_POINTS mmps;
 235
 236                     mmpsize = sizeof(MOUNTMGR_MOUNT_POINT) + path.Length;
 237
 238                     mmp = (MOUNTMGR_MOUNT_POINT*)malloc(mmpsize);
 239
 240                     RtlZeroMemory(mmp, sizeof(MOUNTMGR_MOUNT_POINT));
 241                     mmp->DeviceNameOffset = sizeof(MOUNTMGR_MOUNT_POINT);
 242                     mmp->DeviceNameLength = path.Length;
 243                     RtlCopyMemory(&mmp[1], path.Buffer, path.Length);
 244
 245                     Status = NtDeviceIoControlFile(mountmgr, NULL, NULL, NULL, &iosb, IOCTL_MOUNTMGR_QUERY_POINTS,
 246                                                    mmp, mmpsize, &mmps, sizeof(MOUNTMGR_MOUNT_POINTS));
 247                     if (NT_SUCCESS(Status) || Status == STATUS_BUFFER_OVERFLOW) {
 248                         MOUNTMGR_MOUNT_POINTS* mmps2;
 249
 250                         mmps2 = (MOUNTMGR_MOUNT_POINTS*)malloc(mmps.Size);
 251
 252                         Status = NtDeviceIoControlFile(mountmgr, NULL, NULL, NULL, &iosb, IOCTL_MOUNTMGR_QUERY_POINTS,
 253                                                     mmp, mmpsize, mmps2, mmps.Size);
 254
 255                         if (NT_SUCCESS(Status)) {
 256                             ULONG i;
 257
 258                             for (i = 0; i < mmps2->NumberOfMountPoints; i++) {
 259                                 WCHAR* symlink = (WCHAR*)((UINT8*)mmps2 + mmps2->MountPoints[i].SymbolicLinkNameOffset);
 260
 261                                 if (mmps2->MountPoints[i].SymbolicLinkNameLength == 0x1c &&
 262                                     RtlCompareMemory(symlink, dosdevices, wcslen(dosdevices) * sizeof(WCHAR)) == wcslen(dosdevices) * sizeof(WCHAR) &&
 263                                     symlink[13] == ':'
 264                                 ) {
 265                                     WCHAR dr[3];
 266
 267                                     dr[0] = symlink[12];
 268                                     dr[1] = ':';
 269                                     dr[2] = 0;
 270
 271                                     dev.drive = dr;
 272                                     break;
 273                                 }
 274                             }
 275                         }
 276                     }
 277
 278                     free(mmp);
 279                 }
 280
 281                 if (!dev.is_disk || !dev.has_parts) {
 282                     i = 0;
 283                     while (fs_ident[i].name) {
 284                         if (i == 0 || fs_ident[i].kboff != fs_ident[i-1].kboff) {
 285                             LARGE_INTEGER off;
 286
 287                             off.QuadPart = fs_ident[i].kboff * 1024;
 288                             Status = NtReadFile(file, NULL, NULL, NULL, &iosb, sb, sizeof(sb), &off, NULL);
 289                         }
 290
 291                         if (NT_SUCCESS(Status)) {
 292                             if (RtlCompareMemory(sb + fs_ident[i].sboff, fs_ident[i].magic, fs_ident[i].magiclen) == fs_ident[i].magiclen) {
 293                                 dev.fstype = fs_ident[i].name;
 294
 295                                 if (dev.fstype == L"Btrfs") {
 296                                     superblock* bsb = (superblock*)sb;
 297
 298                                     RtlCopyMemory(&dev.fs_uuid, &bsb->uuid, sizeof(BTRFS_UUID));
 299                                     RtlCopyMemory(&dev.dev_uuid, &bsb->dev_item.device_uuid, sizeof(BTRFS_UUID));
 300                                 }
 301
 302                                 break;
 303                             }
 304                         }
 305
 306                         i++;
 307                     }
 308
 309                     if (dev.fstype == L"Btrfs" && RtlCompareMemory(guid, &GUID_DEVINTERFACE_DISK, sizeof(GUID)) != sizeof(GUID)) {
 310                         std::wstring name;
 311                         std::wstring pref = L"\\Device\\Btrfs{";
 312
 313                         name = get_mountdev_name(file);
 314
 315                         if (name.length() > pref.length() && RtlCompareMemory(name.c_str(), pref.c_str(), pref.length() * sizeof(WCHAR)) == pref.length() * sizeof(WCHAR))
 316                             dev.ignore = TRUE;
 317                     }
 318                 }
 319
 320                 device_list->push_back(dev);
 321
 322 nextitem:
 323                 NtClose(file);
 324             }
 325
 326 nextitem2:
 327             free(detail);
 328
 329             index++;
 330         } while (SetupDiEnumDeviceInterfaces(h, NULL, guid, index, &did));
 331
 332         SetupDiDestroyDeviceInfoList(h);
 333     } else {
 334         ShowError(hwnd, GetLastError());
 335         return;
 336     }
 337 }
 338
 339 static bool sort_devices(device i, device j) {
 340     if (i.disk_num < j.disk_num)
 341         return true;
 342
 343     if (i.disk_num == j.disk_num && i.part_num < j.part_num)
 344         return true;
 345
 346     return false;
 347 }
 348
 349 void BtrfsDeviceAdd::populate_device_tree(HWND tree) {
 350     HWND hwnd = GetParent(tree);
 351     unsigned int i;
 352     ULONG last_disk_num = 0xffffffff;
 353     HTREEITEM diskitem;
 354     NTSTATUS Status;
 355     OBJECT_ATTRIBUTES attr;
 356     UNICODE_STRING us;
 357     IO_STATUS_BLOCK iosb;
 358     HANDLE mountmgr, btrfsh;
 359     btrfs_filesystem* bfs = NULL;
 360
 361     static WCHAR btrfs[] = L"\\Btrfs";
 362
 363     device_list.clear();
 364
 365     RtlInitUnicodeString(&us, MOUNTMGR_DEVICE_NAME);
 366     InitializeObjectAttributes(&attr, &us, 0, NULL, NULL);
 367
 368     Status = NtOpenFile(&mountmgr, FILE_GENERIC_READ | FILE_GENERIC_WRITE, &attr, &iosb,
 369                         FILE_SHARE_READ, FILE_SYNCHRONOUS_IO_ALERT);
 370     if (!NT_SUCCESS(Status)) {
 371         MessageBoxW(hwnd, L"Could not get handle to mount manager.", L"Error", MB_ICONERROR);
 372         return;
 373     }
 374
 375     us.Length = us.MaximumLength = wcslen(btrfs) * sizeof(WCHAR);
 376     us.Buffer = btrfs;
 377
 378     InitializeObjectAttributes(&attr, &us, 0, NULL, NULL);
 379
 380     Status = NtOpenFile(&btrfsh, SYNCHRONIZE | FILE_READ_ATTRIBUTES, &attr, &iosb,
 381                         FILE_SHARE_READ | FILE_SHARE_WRITE, FILE_SYNCHRONOUS_IO_ALERT);
 382     if (NT_SUCCESS(Status)) {
 383         ULONG bfssize = 0;
 384
 385         do {
 386             bfssize += 1024;
 387
 388             if (bfs) free(bfs);
 389             bfs = (btrfs_filesystem*)malloc(bfssize);
 390
 391             Status = NtDeviceIoControlFile(btrfsh, NULL, NULL, NULL, &iosb, IOCTL_BTRFS_QUERY_FILESYSTEMS, NULL, 0, bfs, bfssize);
 392             if (!NT_SUCCESS(Status) && Status != STATUS_BUFFER_OVERFLOW) {
 393                 free(bfs);
 394                 bfs = NULL;
 395                 break;
 396             }
 397         } while (Status == STATUS_BUFFER_OVERFLOW);
 398
 399         if (bfs && bfs->num_devices == 0) { // no mounted filesystems found
 400             free(bfs);
 401             bfs = NULL;
 402         }
 403     }
 404     NtClose(btrfsh);
 405
 406     find_devices(hwnd, &GUID_DEVINTERFACE_DISK, mountmgr, &device_list);
 407     find_devices(hwnd, &GUID_DEVINTERFACE_VOLUME, mountmgr, &device_list);
 408     find_devices(hwnd, &GUID_DEVINTERFACE_HIDDEN_VOLUME, mountmgr, &device_list);
 409
 410     NtClose(mountmgr);
 411
 412     std::sort(device_list.begin(), device_list.end(), sort_devices);
 413
 414     for (i = 0; i < device_list.size(); i++) {
 415         if (!device_list[i].ignore) {
 416             TVINSERTSTRUCTW tis;
 417             HTREEITEM item;
 418             std::wstring name;
 419             WCHAR size[255];
 420
 421             if (device_list[i].disk_num != 0xffffffff && device_list[i].disk_num == last_disk_num)
 422                 tis.hParent = diskitem;
 423             else
 424                 tis.hParent = TVI_ROOT;
 425
 426             tis.hInsertAfter = TVI_LAST;
 427             tis.itemex.mask = TVIF_TEXT | TVIF_STATE | TVIF_PARAM;
 428             tis.itemex.state = TVIS_EXPANDED;
 429             tis.itemex.stateMask = TVIS_EXPANDED;
 430
 431             if (device_list[i].disk_num != 0xffffffff) {
 432                 WCHAR t[255], u[255];
 433
 434                 if (!LoadStringW(module, device_list[i].part_num != 0 ? IDS_PARTITION : IDS_DISK_NUM, t, sizeof(t) / sizeof(WCHAR))) {
 435                     ShowError(hwnd, GetLastError());
 436                     return;
 437                 }
 438
 439                 if (StringCchPrintfW(u, sizeof(u) / sizeof(WCHAR), t, device_list[i].part_num != 0 ? device_list[i].part_num : device_list[i].disk_num) == STRSAFE_E_INSUFFICIENT_BUFFER)
 440                     return;
 441
 442                 name = u;
 443             } else
 444                 name = device_list[i].pnp_name;
 445
 446             // match child Btrfs devices to their parent
 447             if (bfs && device_list[i].drive == L"" && device_list[i].fstype == L"Btrfs") {
 448                 btrfs_filesystem* bfs2 = bfs;
 449
 450                 while (TRUE) {
 451                     if (RtlCompareMemory(&bfs2->uuid, &device_list[i].fs_uuid, sizeof(BTRFS_UUID)) == sizeof(BTRFS_UUID)) {
 452                         ULONG j, k;
 453                         btrfs_filesystem_device* dev;
 454
 455                         for (j = 0; j < bfs2->num_devices; j++) {
 456                             if (j == 0)
 457                                 dev = &bfs2->device;
 458                             else
 459                                 dev = (btrfs_filesystem_device*)((UINT8*)dev + offsetof(btrfs_filesystem_device, name[0]) + dev->name_length);
 460
 461                             if (RtlCompareMemory(&device_list[i].dev_uuid, &device_list[i].dev_uuid, sizeof(BTRFS_UUID)) == sizeof(BTRFS_UUID)) {
 462                                 for (k = 0; k < device_list.size(); k++) {
 463                                     if (k != i && device_list[k].fstype == L"Btrfs" && device_list[k].drive != L"" &&
 464                                         RtlCompareMemory(&device_list[k].fs_uuid, &device_list[i].fs_uuid, sizeof(BTRFS_UUID)) == sizeof(BTRFS_UUID)) {
 465                                         device_list[i].drive = device_list[k].drive;
 466                                         break;
 467                                     }
 468                                 }
 469
 470                                 device_list[i].multi_device = bfs2->num_devices > 1;
 471
 472                                 break;
 473                             }
 474                         }
 475
 476                         break;
 477                     }
 478
 479                     if (bfs2->next_entry != 0)
 480                         bfs2 = (btrfs_filesystem*)((UINT8*)bfs2 + bfs2->next_entry);
 481                     else
 482                         break;
 483                 }
 484             }
 485
 486             name += L" (";
 487
 488             if (device_list[i].friendly_name != L"") {
 489                 name += device_list[i].friendly_name;
 490                 name += L", ";
 491             }
 492
 493             if (device_list[i].drive != L"") {
 494                 name += device_list[i].drive;
 495                 name += L", ";
 496             }
 497
 498             if (device_list[i].fstype != L"") {
 499                 name += device_list[i].fstype;
 500                 name += L", ";
 501             }
 502
 503             format_size(device_list[i].size, size, sizeof(size) / sizeof(WCHAR), FALSE);
 504             name += size;
 505
 506             name += L")";
 507
 508             tis.itemex.pszText = (WCHAR*)name.c_str();
 509             tis.itemex.cchTextMax = name.length();
 510             tis.itemex.lParam = (LPARAM)&device_list[i];
 511
 512             item = (HTREEITEM)SendMessageW(tree, TVM_INSERTITEMW, 0, (LPARAM)&tis);
 513             if (!item) {
 514                 MessageBoxW(hwnd, L"TVM_INSERTITEM failed", L"Error", MB_ICONERROR);
 515                 return;
 516             }
 517
 518             if (device_list[i].part_num == 0) {
 519                 diskitem = item;
 520                 last_disk_num = device_list[i].disk_num;
 521             }
 522         }
 523     }
 524 }
 525
 526 void BtrfsDeviceAdd::AddDevice(HWND hwndDlg) {
 527     WCHAR mess[255], title[255];
 528     NTSTATUS Status;
 529     UNICODE_STRING vn;
 530     OBJECT_ATTRIBUTES attr;
 531     IO_STATUS_BLOCK iosb;
 532     HANDLE h, h2;
 533
 534     if (!sel) {
 535         EndDialog(hwndDlg, 0);
 536         return;
 537     }
 538
 539     if (sel->fstype != L"") {
 540         WCHAR s[255];
 541
 542         if (!LoadStringW(module, IDS_ADD_DEVICE_CONFIRMATION_FS, s, sizeof(s) / sizeof(WCHAR))) {
 543             ShowError(hwndDlg, GetLastError());
 544             return;
 545         }
 546
 547         if (StringCchPrintfW(mess, sizeof(mess) / sizeof(WCHAR), s, sel->fstype.c_str()) == STRSAFE_E_INSUFFICIENT_BUFFER)
 548             return;
 549     } else {
 550         if (!LoadStringW(module, IDS_ADD_DEVICE_CONFIRMATION, mess, sizeof(mess) / sizeof(WCHAR))) {
 551             ShowError(hwndDlg, GetLastError());
 552             return;
 553         }
 554     }
 555
 556     if (!LoadStringW(module, IDS_CONFIRMATION_TITLE, title, sizeof(title) / sizeof(WCHAR))) {
 557         ShowError(hwndDlg, GetLastError());
 558         return;
 559     }
 560
 561     if (MessageBoxW(hwndDlg, mess, title, MB_YESNO) != IDYES)
 562         return;
 563
 564     h = CreateFileW(cmdline, FILE_TRAVERSE | FILE_READ_ATTRIBUTES, FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, NULL,
 565                     OPEN_EXISTING, FILE_FLAG_BACKUP_SEMANTICS | FILE_FLAG_OPEN_REPARSE_POINT, NULL);
 566
 567     if (h == INVALID_HANDLE_VALUE) {
 568         ShowError(hwndDlg, GetLastError());
 569         return;
 570     }
 571
 572     vn.Length = vn.MaximumLength = sel->pnp_name.length() * sizeof(WCHAR);
 573     vn.Buffer = (WCHAR*)sel->pnp_name.c_str();
 574
 575     InitializeObjectAttributes(&attr, &vn, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, NULL, NULL);
 576
 577     Status = NtOpenFile(&h2, FILE_GENERIC_READ | FILE_GENERIC_WRITE, &attr, &iosb, FILE_SHARE_READ, FILE_SYNCHRONOUS_IO_ALERT);
 578     if (!NT_SUCCESS(Status)) {
 579         ShowNtStatusError(hwndDlg, Status);
 580         CloseHandle(h);
 581         return;
 582     }
 583
 584     if (!sel->is_disk) {
 585         Status = NtFsControlFile(h2, NULL, NULL, NULL, &iosb, FSCTL_LOCK_VOLUME, NULL, 0, NULL, 0);
 586         if (!NT_SUCCESS(Status)) {
 587             WCHAR t[255], u[255];
 588
 589             if (!LoadStringW(module, IDS_LOCK_FAILED, t, sizeof(t) / sizeof(WCHAR))) {
 590                 ShowError(hwnd, GetLastError());
 591                 return;
 592             }
 593
 594             if (StringCchPrintfW(u, sizeof(u) / sizeof(WCHAR), t, Status) == STRSAFE_E_INSUFFICIENT_BUFFER)
 595                 return;
 596
 597             if (!LoadStringW(module, IDS_ERROR, title, sizeof(title) / sizeof(WCHAR))) {
 598                 ShowError(hwndDlg, GetLastError());
 599                 return;
 600             }
 601
 602             MessageBoxW(hwndDlg, u, title, MB_ICONERROR);
 603
 604             NtClose(h2);
 605             CloseHandle(h);
 606             return;
 607         }
 608     }
 609
 610     Status = NtFsControlFile(h, NULL, NULL, NULL, &iosb, FSCTL_BTRFS_ADD_DEVICE, &h2, sizeof(HANDLE), NULL, 0);
 611     if (!NT_SUCCESS(Status)) {
 612         ShowNtStatusError(hwndDlg, Status);
 613         NtClose(h2);
 614         CloseHandle(h);
 615         return;
 616     }
 617
 618     if (!sel->is_disk) {
 619         Status = NtFsControlFile(h2, NULL, NULL, NULL, &iosb, FSCTL_DISMOUNT_VOLUME, NULL, 0, NULL, 0);
 620         if (!NT_SUCCESS(Status))
 621             ShowNtStatusError(hwndDlg, Status);
 622
 623         Status = NtFsControlFile(h2, NULL, NULL, NULL, &iosb, FSCTL_UNLOCK_VOLUME, NULL, 0, NULL, 0);
 624         if (!NT_SUCCESS(Status))
 625             ShowNtStatusError(hwndDlg, Status);
 626     }
 627
 628     NtClose(h2);
 629     CloseHandle(h);
 630
 631     EndDialog(hwndDlg, 0);
 632 }
 633
 634 INT_PTR CALLBACK BtrfsDeviceAdd::DeviceAddDlgProc(HWND hwndDlg, UINT uMsg, WPARAM wParam, LPARAM lParam) {
 635     switch (uMsg) {
 636         case WM_INITDIALOG:
 637         {
 638             EnableThemeDialogTexture(hwndDlg, ETDT_ENABLETAB);
 639             populate_device_tree(GetDlgItem(hwndDlg, IDC_DEVICE_TREE));
 640             EnableWindow(GetDlgItem(hwndDlg, IDOK), FALSE);
 641             break;
 642         }
 643
 644         case WM_COMMAND:
 645             switch (HIWORD(wParam)) {
 646                 case BN_CLICKED:
 647                     switch (LOWORD(wParam)) {
 648                         case IDOK:
 649                             AddDevice(hwndDlg);
 650                         return TRUE;
 651
 652                         case IDCANCEL:
 653                             EndDialog(hwndDlg, 0);
 654                         return TRUE;
 655                     }
 656                 break;
 657             }
 658         break;
 659
 660         case WM_NOTIFY:
 661             switch (((LPNMHDR)lParam)->code) {
 662                 case TVN_SELCHANGEDW:
 663                 {
 664                     NMTREEVIEWW* nmtv = (NMTREEVIEWW*)lParam;
 665                     TVITEMW tvi;
 666                     BOOL enable = FALSE;
 667
 668                     RtlZeroMemory(&tvi, sizeof(TVITEMW));
 669                     tvi.hItem = nmtv->itemNew.hItem;
 670                     tvi.mask = TVIF_PARAM | TVIF_HANDLE;
 671
 672                     if (SendMessageW(GetDlgItem(hwndDlg, IDC_DEVICE_TREE), TVM_GETITEMW, 0, (LPARAM)&tvi))
 673                         sel = tvi.lParam == 0 ? NULL : (device*)tvi.lParam;
 674                     else
 675                         sel = NULL;
 676
 677                     if (sel)
 678                         enable = (!sel->is_disk || !sel->has_parts) && !sel->multi_device;
 679
 680                     EnableWindow(GetDlgItem(hwndDlg, IDOK), enable);
 681                     break;
 682                 }
 683             }
 684         break;
 685     }
 686
 687     return FALSE;
 688 }
 689
 690 static INT_PTR CALLBACK stub_DeviceAddDlgProc(HWND hwndDlg, UINT uMsg, WPARAM wParam, LPARAM lParam) {
 691     BtrfsDeviceAdd* bda;
 692
 693     if (uMsg == WM_INITDIALOG) {
 694         SetWindowLongPtr(hwndDlg, GWLP_USERDATA, (LONG_PTR)lParam);
 695         bda = (BtrfsDeviceAdd*)lParam;
 696     } else {
 697         bda = (BtrfsDeviceAdd*)GetWindowLongPtr(hwndDlg, GWLP_USERDATA);
 698     }
 699
 700     if (bda)
 701         return bda->DeviceAddDlgProc(hwndDlg, uMsg, wParam, lParam);
 702     else
 703         return FALSE;
 704 }
 705
 706 void BtrfsDeviceAdd::ShowDialog() {
 707     DialogBoxParamW(module, MAKEINTRESOURCEW(IDD_DEVICE_ADD), hwnd, stub_DeviceAddDlgProc, (LPARAM)this);
 708 }
 709
 710 BtrfsDeviceAdd::BtrfsDeviceAdd(HINSTANCE hinst, HWND hwnd, WCHAR* cmdline) {
 711     this->hinst = hinst;
 712     this->hwnd = hwnd;
 713     this->cmdline = cmdline;
 714
 715     sel = NULL;
 716 }
 717
 718 void BtrfsDeviceResize::do_resize(HWND hwndDlg) {
 719     HANDLE h;
 720     NTSTATUS Status;
 721     IO_STATUS_BLOCK iosb;
 722     btrfs_resize br;
 723
 724     h = CreateFileW(fn, FILE_TRAVERSE | FILE_READ_ATTRIBUTES, FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, NULL,
 725                     OPEN_EXISTING, FILE_FLAG_BACKUP_SEMANTICS | FILE_FLAG_OPEN_REPARSE_POINT, NULL);
 726
 727     if (h == INVALID_HANDLE_VALUE) {
 728         ShowError(hwndDlg, GetLastError());
 729         return;
 730     }
 731
 732     br.device = dev_id;
 733     br.size = new_size;
 734
 735     Status = NtFsControlFile(h, NULL, NULL, NULL, &iosb, FSCTL_BTRFS_RESIZE, &br, sizeof(btrfs_resize), NULL, 0);
 736
 737     if (Status != STATUS_MORE_PROCESSING_REQUIRED && !NT_SUCCESS(Status)) {
 738         ShowNtStatusError(hwndDlg, Status);
 739         CloseHandle(h);
 740         return;
 741     }
 742
 743     CloseHandle(h);
 744
 745     if (Status != STATUS_MORE_PROCESSING_REQUIRED) {
 746         WCHAR s[255], t[255], u[255];
 747
 748         LoadStringW(module, IDS_RESIZE_SUCCESSFUL, s, sizeof(s) / sizeof(WCHAR));
 749         format_size(new_size, u, sizeof(u) / sizeof(WCHAR), TRUE);
 750         StringCchPrintfW(t, sizeof(t) / sizeof(WCHAR), s, dev_id, u);
 751         MessageBoxW(hwndDlg, t, L"", MB_OK);
 752
 753         EndDialog(hwndDlg, 0);
 754     } else {
 755         BtrfsBalance* bb;
 756         HWND par;
 757
 758         par = GetParent(hwndDlg);
 759         EndDialog(hwndDlg, 0);
 760
 761         bb = new BtrfsBalance(fn, FALSE, TRUE);
 762
 763         bb->ShowBalance(par);
 764
 765         delete bb;
 766     }
 767 }
 768
 769 INT_PTR CALLBACK BtrfsDeviceResize::DeviceResizeDlgProc(HWND hwndDlg, UINT uMsg, WPARAM wParam, LPARAM lParam) {
 770     switch (uMsg) {
 771         case WM_INITDIALOG:
 772         {
 773             HANDLE h;
 774             WCHAR s[255], t[255], u[255];
 775
 776             EnableThemeDialogTexture(hwndDlg, ETDT_ENABLETAB);
 777
 778             GetDlgItemTextW(hwndDlg, IDC_RESIZE_DEVICE_ID, s, sizeof(s) / sizeof(WCHAR));
 779             StringCchPrintfW(t, sizeof(t) / sizeof(WCHAR), s, dev_id);
 780             SetDlgItemTextW(hwndDlg, IDC_RESIZE_DEVICE_ID, t);
 781
 782             h = CreateFileW(fn, FILE_TRAVERSE | FILE_READ_ATTRIBUTES, FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, NULL,
 783                             OPEN_EXISTING, FILE_FLAG_BACKUP_SEMANTICS | FILE_FLAG_OPEN_REPARSE_POINT, NULL);
 784
 785             if (h != INVALID_HANDLE_VALUE) {
 786                 NTSTATUS Status;
 787                 IO_STATUS_BLOCK iosb;
 788                 btrfs_device *devices, *bd;
 789                 ULONG devsize;
 790                 BOOL found = FALSE;
 791                 HWND slider;
 792
 793                 devsize = 1024;
 794                 devices = (btrfs_device*)malloc(devsize);
 795
 796                 while (TRUE) {
 797                     Status = NtFsControlFile(h, NULL, NULL, NULL, &iosb, FSCTL_BTRFS_GET_DEVICES, NULL, 0, devices, devsize);
 798                     if (Status == STATUS_BUFFER_OVERFLOW) {
 799                         devsize += 1024;
 800
 801                         free(devices);
 802                         devices = (btrfs_device*)malloc(devsize);
 803                     } else
 804                         break;
 805                 }
 806
 807                 if (!NT_SUCCESS(Status)) {
 808                     free(devices);
 809                     CloseHandle(h);
 810                     return FALSE;
 811                 }
 812
 813                 bd = devices;
 814
 815                 while (TRUE) {
 816                     if (bd->dev_id == dev_id) {
 817                         memcpy(&dev_info, bd, sizeof(btrfs_device));
 818                         found = TRUE;
 819                         break;
 820                     }
 821
 822                     if (bd->next_entry > 0)
 823                         bd = (btrfs_device*)((UINT8*)bd + bd->next_entry);
 824                     else
 825                         break;
 826                 }
 827
 828                 if (!found) {
 829                     free(devices);
 830                     CloseHandle(h);
 831                     return FALSE;
 832                 }
 833
 834                 free(devices);
 835                 CloseHandle(h);
 836
 837                 GetDlgItemTextW(hwndDlg, IDC_RESIZE_CURSIZE, s, sizeof(s) / sizeof(WCHAR));
 838                 format_size(dev_info.size, u, sizeof(u) / sizeof(WCHAR), TRUE);
 839                 StringCchPrintfW(t, sizeof(t) / sizeof(WCHAR), s, u);
 840                 SetDlgItemTextW(hwndDlg, IDC_RESIZE_CURSIZE, t);
 841
 842                 new_size = dev_info.size;
 843
 844                 GetDlgItemTextW(hwndDlg, IDC_RESIZE_NEWSIZE, new_size_text, sizeof(new_size_text) / sizeof(WCHAR));
 845                 StringCchPrintfW(t, sizeof(t) / sizeof(WCHAR), new_size_text, u);
 846                 SetDlgItemTextW(hwndDlg, IDC_RESIZE_NEWSIZE, t);
 847
 848                 slider = GetDlgItem(hwndDlg, IDC_RESIZE_SLIDER);
 849                 SendMessageW(slider, TBM_SETRANGEMIN, FALSE, 0);
 850                 SendMessageW(slider, TBM_SETRANGEMAX, FALSE, dev_info.max_size / 1048576);
 851                 SendMessageW(slider, TBM_SETPOS, TRUE, new_size / 1048576);
 852             } else
 853                 return FALSE;
 854
 855             break;
 856         }
 857
 858         case WM_COMMAND:
 859             switch (HIWORD(wParam)) {
 860                 case BN_CLICKED:
 861                     switch (LOWORD(wParam)) {
 862                         case IDOK:
 863                             do_resize(hwndDlg);
 864                             return TRUE;
 865
 866                         case IDCANCEL:
 867                             EndDialog(hwndDlg, 0);
 868                             return TRUE;
 869                     }
 870                 break;
 871             }
 872         break;
 873
 874         case WM_HSCROLL:
 875         {
 876             WCHAR t[255], u[255];
 877
 878             new_size = UInt32x32To64(SendMessageW(GetDlgItem(hwndDlg, IDC_RESIZE_SLIDER), TBM_GETPOS, 0, 0), 1048576);
 879
 880             format_size(new_size, u, sizeof(u) / sizeof(WCHAR), TRUE);
 881             StringCchPrintfW(t, sizeof(t) / sizeof(WCHAR), new_size_text, u);
 882             SetDlgItemTextW(hwndDlg, IDC_RESIZE_NEWSIZE, t);
 883
 884             EnableWindow(GetDlgItem(hwndDlg, IDOK), new_size > 0 ? TRUE : FALSE);
 885
 886             break;
 887         }
 888     }
 889
 890     return FALSE;
 891 }
 892
 893 static INT_PTR CALLBACK stub_DeviceResizeDlgProc(HWND hwndDlg, UINT uMsg, WPARAM wParam, LPARAM lParam) {
 894     BtrfsDeviceResize* bdr;
 895
 896     if (uMsg == WM_INITDIALOG) {
 897         SetWindowLongPtr(hwndDlg, GWLP_USERDATA, (LONG_PTR)lParam);
 898         bdr = (BtrfsDeviceResize*)lParam;
 899     } else
 900         bdr = (BtrfsDeviceResize*)GetWindowLongPtr(hwndDlg, GWLP_USERDATA);
 901
 902     if (bdr)
 903         return bdr->DeviceResizeDlgProc(hwndDlg, uMsg, wParam, lParam);
 904     else
 905         return FALSE;
 906 }
 907
 908 void BtrfsDeviceResize::ShowDialog(HWND hwnd, WCHAR* fn, UINT64 dev_id) {
 909     this->dev_id = dev_id;
 910     wcscpy(this->fn, fn);
 911
 912     DialogBoxParamW(module, MAKEINTRESOURCEW(IDD_RESIZE), hwnd, stub_DeviceResizeDlgProc, (LPARAM)this);
 913 }
 914
 915 #ifdef __cplusplus
 916 extern "C" {
 917 #endif
 918
 919 void CALLBACK AddDeviceW(HWND hwnd, HINSTANCE hinst, LPWSTR lpszCmdLine, int nCmdShow) {
 920     HANDLE token;
 921     TOKEN_PRIVILEGES tp;
 922     LUID luid;
 923     BtrfsDeviceAdd* bda;
 924
 925     set_dpi_aware();
 926
 927     if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &token)) {
 928         ShowError(hwnd, GetLastError());
 929         return;
 930     }
 931
 932     if (!LookupPrivilegeValueW(NULL, L"SeManageVolumePrivilege", &luid)) {
 933         ShowError(hwnd, GetLastError());
 934         goto end;
 935     }
 936
 937     tp.PrivilegeCount = 1;
 938     tp.Privileges[0].Luid = luid;
 939     tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
 940
 941     if (!AdjustTokenPrivileges(token, FALSE, &tp, sizeof(TOKEN_PRIVILEGES), NULL, NULL)) {
 942         ShowError(hwnd, GetLastError());
 943         goto end;
 944     }
 945
 946     bda = new BtrfsDeviceAdd(hinst, hwnd, lpszCmdLine);
 947     bda->ShowDialog();
 948     delete bda;
 949
 950 end:
 951     CloseHandle(token);
 952 }
 953
 954 void CALLBACK RemoveDeviceW(HWND hwnd, HINSTANCE hinst, LPWSTR lpszCmdLine, int nCmdShow) {
 955     WCHAR *s, *vol, *dev;
 956     UINT64 devid;
 957     HANDLE h, token;
 958     TOKEN_PRIVILEGES tp;
 959     LUID luid;
 960     NTSTATUS Status;
 961     IO_STATUS_BLOCK iosb;
 962     BtrfsBalance* bb;
 963
 964     set_dpi_aware();
 965
 966     if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &token)) {
 967         ShowError(hwnd, GetLastError());
 968         return;
 969     }
 970
 971     if (!LookupPrivilegeValueW(NULL, L"SeManageVolumePrivilege", &luid)) {
 972         ShowError(hwnd, GetLastError());
 973         goto end;
 974     }
 975
 976     tp.PrivilegeCount = 1;
 977     tp.Privileges[0].Luid = luid;
 978     tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
 979
 980     if (!AdjustTokenPrivileges(token, FALSE, &tp, sizeof(TOKEN_PRIVILEGES), NULL, NULL)) {
 981         ShowError(hwnd, GetLastError());
 982         goto end;
 983     }
 984
 985     s = wcsstr(lpszCmdLine, L"|");
 986     if (!s)
 987         goto end;
 988
 989     s[0] = 0;
 990
 991     vol = lpszCmdLine;
 992     dev = &s[1];
 993
 994     devid = _wtoi(dev);
 995     if (devid == 0)
 996         goto end;
 997
 998     h = CreateFileW(vol, FILE_TRAVERSE | FILE_READ_ATTRIBUTES, FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, NULL,
 999                     OPEN_EXISTING, FILE_FLAG_BACKUP_SEMANTICS | FILE_FLAG_OPEN_REPARSE_POINT, NULL);
1000
1001     if (h == INVALID_HANDLE_VALUE) {
1002         ShowError(hwnd, GetLastError());
1003         goto end;
1004     }
1005
1006     Status = NtFsControlFile(h, NULL, NULL, NULL, &iosb, FSCTL_BTRFS_REMOVE_DEVICE, &devid, sizeof(UINT64), NULL, 0);
1007     if (!NT_SUCCESS(Status)) {
1008         if (Status == STATUS_CANNOT_DELETE)
1009             ShowStringError(hwnd, IDS_CANNOT_REMOVE_RAID);
1010         else
1011             ShowNtStatusError(hwnd, Status);
1012
1013         CloseHandle(h);
1014         goto end;
1015     }
1016
1017     CloseHandle(h);
1018
1019     bb = new BtrfsBalance(vol, TRUE);
1020
1021     bb->ShowBalance(hwnd);
1022
1023     delete bb;
1024
1025 end:
1026     CloseHandle(token);
1027 }
1028
1029 void CALLBACK ResizeDeviceW(HWND hwnd, HINSTANCE hinst, LPWSTR lpszCmdLine, int nCmdShow) {
1030     WCHAR *s, *vol, *dev;
1031     UINT64 devid;
1032     HANDLE token;
1033     TOKEN_PRIVILEGES tp;
1034     LUID luid;
1035     BtrfsDeviceResize* bdr;
1036
1037     set_dpi_aware();
1038
1039     if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &token)) {
1040         ShowError(hwnd, GetLastError());
1041         return;
1042     }
1043
1044     if (!LookupPrivilegeValueW(NULL, L"SeManageVolumePrivilege", &luid)) {
1045         ShowError(hwnd, GetLastError());
1046         goto end;
1047     }
1048
1049     tp.PrivilegeCount = 1;
1050     tp.Privileges[0].Luid = luid;
1051     tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
1052
1053     if (!AdjustTokenPrivileges(token, FALSE, &tp, sizeof(TOKEN_PRIVILEGES), NULL, NULL)) {
1054         ShowError(hwnd, GetLastError());
1055         goto end;
1056     }
1057
1058     s = wcsstr(lpszCmdLine, L"|");
1059     if (!s)
1060         goto end;
1061
1062     s[0] = 0;
1063
1064     vol = lpszCmdLine;
1065     dev = &s[1];
1066
1067     devid = _wtoi(dev);
1068     if (devid == 0)
1069         goto end;
1070
1071     bdr = new BtrfsDeviceResize;
1072     bdr->ShowDialog(hwnd, vol, devid);
1073     delete bdr;
1074
1075 end:
1076     CloseHandle(token);
1077 }
1078
1079 #ifdef __cplusplus
1080 }
1081 #endif