ubase

suckless linux base utils
git clone git://git.2f30.org/ubase.git
Log | Files | Refs | README | LICENSE

dd.c (6564B)


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