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