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