micro-utils/src/dd.c

#include <fcntl.h>
#include <stdio.h>
#include <signal.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include "human.h"

static off_t infull, inpart;
static off_t outfull, outpart;
static off_t tbytes;

static void summary(void) {
	fprintf(stderr, "%jd+%jd records in\n", infull, inpart);
	fprintf(stderr, "%jd+%jd records out\n", outfull, outpart);

	fprintf(stderr, "%s total bytes copied\n", mu_humansize(tbytes, 1024));
}

static int openfile(int flag, char *path, int mode) {
	if (!strcmp(path, "-")) {
		if (flag)
			return STDOUT_FILENO;

		return STDIN_FILENO;
	}

	int fd = open(path, mode, 0666);
	if (fd < 0) {
		fprintf(stderr, "dd: %s: %s\n", path, strerror(errno));
		exit(1);
	}

	return fd;
}

static off_t strtonum(char *str) {
	char *p = NULL;
	off_t res = strtoll(str, &p, 0);
	if (str != p) {
		if (!strcmp(p, "b"))
			res *= 512;

		else if (!strcmp(p, "m"))
			res *= 1000000;

		else if (!strcmp(p, "M"))
			res *= 1048576;

		else if (!strcmp(p, "K"))
			res *= 1024;

		else if (!strcmp(p, "k"))
			res *= 1000;

		else if (!strcmp(p, "g"))
			res *= 1000000000;

		else if (!strcmp(p, "G"))
			res *= 1073741824;
	}

	if (res < 0)
		res = 0;

	return res;
}

static int copy(int fd, void *buf, off_t len, off_t max) {
	off_t n = write(fd, buf, len);
	if (n < 0) {
		fprintf(stderr, "dd: %s\n", strerror(errno));
		return 1;
	}

	else if (n == max)
		outfull++;

	else if (n == len)
		outpart++;

	tbytes += n;
	return 0;
}

int main(int argc, char **argv) {
	signal(SIGPIPE, SIG_IGN);

	char *ifp = "-";
	char *ofp = "-";

	off_t count = -1;
	off_t skip = 0;
	off_t seek = 0;

	off_t bs = 0;
	off_t ibs = 512;
	off_t obs = 512;

	int ofd = STDOUT_FILENO;
	int ifd = STDIN_FILENO;

	char notrunc = 0;
	char fsync_flag = 0;

	/* Return value */
	int ret = 1;

	for (int i = 1; i < argc; i++) {
		char *arg = argv[i];

		char *val = strchr(arg, '=');
		if (val == NULL) {
			puts("dd\n\tif=InputFile\n\tof=OutputFile\n\tbs=ibs and obs\n\tibs=Input buffer size\n\tobs=Output buffer size\n\tseek=Skip N obs-sized output blocks\n\tskip=Skip N ibs-sized output blocks\n\tcount=Copy only N input blocks\n\tnconv=notrunc Don't truncate output file\n\tconv=fsync Physically write data out before finishing\n\nN and BYTES may be followed by the following multiplicative\nsuffixes: w=2, b=512, k=1000, K=1024, m=1000*1000,\nM=1024*1024, g=1000*1000*1000, G=1024*1024*1024");
			return 1;
		}

		*val = '\0';
		val++;

		/* Get value */
		if (!strcmp(arg, "if"))
			ifp = val;

		else if (!strcmp(arg, "of"))
			ofp = val;

		else if (!strcmp(arg, "seek"))
			seek = strtonum(val);

		else if (!strcmp(arg, "skip"))
			skip = strtonum(val);

		else if (!strcmp(arg, "count"))
			count = strtonum(val);

		else if (!strcmp(arg, "bs"))
			bs = strtonum(val);

		else if (!strcmp(arg, "ibs"))
			ibs = strtonum(val);

		else if (!strcmp(arg, "obs"))
			obs = strtonum(val);

		char *ptr = strstr(arg, "conv");
		if (ptr) {
			if (strstr(val, "notrunc"))
				notrunc = 1;

			if (strstr(val, "fsync"))
				fsync_flag = 1;
		}
	}

	if (bs) {
		ibs = bs;
		obs = bs;
	}

	/* Make input ibuffer */
	char *ibuf = malloc(ibs);
	if (ibuf == NULL) {
		fprintf(stderr, "dd: %s\n", strerror(errno));
		return 1;
	}

	char *obuf = NULL;
	if (ibs != obs) {
		obuf = malloc(obs);
		if (obuf == NULL) {
			free(ibuf);
			fprintf(stderr, "dd: %s\n", strerror(errno));
			return 1;
		}
	}

	/* Open files. Input */
	ifd = openfile(0, ifp, O_RDONLY);
	if (skip) {
		if (lseek(ifd, skip * ibs, SEEK_CUR) < 0)
			goto CLOSE;
	}

	/* Output */
	int oflag = O_WRONLY | O_CREAT;
	if (seek && !notrunc)
		oflag |= O_TRUNC;

	ofd = openfile(1, ofp, oflag);
	if (seek) {
		if (!notrunc && ftruncate(ofd, seek * ibs) < 0) {
			fprintf(stderr, "dd: %s\n", strerror(errno));
			goto CLOSE;
		}

		if (lseek(ofd, seek * ibs, SEEK_SET) < 0) {
			fprintf(stderr, "dd: %s\n", strerror(errno));
			goto CLOSE;
		}
	}

	/* dd */
	off_t opos = 0;

	while (1) {
		if (count == infull + inpart)
			break;

		off_t n = read(ifd, ibuf, ibs);
		if (n <= 0)
			break;

		else if (ibs == n)
			infull++;

		else
			inpart++;

		if (ibs == obs) {
			if (copy(ofd, ibuf, n, ibs))
				goto CLOSE;
		}

		else {
			char *tmp = ibuf;
			while (n) {
				off_t i = obs - opos;
				if (i > n)
					i = n;

				memcpy(obuf + opos, tmp, i);

				n -= i;
				tmp += i;
				opos += i;
				if (opos == obs) {
					if (copy(ofd, obuf, obs, obs))
						goto CLOSE;

					opos = 0;
				}
			}
		}
	}

	if (opos != 0) {
		if (copy(ofd, obuf, obs, obs))
			goto CLOSE;
	}

	/* End */
	if (fsync_flag)
		fsync(ofd);

	summary();

	ret = 0;

CLOSE:
	free(ibuf);
	if (ibs != obs)
		free(obuf);

	close(ifd);
	close(ofd);
	return ret;
}