1// Copyright 2013 The Chromium Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5// devguid requires Windows.h be imported first.
6#include <windows.h>
7#include <setupapi.h>
8#include <winioctl.h>
9
10#include "base/strings/string_number_conversions.h"
11#include "base/strings/string_util.h"
12#include "base/strings/utf_string_conversions.h"
13#include "base/win/scoped_handle.h"
14#include "chrome/browser/extensions/api/image_writer_private/removable_storage_provider.h"
15
16namespace extensions {
17
18namespace {
19
20bool AddDeviceInfo(HANDLE interface_enumerator,
21                   SP_DEVICE_INTERFACE_DATA* interface_data,
22                   scoped_refptr<StorageDeviceList> device_list) {
23  // Get the required buffer size by calling with a null output buffer.
24  DWORD interface_detail_data_size;
25  BOOL status = SetupDiGetDeviceInterfaceDetail(
26      interface_enumerator,
27      interface_data,
28      NULL,                         // Output buffer.
29      0,                            // Output buffer size.
30      &interface_detail_data_size,  // Receives the buffer size.
31      NULL);                        // Optional DEVINFO_DATA.
32
33  if (status == FALSE && GetLastError() != ERROR_INSUFFICIENT_BUFFER) {
34    PLOG(ERROR) << "SetupDiGetDeviceInterfaceDetail failed";
35    return false;
36  }
37
38  scoped_ptr<char[]> interface_detail_data_buffer(
39      new char[interface_detail_data_size]);
40
41  SP_DEVICE_INTERFACE_DETAIL_DATA* interface_detail_data =
42      reinterpret_cast<SP_DEVICE_INTERFACE_DETAIL_DATA*>(
43          interface_detail_data_buffer.get());
44
45  interface_detail_data->cbSize = sizeof(SP_INTERFACE_DEVICE_DETAIL_DATA);
46
47  status = SetupDiGetDeviceInterfaceDetail(
48      interface_enumerator,
49      interface_data,
50      interface_detail_data, // Output struct.
51      interface_detail_data_size,  // Output struct size.
52      NULL,                        // Receives required size, unneeded.
53      NULL);                       // Optional DEVINFO_Data.
54
55  if (status == FALSE) {
56    PLOG(ERROR) << "SetupDiGetDeviceInterfaceDetail failed";
57    return false;
58  }
59
60  // Open a handle to the device to send DeviceIoControl messages.
61  base::win::ScopedHandle device_handle(CreateFile(
62      interface_detail_data->DevicePath,
63      // Desired access, which is none as we only need metadata.
64      0,
65      // Required to be read + write for devices.
66      FILE_SHARE_READ | FILE_SHARE_WRITE,
67      NULL,           // Optional security attributes.
68      OPEN_EXISTING,  // Devices already exist.
69      0,              // No optional flags.
70      NULL));          // No template file.
71
72  if (!device_handle.IsValid()) {
73    PLOG(ERROR) << "Opening device handle failed.";
74    return false;
75  }
76
77  DISK_GEOMETRY geometry;
78  DWORD bytes_returned;
79  status = DeviceIoControl(
80      device_handle.Get(),           // Device handle.
81      IOCTL_DISK_GET_DRIVE_GEOMETRY, // Flag to request disk size.
82      NULL,                          // Optional additional parameters.
83      0,                             // Optional parameter size.
84      &geometry,                     // output buffer.
85      sizeof(DISK_GEOMETRY),         // output size.
86      &bytes_returned,               // Must be non-null. If overlapped is null,
87                                     // then value is meaningless.
88      NULL);                         // Optional unused overlapped parameter.
89
90  if (status == FALSE) {
91    PLOG(ERROR) << "DeviceIoControl";
92    return false;
93  }
94
95  ULONGLONG disk_capacity = geometry.Cylinders.QuadPart *
96    geometry.TracksPerCylinder *
97    geometry.SectorsPerTrack *
98    geometry.BytesPerSector;
99
100  STORAGE_PROPERTY_QUERY query = STORAGE_PROPERTY_QUERY();
101  query.PropertyId = StorageDeviceProperty;
102  query.QueryType = PropertyStandardQuery;
103
104  scoped_ptr<char[]> output_buf(new char[1024]);
105  status = DeviceIoControl(
106      device_handle.Get(),            // Device handle.
107      IOCTL_STORAGE_QUERY_PROPERTY,   // Flag to request device properties.
108      &query,                         // Query parameters.
109      sizeof(STORAGE_PROPERTY_QUERY), // query parameters size.
110      output_buf.get(),               // output buffer.
111      1024,                           // Size of buffer.
112      &bytes_returned,                // Number of bytes returned.
113                                      // Must not be null.
114      NULL);                          // Optional unused overlapped perameter.
115
116  if (status == FALSE) {
117    PLOG(ERROR) << "Storage property query failed.";
118    return false;
119  }
120
121  STORAGE_DEVICE_DESCRIPTOR* device_descriptor =
122      reinterpret_cast<STORAGE_DEVICE_DESCRIPTOR*>(output_buf.get());
123
124  if (!device_descriptor->RemovableMedia &&
125      !(device_descriptor->BusType == BusTypeUsb)) {
126    // Reject non-removable and non-USB devices.
127    // Return true to indicate success but not add anything to the device list.
128    return true;
129  }
130
131  // Create a drive identifier from the drive number.
132  STORAGE_DEVICE_NUMBER device_number = {0};
133  status = DeviceIoControl(
134      device_handle.Get(),            // Device handle.
135      IOCTL_STORAGE_GET_DEVICE_NUMBER,// Flag to request device number.
136      NULL,                           // Query parameters, should be NULL.
137      0,                              // Query parameters size, should be 0.
138      &device_number,                 // output buffer.
139      sizeof(device_number),          // Size of buffer.
140      &bytes_returned,                // Number of bytes returned.
141      NULL);                          // Optional unused overlapped perameter.
142
143  if (status == FALSE) {
144    PLOG(ERROR) << "Storage device number query failed.";
145    return false;
146  }
147
148  std::string drive_id = "\\\\.\\PhysicalDrive";
149  drive_id.append(base::Uint64ToString(device_number.DeviceNumber));
150
151  linked_ptr<api::image_writer_private::RemovableStorageDevice> device(
152    new api::image_writer_private::RemovableStorageDevice());
153  device->capacity = disk_capacity;
154  device->storage_unit_id = drive_id;
155  device->removable = device_descriptor->RemovableMedia == TRUE;
156
157  if (device_descriptor->VendorIdOffset &&
158      output_buf[device_descriptor->VendorIdOffset]) {
159    device->vendor.assign(output_buf.get() + device_descriptor->VendorIdOffset);
160  }
161
162  std::string product_id;
163  if (device_descriptor->ProductIdOffset &&
164      output_buf[device_descriptor->ProductIdOffset]) {
165    device->model.assign(output_buf.get() + device_descriptor->ProductIdOffset);
166  }
167
168  device_list->data.push_back(device);
169
170  return true;
171}
172
173}  // namespace
174
175bool RemovableStorageProvider::PopulateDeviceList(
176    scoped_refptr<StorageDeviceList> device_list) {
177  HDEVINFO interface_enumerator = SetupDiGetClassDevs(
178      &DiskClassGuid,
179      NULL, // Enumerator.
180      NULL, // Parent window.
181      // Only devices present & interface class.
182      (DIGCF_PRESENT | DIGCF_INTERFACEDEVICE));
183
184  if (interface_enumerator == INVALID_HANDLE_VALUE) {
185    DPLOG(ERROR) << "SetupDiGetClassDevs failed.";
186    return false;
187  }
188
189  DWORD index = 0;
190  SP_DEVICE_INTERFACE_DATA interface_data;
191  interface_data.cbSize = sizeof(SP_INTERFACE_DEVICE_DATA);
192
193  while (SetupDiEnumDeviceInterfaces(
194      interface_enumerator,
195      NULL,                    // Device Info data.
196      &GUID_DEVINTERFACE_DISK, // Only disk devices.
197      index,
198      &interface_data)) {
199    AddDeviceInfo(interface_enumerator, &interface_data, device_list);
200    index++;
201  }
202
203  DWORD error_code = GetLastError();
204
205  if (error_code != ERROR_NO_MORE_ITEMS) {
206    PLOG(ERROR) << "SetupDiEnumDeviceInterfaces failed";
207    SetupDiDestroyDeviceInfoList(interface_enumerator);
208    return false;
209  }
210
211  SetupDiDestroyDeviceInfoList(interface_enumerator);
212  return true;
213}
214
215} // namespace extensions
216