dd.c - ubase - suckless linux base utils
 (HTM) git clone git://git.suckless.org/ubase
 (DIR) Log
 (DIR) Files
 (DIR) Refs
 (DIR) README
 (DIR) 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 }