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
35 using namespace std;
36
37 // Returns the official Windows name for a shortened version of same.
MakeRealName(void)38 void 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
OpenForRead(void)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) {
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....
OpenForWrite(void)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);
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.
Close(void)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
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.
GetBlockSize(void)126 int 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 // In theory, returns the physical block size. In practice, this is only
148 // supported in Linux, as of yet.
149 // TODO: Get this working in Windows.
GetPhysBlockSize(void)150 int DiskIO::GetPhysBlockSize(void) {
151 return 0;
152 } // DiskIO::GetPhysBlockSize()
153
154 // Returns the number of heads, according to the kernel, or 255 if the
155 // correct value can't be determined.
GetNumHeads(void)156 uint32_t DiskIO::GetNumHeads(void) {
157 return UINT32_C(255);
158 } // DiskIO::GetNumHeads();
159
160 // Returns the number of sectors per track, according to the kernel, or 63
161 // if the correct value can't be determined.
GetNumSecsPerTrack(void)162 uint32_t DiskIO::GetNumSecsPerTrack(void) {
163 return UINT32_C(63);
164 } // DiskIO::GetNumSecsPerTrack()
165
166 // Resync disk caches so the OS uses the new partition table. This code varies
167 // a lot from one OS to another.
168 // Returns 1 on success, 0 if the kernel continues to use the old partition table.
DiskSync(void)169 int DiskIO::DiskSync(void) {
170 DWORD i;
171 GET_LENGTH_INFORMATION buf;
172 int retval = 0;
173
174 // If disk isn't open, try to open it....
175 if (!openForWrite) {
176 OpenForWrite();
177 } // if
178
179 if (isOpen) {
180 if (DeviceIoControl(fd, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, &buf, sizeof(buf), &i, NULL) == 0) {
181 cout << "Disk synchronization failed! The computer may use the old partition table\n"
182 << "until you reboot or remove and re-insert the disk!\n";
183 } else {
184 cout << "Disk synchronization succeeded! The computer should now use the new\n"
185 << "partition table.\n";
186 retval = 1;
187 } // if/else
188 } else {
189 cout << "Unable to open the disk for synchronization operation! The computer will\n"
190 << "continue to use the old partition table until you reboot or remove and\n"
191 << "re-insert the disk!\n";
192 } // if (isOpen)
193 return retval;
194 } // DiskIO::DiskSync()
195
196 // Seek to the specified sector. Returns 1 on success, 0 on failure.
Seek(uint64_t sector)197 int DiskIO::Seek(uint64_t sector) {
198 int retval = 1;
199 LARGE_INTEGER seekTo;
200
201 // If disk isn't open, try to open it....
202 if (!isOpen) {
203 retval = OpenForRead();
204 } // if
205
206 if (isOpen) {
207 seekTo.QuadPart = sector * (uint64_t) GetBlockSize();
208 retval = SetFilePointerEx(fd, seekTo, NULL, FILE_BEGIN);
209 if (retval == 0) {
210 errno = GetLastError();
211 cerr << "Error when seeking to " << seekTo.QuadPart << "! Error is " << errno << "\n";
212 retval = 0;
213 } // if
214 } // if
215 return retval;
216 } // DiskIO::Seek()
217
218 // A variant on the standard read() function. Done to work around
219 // limitations in FreeBSD concerning the matching of the sector
220 // size with the number of bytes read.
221 // Returns the number of bytes read into buffer.
Read(void * buffer,int numBytes)222 int DiskIO::Read(void* buffer, int numBytes) {
223 int blockSize = 512, i, numBlocks;
224 char* tempSpace;
225 DWORD retval = 0;
226
227 // If disk isn't open, try to open it....
228 if (!isOpen) {
229 OpenForRead();
230 } // if
231
232 if (isOpen) {
233 // Compute required space and allocate memory
234 blockSize = GetBlockSize();
235 if (numBytes <= blockSize) {
236 numBlocks = 1;
237 tempSpace = new char [blockSize];
238 } else {
239 numBlocks = numBytes / blockSize;
240 if ((numBytes % blockSize) != 0)
241 numBlocks++;
242 tempSpace = new char [numBlocks * blockSize];
243 } // if/else
244 if (tempSpace == NULL) {
245 cerr << "Unable to allocate memory in DiskIO::Read()! Terminating!\n";
246 exit(1);
247 } // if
248
249 // Read the data into temporary space, then copy it to buffer
250 ReadFile(fd, tempSpace, numBlocks * blockSize, &retval, NULL);
251 for (i = 0; i < numBytes; i++) {
252 ((char*) buffer)[i] = tempSpace[i];
253 } // for
254
255 // Adjust the return value, if necessary....
256 if (((numBlocks * blockSize) != numBytes) && (retval > 0))
257 retval = numBytes;
258
259 delete[] tempSpace;
260 } // if (isOpen)
261 return retval;
262 } // DiskIO::Read()
263
264 // A variant on the standard write() function.
265 // Returns the number of bytes written.
Write(void * buffer,int numBytes)266 int DiskIO::Write(void* buffer, int numBytes) {
267 int blockSize = 512, i, numBlocks, retval = 0;
268 char* tempSpace;
269 DWORD numWritten;
270
271 // If disk isn't open, try to open it....
272 if ((!isOpen) || (!openForWrite)) {
273 OpenForWrite();
274 } // if
275
276 if (isOpen) {
277 // Compute required space and allocate memory
278 blockSize = GetBlockSize();
279 if (numBytes <= blockSize) {
280 numBlocks = 1;
281 tempSpace = new char [blockSize];
282 } else {
283 numBlocks = numBytes / blockSize;
284 if ((numBytes % blockSize) != 0) numBlocks++;
285 tempSpace = new char [numBlocks * blockSize];
286 } // if/else
287 if (tempSpace == NULL) {
288 cerr << "Unable to allocate memory in DiskIO::Write()! Terminating!\n";
289 exit(1);
290 } // if
291
292 // Copy the data to my own buffer, then write it
293 for (i = 0; i < numBytes; i++) {
294 tempSpace[i] = ((char*) buffer)[i];
295 } // for
296 for (i = numBytes; i < numBlocks * blockSize; i++) {
297 tempSpace[i] = 0;
298 } // for
299 WriteFile(fd, tempSpace, numBlocks * blockSize, &numWritten, NULL);
300 retval = (int) numWritten;
301
302 // Adjust the return value, if necessary....
303 if (((numBlocks * blockSize) != numBytes) && (retval > 0))
304 retval = numBytes;
305
306 delete[] tempSpace;
307 } // if (isOpen)
308 return retval;
309 } // DiskIO:Write()
310
311 // Returns the size of the disk in blocks.
DiskSize(int * err)312 uint64_t DiskIO::DiskSize(int *err) {
313 uint64_t sectors = 0; // size in sectors
314 DWORD bytes, moreBytes; // low- and high-order bytes of file size
315 GET_LENGTH_INFORMATION buf;
316 DWORD i;
317
318 // If disk isn't open, try to open it....
319 if (!isOpen) {
320 OpenForRead();
321 } // if
322
323 if (isOpen) {
324 // Note to self: I recall testing a simplified version of
325 // this code, similar to what's in the __APPLE__ block,
326 // on Linux, but I had some problems. IIRC, it ran OK on 32-bit
327 // systems but not on 64-bit. Keep this in mind in case of
328 // 32/64-bit issues on MacOS....
329 if (DeviceIoControl(fd, IOCTL_DISK_GET_LENGTH_INFO, NULL, 0, &buf, sizeof(buf), &i, NULL)) {
330 sectors = (uint64_t) buf.Length.QuadPart / GetBlockSize();
331 *err = 0;
332 } else { // doesn't seem to be a disk device; assume it's an image file....
333 bytes = GetFileSize(fd, &moreBytes);
334 sectors = ((uint64_t) bytes + ((uint64_t) moreBytes) * UINT32_MAX) / GetBlockSize();
335 *err = 0;
336 } // if
337 } else {
338 *err = -1;
339 sectors = 0;
340 } // if/else (isOpen)
341
342 return sectors;
343 } // DiskIO::DiskSize()
344