git.mutantstargoat.com Git - instimg/blob - src/rawdisk.c

   1 #include <stdio.h>
   2 #include <stdlib.h>
   3 #include <string.h>
   4 #include <malloc.h>
   5 #include <windows.h>
   6 #include <setupapi.h>
   7 #include <devguid.h>
   8 #include <winioctl.h>
   9 #include "rawdisk.h"
  10
  11 #undef DBG_FAKE_DISKS
  12
  13 static GUID guid_iface_disk = {0x53f56307, 0xb6bf, 0x11d0, {0x94, 0xf2, 0x00, 0xa0, 0xc9, 0x1e, 0xfb, 0x8b}};
  14
  15 #ifndef IOCTL_DISK_UPDATE_PROPERTIES
  16 #define IOCTL_DISK_UPDATE_PROPERTIES CTL_CODE(IOCTL_DISK_BASE,0x0050,METHOD_BUFFERED,FILE_ANY_ACCESS)
  17 #endif
  18
  19 int rawdisk_detect(struct rawdisk_device *disks, int max_disks)
  20 {
  21         int devidx, ifidx, count;
  22         HDEVINFO devset;
  23         SP_DEVINFO_DATA devdata;
  24         SP_DEVICE_INTERFACE_DATA devif;
  25         SP_DEVICE_INTERFACE_DETAIL_DATA_A *devdetail;
  26         DWORD size, regtype;
  27         char devname[1024];
  28
  29         if((devset = SetupDiGetClassDevs(&guid_iface_disk, 0, 0,
  30                         DIGCF_PRESENT | DIGCF_DEVICEINTERFACE)) == INVALID_HANDLE_VALUE) {
  31                 fprintf(stderr, "failed to enumerate devices\n");
  32                 return -1;
  33         }
  34
  35         count = 0;
  36         devidx = 0;
  37         for(;;) {
  38                 memset(&devdata, 0, sizeof devdata);
  39                 devdata.cbSize = sizeof devdata;
  40                 if(!SetupDiEnumDeviceInfo(devset, devidx, &devdata)) {
  41                         if(GetLastError() == ERROR_NO_MORE_ITEMS) break;
  42                         devidx++;
  43                         continue;
  44                 }
  45
  46                 regtype = SPDRP_PHYSICAL_DEVICE_OBJECT_NAME;
  47                 SetupDiGetDeviceRegistryProperty(devset, &devdata, SPDRP_FRIENDLYNAME,
  48                                 &regtype, (unsigned char*)devname, sizeof devname, &size);
  49
  50                 count = 0;
  51                 ifidx = 0;
  52                 for(;;) {
  53                         memset(&devif, 0, sizeof devif);
  54                         devif.cbSize = sizeof devif;
  55                         if(!SetupDiEnumDeviceInterfaces(devset, &devdata, &guid_iface_disk, ifidx, &devif)) {
  56                                 if(GetLastError() == ERROR_NO_MORE_ITEMS) break;
  57                                 ifidx++;
  58                                 continue;
  59                         }
  60
  61                         SetupDiGetDeviceInterfaceDetail(devset, &devif, 0, 0, &size, 0);
  62                         if(!(devdetail = malloc(size))) {
  63                                 fprintf(stderr, "failed to allocate device interface detail buffer (size: %lu)\n", (unsigned long)size);
  64                                 return -1;
  65                         }
  66                         devdetail->cbSize = sizeof *devdetail;
  67                         SetupDiGetDeviceInterfaceDetail(devset, &devif, devdetail, size, 0, 0);
  68
  69
  70                         if(count < max_disks) {
  71                                 disks[count].path = strdup(devdetail->DevicePath);
  72                                 disks[count].name = strdup(devname);
  73                                 if(!disks[count].path || !disks[count].name) {
  74                                         fprintf(stderr, "failed to allocate device strings\n");
  75                                         return -1;
  76                                 }
  77                                 count++;
  78                         }
  79
  80                         free(devdetail);
  81                         ifidx++;
  82                 }
  83
  84                 devidx++;
  85         }
  86
  87 #ifdef DBG_FAKE_DISKS
  88         {
  89                 int i;
  90                 for(i=0; i<3; i++) {
  91                         char buf[32];
  92                         sprintf(buf, "\\\\?\\fake\\disk\\%d", i);
  93                         disks[count].path = strdup(buf);
  94                         sprintf(buf, "FAKE_DISK_%d", i);
  95                         disks[count].name = strdup(buf);
  96                         count++;
  97                 }
  98         }
  99 #endif
 100
 101         return count;
 102 }
 103
 104 int rawdisk_eject(HANDLE hdev)
 105 {
 106         DWORD sz;
 107         return DeviceIoControl(hdev, IOCTL_STORAGE_EJECT_MEDIA, 0, 0, 0, 0, &sz, 0) == 0 ? -1 : 0;
 108 }
 109
 110 int rawdisk_load(HANDLE hdev)
 111 {
 112         DWORD sz;
 113         return DeviceIoControl(hdev, IOCTL_STORAGE_LOAD_MEDIA, 0, 0, 0, 0, &sz, 0) == 0 ? -1 : 0;
 114 }
 115
 116 void rawdisk_refresh(HANDLE hdev)
 117 {
 118         DWORD sz;
 119         char *buf;
 120
 121         /* try to rock the boat so hopefully windows will re-read the disk and it'll show up */
 122         DeviceIoControl(hdev, IOCTL_DISK_UPDATE_PROPERTIES, 0, 0, 0, 0, &sz, 0);
 123
 124         sz = sizeof(DRIVE_LAYOUT_INFORMATION) + sizeof(PARTITION_INFORMATION) * 32;
 125         buf = alloca(sz);
 126
 127         DeviceIoControl(hdev, IOCTL_DISK_GET_DRIVE_LAYOUT, 0, 0, buf, sz, &sz, 0);
 128 }