diskio-unix.cc

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