ucw/ddigger.c

   1 #include <ucw/lib.h>
   2 #include <ucw/io.h>
   3 #include <ucw/opt.h>
   4 #include <stdint.h>
   5 #include <stdio.h>
   6 #include <fcntl.h>
   7 #include <unistd.h>
   8 #include <sys/ioctl.h>
   9 #include <linux/fs.h>
  10
  11 static char *dev_name;
  12 static char *out_name;
  13 static char *stat_name;
  14 static uint block_size = 1048576;
  15 static uint error_skip;
  16 static u64 start_pos;
  17 static u64 end_pos = ~(u64)0;
  18 static int show_status;
  19
  20 struct stat_rec {
  21   u64 pos;
  22   u64 len;
  23 };
  24
  25 static struct stat_rec *status;
  26
  27 #define ASORT_PREFIX(x) ss_##x
  28 #define ASORT_KEY_TYPE struct stat_rec
  29 #define ASORT_LT(x,y) x.pos < y.pos
  30 #include <ucw/sorter/array-simple.h>
  31
  32 static void stat_read(int fd)
  33 {
  34   u64 len = ucw_seek(fd, 0, SEEK_END);
  35   ucw_seek(fd, 0, SEEK_SET);
  36   if (len % sizeof(struct stat_rec))
  37     die("Malformed status file");
  38   if (len > ~0U)
  39     die("Status file too long");
  40
  41   uns items = len / sizeof(struct stat_rec);
  42   GARY_INIT(status, items);
  43   if (careful_read(fd, status, len) <= 0)
  44     die("Error reading status file");
  45
  46   ss_sort(status, items);
  47
  48 #if 0
  49   for (uint i=0; i<items; i++)
  50     printf("%jd %jd\n", (intmax_t) status[i].pos, (intmax_t) status[i].len);
  51 #endif
  52
  53   for (uint i=1; i<items; i++)
  54     if (status[i].pos < status[i-1].pos + status[i-1].len)
  55       die("Overlapping status file entries: [%jd,%jd] with [%jd,%jd]",
  56         (intmax_t) status[i-1].pos,
  57         (intmax_t)(status[i-1].pos + status[i-1].len),
  58         (intmax_t) status[i].pos,
  59         (intmax_t)(status[i].pos + status[i].len));
  60 }
  61
  62 static u64 stat_have(void)
  63 {
  64   u64 have = 0;
  65   for (uint i=0; i<GARY_SIZE(status); i++)
  66     {
  67       u64 x = CLAMP(status[i].pos, start_pos, end_pos);
  68       u64 y = CLAMP(status[i].pos + status[i].len, start_pos, end_pos);
  69       have += y - x;
  70     }
  71   return have;
  72 }
  73
  74 static void stat_show(void)
  75 {
  76   u64 last = 0;
  77   for (uint i=0; i<GARY_SIZE(status); i++)
  78     {
  79       if (status[i].pos > last)
  80         printf("Missing: %ju +%ju\n", (intmax_t) last, (intmax_t)(status[i].pos - last));
  81       last = status[i].pos + status[i].len;
  82     }
  83   printf("End at %ju\n", (intmax_t) last);
  84 }
  85
  86 static struct opt_section options = {
  87   OPT_ITEMS {
  88     OPT_HELP("Usage: ddigger [options] block-device output-file status-file"),
  89     OPT_HELP("   or: ddigger [options] --status - - status-file"),
  90     OPT_HELP(""),
  91     OPT_HELP("Options:"),
  92     OPT_HELP_OPTION,
  93     OPT_STRING(OPT_POSITIONAL(1), NULL, dev_name, OPT_REQUIRED, ""),
  94     OPT_STRING(OPT_POSITIONAL(2), NULL, out_name, OPT_REQUIRED, ""),
  95     OPT_STRING(OPT_POSITIONAL(3), NULL, stat_name, OPT_REQUIRED, ""),
  96     OPT_UINT('b', "block-size", block_size, OPT_REQUIRED_VALUE, "<bytes>\tTransfer block size (default=1M)"),
  97     OPT_UINT('e', "error-skip", error_skip, OPT_REQUIRED_VALUE, "<bytes>\tHow far to skip on error (default=block size)"),
  98     OPT_U64(0, "start", start_pos, OPT_REQUIRED_VALUE, "<bytes>\tStart position (default=beginning of device)"),
  99     OPT_U64(0, "end", end_pos, OPT_REQUIRED_VALUE, "<bytes>\tEnd position (default=end of device)"),
 100     OPT_SWITCH(0, "status", show_status, 1, 0, "\tShow contents of status file"),
 101     OPT_END
 102   }
 103 };
 104
 105 int main(int argc UNUSED, char **argv)
 106 {
 107   opt_parse(&options, argv+1);
 108   if (!error_skip)
 109     error_skip = block_size;
 110
 111   int stat_fd = ucw_open(stat_name, (show_status ? O_RDONLY : O_RDWR | O_CREAT), 0666);
 112   if (stat_fd < 0)
 113     die("Cannot open %s: %m", stat_name);
 114   stat_read(stat_fd);
 115
 116   if (show_status)
 117     {
 118       stat_show();
 119       return 0;
 120     }
 121
 122   int dev_fd = ucw_open(dev_name, O_RDONLY | O_DIRECT);
 123   if (dev_fd < 0)
 124     die("Cannot open block device %s: %m", dev_name);
 125
 126   u64 dev_size;
 127   if (ioctl(dev_fd, BLKGETSIZE64, &dev_size) < 0)
 128     die("BLKGETSIZE64: %m");
 129   start_pos = MIN(dev_size, start_pos);
 130   end_pos = MIN(dev_size, end_pos);
 131
 132   int out_fd = ucw_open(out_name, O_RDWR | O_CREAT, 0666);
 133   if (out_fd < 0)
 134     die("Cannot open %s: %m", out_name);
 135
 136   if (ucw_ftruncate(out_fd, dev_size) < 0)
 137     die("Cannot resize %s: %m", out_name);
 138
 139   u64 remains = (end_pos - start_pos) + stat_have();
 140   ucw_seek(stat_fd, 0, SEEK_END);
 141   *GARY_PUSH(status) = (struct stat_rec) { .pos = end_pos, .len = 0 };
 142
 143   byte *buf = big_alloc(MAX(block_size, 4096));
 144   u64 pos = start_pos;
 145   u64 have = 0;
 146   uint stat_i = 0;
 147   uint errors = 0;
 148
 149   while (pos < end_pos)
 150     {
 151       while (status[stat_i].pos + status[stat_i].len < pos)
 152         stat_i++;
 153       if (status[stat_i].pos <= pos)
 154         {
 155           pos = status[stat_i].pos + status[stat_i].len;
 156           stat_i++;
 157           continue;
 158         }
 159
 160       u64 len = MIN(block_size, status[stat_i].pos - pos);
 161       if (pos % block_size)
 162         len = MIN(len, block_size - pos % block_size);
 163
 164 #define GAUGE(curr,max) (curr), (max), 100*(double)(curr)/(double)((max)?:1)
 165       fprintf(stderr, "Pos: %zd / %zd (%.2f%%)  Rel: %zd / %zd (%.2f%%)  Have: %zd / %zd (%.2f%%)  Errors: %u\r",
 166         GAUGE(pos, dev_size),
 167         GAUGE(pos-start_pos, end_pos-start_pos),
 168         GAUGE(have, remains),
 169         errors);
 170 #undef GAUGE
 171
 172       if (ucw_seek(dev_fd, pos, SEEK_SET) < 0)
 173         die("lseek: %m");
 174       if (ucw_seek(out_fd, pos, SEEK_SET) < 0)
 175         die("lseek: %m");
 176
 177       ssize_t done = read(dev_fd, buf, len);
 178       if (done < (ssize_t) len)
 179         {
 180           fprintf(stderr, "\n");
 181           // msg(L_INFO, "Error: %d bytes read at position %jd", (int) MAX(done, 0), (intmax_t) pos);
 182           errors++;
 183           pos += error_skip;
 184         }
 185       else
 186         {
 187           if (careful_write(out_fd, buf, len) < 0)
 188             die("Error writing %s: %m", out_name);
 189
 190           struct stat_rec sr = { .pos = pos, .len = len };
 191           if (careful_write(stat_fd, &sr, sizeof(sr)) < 0)
 192             die("Error writing status file: %m");
 193
 194           pos += len;
 195           have += len;
 196         }
 197     }
 198
 199   return 0;
 200 }