]> mj.ucw.cz Git - misc.git/blob - ucw/ddigger.c
3df69815431118909eeb9bd06e2acf40ca4efc7f
[misc.git] / 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 }