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
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.
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....
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.
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.
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.
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.
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.
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...)
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!
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.
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.
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.
427 uint64_t DiskIO::DiskSize(int *err) {
428    uint64_t sectors = 0; // size in sectors
429    off_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, &sectors);
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