srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 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 |
srs5694 | 08bb0da | 2010-02-19 17:19:55 -0500 | [diff] [blame] | 22 | #define S_IRGRP 0 |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 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) { |
srs5694 | e321d44 | 2010-01-29 17:44:04 -0500 | [diff] [blame] | 39 | size_t colonPos; |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 40 | |
| 41 | colonPos = userFilename.find(':', 0); |
| 42 | if ((colonPos != string::npos) && (colonPos <= 3)) { |
| 43 | realFilename = "\\\\.\\physicaldrive"; |
| 44 | realFilename += userFilename.substr(0, colonPos); |
srs5694 | 0a69731 | 2010-01-28 21:10:52 -0500 | [diff] [blame] | 45 | } else { |
| 46 | realFilename = userFilename; |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 47 | } // if/else |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 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) { |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 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); |
srs5694 | 55d9261 | 2010-03-07 22:16:07 -0500 | [diff] [blame] | 67 | cerr << "Problem opening " << realFilename << " for reading!\n"; |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 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); |
srs5694 | 08bb0da | 2010-02-19 17:19:55 -0500 | [diff] [blame] | 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 |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 102 | if (fd == INVALID_HANDLE_VALUE) { |
| 103 | CloseHandle(fd); |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 104 | isOpen = 0; |
| 105 | openForWrite = 0; |
srs5694 | 0a69731 | 2010-01-28 21:10:52 -0500 | [diff] [blame] | 106 | errno = GetLastError(); |
| 107 | } else { |
| 108 | isOpen = 1; |
| 109 | openForWrite = 1; |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 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 |
srs5694 | 55d9261 | 2010-03-07 22:16:07 -0500 | [diff] [blame] | 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. |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 126 | int DiskIO::GetBlockSize(void) { |
srs5694 | 55d9261 | 2010-03-07 22:16:07 -0500 | [diff] [blame] | 127 | DWORD blockSize = 0, retBytes; |
| 128 | DISK_GEOMETRY_EX geom; |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 129 | |
| 130 | // If disk isn't open, try to open it.... |
| 131 | if (!isOpen) { |
| 132 | OpenForRead(); |
| 133 | } // if |
| 134 | |
| 135 | if (isOpen) { |
srs5694 | 699941e | 2011-03-21 21:33:57 -0400 | [diff] [blame] | 136 | if (DeviceIoControl(fd, IOCTL_DISK_GET_DRIVE_GEOMETRY_EX, NULL, 0, |
| 137 | &geom, sizeof(geom), &retBytes, NULL)) { |
srs5694 | 55d9261 | 2010-03-07 22:16:07 -0500 | [diff] [blame] | 138 | blockSize = geom.Geometry.BytesPerSector; |
srs5694 | 699941e | 2011-03-21 21:33:57 -0400 | [diff] [blame] | 139 | } else { // was probably an ordinary file; set default value.... |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 140 | blockSize = SECTOR_SIZE; |
srs5694 | 699941e | 2011-03-21 21:33:57 -0400 | [diff] [blame] | 141 | } // if/else |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 142 | } // if (isOpen) |
| 143 | |
| 144 | return (blockSize); |
| 145 | } // DiskIO::GetBlockSize() |
| 146 | |
srs5694 | bf8950c | 2011-03-12 01:23:12 -0500 | [diff] [blame] | 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 | |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 159 | // Resync disk caches so the OS uses the new partition table. This code varies |
| 160 | // a lot from one OS to another. |
srs5694 | a17fe69 | 2011-09-10 20:30:20 -0400 | [diff] [blame] | 161 | // Returns 1 on success, 0 if the kernel continues to use the old partition table. |
| 162 | int DiskIO::DiskSync(void) { |
srs5694 | 0a69731 | 2010-01-28 21:10:52 -0500 | [diff] [blame] | 163 | DWORD i; |
| 164 | GET_LENGTH_INFORMATION buf; |
srs5694 | a17fe69 | 2011-09-10 20:30:20 -0400 | [diff] [blame] | 165 | int retval = 0; |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 166 | |
| 167 | // If disk isn't open, try to open it.... |
srs5694 | 0a69731 | 2010-01-28 21:10:52 -0500 | [diff] [blame] | 168 | if (!openForWrite) { |
| 169 | OpenForWrite(); |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 170 | } // if |
| 171 | |
| 172 | if (isOpen) { |
srs5694 | 0a69731 | 2010-01-28 21:10:52 -0500 | [diff] [blame] | 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"; |
srs5694 | a17fe69 | 2011-09-10 20:30:20 -0400 | [diff] [blame] | 179 | retval = 1; |
srs5694 | 0a69731 | 2010-01-28 21:10:52 -0500 | [diff] [blame] | 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"; |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 185 | } // if (isOpen) |
srs5694 | a17fe69 | 2011-09-10 20:30:20 -0400 | [diff] [blame] | 186 | return retval; |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 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; |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 193 | |
| 194 | // If disk isn't open, try to open it.... |
| 195 | if (!isOpen) { |
| 196 | retval = OpenForRead(); |
| 197 | } // if |
| 198 | |
| 199 | if (isOpen) { |
srs5694 | 0a69731 | 2010-01-28 21:10:52 -0500 | [diff] [blame] | 200 | seekTo.QuadPart = sector * (uint64_t) GetBlockSize(); |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 201 | retval = SetFilePointerEx(fd, seekTo, NULL, FILE_BEGIN); |
| 202 | if (retval == 0) { |
| 203 | errno = GetLastError(); |
srs5694 | 0a69731 | 2010-01-28 21:10:52 -0500 | [diff] [blame] | 204 | cerr << "Error when seeking to " << seekTo.QuadPart << "! Error is " << errno << "\n"; |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 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; |
srs5694 | cb76c67 | 2010-02-11 22:22:22 -0500 | [diff] [blame] | 230 | tempSpace = new char [blockSize]; |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 231 | } else { |
| 232 | numBlocks = numBytes / blockSize; |
srs5694 | cb76c67 | 2010-02-11 22:22:22 -0500 | [diff] [blame] | 233 | if ((numBytes % blockSize) != 0) |
| 234 | numBlocks++; |
| 235 | tempSpace = new char [numBlocks * blockSize]; |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 236 | } // if/else |
srs5694 | 6aae2a9 | 2011-06-10 01:16:51 -0400 | [diff] [blame] | 237 | if (tempSpace == NULL) { |
| 238 | cerr << "Unable to allocate memory in DiskIO::Read()! Terminating!\n"; |
| 239 | exit(1); |
| 240 | } // if |
srs5694 | 3488294 | 2012-03-23 12:49:15 -0400 | [diff] [blame] | 241 | |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 242 | // Read the data into temporary space, then copy it to buffer |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 243 | ReadFile(fd, tempSpace, numBlocks * blockSize, &retval, NULL); |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 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 | |
srs5694 | cb76c67 | 2010-02-11 22:22:22 -0500 | [diff] [blame] | 252 | delete[] tempSpace; |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 253 | } // if (isOpen) |
| 254 | return retval; |
| 255 | } // DiskIO::Read() |
| 256 | |
srs5694 | 4307ef2 | 2012-05-30 12:30:48 -0400 | [diff] [blame] | 257 | // A variant on the standard write() function. |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 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; |
srs5694 | cb76c67 | 2010-02-11 22:22:22 -0500 | [diff] [blame] | 274 | tempSpace = new char [blockSize]; |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 275 | } else { |
| 276 | numBlocks = numBytes / blockSize; |
| 277 | if ((numBytes % blockSize) != 0) numBlocks++; |
srs5694 | cb76c67 | 2010-02-11 22:22:22 -0500 | [diff] [blame] | 278 | tempSpace = new char [numBlocks * blockSize]; |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 279 | } // if/else |
srs5694 | 6aae2a9 | 2011-06-10 01:16:51 -0400 | [diff] [blame] | 280 | if (tempSpace == NULL) { |
| 281 | cerr << "Unable to allocate memory in DiskIO::Write()! Terminating!\n"; |
| 282 | exit(1); |
| 283 | } // if |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 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 |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 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 | |
srs5694 | cb76c67 | 2010-02-11 22:22:22 -0500 | [diff] [blame] | 299 | delete[] tempSpace; |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 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 |
srs5694 | 0a69731 | 2010-01-28 21:10:52 -0500 | [diff] [blame] | 307 | DWORD bytes, moreBytes; // low- and high-order bytes of file size |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 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.... |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 322 | if (DeviceIoControl(fd, IOCTL_DISK_GET_LENGTH_INFO, NULL, 0, &buf, sizeof(buf), &i, NULL)) { |
| 323 | sectors = (uint64_t) buf.Length.QuadPart / GetBlockSize(); |
srs5694 | 0a69731 | 2010-01-28 21:10:52 -0500 | [diff] [blame] | 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) |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 334 | |
srs5694 | add79a6 | 2010-01-26 15:59:58 -0500 | [diff] [blame] | 335 | return sectors; |
| 336 | } // DiskIO::DiskSize() |