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