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, §ors); 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