• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 // Returns the number of heads, according to the kernel, or 255 if the
148 // correct value can't be determined.
GetNumHeads(void)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.
GetNumSecsPerTrack(void)155 uint32_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.
DiskSync(void)162 int 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.
Seek(uint64_t sector)190 int 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.
Read(void * buffer,int numBytes)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;
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.
Write(void * buffer,int numBytes)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;
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.
DiskSize(int * err)305 uint64_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