1// Copyright 2017 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5#include <dirent.h>
6#include <errno.h>
7#include <fcntl.h>
8#include <limits.h>
9#include <stdarg.h>
10#include <stdbool.h>
11#include <stdio.h>
12#include <stdlib.h>
13#include <string.h>
14#include <sys/stat.h>
15#include <unistd.h>
16
17int usage(void) {
18    fprintf(stderr, "usage: dd [OPTIONS]\n");
19    fprintf(stderr, "dd can be used to convert and copy files\n");
20    fprintf(stderr, " bs=BYTES  : read and write BYTES bytes at a time\n");
21    fprintf(stderr, " count=N   : copy only N input blocks\n");
22    fprintf(stderr, " ibs=BYTES : read BYTES bytes at a time (default: 512)\n");
23    fprintf(stderr, " if=FILE   : read from FILE instead of stdin\n");
24    fprintf(stderr, " obs=BYTES : write BYTES bytes at a time (default: 512)\n");
25    fprintf(stderr, " of=FILE   : write to FILE instead of stdout\n");
26    fprintf(stderr, " seek=N    : skip N obs-sized blocks at start of output\n");
27    fprintf(stderr, " skip=N    : skip N ibs-sized blocks at start of input\n");
28    fprintf(stderr, " N and BYTES may be followed by the following multiplicitive\n"
29                    " suffixes: c = 1, w = 2, b = 512, kB = 1000, K = 1024,\n"
30                    "           MB = 1000 * 1000, M = 1024 * 1024, xM = M,\n"
31                    "           GB = 1000 * 1000 * 1000, G = 1024 * 1024 * 1024\n");
32    fprintf(stderr, " --help : Show this help message\n");
33    return -1;
34}
35
36// Returns "true" if the argument matches the prefix.
37// In this case, moves the argument past the prefix.
38bool prefix_match(const char** arg, const char* prefix) {
39    if (!strncmp(*arg, prefix, strlen(prefix))) {
40        *arg += strlen(prefix);
41        return true;
42    }
43    return false;
44}
45
46#define MAYBE_MULTIPLY_SUFFIX(str, out, suffix, value) \
47        if (!strcmp((str), (suffix))) {                \
48            (out) *= (value);                          \
49            return 0;                                  \
50        }
51
52// Parse the formatted size string from |s|, and place
53// the result in |out|.
54//
55// Returns 0 on success.
56int parse_size(const char* s, size_t* out) {
57    char* endptr;
58    if (!(*s >= '0' && *s <= '9')) {
59        goto done;
60    }
61    *out = strtol(s, &endptr, 10);
62    if (*endptr == '\0') {
63        return 0;
64    }
65
66    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "G", 1UL << 30);
67    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "GB", 1000 * 1000 * 1000UL);
68    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "xM", 1UL << 20);
69    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "M", 1UL << 20);
70    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "MB", 1000 * 1000UL);
71    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "K", 1UL << 10);
72    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "kB", 1000UL);
73    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "b", 512UL);
74    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "w", 2UL);
75    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "c", 1UL);
76
77done:
78    fprintf(stderr, "Couldn't parse size string: %s\n", s);
79    return -1;
80}
81
82typedef struct {
83    bool use_count;
84    size_t count;
85    size_t input_bs;
86    size_t input_skip;
87    size_t output_bs;
88    size_t output_seek;
89    char input[PATH_MAX];
90    char output[PATH_MAX];
91} dd_options_t;
92
93int parse_args(int argc, const char** argv, dd_options_t* options) {
94    while (argc > 1) {
95        const char* arg = argv[1];
96        if (prefix_match(&arg, "bs=")) {
97            size_t size;
98            if (parse_size(arg, &size)) {
99                return usage();
100            }
101            options->input_bs = size;
102            options->output_bs = size;
103        } else if (prefix_match(&arg, "count=")) {
104            if (parse_size(arg, &options->count)) {
105                return usage();
106            }
107            options->use_count = true;
108        } else if (prefix_match(&arg, "ibs=")) {
109            if (parse_size(arg, &options->input_bs)) {
110                return usage();
111            }
112        } else if (prefix_match(&arg, "obs=")) {
113            if (parse_size(arg, &options->output_bs)) {
114                return usage();
115            }
116        } else if (prefix_match(&arg, "seek=")) {
117            if (parse_size(arg, &options->output_seek)) {
118                return usage();
119            }
120        } else if (prefix_match(&arg, "skip=")) {
121            if (parse_size(arg, &options->input_skip)) {
122                return usage();
123            }
124        } else if (prefix_match(&arg, "if=")) {
125            strncpy(options->input, arg, PATH_MAX);
126            options->input[PATH_MAX - 1] = '\0';
127        } else if (prefix_match(&arg, "of=")) {
128            strncpy(options->output, arg, PATH_MAX);
129            options->output[PATH_MAX - 1] = '\0';
130        } else {
131            return usage();
132        }
133        argc--;
134        argv++;
135    }
136    return 0;
137}
138
139#define MIN(x,y) ((x) < (y) ? (x) : (y))
140#define MAX(x,y) ((x) < (y) ? (y) : (x))
141
142int main(int argc, const char** argv) {
143    dd_options_t options;
144    memset(&options, 0, sizeof(dd_options_t));
145    options.input_bs = 512;
146    options.output_bs = 512;
147    int r;
148    if ((r = parse_args(argc, argv, &options))) {
149        return r;
150    }
151
152    if (options.input_bs == 0 || options.output_bs == 0) {
153        fprintf(stderr, "block sizes must be greater than zero\n");
154        return -1;
155    }
156
157    options.input_skip *= options.input_bs;
158    options.output_seek *= options.output_bs;
159
160    // Input and output fds
161    int in = -1;
162    int out = -1;
163    // Buffer to contain partially read data
164    char* buf = NULL;
165    // Return code
166    r = -1;
167    // Number of full records copied to/from target
168    size_t records_in = 0;
169    size_t records_out = 0;
170    // Size of remaining "partial" transfer from input / to output.
171    size_t record_in_partial = 0;
172    size_t record_out_partial = 0;
173
174    if (*options.input == '\0') {
175        in = STDIN_FILENO;
176    } else {
177        in = open(options.input, O_RDONLY);
178        if (in < 0) {
179            fprintf(stderr, "Couldn't open input file %s : %d\n", options.input, errno);
180            goto done;
181        }
182    }
183
184    if (*options.output == '\0') {
185        out = STDOUT_FILENO;
186    } else {
187        out = open(options.output, O_WRONLY | O_CREAT);
188        if (out < 0) {
189            fprintf(stderr, "Couldn't open output file %s : %d\n", options.output, errno);
190            goto done;
191        }
192    }
193
194    buf = malloc(MAX(options.output_bs, options.input_bs));
195    if (buf == NULL) {
196        fprintf(stderr, "No memory\n");
197        goto done;
198    }
199
200    if (options.input_skip != 0) {
201        // Try seeking first; if that doesn't work, try reading to an input buffer.
202        if (lseek(in, options.input_skip, SEEK_SET) != (off_t) options.input_skip) {
203            while (options.input_skip) {
204                if (read(in, buf, options.input_bs) != (ssize_t) options.input_bs) {
205                    fprintf(stderr, "Couldn't read from input\n");
206                    goto done;
207                }
208                options.input_skip -= options.input_bs;
209            }
210        }
211    }
212
213    if (options.output_seek != 0) {
214        if (lseek(out, options.output_seek, SEEK_SET) != (off_t) options.output_seek) {
215            fprintf(stderr, "Failed to seek on output\n");
216            goto done;
217        }
218    }
219
220    if (MAX(options.input_bs, options.output_bs) %
221        MIN(options.input_bs, options.output_bs) != 0) {
222        // TODO(smklein): Implement this case, rather than returning an error
223        fprintf(stderr, "Input and output block sizes must be multiples\n");
224        goto done;
225    }
226
227    bool terminating = false;
228    size_t rlen = 0;
229    while (true) {
230        if (options.use_count && !options.count) {
231            r = 0;
232            goto done;
233        }
234
235        // Read as much as we can (up to input_bs) into our target buffer
236        ssize_t rout;
237        if ((rout = read(in, buf, options.input_bs)) != (ssize_t) options.input_bs) {
238            terminating = true;
239        }
240        if (rout == (ssize_t) options.input_bs) {
241            records_in++;
242        } else if (rout > 0) {
243            record_in_partial = rout;
244        }
245        if (rout > 0) {
246            rlen += rout;
247        }
248        if (options.use_count) {
249            --options.count;
250            if (options.count == 0) {
251                terminating = true;
252            }
253        }
254
255        // If we can (or should, due to impending termination), dump the read
256        // buffer into the output file.
257        if (rlen >= options.output_bs || terminating) {
258            size_t off = 0;
259            while (off != rlen) {
260                size_t wlen = MIN(options.output_bs, rlen - off);
261                if (write(out, buf + off, wlen) != (ssize_t) wlen) {
262                    fprintf(stderr, "Couldn't write %zu bytes to output\n", wlen);
263                    goto done;
264                }
265                if (wlen == options.output_bs) {
266                    records_out++;
267                } else {
268                    record_out_partial = wlen;
269                }
270                off += wlen;
271            }
272            rlen = 0;
273        }
274
275        if (terminating) {
276            r = 0;
277            goto done;
278        }
279    }
280
281done:
282    printf("%zu+%u records in\n", records_in, record_in_partial ? 1 : 0);
283    printf("%zu+%u records out\n", records_out, record_out_partial ? 1 : 0);
284    printf("%zu bytes copied\n", records_out * options.output_bs + record_out_partial);
285
286    if (in != -1) {
287        close(in);
288    }
289    if (out != -1) {
290        close(out);
291    }
292    free(buf);
293    return r;
294}
295