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