xref: /reactos/dll/shellext/shellbtrfs/devices.cpp (revision b826992a)
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 
get_mountdev_name(const nt_handle & h)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 
find_devices(HWND hwnd,const GUID * guid,const mountmgr & mm,vector<device> & device_list)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
984