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 }