Home | History | Annotate | Download | only in gptfdisk 
      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 (at) 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()
    337