Home | History | Annotate | Download | only in gptfdisk
      1 //
      2 // C++ Interface: diskio (Unix components [Linux, FreeBSD, Mac OS X])
      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 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 #ifndef __STDC_CONSTANT_MACROS
     17 #define __STDC_CONSTANT_MACROS
     18 #endif
     19 
     20 #include <sys/ioctl.h>
     21 #include <string.h>
     22 #include <string>
     23 #include <stdint.h>
     24 #include <unistd.h>
     25 #include <errno.h>
     26 #include <fcntl.h>
     27 #include <sys/stat.h>
     28 #include <unistd.h>
     29 
     30 #ifdef __linux__
     31 #include "linux/hdreg.h"
     32 #endif
     33 
     34 #include <iostream>
     35 
     36 #include "diskio.h"
     37 
     38 using namespace std;
     39 
     40 // Returns the official "real" name for a shortened version of same.
     41 // Trivial here; more important in Windows
     42 void DiskIO::MakeRealName(void) {
     43    realFilename = userFilename;
     44 } // DiskIO::MakeRealName()
     45 
     46 // Open the currently on-record file for reading. Returns 1 if the file is
     47 // already open or is opened by this call, 0 if opening the file doesn't
     48 // work.
     49 int DiskIO::OpenForRead(void) {
     50    int shouldOpen = 1;
     51    struct stat64 st;
     52 
     53    if (isOpen) { // file is already open
     54       if (openForWrite) {
     55          Close();
     56       } else {
     57          shouldOpen = 0;
     58       } // if/else
     59    } // if
     60 
     61    if (shouldOpen) {
     62       fd = open(realFilename.c_str(), O_RDONLY);
     63       if (fd == -1) {
     64          cerr << "Problem opening " << realFilename << " for reading! Error is " << errno << ".\n";
     65          if (errno == EACCES) // User is probably not running as root
     66             cerr << "You must run this program as root or use sudo!\n";
     67          if (errno == ENOENT)
     68             cerr << "The specified file does not exist!\n";
     69          realFilename = "";
     70          userFilename = "";
     71          isOpen = 0;
     72          openForWrite = 0;
     73       } else {
     74          isOpen = 0;
     75          openForWrite = 0;
     76          if (fstat64(fd, &st) == 0) {
     77             if (S_ISDIR(st.st_mode))
     78                cerr << "The specified path is a directory!\n";
     79 #if !defined(__FreeBSD__) && !defined(__APPLE__)
     80             else if (S_ISCHR(st.st_mode))
     81                cerr << "The specified path is a character device!\n";
     82 #endif
     83             else if (S_ISFIFO(st.st_mode))
     84                cerr << "The specified path is a FIFO!\n";
     85             else if (S_ISSOCK(st.st_mode))
     86                cerr << "The specified path is a socket!\n";
     87             else
     88                isOpen = 1;
     89          } // if (fstat64()...)
     90       } // if/else
     91    } // if
     92 
     93    return isOpen;
     94 } // DiskIO::OpenForRead(void)
     95 
     96 // An extended file-open function. This includes some system-specific checks.
     97 // Returns 1 if the file is open, 0 otherwise....
     98 int DiskIO::OpenForWrite(void) {
     99    if ((isOpen) && (openForWrite))
    100       return 1;
    101 
    102    // Close the disk, in case it's already open for reading only....
    103    Close();
    104 
    105    // try to open the device; may fail....
    106    fd = open(realFilename.c_str(), O_WRONLY | O_CREAT, S_IWUSR | S_IRUSR | S_IRGRP | S_IROTH);
    107 #ifdef __APPLE__
    108    // MacOS X requires a shared lock under some circumstances....
    109    if (fd < 0) {
    110       cerr << "Warning: Devices opened with shared lock will not have their\npartition table automatically reloaded!\n";
    111       fd = open(realFilename.c_str(), O_WRONLY | O_SHLOCK);
    112    } // if
    113 #endif
    114    if (fd >= 0) {
    115       isOpen = 1;
    116       openForWrite = 1;
    117    } else {
    118       isOpen = 0;
    119       openForWrite = 0;
    120    } // if/else
    121    return isOpen;
    122 } // DiskIO::OpenForWrite(void)
    123 
    124 // Close the disk device. Note that this does NOT erase the stored filenames,
    125 // so the file can be re-opened without specifying the filename.
    126 void DiskIO::Close(void) {
    127    if (isOpen)
    128       if (close(fd) < 0)
    129          cerr << "Warning! Problem closing file!\n";
    130    isOpen = 0;
    131    openForWrite = 0;
    132 } // DiskIO::Close()
    133 
    134 // Returns block size of device pointed to by fd file descriptor. If the ioctl
    135 // returns an error condition, print a warning but return a value of SECTOR_SIZE
    136 // (512). If the disk can't be opened at all, return a value of 0.
    137 int DiskIO::GetBlockSize(void) {
    138    int err = -1, blockSize = 0;
    139 #ifdef __sun__
    140    struct dk_minfo minfo;
    141 #endif
    142 
    143    // If disk isn't open, try to open it....
    144    if (!isOpen) {
    145       OpenForRead();
    146    } // if
    147 
    148    if (isOpen) {
    149 #ifdef __APPLE__
    150       err = ioctl(fd, DKIOCGETBLOCKSIZE, &blockSize);
    151 #endif
    152 #ifdef __sun__
    153       err = ioctl(fd, DKIOCGMEDIAINFO, &minfo);
    154       if (err == 0)
    155           blockSize = minfo.dki_lbsize;
    156 #endif
    157 #if defined (__FreeBSD__) || defined (__FreeBSD_kernel__)
    158       err = ioctl(fd, DIOCGSECTORSIZE, &blockSize);
    159 #endif
    160 #ifdef __linux__
    161       err = ioctl(fd, BLKSSZGET, &blockSize);
    162 #endif
    163 
    164       if (err == -1) {
    165          blockSize = SECTOR_SIZE;
    166          // ENOTTY = inappropriate ioctl; probably being called on a disk image
    167          // file, so don't display the warning message....
    168          // 32-bit code returns EINVAL, I don't know why. I know I'm treading on
    169          // thin ice here, but it should be OK in all but very weird cases....
    170          if ((errno != ENOTTY) && (errno != EINVAL)) {
    171             cerr << "\aError " << errno << " when determining sector size! Setting sector size to "
    172                  << SECTOR_SIZE << "\n";
    173             cout << "Disk device is " << realFilename << "\n";
    174          } // if
    175       } // if (err == -1)
    176    } // if (isOpen)
    177 
    178    return (blockSize);
    179 } // DiskIO::GetBlockSize()
    180 
    181 // Returns the number of heads, according to the kernel, or 255 if the
    182 // correct value can't be determined.
    183 uint32_t DiskIO::GetNumHeads(void) {
    184    uint32_t numHeads = 255;
    185 
    186 #ifdef HDIO_GETGEO
    187    struct hd_geometry geometry;
    188 
    189    // If disk isn't open, try to open it....
    190    if (!isOpen)
    191       OpenForRead();
    192 
    193    if (!ioctl(fd, HDIO_GETGEO, &geometry))
    194       numHeads = (uint32_t) geometry.heads;
    195 #endif
    196    return numHeads;
    197 } // DiskIO::GetNumHeads();
    198 
    199 // Returns the number of sectors per track, according to the kernel, or 63
    200 // if the correct value can't be determined.
    201 uint32_t DiskIO::GetNumSecsPerTrack(void) {
    202    uint32_t numSecs = 63;
    203 
    204    #ifdef HDIO_GETGEO
    205    struct hd_geometry geometry;
    206 
    207    // If disk isn't open, try to open it....
    208    if (!isOpen)
    209       OpenForRead();
    210 
    211    if (!ioctl(fd, HDIO_GETGEO, &geometry))
    212       numSecs = (uint32_t) geometry.sectors;
    213    #endif
    214    return numSecs;
    215 } // DiskIO::GetNumSecsPerTrack()
    216 
    217 // Resync disk caches so the OS uses the new partition table. This code varies
    218 // a lot from one OS to another.
    219 // Returns 1 on success, 0 if the kernel continues to use the old partition table.
    220 // (Note that for most OSes, the default of 0 is returned because I've not yet
    221 // looked into how to test for success in the underlying system calls...)
    222 int DiskIO::DiskSync(void) {
    223    int i, retval = 0, platformFound = 0;
    224 
    225    // If disk isn't open, try to open it....
    226    if (!isOpen) {
    227       OpenForRead();
    228    } // if
    229 
    230    if (isOpen) {
    231       sync();
    232 #if defined(__APPLE__) || defined(__sun__)
    233       cout << "Warning: The kernel may continue to use old or deleted partitions.\n"
    234            << "You should reboot or remove the drive.\n";
    235                /* don't know if this helps
    236                * it definitely will get things on disk though:
    237                * http://topiks.org/mac-os-x/0321278542/ch12lev1sec8.html */
    238 #ifdef __sun__
    239       i = ioctl(fd, DKIOCFLUSHWRITECACHE);
    240 #else
    241       i = ioctl(fd, DKIOCSYNCHRONIZECACHE);
    242 #endif
    243       platformFound++;
    244 #endif
    245 #if defined (__FreeBSD__) || defined (__FreeBSD_kernel__)
    246       sleep(2);
    247       i = ioctl(fd, DIOCGFLUSH);
    248       cout << "Warning: The kernel may continue to use old or deleted partitions.\n"
    249            << "You should reboot or remove the drive.\n";
    250       platformFound++;
    251 #endif
    252 #ifdef __linux__
    253       sleep(1); // Theoretically unnecessary, but ioctl() fails sometimes if omitted....
    254       fsync(fd);
    255       i = ioctl(fd, BLKRRPART);
    256       if (i) {
    257          cout << "Warning: The kernel is still using the old partition table.\n"
    258               << "The new table will be used at the next reboot.\n";
    259       } else {
    260          retval = 1;
    261       } // if/else
    262       platformFound++;
    263 #endif
    264       if (platformFound == 0)
    265          cerr << "Warning: Platform not recognized!\n";
    266       if (platformFound > 1)
    267          cerr << "\nWarning: We seem to be running on multiple platforms!\n";
    268    } // if (isOpen)
    269    return retval;
    270 } // DiskIO::DiskSync()
    271 
    272 // Seek to the specified sector. Returns 1 on success, 0 on failure.
    273 // Note that seeking beyond the end of the file is NOT detected as a failure!
    274 int DiskIO::Seek(uint64_t sector) {
    275    int retval = 1;
    276    off64_t seekTo, sought;
    277 
    278    // If disk isn't open, try to open it....
    279    if (!isOpen) {
    280       retval = OpenForRead();
    281    } // if
    282 
    283    if (isOpen) {
    284       seekTo = sector * (uint64_t) GetBlockSize();
    285       sought = lseek64(fd, seekTo, SEEK_SET);
    286       if (sought != seekTo) {
    287          retval = 0;
    288       } // if
    289    } // if
    290    return retval;
    291 } // DiskIO::Seek()
    292 
    293 // A variant on the standard read() function. Done to work around
    294 // limitations in FreeBSD concerning the matching of the sector
    295 // size with the number of bytes read.
    296 // Returns the number of bytes read into buffer.
    297 int DiskIO::Read(void* buffer, int numBytes) {
    298    int blockSize, numBlocks, retval = 0;
    299    char* tempSpace;
    300 
    301    // If disk isn't open, try to open it....
    302    if (!isOpen) {
    303       OpenForRead();
    304    } // if
    305 
    306    if (isOpen) {
    307       // Compute required space and allocate memory
    308       blockSize = GetBlockSize();
    309       if (numBytes <= blockSize) {
    310          numBlocks = 1;
    311          tempSpace = new char [blockSize];
    312       } else {
    313          numBlocks = numBytes / blockSize;
    314          if ((numBytes % blockSize) != 0)
    315             numBlocks++;
    316          tempSpace = new char [numBlocks * blockSize];
    317       } // if/else
    318       if (tempSpace == NULL) {
    319          cerr << "Unable to allocate memory in DiskIO::Read()! Terminating!\n";
    320          exit(1);
    321       } // if
    322 
    323       // Read the data into temporary space, then copy it to buffer
    324       retval = read(fd, tempSpace, numBlocks * blockSize);
    325       memcpy(buffer, tempSpace, numBytes);
    326 
    327       // Adjust the return value, if necessary....
    328       if (((numBlocks * blockSize) != numBytes) && (retval > 0))
    329          retval = numBytes;
    330 
    331       delete[] tempSpace;
    332    } // if (isOpen)
    333    return retval;
    334 } // DiskIO::Read()
    335 
    336 // A variant on the standard write() function. Done to work around
    337 // limitations in FreeBSD concerning the matching of the sector
    338 // size with the number of bytes read.
    339 // Returns the number of bytes written.
    340 int DiskIO::Write(void* buffer, int numBytes) {
    341    int blockSize = 512, i, numBlocks, retval = 0;
    342    char* tempSpace;
    343 
    344    // If disk isn't open, try to open it....
    345    if ((!isOpen) || (!openForWrite)) {
    346       OpenForWrite();
    347    } // if
    348 
    349    if (isOpen) {
    350       // Compute required space and allocate memory
    351       blockSize = GetBlockSize();
    352       if (numBytes <= blockSize) {
    353          numBlocks = 1;
    354          tempSpace = new char [blockSize];
    355       } else {
    356          numBlocks = numBytes / blockSize;
    357          if ((numBytes % blockSize) != 0) numBlocks++;
    358          tempSpace = new char [numBlocks * blockSize];
    359       } // if/else
    360       if (tempSpace == NULL) {
    361          cerr << "Unable to allocate memory in DiskIO::Write()! Terminating!\n";
    362          exit(1);
    363       } // if
    364 
    365       // Copy the data to my own buffer, then write it
    366       memcpy(tempSpace, buffer, numBytes);
    367       for (i = numBytes; i < numBlocks * blockSize; i++) {
    368          tempSpace[i] = 0;
    369       } // for
    370       retval = write(fd, tempSpace, numBlocks * blockSize);
    371 
    372       // Adjust the return value, if necessary....
    373       if (((numBlocks * blockSize) != numBytes) && (retval > 0))
    374          retval = numBytes;
    375 
    376       delete[] tempSpace;
    377    } // if (isOpen)
    378    return retval;
    379 } // DiskIO:Write()
    380 
    381 /**************************************************************************************
    382  *                                                                                    *
    383  * Below functions are lifted from various sources, as documented in comments before  *
    384  * each one.                                                                          *
    385  *                                                                                    *
    386  **************************************************************************************/
    387 
    388 // The disksize function is taken from the Linux fdisk code and modified
    389 // greatly since then to enable FreeBSD and MacOS support, as well as to
    390 // return correct values for disk image files.
    391 uint64_t DiskIO::DiskSize(int *err) {
    392    uint64_t sectors = 0; // size in sectors
    393    off_t bytes = 0; // size in bytes
    394    struct stat64 st;
    395    int platformFound = 0;
    396 #ifdef __sun__
    397    struct dk_minfo minfo;
    398 #endif
    399 
    400    // If disk isn't open, try to open it....
    401    if (!isOpen) {
    402       OpenForRead();
    403    } // if
    404 
    405    if (isOpen) {
    406       // Note to self: I recall testing a simplified version of
    407       // this code, similar to what's in the __APPLE__ block,
    408       // on Linux, but I had some problems. IIRC, it ran OK on 32-bit
    409       // systems but not on 64-bit. Keep this in mind in case of
    410       // 32/64-bit issues on MacOS....
    411 #ifdef __APPLE__
    412       *err = ioctl(fd, DKIOCGETBLOCKCOUNT, &sectors);
    413       platformFound++;
    414 #endif
    415 #ifdef __sun__
    416       *err = ioctl(fd, DKIOCGMEDIAINFO, &minfo);
    417       if (*err == 0)
    418           sectors = minfo.dki_capacity;
    419       platformFound++;
    420 #endif
    421 #if defined (__FreeBSD__) || defined (__FreeBSD_kernel__)
    422       *err = ioctl(fd, DIOCGMEDIASIZE, &bytes);
    423       long long b = GetBlockSize();
    424       sectors = bytes / b;
    425       platformFound++;
    426 #endif
    427 #ifdef __linux__
    428       long sz;
    429       long long b;
    430       *err = ioctl(fd, BLKGETSIZE, &sz);
    431       if (*err) {
    432          sectors = sz = 0;
    433       } // if
    434       if ((!*err) || (errno == EFBIG)) {
    435          *err = ioctl(fd, BLKGETSIZE64, &b);
    436          if (*err || b == 0 || b == sz)
    437             sectors = sz;
    438          else
    439             sectors = (b >> 9);
    440       } // if
    441       // Unintuitively, the above returns values in 512-byte blocks, no
    442       // matter what the underlying device's block size. Correct for this....
    443       sectors /= (GetBlockSize() / 512);
    444       platformFound++;
    445 #endif
    446       if (platformFound != 1)
    447          cerr << "Warning! We seem to be running on no known platform!\n";
    448 
    449       // The above methods have failed, so let's assume it's a regular
    450       // file (a QEMU image, dd backup, or what have you) and see what
    451       // fstat() gives us....
    452       if ((sectors == 0) || (*err == -1)) {
    453          if (fstat64(fd, &st) == 0) {
    454             bytes = st.st_size;
    455             if ((bytes % UINT64_C(512)) != 0)
    456                cerr << "Warning: File size is not a multiple of 512 bytes!"
    457                     << " Misbehavior is likely!\n\a";
    458             sectors = bytes / UINT64_C(512);
    459          } // if
    460       } // if
    461    } // if (isOpen)
    462    return sectors;
    463 } // DiskIO::DiskSize()
    464