A digging dd

diff --git a/ucw/ddigger.c b/ucw/ddigger.c
new file mode 100644 (file)
index 0000000..3df6981
--- /dev/null
+++ b/ucw/ddigger.c
@@ -0,0 +1,179 @@
+#include <ucw/lib.h>
+#include <ucw/io.h>
+#include <ucw/opt.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <fcntl.h>
+#include <unistd.h>
+#include <sys/ioctl.h>
+#include <linux/fs.h>
+
+static char *dev_name;
+static char *out_name;
+static char *stat_name;
+static uint block_size = 1048576;
+static uint error_skip;
+static u64 start_pos;
+static u64 end_pos = ~(u64)0;
+
+struct stat_rec {
+  u64 pos;
+  u64 len;
+};
+
+static struct stat_rec *status;
+
+#define ASORT_PREFIX(x) ss_##x
+#define ASORT_KEY_TYPE struct stat_rec
+#define ASORT_LT(x,y) x.pos < y.pos
+#include <ucw/sorter/array-simple.h>
+
+static void stat_read(int fd)
+{
+  u64 len = ucw_seek(fd, 0, SEEK_END);
+  ucw_seek(fd, 0, SEEK_SET);
+  if (len % sizeof(struct stat_rec))
+    die("Malformed status file");
+  if (len > ~0U)
+    die("Status file too long");
+
+  uns items = len / sizeof(struct stat_rec);
+  GARY_INIT(status, items);
+  if (careful_read(fd, status, len) <= 0)
+    die("Error reading status file");
+
+  ss_sort(status, items);
+
+#if 0
+  for (uint i=0; i<items; i++)
+    printf("%jd %jd\n", (intmax_t) status[i].pos, (intmax_t) status[i].len);
+#endif
+
+  for (uint i=1; i<items; i++)
+    if (status[i].pos < status[i-1].pos + status[i-1].len)
+      die("Overlapping status file entries: [%jd,%jd] with [%jd,%jd]",
+       (intmax_t) status[i-1].pos,
+       (intmax_t)(status[i-1].pos + status[i-1].len),
+       (intmax_t) status[i].pos,
+       (intmax_t)(status[i].pos + status[i].len));
+}
+
+static u64 stat_have(void)
+{
+  u64 have = 0;
+  for (uint i=0; i<GARY_SIZE(status); i++)
+    {
+      u64 x = CLAMP(status[i].pos, start_pos, end_pos);
+      u64 y = CLAMP(status[i].pos + status[i].len, start_pos, end_pos);
+      have += y - x;
+    }
+  return have;
+}
+
+static struct opt_section options = {
+  OPT_ITEMS {
+    OPT_HELP("Usage: ddigger [options] block-device output-file status-file"),
+    OPT_HELP(""),
+    OPT_HELP("Options:"),
+    OPT_HELP_OPTION,
+    OPT_STRING(OPT_POSITIONAL(1), NULL, dev_name, OPT_REQUIRED, ""),
+    OPT_STRING(OPT_POSITIONAL(2), NULL, out_name, OPT_REQUIRED, ""),
+    OPT_STRING(OPT_POSITIONAL(3), NULL, stat_name, OPT_REQUIRED, ""),
+    OPT_UINT('b', "block-size", block_size, OPT_REQUIRED_VALUE, "<bytes>\tTransfer block size (default=1M)"),
+    OPT_UINT('e', "error-skip", error_skip, OPT_REQUIRED_VALUE, "<bytes>\tHow far to skip on error (default=block size)"),
+    OPT_U64(0, "start", start_pos, OPT_REQUIRED_VALUE, "<bytes>\tStart position (default=beginning of device)"),
+    OPT_U64(0, "end", end_pos, OPT_REQUIRED_VALUE, "<bytes>\tEnd position (default=end of device)"),
+    OPT_END
+  }
+};
+
+int main(int argc UNUSED, char **argv)
+{
+  opt_parse(&options, argv+1);
+  if (!error_skip)
+    error_skip = block_size;
+
+  int dev_fd = ucw_open(dev_name, O_RDONLY | O_DIRECT);
+  if (dev_fd < 0)
+    die("Cannot open block device %s: %m", dev_name);
+
+  u64 dev_size;
+  if (ioctl(dev_fd, BLKGETSIZE64, &dev_size) < 0)
+    die("BLKGETSIZE64: %m");
+  start_pos = MIN(dev_size, start_pos);
+  end_pos = MIN(dev_size, end_pos);
+
+  int out_fd = ucw_open(out_name, O_RDWR | O_CREAT, 0666);
+  if (out_fd < 0)
+    die("Cannot open %s: %m", out_name);
+
+  if (ucw_ftruncate(out_fd, dev_size) < 0)
+    die("Cannot resize %s: %m", out_name);
+
+  int stat_fd = ucw_open(stat_name, O_RDWR | O_CREAT, 0666);
+  if (stat_fd < 0)
+    die("Cannot open %s: %m", stat_name);
+  stat_read(stat_fd);
+  u64 remains = (end_pos - start_pos) + stat_have();
+
+  ucw_seek(stat_fd, 0, SEEK_END);
+  *GARY_PUSH(status) = (struct stat_rec) { .pos = end_pos, .len = 0 };
+
+  byte *buf = big_alloc(MAX(block_size, 4096));
+  u64 pos = start_pos;
+  u64 have = 0;
+  uint stat_i = 0;
+  uint errors = 0;
+
+  while (pos < end_pos)
+    {
+      while (status[stat_i].pos + status[stat_i].len < pos)
+       stat_i++;
+      if (status[stat_i].pos <= pos)
+       {
+         pos = status[stat_i].pos + status[stat_i].len;
+         stat_i++;
+         continue;
+       }
+
+      u64 len = MIN(block_size, status[stat_i].pos - pos);
+      if (pos % block_size)
+       len = MIN(len, block_size - pos % block_size);
+
+#define GAUGE(curr,max) (curr), (max), 100*(double)(curr)/(double)((max)?:1)
+      fprintf(stderr, "Pos: %zd / %zd (%.2f%%)  Rel: %zd / %zd (%.2f%%)  Have: %zd / %zd (%.2f%%)  Errors: %u\r",
+       GAUGE(pos, dev_size),
+       GAUGE(pos-start_pos, end_pos-start_pos),
+       GAUGE(have, remains),
+       errors);
+#undef GAUGE
+
+      if (ucw_seek(dev_fd, pos, SEEK_SET) < 0)
+       die("lseek: %m");
+      if (ucw_seek(out_fd, pos, SEEK_SET) < 0)
+       die("lseek: %m");
+
+      ssize_t done = read(dev_fd, buf, len);
+      if (done < (ssize_t) len)
+       {
+         fprintf(stderr, "\n");
+         // msg(L_INFO, "Error: %d bytes read at position %jd", (int) MAX(done, 0), (intmax_t) pos);
+         errors++;
+         pos += error_skip;
+       }
+      else
+       {
+         if (careful_write(out_fd, buf, len) < 0)
+           die("Error writing %s: %m", out_name);
+
+         struct stat_rec sr = { .pos = pos, .len = len };
+         if (careful_write(stat_fd, &sr, sizeof(sr)) < 0)
+           die("Error writing status file: %m");
+
+         pos += len;
+         have += len;
+       }
+    }
+
+  return 0;
+}