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