1 //
2 // C++ Interface: diskio (Unix components [Linux, FreeBSD, Mac OS X])
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 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 <sys/ioctl.h>
21 #include <string.h>
22 #include <string>
23 #include <stdint.h>
24 #include <unistd.h>
25 #include <errno.h>
26 #include <fcntl.h>
27 #include <sys/stat.h>
28 #include <unistd.h>
29
30 #ifdef __linux__
31 #include "linux/hdreg.h"
32 #endif
33
34 #include <iostream>
35
36 #include "diskio.h"
37
38 using namespace std;
39
40 // Returns the official "real" name for a shortened version of same.
41 // Trivial here; more important in Windows
MakeRealName(void)42 void DiskIO::MakeRealName(void) {
43 realFilename = userFilename;
44 } // DiskIO::MakeRealName()
45
46 // Open the currently on-record file for reading. Returns 1 if the file is
47 // already open or is opened by this call, 0 if opening the file doesn't
48 // work.
OpenForRead(void)49 int DiskIO::OpenForRead(void) {
50 int shouldOpen = 1;
51 struct stat64 st;
52
53 if (isOpen) { // file is already open
54 if (openForWrite) {
55 Close();
56 } else {
57 shouldOpen = 0;
58 } // if/else
59 } // if
60
61 if (shouldOpen) {
62 fd = open(realFilename.c_str(), O_RDONLY);
63 if (fd == -1) {
64 cerr << "Problem opening " << realFilename << " for reading! Error is " << errno << ".\n";
65 if (errno == EACCES) // User is probably not running as root
66 cerr << "You must run this program as root or use sudo!\n";
67 if (errno == ENOENT)
68 cerr << "The specified file does not exist!\n";
69 realFilename = "";
70 userFilename = "";
71 isOpen = 0;
72 openForWrite = 0;
73 } else {
74 isOpen = 0;
75 openForWrite = 0;
76 if (fstat64(fd, &st) == 0) {
77 if (S_ISDIR(st.st_mode))
78 cerr << "The specified path is a directory!\n";
79 #if !defined(__FreeBSD__) && !defined(__APPLE__)
80 else if (S_ISCHR(st.st_mode))
81 cerr << "The specified path is a character device!\n";
82 #endif
83 else if (S_ISFIFO(st.st_mode))
84 cerr << "The specified path is a FIFO!\n";
85 else if (S_ISSOCK(st.st_mode))
86 cerr << "The specified path is a socket!\n";
87 else
88 isOpen = 1;
89 } // if (fstat64()...)
90 } // if/else
91 } // if
92
93 return isOpen;
94 } // DiskIO::OpenForRead(void)
95
96 // An extended file-open function. This includes some system-specific checks.
97 // Returns 1 if the file is open, 0 otherwise....
OpenForWrite(void)98 int DiskIO::OpenForWrite(void) {
99 if ((isOpen) && (openForWrite))
100 return 1;
101
102 // Close the disk, in case it's already open for reading only....
103 Close();
104
105 // try to open the device; may fail....
106 fd = open(realFilename.c_str(), O_WRONLY | O_CREAT, S_IWUSR | S_IRUSR | S_IRGRP | S_IROTH);
107 #ifdef __APPLE__
108 // MacOS X requires a shared lock under some circumstances....
109 if (fd < 0) {
110 cerr << "Warning: Devices opened with shared lock will not have their\npartition table automatically reloaded!\n";
111 fd = open(realFilename.c_str(), O_WRONLY | O_SHLOCK);
112 } // if
113 #endif
114 if (fd >= 0) {
115 isOpen = 1;
116 openForWrite = 1;
117 } else {
118 isOpen = 0;
119 openForWrite = 0;
120 } // if/else
121 return isOpen;
122 } // DiskIO::OpenForWrite(void)
123
124 // Close the disk device. Note that this does NOT erase the stored filenames,
125 // so the file can be re-opened without specifying the filename.
Close(void)126 void DiskIO::Close(void) {
127 if (isOpen)
128 if (close(fd) < 0)
129 cerr << "Warning! Problem closing file!\n";
130 isOpen = 0;
131 openForWrite = 0;
132 } // DiskIO::Close()
133
134 // Returns block size of device pointed to by fd file descriptor. If the ioctl
135 // returns an error condition, print a warning but return a value of SECTOR_SIZE
136 // (512). If the disk can't be opened at all, return a value of 0.
GetBlockSize(void)137 int DiskIO::GetBlockSize(void) {
138 int err = -1, blockSize = 0;
139 #ifdef __sun__
140 struct dk_minfo minfo;
141 #endif
142
143 // If disk isn't open, try to open it....
144 if (!isOpen) {
145 OpenForRead();
146 } // if
147
148 if (isOpen) {
149 #ifdef __APPLE__
150 err = ioctl(fd, DKIOCGETBLOCKSIZE, &blockSize);
151 #endif
152 #ifdef __sun__
153 err = ioctl(fd, DKIOCGMEDIAINFO, &minfo);
154 if (err == 0)
155 blockSize = minfo.dki_lbsize;
156 #endif
157 #if defined (__FreeBSD__) || defined (__FreeBSD_kernel__)
158 err = ioctl(fd, DIOCGSECTORSIZE, &blockSize);
159 #endif
160 #ifdef __linux__
161 err = ioctl(fd, BLKSSZGET, &blockSize);
162 #endif
163
164 if (err == -1) {
165 blockSize = SECTOR_SIZE;
166 // ENOTTY = inappropriate ioctl; probably being called on a disk image
167 // file, so don't display the warning message....
168 // 32-bit code returns EINVAL, I don't know why. I know I'm treading on
169 // thin ice here, but it should be OK in all but very weird cases....
170 if ((errno != ENOTTY) && (errno != EINVAL)) {
171 cerr << "\aError " << errno << " when determining sector size! Setting sector size to "
172 << SECTOR_SIZE << "\n";
173 cout << "Disk device is " << realFilename << "\n";
174 } // if
175 } // if (err == -1)
176 } // if (isOpen)
177
178 return (blockSize);
179 } // DiskIO::GetBlockSize()
180
181 // Returns the number of heads, according to the kernel, or 255 if the
182 // correct value can't be determined.
GetNumHeads(void)183 uint32_t DiskIO::GetNumHeads(void) {
184 uint32_t numHeads = 255;
185
186 #ifdef HDIO_GETGEO
187 struct hd_geometry geometry;
188
189 // If disk isn't open, try to open it....
190 if (!isOpen)
191 OpenForRead();
192
193 if (!ioctl(fd, HDIO_GETGEO, &geometry))
194 numHeads = (uint32_t) geometry.heads;
195 #endif
196 return numHeads;
197 } // DiskIO::GetNumHeads();
198
199 // Returns the number of sectors per track, according to the kernel, or 63
200 // if the correct value can't be determined.
GetNumSecsPerTrack(void)201 uint32_t DiskIO::GetNumSecsPerTrack(void) {
202 uint32_t numSecs = 63;
203
204 #ifdef HDIO_GETGEO
205 struct hd_geometry geometry;
206
207 // If disk isn't open, try to open it....
208 if (!isOpen)
209 OpenForRead();
210
211 if (!ioctl(fd, HDIO_GETGEO, &geometry))
212 numSecs = (uint32_t) geometry.sectors;
213 #endif
214 return numSecs;
215 } // DiskIO::GetNumSecsPerTrack()
216
217 // Resync disk caches so the OS uses the new partition table. This code varies
218 // a lot from one OS to another.
219 // Returns 1 on success, 0 if the kernel continues to use the old partition table.
220 // (Note that for most OSes, the default of 0 is returned because I've not yet
221 // looked into how to test for success in the underlying system calls...)
DiskSync(void)222 int DiskIO::DiskSync(void) {
223 int i, retval = 0, platformFound = 0;
224
225 // If disk isn't open, try to open it....
226 if (!isOpen) {
227 OpenForRead();
228 } // if
229
230 if (isOpen) {
231 sync();
232 #if defined(__APPLE__) || defined(__sun__)
233 cout << "Warning: The kernel may continue to use old or deleted partitions.\n"
234 << "You should reboot or remove the drive.\n";
235 /* don't know if this helps
236 * it definitely will get things on disk though:
237 * http://topiks.org/mac-os-x/0321278542/ch12lev1sec8.html */
238 #ifdef __sun__
239 i = ioctl(fd, DKIOCFLUSHWRITECACHE);
240 #else
241 i = ioctl(fd, DKIOCSYNCHRONIZECACHE);
242 #endif
243 platformFound++;
244 #endif
245 #if defined (__FreeBSD__) || defined (__FreeBSD_kernel__)
246 sleep(2);
247 i = ioctl(fd, DIOCGFLUSH);
248 cout << "Warning: The kernel may continue to use old or deleted partitions.\n"
249 << "You should reboot or remove the drive.\n";
250 platformFound++;
251 #endif
252 #ifdef __linux__
253 sleep(1); // Theoretically unnecessary, but ioctl() fails sometimes if omitted....
254 fsync(fd);
255 i = ioctl(fd, BLKRRPART);
256 if (i) {
257 cout << "Warning: The kernel is still using the old partition table.\n"
258 << "The new table will be used at the next reboot.\n";
259 } else {
260 retval = 1;
261 } // if/else
262 platformFound++;
263 #endif
264 if (platformFound == 0)
265 cerr << "Warning: Platform not recognized!\n";
266 if (platformFound > 1)
267 cerr << "\nWarning: We seem to be running on multiple platforms!\n";
268 } // if (isOpen)
269 return retval;
270 } // DiskIO::DiskSync()
271
272 // Seek to the specified sector. Returns 1 on success, 0 on failure.
273 // Note that seeking beyond the end of the file is NOT detected as a failure!
Seek(uint64_t sector)274 int DiskIO::Seek(uint64_t sector) {
275 int retval = 1;
276 off64_t seekTo, sought;
277
278 // If disk isn't open, try to open it....
279 if (!isOpen) {
280 retval = OpenForRead();
281 } // if
282
283 if (isOpen) {
284 seekTo = sector * (uint64_t) GetBlockSize();
285 sought = lseek64(fd, seekTo, SEEK_SET);
286 if (sought != seekTo) {
287 retval = 0;
288 } // if
289 } // if
290 return retval;
291 } // DiskIO::Seek()
292
293 // A variant on the standard read() function. Done to work around
294 // limitations in FreeBSD concerning the matching of the sector
295 // size with the number of bytes read.
296 // Returns the number of bytes read into buffer.
Read(void * buffer,int numBytes)297 int DiskIO::Read(void* buffer, int numBytes) {
298 int blockSize, numBlocks, retval = 0;
299 char* tempSpace;
300
301 // If disk isn't open, try to open it....
302 if (!isOpen) {
303 OpenForRead();
304 } // if
305
306 if (isOpen) {
307 // Compute required space and allocate memory
308 blockSize = GetBlockSize();
309 if (numBytes <= blockSize) {
310 numBlocks = 1;
311 tempSpace = new char [blockSize];
312 } else {
313 numBlocks = numBytes / blockSize;
314 if ((numBytes % blockSize) != 0)
315 numBlocks++;
316 tempSpace = new char [numBlocks * blockSize];
317 } // if/else
318 if (tempSpace == NULL) {
319 cerr << "Unable to allocate memory in DiskIO::Read()! Terminating!\n";
320 exit(1);
321 } // if
322
323 // Read the data into temporary space, then copy it to buffer
324 retval = read(fd, tempSpace, numBlocks * blockSize);
325 memcpy(buffer, tempSpace, numBytes);
326
327 // Adjust the return value, if necessary....
328 if (((numBlocks * blockSize) != numBytes) && (retval > 0))
329 retval = numBytes;
330
331 delete[] tempSpace;
332 } // if (isOpen)
333 return retval;
334 } // DiskIO::Read()
335
336 // A variant on the standard write() function. Done to work around
337 // limitations in FreeBSD concerning the matching of the sector
338 // size with the number of bytes read.
339 // Returns the number of bytes written.
Write(void * buffer,int numBytes)340 int DiskIO::Write(void* buffer, int numBytes) {
341 int blockSize = 512, i, numBlocks, retval = 0;
342 char* tempSpace;
343
344 // If disk isn't open, try to open it....
345 if ((!isOpen) || (!openForWrite)) {
346 OpenForWrite();
347 } // if
348
349 if (isOpen) {
350 // Compute required space and allocate memory
351 blockSize = GetBlockSize();
352 if (numBytes <= blockSize) {
353 numBlocks = 1;
354 tempSpace = new char [blockSize];
355 } else {
356 numBlocks = numBytes / blockSize;
357 if ((numBytes % blockSize) != 0) numBlocks++;
358 tempSpace = new char [numBlocks * blockSize];
359 } // if/else
360 if (tempSpace == NULL) {
361 cerr << "Unable to allocate memory in DiskIO::Write()! Terminating!\n";
362 exit(1);
363 } // if
364
365 // Copy the data to my own buffer, then write it
366 memcpy(tempSpace, buffer, numBytes);
367 for (i = numBytes; i < numBlocks * blockSize; i++) {
368 tempSpace[i] = 0;
369 } // for
370 retval = write(fd, tempSpace, numBlocks * blockSize);
371
372 // Adjust the return value, if necessary....
373 if (((numBlocks * blockSize) != numBytes) && (retval > 0))
374 retval = numBytes;
375
376 delete[] tempSpace;
377 } // if (isOpen)
378 return retval;
379 } // DiskIO:Write()
380
381 /**************************************************************************************
382 * *
383 * Below functions are lifted from various sources, as documented in comments before *
384 * each one. *
385 * *
386 **************************************************************************************/
387
388 // The disksize function is taken from the Linux fdisk code and modified
389 // greatly since then to enable FreeBSD and MacOS support, as well as to
390 // return correct values for disk image files.
DiskSize(int * err)391 uint64_t DiskIO::DiskSize(int *err) {
392 uint64_t sectors = 0; // size in sectors
393 off_t bytes = 0; // size in bytes
394 struct stat64 st;
395 int platformFound = 0;
396 #ifdef __sun__
397 struct dk_minfo minfo;
398 #endif
399
400 // If disk isn't open, try to open it....
401 if (!isOpen) {
402 OpenForRead();
403 } // if
404
405 if (isOpen) {
406 // Note to self: I recall testing a simplified version of
407 // this code, similar to what's in the __APPLE__ block,
408 // on Linux, but I had some problems. IIRC, it ran OK on 32-bit
409 // systems but not on 64-bit. Keep this in mind in case of
410 // 32/64-bit issues on MacOS....
411 #ifdef __APPLE__
412 *err = ioctl(fd, DKIOCGETBLOCKCOUNT, §ors);
413 platformFound++;
414 #endif
415 #ifdef __sun__
416 *err = ioctl(fd, DKIOCGMEDIAINFO, &minfo);
417 if (*err == 0)
418 sectors = minfo.dki_capacity;
419 platformFound++;
420 #endif
421 #if defined (__FreeBSD__) || defined (__FreeBSD_kernel__)
422 *err = ioctl(fd, DIOCGMEDIASIZE, &bytes);
423 long long b = GetBlockSize();
424 sectors = bytes / b;
425 platformFound++;
426 #endif
427 #ifdef __linux__
428 long sz;
429 long long b;
430 *err = ioctl(fd, BLKGETSIZE, &sz);
431 if (*err) {
432 sectors = sz = 0;
433 } // if
434 if ((!*err) || (errno == EFBIG)) {
435 *err = ioctl(fd, BLKGETSIZE64, &b);
436 if (*err || b == 0 || b == sz)
437 sectors = sz;
438 else
439 sectors = (b >> 9);
440 } // if
441 // Unintuitively, the above returns values in 512-byte blocks, no
442 // matter what the underlying device's block size. Correct for this....
443 sectors /= (GetBlockSize() / 512);
444 platformFound++;
445 #endif
446 if (platformFound != 1)
447 cerr << "Warning! We seem to be running on no known platform!\n";
448
449 // The above methods have failed, so let's assume it's a regular
450 // file (a QEMU image, dd backup, or what have you) and see what
451 // fstat() gives us....
452 if ((sectors == 0) || (*err == -1)) {
453 if (fstat64(fd, &st) == 0) {
454 bytes = st.st_size;
455 if ((bytes % UINT64_C(512)) != 0)
456 cerr << "Warning: File size is not a multiple of 512 bytes!"
457 << " Misbehavior is likely!\n\a";
458 sectors = bytes / UINT64_C(512);
459 } // if
460 } // if
461 } // if (isOpen)
462 return sectors;
463 } // DiskIO::DiskSize()
464