File: diskio-windows.cc

package info (click to toggle)
gdisk 1.0.5-1
  • links: PTS, VCS
  • area: main
  • in suites: bullseye, sid
  • size: 940 kB
  • sloc: cpp: 8,692; sh: 244; makefile: 48
file content (343 lines) | stat: -rw-r--r-- 10,905 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
//
// C++ Interface: diskio (Windows-specific components)
//
// Description: Class to handle low-level disk I/O for GPT fdisk
//
//
// Author: Rod Smith <rodsmith@rodsbooks.com>, (C) 2009
//
// Copyright: See COPYING file that comes with this distribution
//
//
// This program is copyright (c) 2009, 2010 by Roderick W. Smith. It is distributed
// under the terms of the GNU GPL version 2, as detailed in the COPYING file.

#define __STDC_LIMIT_MACROS
#define __STDC_CONSTANT_MACROS

#include <windows.h>
#include <winioctl.h>
#define fstat64 fstat
#define stat64 stat
#define S_IRGRP 0
#define S_IROTH 0
#include <stdio.h>
#include <string>
#include <stdint.h>
#include <errno.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <iostream>

#include "support.h"
#include "diskio.h"

using namespace std;

// Returns the official Windows name for a shortened version of same.
void DiskIO::MakeRealName(void) {
   size_t colonPos;

   colonPos = userFilename.find(':', 0);
   if ((colonPos != string::npos) && (colonPos <= 3)) {
      realFilename = "\\\\.\\physicaldrive";
      realFilename += userFilename.substr(0, colonPos);
   } else {
      realFilename = userFilename;
   } // if/else
} // DiskIO::MakeRealName()

// Open the currently on-record file for reading
int DiskIO::OpenForRead(void) {
   int shouldOpen = 1;

   if (isOpen) { // file is already open
      if (openForWrite) {
         Close();
      } else {
         shouldOpen = 0;
      } // if/else
   } // if

   if (shouldOpen) {
      fd = CreateFile(realFilename.c_str(),GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE,
                      NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
      if (fd == INVALID_HANDLE_VALUE) {
         CloseHandle(fd);
         cerr << "Problem opening " << realFilename << " for reading!\n";
         realFilename = "";
         userFilename = "";
         isOpen = 0;
         openForWrite = 0;
      } else {
         isOpen = 1;
         openForWrite = 0;
      } // if/else
   } // if

   return isOpen;
} // DiskIO::OpenForRead(void)

// An extended file-open function. This includes some system-specific checks.
// Returns 1 if the file is open, 0 otherwise....
int DiskIO::OpenForWrite(void) {
   if ((isOpen) && (openForWrite))
      return 1;

   // Close the disk, in case it's already open for reading only....
   Close();

   // try to open the device; may fail....
   fd = CreateFile(realFilename.c_str(), GENERIC_READ | GENERIC_WRITE,
                   FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING,
                   FILE_ATTRIBUTE_NORMAL, NULL);
   // Preceding call can fail when creating backup files; if so, try
   // again with different option...
   if (fd == INVALID_HANDLE_VALUE) {
      CloseHandle(fd);
      fd = CreateFile(realFilename.c_str(), GENERIC_READ | GENERIC_WRITE,
                      FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_ALWAYS,
                      FILE_ATTRIBUTE_NORMAL, NULL);
   } // if
   if (fd == INVALID_HANDLE_VALUE) {
      CloseHandle(fd);
      isOpen = 0;
      openForWrite = 0;
      errno = GetLastError();
   } else {
      isOpen = 1;
      openForWrite = 1;
   } // if/else
   return isOpen;
} // DiskIO::OpenForWrite(void)

// Close the disk device. Note that this does NOT erase the stored filenames,
// so the file can be re-opened without specifying the filename.
void DiskIO::Close(void) {
   if (isOpen)
      CloseHandle(fd);
   isOpen = 0;
   openForWrite = 0;
} // DiskIO::Close()

// Returns block size of device pointed to by fd file descriptor. If the ioctl
// returns an error condition, assume it's a disk file and return a value of
// SECTOR_SIZE (512). If the disk can't be opened at all, return a value of 0.
int DiskIO::GetBlockSize(void) {
   DWORD blockSize = 0, retBytes;
   DISK_GEOMETRY_EX geom;

   // If disk isn't open, try to open it....
   if (!isOpen) {
      OpenForRead();
   } // if

   if (isOpen) {
      if (DeviceIoControl(fd, IOCTL_DISK_GET_DRIVE_GEOMETRY_EX, NULL, 0,
                          &geom, sizeof(geom), &retBytes, NULL)) {
         blockSize = geom.Geometry.BytesPerSector;
      } else { // was probably an ordinary file; set default value....
         blockSize = SECTOR_SIZE;
      } // if/else
   } // if (isOpen)

   return (blockSize);
} // DiskIO::GetBlockSize()

// In theory, returns the physical block size. In practice, this is only
// supported in Linux, as of yet.
// TODO: Get this working in Windows.
int DiskIO::GetPhysBlockSize(void) {
   return 0;
} // DiskIO::GetPhysBlockSize()

// Returns the number of heads, according to the kernel, or 255 if the
// correct value can't be determined.
uint32_t DiskIO::GetNumHeads(void) {
   return UINT32_C(255);
} // DiskIO::GetNumHeads();

// Returns the number of sectors per track, according to the kernel, or 63
// if the correct value can't be determined.
uint32_t DiskIO::GetNumSecsPerTrack(void) {
   return UINT32_C(63);
} // DiskIO::GetNumSecsPerTrack()

// Resync disk caches so the OS uses the new partition table. This code varies
// a lot from one OS to another.
// Returns 1 on success, 0 if the kernel continues to use the old partition table.
int DiskIO::DiskSync(void) {
   DWORD i;
   GET_LENGTH_INFORMATION buf;
   int retval = 0;

   // If disk isn't open, try to open it....
   if (!openForWrite) {
      OpenForWrite();
   } // if

   if (isOpen) {
      if (DeviceIoControl(fd, IOCTL_DISK_UPDATE_PROPERTIES, NULL, 0, &buf, sizeof(buf), &i, NULL) == 0) {
         cout << "Disk synchronization failed! The computer may use the old partition table\n"
              << "until you reboot or remove and re-insert the disk!\n";
      } else {
         cout << "Disk synchronization succeeded! The computer should now use the new\n"
              << "partition table.\n";
         retval = 1;
      } // if/else
   } else {
      cout << "Unable to open the disk for synchronization operation! The computer will\n"
           << "continue to use the old partition table until you reboot or remove and\n"
           << "re-insert the disk!\n";
   } // if (isOpen)
   return retval;
} // DiskIO::DiskSync()

// Seek to the specified sector. Returns 1 on success, 0 on failure.
int DiskIO::Seek(uint64_t sector) {
   int retval = 1;
   LARGE_INTEGER seekTo;

   // If disk isn't open, try to open it....
   if (!isOpen) {
      retval = OpenForRead();
   } // if

   if (isOpen) {
      seekTo.QuadPart = sector * (uint64_t) GetBlockSize();
      retval = SetFilePointerEx(fd, seekTo, NULL, FILE_BEGIN);
      if (retval == 0) {
         errno = GetLastError();
         cerr << "Error when seeking to " << seekTo.QuadPart << "! Error is " << errno << "\n";
         retval = 0;
      } // if
   } // if
   return retval;
} // DiskIO::Seek()

// A variant on the standard read() function. Done to work around
// limitations in FreeBSD concerning the matching of the sector
// size with the number of bytes read.
// Returns the number of bytes read into buffer.
int DiskIO::Read(void* buffer, int numBytes) {
   int blockSize = 512, i, numBlocks;
   char* tempSpace;
   DWORD retval = 0;

   // If disk isn't open, try to open it....
   if (!isOpen) {
      OpenForRead();
   } // if

   if (isOpen) {
      // Compute required space and allocate memory
      blockSize = GetBlockSize();
      if (numBytes <= blockSize) {
         numBlocks = 1;
         tempSpace = new char [blockSize];
      } else {
         numBlocks = numBytes / blockSize;
         if ((numBytes % blockSize) != 0)
            numBlocks++;
         tempSpace = new char [numBlocks * blockSize];
      } // if/else
      if (tempSpace == NULL) {
         cerr << "Unable to allocate memory in DiskIO::Read()! Terminating!\n";
         exit(1);
      } // if

      // Read the data into temporary space, then copy it to buffer
      ReadFile(fd, tempSpace, numBlocks * blockSize, &retval, NULL);
      for (i = 0; i < numBytes; i++) {
         ((char*) buffer)[i] = tempSpace[i];
      } // for

      // Adjust the return value, if necessary....
      if (((numBlocks * blockSize) != numBytes) && (retval > 0))
         retval = numBytes;

      delete[] tempSpace;
   } // if (isOpen)
   return retval;
} // DiskIO::Read()

// A variant on the standard write() function.
// Returns the number of bytes written.
int DiskIO::Write(void* buffer, int numBytes) {
   int blockSize = 512, i, numBlocks, retval = 0;
   char* tempSpace;
   DWORD numWritten;

   // If disk isn't open, try to open it....
   if ((!isOpen) || (!openForWrite)) {
      OpenForWrite();
   } // if

   if (isOpen) {
      // Compute required space and allocate memory
      blockSize = GetBlockSize();
      if (numBytes <= blockSize) {
         numBlocks = 1;
         tempSpace = new char [blockSize];
      } else {
         numBlocks = numBytes / blockSize;
         if ((numBytes % blockSize) != 0) numBlocks++;
         tempSpace = new char [numBlocks * blockSize];
      } // if/else
      if (tempSpace == NULL) {
         cerr << "Unable to allocate memory in DiskIO::Write()! Terminating!\n";
         exit(1);
      } // if

      // Copy the data to my own buffer, then write it
      for (i = 0; i < numBytes; i++) {
         tempSpace[i] = ((char*) buffer)[i];
      } // for
      for (i = numBytes; i < numBlocks * blockSize; i++) {
         tempSpace[i] = 0;
      } // for
      WriteFile(fd, tempSpace, numBlocks * blockSize, &numWritten, NULL);
      retval = (int) numWritten;

      // Adjust the return value, if necessary....
      if (((numBlocks * blockSize) != numBytes) && (retval > 0))
         retval = numBytes;

      delete[] tempSpace;
   } // if (isOpen)
   return retval;
} // DiskIO:Write()

// Returns the size of the disk in blocks.
uint64_t DiskIO::DiskSize(int *err) {
   uint64_t sectors = 0; // size in sectors
   DWORD bytes, moreBytes; // low- and high-order bytes of file size
   GET_LENGTH_INFORMATION buf;
   DWORD i;

   // If disk isn't open, try to open it....
   if (!isOpen) {
      OpenForRead();
   } // if

   if (isOpen) {
      // Note to self: I recall testing a simplified version of
      // this code, similar to what's in the __APPLE__ block,
      // on Linux, but I had some problems. IIRC, it ran OK on 32-bit
      // systems but not on 64-bit. Keep this in mind in case of
      // 32/64-bit issues on MacOS....
      if (DeviceIoControl(fd, IOCTL_DISK_GET_LENGTH_INFO, NULL, 0, &buf, sizeof(buf), &i, NULL)) {
         sectors = (uint64_t) buf.Length.QuadPart / GetBlockSize();
         *err = 0;
      } else { // doesn't seem to be a disk device; assume it's an image file....
         bytes = GetFileSize(fd, &moreBytes);
         sectors = ((uint64_t) bytes + ((uint64_t) moreBytes) * UINT32_MAX) / GetBlockSize();
         *err = 0;
      } // if
   } else {
      *err = -1;
      sectors = 0;
   } // if/else (isOpen)

   return sectors;
} // DiskIO::DiskSize()