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