text: // Copyright 2014 The Chromium Authors. All rights reserved.
  // Use of this source code is governed by a BSD-style license that can be
  // found in the LICENSE file.
  
  #include <windows.h>
  #include <setupapi.h>
  #include <winioctl.h>
  
  #include "chrome/utility/image_writer/error_messages.h"
  #include "chrome/utility/image_writer/image_writer.h"
  
  namespace image_writer {
  
  const size_t kStorageQueryBufferSize = 1024;
  
  bool ImageWriter::IsValidDevice() {
    base::win::ScopedHandle device_handle(
        CreateFile(device_path_.value().c_str(),
                   GENERIC_READ | GENERIC_WRITE,
                   FILE_SHARE_READ | FILE_SHARE_WRITE,
                   NULL,
                   OPEN_EXISTING,
                   FILE_FLAG_NO_BUFFERING | FILE_FLAG_WRITE_THROUGH,
                   NULL));
    if (!device_handle.IsValid()) {
      Error(error::kOpenDevice);
      return false;
    }
  
    STORAGE_PROPERTY_QUERY query = STORAGE_PROPERTY_QUERY();
    query.PropertyId = StorageDeviceProperty;
    query.QueryType = PropertyStandardQuery;
    DWORD bytes_returned;
  
    scoped_ptr<char[]> output_buf(new char[kStorageQueryBufferSize]);
    BOOL status = DeviceIoControl(
        device_handle.Get(),             // Device handle.
        IOCTL_STORAGE_QUERY_PROPERTY,    // Flag to request device properties.
        &query,                          // Query parameters.
        sizeof(STORAGE_PROPERTY_QUERY),  // query parameters size.
        output_buf.get(),                // output buffer.
        kStorageQueryBufferSize,         // Size of buffer.
        &bytes_returned,                 // Number of bytes returned.
                                         // Must not be null.
        NULL);                           // Optional unused overlapped perameter.
  
    if (!status) {
      PLOG(ERROR) << "Storage property query failed";
      return false;
    }
  
    STORAGE_DEVICE_DESCRIPTOR* device_descriptor =
        reinterpret_cast<STORAGE_DEVICE_DESCRIPTOR*>(output_buf.get());
  
    return device_descriptor->RemovableMedia == TRUE ||
           device_descriptor->BusType == BusTypeUsb;
  }
  
  bool ImageWriter::OpenDevice() {
    // Windows requires that device files be opened with FILE_FLAG_NO_BUFFERING
    // and FILE_FLAG_WRITE_THROUGH.  These two flags are not part of base::File.
    device_file_ =
        base::File(CreateFile(device_path_.value().c_str(),
                              GENERIC_READ | GENERIC_WRITE,
                              FILE_SHARE_READ | FILE_SHARE_WRITE,
                              NULL,
                              OPEN_EXISTING,
                              FILE_FLAG_NO_BUFFERING | FILE_FLAG_WRITE_THROUGH,
                              NULL));
    return device_file_.IsValid();
  }
  
  void ImageWriter::UnmountVolumes(const base::Closure& continuation) {
    if (!InitializeFiles()) {
      return;
    }
  
    STORAGE_DEVICE_NUMBER sdn = {0};
    DWORD bytes_returned;
  
    BOOL status = DeviceIoControl(
        device_file_.GetPlatformFile(),
        IOCTL_STORAGE_GET_DEVICE_NUMBER,
        NULL,             // Unused, must be NULL.
        0,                // Unused, must be 0.
        &sdn,             // An input buffer to hold the STORAGE_DEVICE_NUMBER
        sizeof(sdn),      // The size of the input buffer.
        &bytes_returned,  // the actual number of bytes returned.
        NULL);            // Unused overlap.
    if (!status) {
      PLOG(ERROR) << "Unable to get device number.";
      return;
    }
  
    ULONG device_number = sdn.DeviceNumber;
  
    TCHAR volume_path[MAX_PATH + 1];
    HANDLE volume_finder = FindFirstVolume(volume_path, MAX_PATH + 1);
    if (volume_finder == INVALID_HANDLE_VALUE) {
      return;
    }
  
    HANDLE volume_handle;
    bool first_volume = true;
    bool success = true;
  
    while (first_volume ||
           FindNextVolume(volume_finder, volume_path, MAX_PATH + 1)) {
      first_volume = false;
  
      size_t length = wcsnlen(volume_path, MAX_PATH + 1);
      if (length < 1) {
        continue;
      }
      volume_path[length - 1] = L'\0';
  
      volume_handle = CreateFile(volume_path,
                                 GENERIC_READ | GENERIC_WRITE,
                                 FILE_SHARE_READ | FILE_SHARE_WRITE,
                                 NULL,
                                 OPEN_EXISTING,
                                 0,
                                 NULL);
      if (volume_handle == INVALID_HANDLE_VALUE) {
        PLOG(ERROR) << "Opening volume handle failed.";
        success = false;
        break;
      }
  
      volume_handles_.push_back(volume_handle);
  
      VOLUME_DISK_EXTENTS disk_extents = {0};
      status = DeviceIoControl(volume_handle,
                               IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
                               NULL,
                               0,
                               &disk_extents,
                               sizeof(disk_extents),
                               &bytes_returned,
                               NULL);
  
      if (!status) {
        DWORD error = GetLastError();
        if (error == ERROR_MORE_DATA || error == ERROR_INVALID_FUNCTION ||
            error == ERROR_NOT_READY) {
          continue;
        } else {
          PLOG(ERROR) << "Unable to get volume disk extents.";
          success = false;
          break;
        }
      }
  
      if (disk_extents.NumberOfDiskExtents != 1 ||
          disk_extents.Extents[0].DiskNumber != device_number) {
        continue;
      }
  
      status = DeviceIoControl(volume_handle,
                               FSCTL_LOCK_VOLUME,
                               NULL,
                               0,
                               NULL,
                               0,
                               &bytes_returned,
                               NULL);
      if (!status) {
        PLOG(ERROR) << "Unable to lock volume.";
        success = false;
        break;
      }
  
      status = DeviceIoControl(volume_handle,
                               FSCTL_DISMOUNT_VOLUME,
                               NULL,
                               0,
                               NULL,
                               0,
                               &bytes_returned,
                               NULL);
      if (!status) {
        DWORD error = GetLastError();
        if (error != ERROR_NOT_SUPPORTED) {
          PLOG(ERROR) << "Unable to dismount volume.";
          success = false;
          break;
        }
      }
    }
  
    if (volume_finder != INVALID_HANDLE_VALUE) {
      FindVolumeClose(volume_finder);
    }
  
    if (success)
      continuation.Run();
  }
  
  }  // namespace image_writer
  