harakit/src/dj.c

/*
 * Copyright (c) 2024 DTB <trinity@trinity.moe>
 * SPDX-License-Identifier: AGPL-3.0-or-later
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU Affero General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or (at your option) any
 * later version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
 * details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program. If not, see https://www.gnu.org/licenses/.
 */

#include <assert.h> /* assert(3) */
#include <errno.h> /* errno */
#include <fcntl.h> /* open(2) */
#include <stdio.h> /* fprintf(3), stderr */
#include <stdlib.h> /* malloc(3), strtol(3), size_t */
#include <string.h> /* memcpy(3), memmove(3), memset(3) */
#if !defined EX_OK || !defined EX_OSERR || !defined EX_USAGE
#	include <sysexits.h>
#endif
#include <unistd.h> /* close(2), getopt(3), lseek(2), read(2), write(2),
                     * optarg, optind, STDIN_FILENO, STDOUT_FILENO */
#include <sys/stat.h> /* S_IRGRP, S_IROTH, S_IRUSR, S_IWGRP, S_IWOTH,
                         S_IWUSR */
extern int errno;

char *program_name = "dj";

/* dj uses two structures that respectively correspond to the reading and
 * writing ends of its jockeyed "pipe". User-configurable members are noted
 * with their relevant options. */
struct Io{
	int bs;       	/* buffer size (-bB) */
	size_t bufuse;	/* buffer usage */
	char *buf;		/* buffer */
	size_t bytes;	/* bytes processed */
	int fd;			/* file descriptor */
	int fl;			/* file opening flags */
	char *fn;		/* file name (may be stdin_name or stdout_name) (-io) */
	size_t prec;	/* partial records processed */
	size_t rec; 	/* records processed */
	long seek;		/* bytes to seek/skip (will be 0 after skippage) (-sS) */
}; 

/* To be assigned to main:fmt and used with printio(). */
static char *fmt_asv   = "%d\037%d\036%d\037%d\035%d\036%d\034";
static char *fmt_human = "%d+%d > %d+%d; %d > %d\n";

static char *stdin_name = "<stdin>";
static char *stdout_name = "<stdout>";

static int read_flags = O_RDONLY; /* Consistent with Busybox dd(1). */
static int write_flags = O_WRONLY | O_CREAT;

#define MIN(a, b) (((a) < (b)) ? (a) : (b))

/* Macro to check if fd is stdin or stdout */
#define fdisstd(fd) ((fd) == STDIN_FILENO || (fd) == STDOUT_FILENO)

/* Reads io->bs bytes from *io's file descriptor into io->buf, storing the
 * number of read bytes in io->bufuse and updating io->bytes. If the buf isn't
 * saturated but is still read into, io->prec will be incremented. If the buf
 * is saturated, io->rec will be incremented. If io->bufuse is 0, errno will
 * probably be set. Returns io. */
static struct Io *
Io_read(struct Io *io){

	assert(io->bs > 0);

	io->bytes += (io->bufuse = read(io->fd, io->buf, io->bs));

	assert(io->bufuse <= io->bs);

	io->prec += (0 < io->bufuse && io->bufuse < io->bs);
	io->rec  += (io->bufuse == io->bs);

	return io;
}

/* Writes io->bufuse (>0) units from io->buf to io->fd, permuting any unwritten
 * bytes to the start of io->buf and updating io->bufuse. If io->bufuse doesn't
 * change, errno will probably be set. Returns io. */
static struct Io *
Io_write(struct Io *io){
	int t;

	assert(io->bufuse > 0);
	assert(io->bufuse <= io->bs);

	if((t = write(io->fd, io->buf, io->bufuse)) > 0)
		memmove(io->buf, io->buf + t, (io->bufuse -= t));

	io->bytes += t;
	io->prec  += (t > 0 && io->bufuse > 0);
	io->rec   += (t > 0 && io->bufuse == 0);

	return io;
}

/* Prints an error message suitable for the event of an operating system error,
 * with the error itself to be described in the string s. */
static int
oserr(char *s){

	fprintf(stderr, "%s: %s: %s\n", program_name, s, strerror(errno));

	return EX_OSERR;
}

/* Prints statistics regarding the use of dj, particularly partially and
 * completely read and written records. */
static void
fprintio(FILE *stream, char *fmt, struct Io io[2]){

	fprintf(stream, fmt,
		io[0].rec, io[0].prec, io[1].rec, io[1].prec,
		io[0].bytes, io[1].bytes);

	return;
}

/* Parses the string s to an integer, returning either the integer or in the
 * case of an error a negative integer. This is used for argument parsing
 * (e.g. -B [int]) in dj and no negative integer would be valid anyway. */
static long
parse(char *s){
	long r;

	errno = 0;
	r = strtol(s, &s, 0);
	return (*s == '\0' /* no chars left unparsed */ && errno == 0)
		? r
		: -1;
}

static int
usage(char *s){

	fprintf(stderr, "Usage: %s [-Hn] [-a byte] [-c count]\n"
		"\t[-i file] [-b block_size] [-s offset]\n"
		"\t[-o file] [-B block_size] [-S offset]\n",
		program_name);

	return EX_USAGE;
}

int main(int argc, char *argv[]){
 	int align;    /* low 8b used, negative if no alignment is being done */
	int count;    /* 0 if dj(1) runs until no more reads are possible */
	char *fmt;    /* == fmt_asv (default) or fmt_human (-H) */
	size_t i;     /* side of io being modified */
	char noerror; /* 0=exits (default) 1=retries on partial reads or writes */
	struct Io io[2];

	/* Set defaults. */
	align = -1;
	count = 0;
	fmt = fmt_asv;
	noerror = 0;
	for(i = 0; i < 2; ++i){
		io[i].bs = 1024 /* 1 KiB */; /* GNU dd(1) default; POSIX says 512B */
		io[i].bufuse = 0;
		io[i].bytes = 0;
		io[i].fd = i == 0 ? STDIN_FILENO : STDOUT_FILENO;
		io[i].fn = i == 0 ? stdin_name   : stdout_name;
		io[i].fl = i == 0 ? read_flags   : write_flags;
		io[i].prec = 0;
		io[i].rec = 0;
		io[i].seek = 0;
	}

	if(argc > 0){
		int c;

		program_name = argv[0];
		while((c = getopt(argc, argv, ":a:b:B:c:i:hHns:S:o:")) != -1)
			switch(c){
			case 'i': case 'o': i = (c == 'o');
				if(optarg[0] == '-' && optarg[1] == '\0'){ /* optarg == "-" */
					io[i].fd = i == 0 ? STDIN_FILENO : STDOUT_FILENO;
					io[i].fn = i == 0 ? stdin_name   : stdout_name;
					break;
				}else{
					int fd;

					if((fd = open(optarg, io[i].fl, /* touch(1p) flags */
									S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP
										| S_IROTH | S_IWOTH))
								!= -1
							&& (fdisstd(io[i].fd) || close(io[i].fd) == 0)){
						io[i].fd = fd;
						io[i].fn = optarg;
						break;
					}
				}
				return oserr(optarg);
			case 'n': noerror = 1;     break;
			case 'H': fmt = fmt_human; break;
			case 'a':
				if(optarg[0] == '\0' || optarg[1] == '\0'){
					align = optarg[0];
					break;
				}
				/* FALLTHROUGH */
			case 'c': case 'b': case 's': case 'B': case 'S': /* numbers */
				if(c == 'c' && (count = parse(optarg)) >= 0)
					break;
				i = (c >= 'A' && c <= 'Z');
				c |= 0x20 /* 0b 0010 0000 */; /* (ASCII) make lowercase */
				if((c == 'b' && (io[i].bs = parse(optarg)) > 0)
						|| (c == 's' && (io[i].seek = parse(optarg)) >= 0))
					break;
				/* FALLTHROUGH */
			default:
				return usage(program_name);
			}
	}

	assert(io->fd != STDIN_FILENO  || io->fl == read_flags);
	assert(io->fd != STDOUT_FILENO || io->fl == write_flags);

	if(argc > optind)
		return usage(program_name);

	for(i = 0; i < 2; ++i){
		/* buffer allocation */
		if((io[i].buf = malloc(io[i].bs * (sizeof *(io[i].buf)))) == NULL){
			fprintf(stderr, "%s: Failed to allocate %d bytes\n",
				program_name, io[i].bs);
			return EX_OSERR;
		}
		/* easy seeking */
		if(!fdisstd(io[i].fd) && lseek(io[i].fd, io[i].seek, SEEK_SET) != -1)
			io[i].seek = 0;
	}

	/* hard skipping */
	if(io[0].seek > 0){
		do{
			if((io[0].bufuse = read(
						io[0].fd, io[0].buf, MIN(io[0].bs, io[0].seek)))
					== 0)
				/* second chance */
				io->bufuse = read(
					io[0].fd, io[0].buf, MIN(io[0].bs, io[0].seek));
		}while((io[0].seek -= io[0].bufuse) > 0 && io[0].bufuse != 0);
		io[0].bufuse = 0;
	}

	/* hard seeking */
	if(io[1].seek > 0){
		memset(io[1].buf, '\0', io[1].bs);
		/* We're going to cheat and use bufuse as the retval for write(2),
		 * which is fine because it'll be zeroed as this function returns
		 * anyway. */
		do{
			if((io[1].bufuse = write(
						io[1].fd, io[1].buf, MIN(io[1].bs, io[1].seek)))
					== 0)
				/* second chance */
				io[1].bufuse = write(
					io[1].fd, io[1].buf, MIN(io[1].bs, io[1].seek));
		}while((io[1].seek -= io[1].bufuse) > 0 && io[1].bufuse != 0);
		io[1].bufuse = 0;
	}

	/* Sought bytes aren't counted in the statistics because successful seeking
	 * is guaranteed here. */
	for(i = 0; i < 2; ++i)
		if(io[i].seek > 0)
			return oserr(io[i].fn);

	do{	/* read */
		if(Io_read(&io[0])->bufuse == 0 && !noerror)
			Io_read(&io[0]); /* second chance */
		if(io[0].bufuse == 0) /* that's all she wrote */
			break;

		if(io[0].bufuse < io[0].bs){
			fprintf(stderr, "%s: Partial read:\n\t", program_name);
			fprintio(stderr, fmt, io);
			if(!noerror)
				count = 1;
			if(align >= 0){
				/* fill the rest of the ibuf with padding */
				memset(io[0].buf + io[0].bufuse, align,
					io[0].bs - io[0].bufuse);
				io->bufuse = io->bs;
			}
		}

		/* write */
		do{
			int t;

			if(io[0].bs <= io[1].bs){
				int n;

				/* saturate obuf */
				memcpy(io[1].buf, io[0].buf,
				       (io[1].bufuse = (n = MIN(io[0].bufuse, io[1].bs))));
				/* permute the copied units out of ibuf */
				memmove(io[0].buf, io[0].buf + n, (io[0].bufuse -= n));
			}else /* if(io[0].bs < io[1].bs) */ {
				int n;

				/* drain what we can from ibuf */
				memcpy(io[1].buf + io[1].bufuse, io[0].buf,
				       (n = MIN(io[0].bufuse, io[1].bs - io[1].bufuse)));
				io[1].bufuse += n;
				/* permute out the copied units */
				memmove(io[0].buf, io[0].buf + n, io[0].bs - n);
				io[0].bufuse -= n;

				if(io[0].bs + io[1].bufuse <= io[1].bs && count != 1)
					continue; /* obuf not saturated - we could write more */
			}

			t = io[1].bufuse;
			if(Io_write(&io[1])->bufuse == t && !noerror)
				Io_write(&io[1]); /* second chance */
			if(io[1].bufuse == t){ /* no more love */
				count = 1;
				break;
			}

			if(0 < io[1].bufuse && io[1].bufuse < t){
				fprintf(stderr, "%s: Partial write:\n\t", program_name);
				fprintio(stderr, fmt, io);
				if(!noerror)
					count = 1;
			}
		}while(io[0].bufuse > 0);
	}while(count == 0 || --count > 0);

	fprintio(stderr, fmt, io);

	return EX_OK;
}