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