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 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 35using namespace std; 36 37// Returns the official Windows name for a shortened version of same. 38void 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 51int 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.... 83int 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. 116void 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. 126int 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. 149uint32_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. 155uint32_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. 162int 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. 190int 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. 215int 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. 259int 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. 305uint64_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