xref: /third_party/gptfdisk/diskio-unix.cc (revision cf200d32) 
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
40using 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
52void 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.
59int 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....
120int 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.
148void 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.
159int 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.
206int 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.
226uint32_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.
244uint32_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...)
265int 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!
318int 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.
341int 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.
384int 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.
435uint64_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, &sectors);
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