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