morpheus-base

morpheus base system
git clone git://git.2f30.org/morpheus-base
Log | Files | Refs

dd.c (6517B)


      1 /* (C) 2011-2012 Sebastian Krahmer all rights reserved.
      2  *
      3  * Optimized dd, to speed up backups etc.
      4  *
      5  * Permission has been granted to release this code under MIT/X.
      6  * The original code is at https://github.com/stealth/odd.  This
      7  * version of the code has been modified by sin@2f30.org.
      8  */
      9 #include <sys/ioctl.h>
     10 #include <sys/mount.h>
     11 #include <sys/select.h>
     12 #include <sys/stat.h>
     13 #include <sys/time.h>
     14 #include <sys/types.h>
     15 #include <sys/vfs.h>
     16 
     17 #include <errno.h>
     18 #include <fcntl.h>
     19 #include <inttypes.h>
     20 #include <signal.h>
     21 #include <stdio.h>
     22 #include <stdlib.h>
     23 #include <stdint.h>
     24 #include <string.h>
     25 #include <time.h>
     26 #include <unistd.h>
     27 
     28 #include "util.h"
     29 
     30 struct dd_config {
     31 	const char *in, *out;
     32 	uint64_t skip, seek, count, b_in, b_out, rec_in, rec_out;
     33 	off_t fsize;
     34 	blksize_t bs;
     35 	char quiet, nosync, direct;
     36 	int saved_errno;
     37 	time_t t_start, t_end;
     38 };
     39 
     40 static int sigint = 0;
     41 
     42 static int
     43 prepare_copy(struct dd_config *ddc, int *ifd, int *ofd)
     44 {
     45 	struct stat st;
     46 	int fli = O_RDONLY|O_LARGEFILE|O_NOCTTY, flo = O_WRONLY|O_LARGEFILE|O_NOATIME|O_NOCTTY;
     47 	uid_t euid = 0;
     48 	long pagesize;
     49 
     50 	if (ddc->direct) {
     51 		fli |= O_DIRECT;
     52 		flo |= O_DIRECT;
     53 	}
     54 
     55 	if (stat(ddc->in, &st) < 0) {
     56 		ddc->saved_errno = errno;
     57 		return -1;
     58 	}
     59 
     60 	euid = geteuid();
     61 
     62 	if (!euid || st.st_uid == euid)
     63 		fli |= O_NOATIME;
     64 
     65 	if ((*ifd = open(ddc->in, fli)) < 0) {
     66 		ddc->saved_errno = errno;
     67 		return -1;
     68 	}
     69 
     70 	ddc->fsize = st.st_size;
     71 
     72 	/* If "bsize" is not given, use optimum of both FS' */
     73 	if (!ddc->bs) {
     74 		struct statfs fst;
     75 		memset(&fst, 0, sizeof(fst));
     76 		pagesize = sysconf(_SC_PAGESIZE);
     77 		if (pagesize <= 0)
     78 			pagesize = 0x1000;
     79 		if (statfs(ddc->out, &fst) < 0 || fst.f_bsize == 0)
     80 			fst.f_bsize = pagesize;
     81 		if ((unsigned long)fst.f_bsize > (unsigned long)st.st_blksize)
     82 			ddc->bs = fst.f_bsize;
     83 		else
     84 			ddc->bs = st.st_blksize;
     85 		if (ddc->bs == 0)
     86 			ddc->bs = pagesize;
     87 	}
     88 
     89 	/* check if device or regular file */
     90 	if (!S_ISREG(st.st_mode)) {
     91 		if (S_ISBLK(st.st_mode)) {
     92 			if (ioctl(*ifd, BLKGETSIZE64, &ddc->fsize) < 0) {
     93 				ddc->saved_errno = errno;
     94 				close(*ifd);
     95 				return -1;
     96 			}
     97 		} else {
     98 			ddc->fsize = (off_t)-1;
     99 			if (ddc->count)
    100 				ddc->fsize = ddc->count*ddc->bs;
    101 		}
    102 	}
    103 
    104 	/* skip and seek are in block items */
    105 	ddc->skip *= ddc->bs;
    106 	ddc->seek *= ddc->bs;
    107 
    108 	/* skip more bytes than are inside source file? */
    109 	if (ddc->fsize != (off_t)-1 && ddc->skip >= (uint64_t)ddc->fsize) {
    110 		ddc->saved_errno = EINVAL;
    111 		close(*ifd);
    112 		return -1;
    113 	}
    114 
    115 	if (!ddc->seek)
    116 		flo |= O_CREAT|O_TRUNC;
    117 
    118 	if ((*ofd = open(ddc->out, flo, st.st_mode)) < 0) {
    119 		ddc->saved_errno = errno;
    120 		close(*ifd);
    121 		return -1;
    122 	}
    123 
    124 	lseek(*ifd, ddc->skip, SEEK_SET);
    125 	lseek(*ofd, ddc->seek, SEEK_SET);
    126 	posix_fadvise(*ifd, ddc->skip, 0, POSIX_FADV_SEQUENTIAL);
    127 	posix_fadvise(*ofd, 0, 0, POSIX_FADV_DONTNEED);
    128 
    129 	/* count is in block items too */
    130 	ddc->count *= ddc->bs;
    131 
    132 	/* If no count is given, its the filesize minus skip offset */
    133 	if (ddc->count == 0)
    134 		ddc->count = ddc->fsize - ddc->skip;
    135 
    136 	return 0;
    137 }
    138 
    139 static int
    140 copy_splice(struct dd_config *ddc)
    141 {
    142 	int ifd, ofd, p[2] = {-1, -1};
    143 	ssize_t r = 0;
    144 	size_t n = 0;
    145 	fd_set rfd, wfd;
    146 
    147 	if (prepare_copy(ddc, &ifd, &ofd) < 0)
    148 		return -1;
    149 	if (pipe(p) < 0) {
    150 		ddc->saved_errno = errno;
    151 		close(ifd); close(ofd);
    152 		close(p[0]); close(p[1]);
    153 		return -1;
    154 	}
    155 #ifdef F_SETPIPE_SZ
    156 	for (n = 29; n >= 20; --n) {
    157 		if (fcntl(p[0], F_SETPIPE_SZ, 1<<n) != -1)
    158 			break;
    159 	}
    160 #endif
    161 	n = ddc->bs;
    162 	for (;ddc->b_out != ddc->count && !sigint;) {
    163 		FD_ZERO(&rfd);
    164 		FD_ZERO(&wfd);
    165 		FD_SET(ifd, &rfd);
    166 		FD_SET(ofd, &wfd);
    167 		r = select(ifd > ofd ? ifd + 1 : ofd + 1, &rfd, &wfd, NULL, NULL);
    168 		if (r < 0) {
    169 			ddc->saved_errno = errno;
    170 			break;
    171 		}
    172 		if (FD_ISSET(ifd, &rfd) == 1 && FD_ISSET(ofd, &wfd) == 1) {
    173 			if (n > ddc->count - ddc->b_out)
    174 				n = ddc->count - ddc->b_out;
    175 			r = splice(ifd, NULL, p[1], NULL, n, SPLICE_F_MORE);
    176 			if (r <= 0) {
    177 				ddc->saved_errno = errno;
    178 				break;
    179 			}
    180 			++ddc->rec_in;
    181 			r = splice(p[0], NULL, ofd, NULL, r, SPLICE_F_MORE);
    182 			if (r <= 0) {
    183 				ddc->saved_errno = errno;
    184 				break;
    185 			}
    186 			ddc->b_out += r;
    187 			++ddc->rec_out;
    188 		}
    189 	}
    190 	close(ifd);
    191 	close(ofd);
    192 	close(p[0]);
    193 	close(p[1]);
    194 	if (r < 0)
    195 		return -1;
    196 	return 0;
    197 }
    198 
    199 static int
    200 copy(struct dd_config *ddc)
    201 {
    202 	int r = 0;
    203 
    204 	ddc->t_start = time(NULL);
    205 
    206 	r = copy_splice(ddc);
    207 	ddc->t_end = time(NULL);
    208 
    209 	/* avoid div by zero */
    210 	if (ddc->t_start == ddc->t_end)
    211 		++ddc->t_end;
    212 	return r;
    213 }
    214 
    215 static void
    216 usage(void)
    217 {
    218 	eprintf("usage: %s [-h] [if=infile] [of=outfile] [bs[=N]] [seek=N] [skip=N] [count=N] [direct] [quiet] [nosync]\n", argv0);
    219 }
    220 
    221 static void
    222 print_stat(const struct dd_config *ddc)
    223 {
    224 	if (ddc->quiet)
    225 		return;
    226 
    227 	fprintf(stderr, "%"PRIu64" records in\n", ddc->rec_in);
    228 	fprintf(stderr, "%"PRIu64" records out\n", ddc->rec_out);
    229 	fprintf(stderr, "%"PRIu64" bytes (%"PRIu64" MB) copied", ddc->b_out,
    230 		ddc->b_out/(1<<20));
    231 	fprintf(stderr, ", %lu s, %f MB/s\n",
    232 		(unsigned long)ddc->t_end - ddc->t_start,
    233 		((double)(ddc->b_out/(1<<20)))/(ddc->t_end - ddc->t_start));
    234 }
    235 
    236 static void
    237 sig_int(int unused)
    238 {
    239 	(void) unused;
    240 	fprintf(stderr, "SIGINT! Aborting ...\n");
    241 	sigint = 1;
    242 }
    243 
    244 int
    245 main(int argc, char *argv[])
    246 {
    247 	int i = 0;
    248 	char buf[1024];
    249 	struct dd_config config;
    250 
    251 	argv0 = argv[0];
    252 	memset(&config, 0, sizeof(config));
    253 	config.bs = 1<<16;
    254 	config.in = "/dev/stdin";
    255 	config.out = "/dev/stdout";
    256 
    257 	/* emulate 'dd' argument parsing */
    258 	for (i = 1; i < argc; ++i) {
    259 		memset(buf, 0, sizeof(buf));
    260 		if (sscanf(argv[i], "if=%1023s", buf) == 1)
    261 			config.in = strdup(buf);
    262 		else if (sscanf(argv[i], "of=%1023s", buf) == 1)
    263 			config.out = strdup(buf);
    264 		else if (sscanf(argv[i], "skip=%1023s", buf) == 1)
    265 			config.skip = estrtoul(buf, 0);
    266 		else if (sscanf(argv[i], "seek=%1023s", buf) == 1)
    267 			config.seek = estrtoul(buf, 0);
    268 		else if (sscanf(argv[i], "count=%1023s", buf) == 1)
    269 			config.count = estrtoul(buf, 0);
    270 		else if (strcmp(argv[i], "direct") == 0)
    271 			config.direct = 1;
    272 		else if (sscanf(argv[i], "bs=%1023s", buf) == 1)
    273 			config.bs = estrtoul(buf, 0);
    274 		else if (strcmp(argv[i], "bs") == 0)
    275 			config.bs = 0;
    276 		else if (strcmp(argv[i], "quiet") == 0)
    277 			config.quiet = 1;
    278 		else if (strcmp(argv[i], "nosync") == 0)
    279 			config.nosync = 1;
    280 		else if (strcmp(argv[i], "-h") == 0)
    281 			usage();
    282 	}
    283 
    284 	if (!config.in || !config.out)
    285 		usage();
    286 
    287 	signal(SIGPIPE, SIG_IGN);
    288 	signal(SIGINT, sig_int);
    289 
    290 	if (copy(&config) < 0)
    291 		weprintf("copy:");
    292 	print_stat(&config);
    293 
    294 	if (config.nosync == 0)
    295 		sync(); sync();
    296 	return config.saved_errno;
    297 }