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 #ifndef __STDC_CONSTANT_MACROS
17 #define __STDC_CONSTANT_MACROS
18 #endif
19
20 #include <windows.h>
21 #include <winioctl.h>
22 #define fstat64 fstat
23 #define stat64 stat
24 #define S_IRGRP 0
25 #define S_IROTH 0
26 #include <stdio.h>
27 #include <string>
28 #include <stdint.h>
29 #include <errno.h>
30 #include <fcntl.h>
31 #include <sys/stat.h>
32 #include <iostream>
33
34 #include "support.h"
35 #include "diskio.h"
36
37 using namespace std;
38
39 // Returns the official Windows name for a shortened version of same.
MakeRealName(void)40 void DiskIO::MakeRealName(void) {
41 size_t colonPos;
42
43 colonPos = userFilename.find(':', 0);
44 if ((colonPos != string::npos) && (colonPos <= 3)) {
45 realFilename = "\\\\.\\physicaldrive";
46 realFilename += userFilename.substr(0, colonPos);
47 } else {
48 realFilename = userFilename;
49 } // if/else
50 } // DiskIO::MakeRealName()
51
52 // Open the currently on-record file for reading
OpenForRead(void)53 int DiskIO::OpenForRead(void) {
54 int shouldOpen = 1;
55
56 if (isOpen) { // file is already open
57 if (openForWrite) {
58 Close();
59 } else {
60 shouldOpen = 0;
61 } // if/else
62 } // if
63
64 if (shouldOpen) {
65 fd = CreateFile(realFilename.c_str(),GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE,
66 NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
67 if (fd == INVALID_HANDLE_VALUE) {
68 cerr << "Problem opening " << realFilename << " for reading!\n";
69 realFilename = "";
70 userFilename = "";
71 isOpen = 0;
72 openForWrite = 0;
73 } else {
74 isOpen = 1;
75 openForWrite = 0;
76 } // if/else
77 } // if
78
79 return isOpen;
80 } // DiskIO::OpenForRead(void)
81
82 // An extended file-open function. This includes some system-specific checks.
83 // Returns 1 if the file is open, 0 otherwise....
OpenForWrite(void)84 int DiskIO::OpenForWrite(void) {
85 if ((isOpen) && (openForWrite))
86 return 1;
87
88 // Close the disk, in case it's already open for reading only....
89 Close();
90
91 // try to open the device; may fail....
92 fd = CreateFile(realFilename.c_str(), GENERIC_READ | GENERIC_WRITE,
93 FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING,
94 FILE_ATTRIBUTE_NORMAL, NULL);
95 // Preceding call can fail when creating backup files; if so, try
96 // again with different option...
97 if (fd == INVALID_HANDLE_VALUE) {
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 isOpen = 0;
104 openForWrite = 0;
105 errno = GetLastError();
106 } else {
107 isOpen = 1;
108 openForWrite = 1;
109 } // if/else
110 return isOpen;
111 } // DiskIO::OpenForWrite(void)
112
113 // Close the disk device. Note that this does NOT erase the stored filenames,
114 // so the file can be re-opened without specifying the filename.
Close(void)115 void DiskIO::Close(void) {
116 if (isOpen) {
117 CloseHandle(fd);
118 fd = INVALID_HANDLE_VALUE;
119 }
120 isOpen = 0;
121 openForWrite = 0;
122 } // DiskIO::Close()
123
124 // Returns block size of device pointed to by fd file descriptor. If the ioctl
125 // returns an error condition, assume it's a disk file and return a value of
126 // SECTOR_SIZE (512). If the disk can't be opened at all, return a value of 0.
GetBlockSize(void)127 int DiskIO::GetBlockSize(void) {
128 DWORD blockSize = 0, retBytes;
129 DISK_GEOMETRY_EX geom;
130
131 // If disk isn't open, try to open it....
132 if (!isOpen) {
133 OpenForRead();
134 } // if
135
136 if (isOpen) {
137 if (DeviceIoControl(fd, IOCTL_DISK_GET_DRIVE_GEOMETRY_EX, NULL, 0,
138 &geom, sizeof(geom), &retBytes, NULL)) {
139 blockSize = geom.Geometry.BytesPerSector;
140 } else { // was probably an ordinary file; set default value....
141 blockSize = SECTOR_SIZE;
142 } // if/else
143 } // if (isOpen)
144
145 return (blockSize);
146 } // DiskIO::GetBlockSize()
147
148 // In theory, returns the physical block size. In practice, this is only
149 // supported in Linux, as of yet.
150 // TODO: Get this working in Windows.
GetPhysBlockSize(void)151 int DiskIO::GetPhysBlockSize(void) {
152 return 0;
153 } // DiskIO::GetPhysBlockSize()
154
155 // Returns the number of heads, according to the kernel, or 255 if the
156 // correct value can't be determined.
GetNumHeads(void)157 uint32_t DiskIO::GetNumHeads(void) {
158 return UINT32_C(255);
159 } // DiskIO::GetNumHeads();
160
161 // Returns the number of sectors per track, according to the kernel, or 63
162 // if the correct value can't be determined.
GetNumSecsPerTrack(void)163 uint32_t DiskIO::GetNumSecsPerTrack(void) {
164 return UINT32_C(63);
165 } // DiskIO::GetNumSecsPerTrack()
166
167 // Resync disk caches so the OS uses the new partition table. This code varies
168 // a lot from one OS to another.
169 // Returns 1 on success, 0 if the kernel continues to use the old partition table.
DiskSync(void)170 int DiskIO::DiskSync(void) {
171 DWORD i;
172 GET_LENGTH_INFORMATION buf;
173 int retval = 0;
174
175 // If disk isn't open, try to open it....
176 if (!openForWrite) {
177 OpenForWrite();
178 } // if
179
180 if (isOpen) {
181 if (DeviceIoControl(fd, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, &buf, sizeof(buf), &i, NULL) == 0) {
182 cout << "Disk synchronization failed! The computer may use the old partition table\n"
183 << "until you reboot or remove and re-insert the disk!\n";
184 } else {
185 cout << "Disk synchronization succeeded! The computer should now use the new\n"
186 << "partition table.\n";
187 retval = 1;
188 } // if/else
189 } else {
190 cout << "Unable to open the disk for synchronization operation! The computer will\n"
191 << "continue to use the old partition table until you reboot or remove and\n"
192 << "re-insert the disk!\n";
193 } // if (isOpen)
194 return retval;
195 } // DiskIO::DiskSync()
196
197 // Seek to the specified sector. Returns 1 on success, 0 on failure.
Seek(uint64_t sector)198 int DiskIO::Seek(uint64_t sector) {
199 int retval = 1;
200 LARGE_INTEGER seekTo;
201
202 // If disk isn't open, try to open it....
203 if (!isOpen) {
204 retval = OpenForRead();
205 } // if
206
207 if (isOpen) {
208 seekTo.QuadPart = sector * (uint64_t) GetBlockSize();
209 retval = SetFilePointerEx(fd, seekTo, NULL, FILE_BEGIN);
210 if (retval == 0) {
211 errno = GetLastError();
212 cerr << "Error when seeking to " << seekTo.QuadPart << "! Error is " << errno << "\n";
213 retval = 0;
214 } // if
215 } // if
216 return retval;
217 } // DiskIO::Seek()
218
219 // A variant on the standard read() function. Done to work around
220 // limitations in FreeBSD concerning the matching of the sector
221 // size with the number of bytes read.
222 // Returns the number of bytes read into buffer.
Read(void* buffer, int numBytes)223 int DiskIO::Read(void* buffer, int numBytes) {
224 int blockSize = 512, i, numBlocks;
225 char* tempSpace;
226 DWORD retval = 0;
227
228 // If disk isn't open, try to open it....
229 if (!isOpen) {
230 OpenForRead();
231 } // if
232
233 if (isOpen) {
234 // Compute required space and allocate memory
235 blockSize = GetBlockSize();
236 if (numBytes <= blockSize) {
237 numBlocks = 1;
238 tempSpace = new char [blockSize];
239 } else {
240 numBlocks = numBytes / blockSize;
241 if ((numBytes % blockSize) != 0)
242 numBlocks++;
243 tempSpace = new char [numBlocks * blockSize];
244 } // if/else
245 if (tempSpace == NULL) {
246 cerr << "Unable to allocate memory in DiskIO::Read()! Terminating!\n";
247 exit(1);
248 } // if
249
250 // Read the data into temporary space, then copy it to buffer
251 ReadFile(fd, tempSpace, numBlocks * blockSize, &retval, NULL);
252 for (i = 0; i < numBytes; i++) {
253 ((char*) buffer)[i] = tempSpace[i];
254 } // for
255
256 // Adjust the return value, if necessary....
257 if (((numBlocks * blockSize) != numBytes) && (retval > 0))
258 retval = numBytes;
259
260 delete[] tempSpace;
261 } // if (isOpen)
262 return retval;
263 } // DiskIO::Read()
264
265 // A variant on the standard write() function.
266 // Returns the number of bytes written.
Write(void* buffer, int numBytes)267 int DiskIO::Write(void* buffer, int numBytes) {
268 int blockSize = 512, i, numBlocks, retval = 0;
269 char* tempSpace;
270 DWORD numWritten;
271
272 // If disk isn't open, try to open it....
273 if ((!isOpen) || (!openForWrite)) {
274 OpenForWrite();
275 } // if
276
277 if (isOpen) {
278 // Compute required space and allocate memory
279 blockSize = GetBlockSize();
280 if (numBytes <= blockSize) {
281 numBlocks = 1;
282 tempSpace = new char [blockSize];
283 } else {
284 numBlocks = numBytes / blockSize;
285 if ((numBytes % blockSize) != 0) numBlocks++;
286 tempSpace = new char [numBlocks * blockSize];
287 } // if/else
288 if (tempSpace == NULL) {
289 cerr << "Unable to allocate memory in DiskIO::Write()! Terminating!\n";
290 exit(1);
291 } // if
292
293 // Copy the data to my own buffer, then write it
294 for (i = 0; i < numBytes; i++) {
295 tempSpace[i] = ((char*) buffer)[i];
296 } // for
297 for (i = numBytes; i < numBlocks * blockSize; i++) {
298 tempSpace[i] = 0;
299 } // for
300 WriteFile(fd, tempSpace, numBlocks * blockSize, &numWritten, NULL);
301 retval = (int) numWritten;
302
303 // Adjust the return value, if necessary....
304 if (((numBlocks * blockSize) != numBytes) && (retval > 0))
305 retval = numBytes;
306
307 delete[] tempSpace;
308 } // if (isOpen)
309 return retval;
310 } // DiskIO:Write()
311
312 // Returns the size of the disk in blocks.
DiskSize(int *err)313 uint64_t DiskIO::DiskSize(int *err) {
314 uint64_t sectors = 0; // size in sectors
315 DWORD bytes, moreBytes; // low- and high-order bytes of file size
316 GET_LENGTH_INFORMATION buf;
317 DWORD i;
318
319 // If disk isn't open, try to open it....
320 if (!isOpen) {
321 OpenForRead();
322 } // if
323
324 if (isOpen) {
325 // Note to self: I recall testing a simplified version of
326 // this code, similar to what's in the __APPLE__ block,
327 // on Linux, but I had some problems. IIRC, it ran OK on 32-bit
328 // systems but not on 64-bit. Keep this in mind in case of
329 // 32/64-bit issues on MacOS....
330 if (DeviceIoControl(fd, IOCTL_DISK_GET_LENGTH_INFO, NULL, 0, &buf, sizeof(buf), &i, NULL)) {
331 sectors = (uint64_t) buf.Length.QuadPart / GetBlockSize();
332 *err = 0;
333 } else { // doesn't seem to be a disk device; assume it's an image file....
334 bytes = GetFileSize(fd, &moreBytes);
335 sectors = ((uint64_t) bytes + ((uint64_t) moreBytes) * UINT32_MAX) / GetBlockSize();
336 *err = 0;
337 } // if
338 } else {
339 *err = -1;
340 sectors = 0;
341 } // if/else (isOpen)
342
343 return sectors;
344 } // DiskIO::DiskSize()
345