diskio-windows.cc

   1 //
   2 // C++ Interface: diskio (Windows-specific components)
   3 //
   4 // Description: Class to handle low-level disk I/O for GPT fdisk
   5 //
   6 //
   7 // Author: Rod Smith <rodsmith@rodsbooks.com>, (C) 2009
   8 //
   9 // Copyright: See COPYING file that comes with this distribution
  10 //
  11 //
  12 // This program is copyright (c) 2009, 2010 by Roderick W. Smith. It is distributed
  13 // under the terms of the GNU GPL version 2, as detailed in the COPYING file.
  14
  15 #define __STDC_LIMIT_MACROS
  16 #define __STDC_CONSTANT_MACROS
  17
  18 #include <windows.h>
  19 #include <winioctl.h>
  20 #define fstat64 fstat
  21 #define stat64 stat
  22 #define S_IRGRP 0
  23 #define S_IROTH 0
  24 #include <stdio.h>
  25 #include <string>
  26 #include <stdint.h>
  27 #include <errno.h>
  28 #include <fcntl.h>
  29 #include <sys/stat.h>
  30 #include <iostream>
  31
  32 #include "support.h"
  33 #include "diskio.h"
  34
  35 using namespace std;
  36
  37 // Returns the official Windows name for a shortened version of same.
  38 void DiskIO::MakeRealName(void) {
  39    size_t colonPos;
  40
  41    colonPos = userFilename.find(':', 0);
  42    if ((colonPos != string::npos) && (colonPos <= 3)) {
  43       realFilename = "\\\\.\\physicaldrive";
  44       realFilename += userFilename.substr(0, colonPos);
  45    } else {
  46       realFilename = userFilename;
  47    } // if/else
  48 } // DiskIO::MakeRealName()
  49
  50 // Open the currently on-record file for reading
  51 int DiskIO::OpenForRead(void) {
  52    int shouldOpen = 1;
  53
  54    if (isOpen) { // file is already open
  55       if (openForWrite) {
  56          Close();
  57       } else {
  58          shouldOpen = 0;
  59       } // if/else
  60    } // if
  61
  62    if (shouldOpen) {
  63       fd = CreateFile(realFilename.c_str(),GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE,
  64                       NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
  65       if (fd == INVALID_HANDLE_VALUE) {
  66          CloseHandle(fd);
  67          cerr << "Problem opening " << realFilename << " for reading!\n";
  68          realFilename = "";
  69          userFilename = "";
  70          isOpen = 0;
  71          openForWrite = 0;
  72       } else {
  73          isOpen = 1;
  74          openForWrite = 0;
  75       } // if/else
  76    } // if
  77
  78    return isOpen;
  79 } // DiskIO::OpenForRead(void)
  80
  81 // An extended file-open function. This includes some system-specific checks.
  82 // Returns 1 if the file is open, 0 otherwise....
  83 int DiskIO::OpenForWrite(void) {
  84    if ((isOpen) && (openForWrite))
  85       return 1;
  86
  87    // Close the disk, in case it's already open for reading only....
  88    Close();
  89
  90    // try to open the device; may fail....
  91    fd = CreateFile(realFilename.c_str(), GENERIC_READ | GENERIC_WRITE,
  92                    FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING,
  93                    FILE_ATTRIBUTE_NORMAL, NULL);
  94    // Preceding call can fail when creating backup files; if so, try
  95    // again with different option...
  96    if (fd == INVALID_HANDLE_VALUE) {
  97       CloseHandle(fd);
  98       fd = CreateFile(realFilename.c_str(), GENERIC_READ | GENERIC_WRITE,
  99                       FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_ALWAYS,
 100                       FILE_ATTRIBUTE_NORMAL, NULL);
 101    } // if
 102    if (fd == INVALID_HANDLE_VALUE) {
 103       CloseHandle(fd);
 104       isOpen = 0;
 105       openForWrite = 0;
 106       errno = GetLastError();
 107    } else {
 108       isOpen = 1;
 109       openForWrite = 1;
 110    } // if/else
 111    return isOpen;
 112 } // DiskIO::OpenForWrite(void)
 113
 114 // Close the disk device. Note that this does NOT erase the stored filenames,
 115 // so the file can be re-opened without specifying the filename.
 116 void DiskIO::Close(void) {
 117    if (isOpen)
 118       CloseHandle(fd);
 119    isOpen = 0;
 120    openForWrite = 0;
 121 } // DiskIO::Close()
 122
 123 // Returns block size of device pointed to by fd file descriptor. If the ioctl
 124 // returns an error condition, assume it's a disk file and return a value of
 125 // SECTOR_SIZE (512). If the disk can't be opened at all, return a value of 0.
 126 int DiskIO::GetBlockSize(void) {
 127    DWORD blockSize = 0, retBytes;
 128    DISK_GEOMETRY_EX geom;
 129
 130    // If disk isn't open, try to open it....
 131    if (!isOpen) {
 132       OpenForRead();
 133    } // if
 134
 135    if (isOpen) {
 136       if (DeviceIoControl(fd, IOCTL_DISK_GET_DRIVE_GEOMETRY_EX, NULL, 0,
 137                           &geom, sizeof(geom), &retBytes, NULL)) {
 138          blockSize = geom.Geometry.BytesPerSector;
 139       } else { // was probably an ordinary file; set default value....
 140          blockSize = SECTOR_SIZE;
 141       } // if/else
 142    } // if (isOpen)
 143
 144    return (blockSize);
 145 } // DiskIO::GetBlockSize()
 146
 147 // Returns the number of heads, according to the kernel, or 255 if the
 148 // correct value can't be determined.
 149 uint32_t DiskIO::GetNumHeads(void) {
 150    return UINT32_C(255);
 151 } // DiskIO::GetNumHeads();
 152
 153 // Returns the number of sectors per track, according to the kernel, or 63
 154 // if the correct value can't be determined.
 155 uint32_t DiskIO::GetNumSecsPerTrack(void) {
 156    return UINT32_C(63);
 157 } // DiskIO::GetNumSecsPerTrack()
 158
 159 // Resync disk caches so the OS uses the new partition table. This code varies
 160 // a lot from one OS to another.
 161 // Returns 1 on success, 0 if the kernel continues to use the old partition table.
 162 int DiskIO::DiskSync(void) {
 163    DWORD i;
 164    GET_LENGTH_INFORMATION buf;
 165    int retval = 0;
 166
 167    // If disk isn't open, try to open it....
 168    if (!openForWrite) {
 169       OpenForWrite();
 170    } // if
 171
 172    if (isOpen) {
 173       if (DeviceIoControl(fd, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, &buf, sizeof(buf), &i, NULL) == 0) {
 174          cout << "Disk synchronization failed! The computer may use the old partition table\n"
 175               << "until you reboot or remove and re-insert the disk!\n";
 176       } else {
 177          cout << "Disk synchronization succeeded! The computer should now use the new\n"
 178               << "partition table.\n";
 179          retval = 1;
 180       } // if/else
 181    } else {
 182       cout << "Unable to open the disk for synchronization operation! The computer will\n"
 183            << "continue to use the old partition table until you reboot or remove and\n"
 184            << "re-insert the disk!\n";
 185    } // if (isOpen)
 186    return retval;
 187 } // DiskIO::DiskSync()
 188
 189 // Seek to the specified sector. Returns 1 on success, 0 on failure.
 190 int DiskIO::Seek(uint64_t sector) {
 191    int retval = 1;
 192    LARGE_INTEGER seekTo;
 193
 194    // If disk isn't open, try to open it....
 195    if (!isOpen) {
 196       retval = OpenForRead();
 197    } // if
 198
 199    if (isOpen) {
 200       seekTo.QuadPart = sector * (uint64_t) GetBlockSize();
 201       retval = SetFilePointerEx(fd, seekTo, NULL, FILE_BEGIN);
 202       if (retval == 0) {
 203          errno = GetLastError();
 204          cerr << "Error when seeking to " << seekTo.QuadPart << "! Error is " << errno << "\n";
 205          retval = 0;
 206       } // if
 207    } // if
 208    return retval;
 209 } // DiskIO::Seek()
 210
 211 // A variant on the standard read() function. Done to work around
 212 // limitations in FreeBSD concerning the matching of the sector
 213 // size with the number of bytes read.
 214 // Returns the number of bytes read into buffer.
 215 int DiskIO::Read(void* buffer, int numBytes) {
 216    int blockSize = 512, i, numBlocks;
 217    char* tempSpace;
 218    DWORD retval = 0;
 219
 220    // If disk isn't open, try to open it....
 221    if (!isOpen) {
 222       OpenForRead();
 223    } // if
 224
 225    if (isOpen) {
 226       // Compute required space and allocate memory
 227       blockSize = GetBlockSize();
 228       if (numBytes <= blockSize) {
 229          numBlocks = 1;
 230          tempSpace = new char [blockSize];
 231       } else {
 232          numBlocks = numBytes / blockSize;
 233          if ((numBytes % blockSize) != 0)
 234             numBlocks++;
 235          tempSpace = new char [numBlocks * blockSize];
 236       } // if/else
 237       if (tempSpace == NULL) {
 238          cerr << "Unable to allocate memory in DiskIO::Read()! Terminating!\n";
 239          exit(1);
 240       } // if
 241
 242       // Read the data into temporary space, then copy it to buffer
 243       ReadFile(fd, tempSpace, numBlocks * blockSize, &retval, NULL);
 244       for (i = 0; i < numBytes; i++) {
 245          ((char*) buffer)[i] = tempSpace[i];
 246       } // for
 247
 248       // Adjust the return value, if necessary....
 249       if (((numBlocks * blockSize) != numBytes) && (retval > 0))
 250          retval = numBytes;
 251
 252       delete[] tempSpace;
 253    } // if (isOpen)
 254    return retval;
 255 } // DiskIO::Read()
 256
 257 // A variant on the standard write() function.
 258 // Returns the number of bytes written.
 259 int DiskIO::Write(void* buffer, int numBytes) {
 260    int blockSize = 512, i, numBlocks, retval = 0;
 261    char* tempSpace;
 262    DWORD numWritten;
 263
 264    // If disk isn't open, try to open it....
 265    if ((!isOpen) || (!openForWrite)) {
 266       OpenForWrite();
 267    } // if
 268
 269    if (isOpen) {
 270       // Compute required space and allocate memory
 271       blockSize = GetBlockSize();
 272       if (numBytes <= blockSize) {
 273          numBlocks = 1;
 274          tempSpace = new char [blockSize];
 275       } else {
 276          numBlocks = numBytes / blockSize;
 277          if ((numBytes % blockSize) != 0) numBlocks++;
 278          tempSpace = new char [numBlocks * blockSize];
 279       } // if/else
 280       if (tempSpace == NULL) {
 281          cerr << "Unable to allocate memory in DiskIO::Write()! Terminating!\n";
 282          exit(1);
 283       } // if
 284
 285       // Copy the data to my own buffer, then write it
 286       for (i = 0; i < numBytes; i++) {
 287          tempSpace[i] = ((char*) buffer)[i];
 288       } // for
 289       for (i = numBytes; i < numBlocks * blockSize; i++) {
 290          tempSpace[i] = 0;
 291       } // for
 292       WriteFile(fd, tempSpace, numBlocks * blockSize, &numWritten, NULL);
 293       retval = (int) numWritten;
 294
 295       // Adjust the return value, if necessary....
 296       if (((numBlocks * blockSize) != numBytes) && (retval > 0))
 297          retval = numBytes;
 298
 299       delete[] tempSpace;
 300    } // if (isOpen)
 301    return retval;
 302 } // DiskIO:Write()
 303
 304 // Returns the size of the disk in blocks.
 305 uint64_t DiskIO::DiskSize(int *err) {
 306    uint64_t sectors = 0; // size in sectors
 307    DWORD bytes, moreBytes; // low- and high-order bytes of file size
 308    GET_LENGTH_INFORMATION buf;
 309    DWORD i;
 310
 311    // If disk isn't open, try to open it....
 312    if (!isOpen) {
 313       OpenForRead();
 314    } // if
 315
 316    if (isOpen) {
 317       // Note to self: I recall testing a simplified version of
 318       // this code, similar to what's in the __APPLE__ block,
 319       // on Linux, but I had some problems. IIRC, it ran OK on 32-bit
 320       // systems but not on 64-bit. Keep this in mind in case of
 321       // 32/64-bit issues on MacOS....
 322       if (DeviceIoControl(fd, IOCTL_DISK_GET_LENGTH_INFO, NULL, 0, &buf, sizeof(buf), &i, NULL)) {
 323          sectors = (uint64_t) buf.Length.QuadPart / GetBlockSize();
 324          *err = 0;
 325       } else { // doesn't seem to be a disk device; assume it's an image file....
 326          bytes = GetFileSize(fd, &moreBytes);
 327          sectors = ((uint64_t) bytes + ((uint64_t) moreBytes) * UINT32_MAX) / GetBlockSize();
 328          *err = 0;
 329       } // if
 330    } else {
 331       *err = -1;
 332       sectors = 0;
 333    } // if/else (isOpen)
 334
 335    return sectors;
 336 } // DiskIO::DiskSize()