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 <ctype.h> /* isupper(3), tolower(3) */
#include <errno.h> /* errno */
#include <fcntl.h> /* open(2) */
#include <stdio.h> /* fprintf(3), stderr */
#include <stdlib.h> /* free(3), malloc(3), strtol(3), size_t */
#include <string.h> /* memcpy(3), memmove(3), memset(3) */
#include <sysexits.h> /* EX_OK, EX_USAGE */
#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;

/* 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 */
	int bytes;		/* bytes processed */
	int fd;			/* file descriptor */
	int fl;			/* file opening flags */
	char *fn;		/* file name (may be stdin_name or stdout_name) (-io) */
	int prec;		/* partial records processed */
	int rec;		/* records processed */
	long seek;		/* bytes to seek/skip (will be 0 after skippage) (-sS) */
}; 

/* To be assigned to main:fmt_output and used with output(). */
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 *program_name = "<no argv[0]>";
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 a std* file, e.g. stdin. */
#define fdisstd(fd) \
		((fd) == STDIN_FILENO \
		|| (fd) == STDOUT_FILENO \
		|| (fd) == STDERR_FILENO)

/* Macro to call the cleanup functions that operate on struct io on the
 * particular io[2] used in main. Error conditions are not checked because this
 * is only used when the program is about to terminate (hence its name). */
#define terminate(io) do{ \
	Io_buffree(&(io)[0]); \
	Io_buffree(&(io)[1]); \
	Io_fdclose(&(io)[0]); \
	Io_fdclose(&(io)[1]); }while(0)

/* Allocates *io's buffer. Returns NULL if unsuccessful. */
static void *
Io_bufalloc(struct Io *io){

	return (io->buf = malloc(io->bs * (sizeof *io->buf)));
}

/* Frees *io's buffer. Returns io. */
static struct Io *
Io_buffree(struct Io *io){

	free(io->buf);

	return io;
}

/* Fills the unused portion of io's buffer with padding, updating io->bufuse.
 * Returns io. */
static struct Io *
Io_bufrpad(struct Io *io, int padding){
	
	memset(io->buf + io->bufuse, padding, io->bs - io->bufuse);
	io->bufuse = io->bs;

	return io;
}

/* Copies from the buffer in src as much as possible to the free space in the
 * dest buffer, removing the copied units from src and permuting the remaining
 * units in the src buffer to the start of the buffer, modifying both the src
 * and dest bufuse and returning dest. */
static struct Io*
Io_bufxapp(struct Io *dest, struct Io *src){
	int n;

	n = MIN(src->bufuse, dest->bs - dest->bufuse);
	memcpy(dest->buf + dest->bufuse, src->buf, n);
	dest->bufuse += n;
	memmove(src->buf, src->buf + n, src->bs - n);
	src->bufuse -= n;

	return dest;
}

/* Copies from the buffer in src to the buffer in dest no more than n units,
 * removing the copied units from src and permuting the remaining units in the
 * src buffer to the start of the buffer, modifying both the src and dest
 * bufuse and returning dest. */
static struct Io*
Io_bufxfer(struct Io *dest, struct Io *src, int n){

	memcpy(dest->buf, src->buf, (dest->bufuse = n));
	memmove(src->buf, src->buf + n, (src->bufuse -= n));

	return dest;
}

/* Closes io->fn and returns -1 on error, otherwise io->fd. */
static int
Io_fdclose(struct Io *io){

	return fdisstd(io->fd)
		? 0
		: close(io->fd);
}

/* Opens io->fn and saves the file descriptor into io->fd. Returns io->fd,
 * which will be -1 if an error occured. */
static int
Io_fdopen(struct Io *io, char *fn){
	int fd;
	
	if((fd = open(fn, io->fl,
					/* these are the flags used by touch(1p) */
					S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP | S_IROTH | S_IWOTH))
				!= -1
			&& Io_fdclose(io) == 0){
		io->fd = fd;
		io->fn = fn;
	}

	return fd;
}

/* Seeks io->seek bytes through *io's file descriptor, (counter-intuitively)
 * returning -1 if successful and a sysexits.h exit code if an unrecoverable
 * error occurred. io->buf will be cleared of useful bytes and io->seek will
 * be set to zero to indicate the seek occurred. */
static int
Io_fdseek(struct Io *io){
	int (*op)(int, void *, size_t);
	
	if(!fdisstd(io->fd) && lseek(io->fd, io->seek, SEEK_SET) != -1)
		return -1;

	/* repeated code to get the condition out of the loop */
	if(io->fl == write_flags){
		memset(io->buf, '\0', io->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->bufuse = write(io->fd, io->buf, MIN(io->bs, io->seek)))
					== 0)
				/* second chance */
				io->bufuse = write(io->fd, io->buf, MIN(io->bs, io->seek));
		}while((io->seek -= io->bufuse) > 0 && io->bufuse != 0);
	}else if(io->fl == read_flags){
		do{
			if((io->bufuse = read(io->fd, io->buf, MIN(io->bs, io->seek)))
					== 0)
				/* second chance */
				io->bufuse = read(io->fd, io->buf, MIN(io->bs, io->seek));
		}while((io->seek -= io->bufuse) > 0 && io->bufuse != 0);
	}else
		return EX_SOFTWARE;

	io->bufuse = 0;

	return -1;
}

/* 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 io->bufuse is
 * 0, errno will probably be set. Returns io. */
static struct Io *
Io_read(struct Io *io){

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

	return io;
}

/* Sets the variables in a struct *io to the defaults. Identifies the read/
 * write ends of the "pipe" by checking io->fl. Returns io. */
static struct Io *
Io_setdefaults(struct Io *io){

	io->bs = 1024 /* bytes; 1 KiB */; /* GNU dd(1) default; POSIX says 512B */
	io->buf = NULL;
	io->bytes = 0;
	io->fd = (io->fl == read_flags) ? STDIN_FILENO : STDOUT_FILENO;
	io->fn = (io->fl == read_flags) ? stdin_name   : stdout_name;
	io->prec = 0;
	io->rec = 0;
	io->seek = 0;

	return io;
}

/* Writes io->bufuse 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;

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

	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
output(struct Io io[2], char *fmt){

	fprintf(stderr, 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(void){

	fprintf(stderr, "Usage: %s (-AHn) (-a [byte]) (-c [count])\n"
		"\t(-i [input file]) (-b [input block size]) (-s [input offset])\n"
		"\t(-o [output file]) (-B [output block size]) (-S [output 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 */
	struct Io ep[2]; /* "engineered pipe"; also "extended play", for the DJ */
	char *fmt_output; /* fmt_asv (default) or fmt_human (-H) */
	char noerror; /* 0=exits (default) 1=retries on partial reads or writes */
	int c;
	int i;

	/* Set defaults. */
	align = -1;
	count = 0;
	fmt_output = fmt_asv;
	noerror = 0;
	ep[0].fl = read_flags;
	Io_setdefaults(&ep[0]);
	ep[1].fl = write_flags;
	Io_setdefaults(&ep[1]);

	if(argc > 0){
		program_name = argv[0];
		while((c = getopt(argc, argv, "a:Ab:B:c:i:hHns:S:o:")) != -1)
			switch(c){
			case 'i': case 'o':
				i = (c == 'o');
				if(optarg[0] == '-' && optarg[1] == '\0'){ /* optarg == "-" */
					ep[i].fd = (i == 0) ? STDIN_FILENO : STDOUT_FILENO;
					ep[i].fn = (i == 0) ? stdin_name   : stdout_name;
					break;
				}else if(Io_fdopen(&ep[i], optarg) != -1)
					break;
				terminate(ep);
				return oserr(optarg);
			case 'A': align = '\0'; break;
			case 'n': noerror = 1; break;
			case 'H': fmt_output = 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':
				if(c == 'c' && (count = parse(optarg)) >= 0)
					break;
				i = isupper(c);
				c = tolower(c);
				if((c == 'b' && (ep[i].bs = parse(optarg)) > 0)
						|| (c == 's' && (ep[i].seek = parse(optarg)) >= 0))
					break;
				/* FALLTHROUGH */
			default:
				terminate(ep);
				return usage();
			}
	}

	if(argc > optind){
		terminate(ep);
		return usage();
	}

	for(i = 0; i <= 1; ++i){
		if(Io_bufalloc(&ep[i]) == NULL){
			fprintf(stderr, "%s: Failed to allocate %d bytes\n",
				program_name, ep[i].bs);
			terminate(ep);
			return EX_OSERR;
		}else if(ep[i].seek > 0)
			switch(Io_fdseek(&ep[i])){
			case EX_OK:
				output(ep, fmt_output);
				terminate(ep);
				return EX_OK;
			}
	}

	do{	/* read */
		Io_read(&ep[0]);
		if(!noerror && ep[0].bufuse == 0)
			Io_read(&ep[0]); /* second chance */
		if(ep[0].bufuse == 0) /* that's all she wrote */
			break;
		else if(ep[0].bufuse < ep[0].bs){
			++ep[0].prec;
			fprintf(stderr, "%s: Partial read:\n\t", program_name);
			output(ep, fmt_output);
			if(!noerror)
				count = 1;
			if(align >= 0)
				Io_bufrpad(&ep[0], align);
		}else
			++ep[0].rec;

		/* write */
		do{	if(ep[1].bs > ep[0].bs){ /* io[1].bs > io[0].bs */
				Io_bufxapp(&ep[1], &ep[0]);
				if(ep[0].bs + ep[1].bufuse <= ep[1].bs && count != 1)
					continue; /* we could write more */
			}else
				Io_bufxfer(&ep[1], &ep[0], MIN(ep[0].bufuse, ep[1].bs));

			c = ep[1].bufuse;
			Io_write(&ep[1]);
			if(!noerror && ep[1].bufuse == c)
				Io_write(&ep[1]); /* second chance */
			if(c == ep[1].bufuse){ /* no more love */
				count = 1;
				break;
			}else if(c > ep[1].bufuse && ep[1].bufuse > 0){
				ep[1].prec += 1;
				fprintf(stderr, "%s: Partial write:\n\t", program_name);
				output(ep, fmt_output);
				if(!noerror)
					count = 1;
			}else if(ep[1].bufuse == 0 && c < ep[1].bs)
				++ep[1].prec;
			else
				++ep[1].rec;
		}while(ep[0].bufuse > 0);
	}while(count == 0 || --count > 0);

	output(ep, fmt_output);
	terminate(ep);

	return EX_OK;
}