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