/*	$OpenBSD: frontend_lpr.c,v 1.2 2019/04/04 19:25:45 eric Exp $	*/

/*
 * Copyright (c) 2017 Eric Faurot <eric@openbsd.org>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>

#include <ctype.h>
#include <errno.h>
#include <limits.h>
#include <netdb.h>
#include <stdarg.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>

#include "lpd.h"
#include "lp.h"

#include "io.h"
#include "log.h"
#include "proc.h"

#define SERVER_TIMEOUT	30000
#define CLIENT_TIMEOUT	 5000

#define	MAXARG	50

#define	F_ZOMBIE	0x1
#define	F_WAITADDRINFO	0x2

#define STATE_READ_COMMAND	0
#define STATE_READ_FILE		1

struct lpr_conn {
	SPLAY_ENTRY(lpr_conn)	 entry;
	uint32_t		 id;
	char			 hostname[NI_MAXHOST];
	struct io		*io;
	int			 state;
	int			 flags;
	int			 recvjob;
	int			 recvcf;
	size_t			 expect;
	FILE			*ofp;	/* output file when receiving data */
	int			 ifd;	/* input file for displayq/rmjob */

	char			*cmd;
	int			 ai_done;
	struct addrinfo		*ai;
	struct io		*iofwd;
};

SPLAY_HEAD(lpr_conn_tree, lpr_conn);

static int lpr_conn_cmp(struct lpr_conn *, struct lpr_conn *);
SPLAY_PROTOTYPE(lpr_conn_tree, lpr_conn, entry, lpr_conn_cmp);

static void lpr_on_allowedhost(struct lpr_conn *, const char *, const char *);
static void lpr_on_recvjob(struct lpr_conn *, int);
static void lpr_on_recvjob_file(struct lpr_conn *, int, size_t, int, int);
static void lpr_on_request(struct lpr_conn *, int, const char *, const char *);
static void lpr_on_getaddrinfo(void *, int, struct addrinfo *);

static void lpr_io_dispatch(struct io *, int, void *);
static int  lpr_readcommand(struct lpr_conn *);
static int  lpr_readfile(struct lpr_conn *);
static int  lpr_parsejobfilter(struct lpr_conn *, struct lp_jobfilter *,
    int, char **);

static void lpr_free(struct lpr_conn *);
static void lpr_close(struct lpr_conn *);
static void lpr_ack(struct lpr_conn *, char);
static void lpr_reply(struct lpr_conn *, const char *);
static void lpr_stream(struct lpr_conn *);
static void lpr_forward(struct lpr_conn *);

static void lpr_iofwd_dispatch(struct io *, int, void *);

static struct lpr_conn_tree conns;

void
lpr_init(void)
{
	SPLAY_INIT(&conns);
}

void
lpr_conn(uint32_t connid, struct listener *l, int sock,
    const struct sockaddr *sa)
{
	struct lpr_conn *conn;

	if ((conn = calloc(1, sizeof(*conn))) == NULL) {
		log_warn("%s: calloc", __func__);
		close(sock);
		frontend_conn_closed(connid);
		return;
	}
	conn->id = connid;
	conn->ifd = -1;
	conn->io = io_new();
	if (conn->io == NULL) {
		log_warn("%s: io_new", __func__);
		free(conn);
		close(sock);
		frontend_conn_closed(connid);
		return;
	}
	SPLAY_INSERT(lpr_conn_tree, &conns, conn);
	io_set_callback(conn->io, lpr_io_dispatch, conn);
	io_set_timeout(conn->io, CLIENT_TIMEOUT);
	io_set_write(conn->io);
	io_attach(conn->io, sock);

	conn->state = STATE_READ_COMMAND;
	m_create(p_engine, IMSG_LPR_ALLOWEDHOST, conn->id, 0, -1);
	m_add_sockaddr(p_engine, sa);
	m_close(p_engine);
}

void
lpr_dispatch_engine(struct imsgproc *proc, struct imsg *imsg)
{
	struct lpr_conn *conn = NULL, key;
	const char *hostname, *reject, *cmd;
	size_t sz;
	int ack, cf = 0;

	key.id = imsg->hdr.peerid;
	if (key.id) {
		conn = SPLAY_FIND(lpr_conn_tree, &conns, &key);
		if (conn == NULL) {
			log_debug("%08x dead-session", key.id);
			if (imsg->fd != -1)
				close(imsg->fd);
			return;
		}
	}

	switch (imsg->hdr.type) {
	case IMSG_LPR_ALLOWEDHOST:
		m_get_string(proc, &hostname);
		m_get_string(proc, &reject);
		m_end(proc);
		lpr_on_allowedhost(conn, hostname, reject);
		break;

	case IMSG_LPR_RECVJOB:
		m_get_int(proc, &ack);
		m_end(proc);
		lpr_on_recvjob(conn, ack);
		break;

	case IMSG_LPR_RECVJOB_CF:
		cf = 1;
	case IMSG_LPR_RECVJOB_DF:
		m_get_int(proc, &ack);
		m_get_size(proc, &sz);
		m_end(proc);
		lpr_on_recvjob_file(conn, ack, sz, cf, imsg->fd);
		break;

	case IMSG_LPR_DISPLAYQ:
	case IMSG_LPR_RMJOB:
		m_get_string(proc, &hostname);
		m_get_string(proc, &cmd);
		m_end(proc);
		lpr_on_request(conn, imsg->fd, hostname, cmd);
		break;

	default:
		fatalx("%s: unexpected imsg %s", __func__,
		    log_fmt_imsgtype(imsg->hdr.type));
	}
}

static void
lpr_on_allowedhost(struct lpr_conn *conn, const char *hostname,
    const char *reject)
{
	strlcpy(conn->hostname, hostname, sizeof(conn->hostname));
	if (reject)
		lpr_reply(conn, reject);
	else
		io_set_read(conn->io);
}

static void
lpr_on_recvjob(struct lpr_conn *conn, int ack)
{
	if (ack == LPR_ACK)
		conn->recvjob = 1;
	else
		log_debug("%08x recvjob failed", conn->id);
	lpr_ack(conn, ack);
}

static void
lpr_on_recvjob_file(struct lpr_conn *conn, int ack, size_t sz, int cf, int fd)
{
	if (ack != LPR_ACK) {
		lpr_ack(conn, ack);
		return;
	}

	if (fd == -1) {
		log_warnx("%s: failed to get fd", __func__);
		lpr_ack(conn, LPR_NACK);
		return;
	}

	conn->ofp = fdopen(fd, "w");
	if (conn->ofp == NULL) {
		log_warn("%s: fdopen", __func__);
		close(fd);
		lpr_ack(conn, LPR_NACK);
		return;
	}

	conn->expect = sz;
	if (cf)
		conn->recvcf = cf;
	conn->state = STATE_READ_FILE;

	lpr_ack(conn, LPR_ACK);
}

static void
lpr_on_request(struct lpr_conn *conn, int fd, const char *hostname,
    const char *cmd)
{
	struct addrinfo hints;

	if (fd == -1) {
		log_warnx("%s: no fd received", __func__);
		lpr_close(conn);
		return;
	}

	log_debug("%08x stream init", conn->id);
	conn->ifd = fd;

	/* Prepare for command forwarding if necessary. */
	if (cmd) {
		log_debug("%08x forwarding to %s: \\%d%s", conn->id, hostname,
		    cmd[0], cmd + 1);
		conn->cmd = strdup(cmd);
		if (conn->cmd == NULL)
			log_warn("%s: strdup", __func__);
		else {
			memset(&hints, 0, sizeof(hints));
			hints.ai_socktype = SOCK_STREAM;
			conn->flags |= F_WAITADDRINFO;
			/*
			 * The callback might run immediatly, so conn->ifd
			 * must be set before, to block lpr_forward().
			 */
			resolver_getaddrinfo(hostname, "printer", &hints,
			    lpr_on_getaddrinfo, conn);
		}
	}

	lpr_stream(conn);
}

static void
lpr_on_getaddrinfo(void *arg, int r, struct addrinfo *ai)
{
	struct lpr_conn *conn = arg;

	conn->flags &= ~F_WAITADDRINFO;
	if (conn->flags & F_ZOMBIE) {
		if (ai)
			freeaddrinfo(ai);
		lpr_free(conn);
	}
	else {
		conn->ai_done = 1;
		conn->ai = ai;
		lpr_forward(conn);
	}
}

static void
lpr_io_dispatch(struct io *io, int evt, void *arg)
{
	struct lpr_conn *conn = arg;
	int r;

	switch (evt) {
	case IO_DATAIN:
		switch(conn->state) {
		case STATE_READ_COMMAND:
			r = lpr_readcommand(conn);
			break;
		case STATE_READ_FILE:
			r = lpr_readfile(conn);
			break;
		default:
			fatal("%s: unexpected state %d", __func__, conn->state);
		}

		if (r == 0)
			io_set_write(conn->io);
		return;

	case IO_LOWAT:
		if (conn->recvjob)
			io_set_read(conn->io);
		else if (conn->ifd != -1)
			lpr_stream(conn);
		else if (conn->cmd == NULL)
			lpr_close(conn);
		return;

	case IO_DISCONNECTED:
		log_debug("%08x disconnected", conn->id);
		/*
		 * Some clients don't wait for the last acknowledgment to close
		 * the session.  So just consider it is closed normally.
		 */
	case IO_CLOSED:
		if (conn->recvcf && conn->state == STATE_READ_COMMAND) {
			/*
			 * Commit the transaction if we received a control file
			 * and the last file was received correctly.
			 */
			m_compose(p_engine, IMSG_LPR_RECVJOB_COMMIT, conn->id,
			    0, -1, NULL, 0);
			conn->recvjob = 0;
		}
		break;

	case IO_TIMEOUT:
		log_debug("%08x timeout", conn->id);
		break;

	case IO_ERROR:
		log_debug("%08x io-error", conn->id);
		break;

	default:
		fatalx("%s: unexpected event %d", __func__, evt);
	}

	lpr_close(conn);
}

static int
lpr_readcommand(struct lpr_conn *conn)
{
	struct lp_jobfilter jf;
	size_t count;
	const char *errstr;
	char *argv[MAXARG], *line;
	int i, argc, cmd;

	line = io_getline(conn->io, NULL);
	if (line == NULL) {
		if (io_datalen(conn->io) >= LPR_MAXCMDLEN) {
			lpr_reply(conn, "Request line too long");
			return 0;
		}
		return -1;
	}

	cmd = line[0];
	line++;

	if (cmd == 0) {
		lpr_reply(conn, "No command");
		return 0;
	}

	log_debug("%08x cmd \\%d", conn->id, cmd);

	/* Parse the command. */
	for (argc = 0; argc < MAXARG; ) {
		argv[argc] = strsep(&line, " \t");
		if (argv[argc] == NULL)
			break;
		if (argv[argc][0] != '\0')
			argc++;
	}
	if (argc == MAXARG) {
		lpr_reply(conn, "Argument list too long");
		return 0;
	}

	if (argc == 0) {
		lpr_reply(conn, "No queue specified");
		return 0;
	}

#define CMD(c)    ((int)(c))
#define SUBCMD(c) (0x100 | (int)(c))

	if (conn->recvjob)
		cmd |= 0x100;
	switch (cmd) {
	case CMD('\1'):	/* PRINT <prn> */
		m_create(p_engine, IMSG_LPR_PRINTJOB, 0, 0, -1);
		m_add_string(p_engine, argv[0]);
		m_close(p_engine);
		lpr_ack(conn, LPR_ACK);
		return 0;

	case CMD('\2'):	/* RECEIVE JOB <prn> */
		m_create(p_engine, IMSG_LPR_RECVJOB, conn->id, 0, -1);
		m_add_string(p_engine, conn->hostname);
		m_add_string(p_engine, argv[0]);
		m_close(p_engine);
		return 0;

	case CMD('\3'):	/* QUEUE STATE SHORT <prn> [job#...] [user..] */
	case CMD('\4'):	/* QUEUE STATE LONG  <prn> [job#...] [user..] */
		if (lpr_parsejobfilter(conn, &jf, argc - 1, argv + 1) == -1)
			return 0;

		m_create(p_engine, IMSG_LPR_DISPLAYQ, conn->id, 0, -1);
		m_add_int(p_engine, (cmd == '\3') ? 0 : 1);
		m_add_string(p_engine, conn->hostname);
		m_add_string(p_engine, argv[0]);
		m_add_int(p_engine, jf.njob);
		for (i = 0; i < jf.njob; i++)
			m_add_int(p_engine, jf.jobs[i]);
		m_add_int(p_engine, jf.nuser);
		for (i = 0; i < jf.nuser; i++)
			m_add_string(p_engine, jf.users[i]);
		m_close(p_engine);
		return 0;

	case CMD('\5'):	/* REMOVE JOBS <prn> <agent> [job#...] [user..] */
		if (argc < 2) {
			lpr_reply(conn, "No agent specified");
			return 0;
		}
		if (lpr_parsejobfilter(conn, &jf, argc - 2, argv + 2) == -1)
			return 0;

		m_create(p_engine, IMSG_LPR_RMJOB, conn->id, 0, -1);
		m_add_string(p_engine, conn->hostname);
		m_add_string(p_engine, argv[0]);
		m_add_string(p_engine, argv[1]);
		m_add_int(p_engine, jf.njob);
		for (i = 0; i < jf.njob; i++)
			m_add_int(p_engine, jf.jobs[i]);
		m_add_int(p_engine, jf.nuser);
		for (i = 0; i < jf.nuser; i++)
			m_add_string(p_engine, jf.users[i]);
		m_close(p_engine);
		return 0;

	case SUBCMD('\1'):	/* ABORT */
		m_compose(p_engine, IMSG_LPR_RECVJOB_CLEAR, conn->id, 0, -1,
		    NULL, 0);
		conn->recvcf = 0;
		lpr_ack(conn, LPR_ACK);
		return 0;

	case SUBCMD('\2'):	/* CONTROL FILE <size> <filename> */
	case SUBCMD('\3'):	/* DATA FILE    <size> <filename> */
		if (argc != 2) {
			log_debug("%08x invalid number of argument", conn->id);
			lpr_ack(conn, LPR_NACK);
			return 0;
		}
		errstr = NULL;
		count = strtonum(argv[0], 1, LPR_MAXFILESIZE, &errstr);
		if (errstr) {
			log_debug("%08x invalid file size: %s", conn->id,
			    strerror(errno));
			lpr_ack(conn, LPR_NACK);
			return 0;
		}

		if (cmd == SUBCMD('\2')) {
			if (conn->recvcf) {
				log_debug("%08x cf file already received",
				    conn->id);
				lpr_ack(conn, LPR_NACK);
				return 0;
			}
			m_create(p_engine, IMSG_LPR_RECVJOB_CF, conn->id, 0,
			    -1);
		}
		else
			m_create(p_engine, IMSG_LPR_RECVJOB_DF, conn->id, 0,
			    -1);
		m_add_size(p_engine, count);
		m_add_string(p_engine, argv[1]);
		m_close(p_engine);
		return 0;

	default:
		if (conn->recvjob)
			lpr_reply(conn, "Protocol error");
		else
			lpr_reply(conn, "Illegal service request");
		return 0;
	}
}

static int
lpr_readfile(struct lpr_conn *conn)
{
	size_t len, w;
	char *data;

	if (conn->expect) {
		/* Read file content. */
		data = io_data(conn->io);
		len = io_datalen(conn->io);
		if (len > conn->expect)
			len = conn->expect;

		log_debug("%08x %zu bytes received", conn->id, len);

		w = fwrite(data, 1, len, conn->ofp);
		if (w != len) {
			log_warnx("%s: fwrite", __func__);
			lpr_close(conn);
			return -1;
		}
		io_drop(conn->io, w);
		conn->expect -= w;
		if (conn->expect)
			return -1;

		fclose(conn->ofp);
		conn->ofp = NULL;

		log_debug("%08x file received", conn->id);
	}

	/* Try to read '\0'. */
	len = io_datalen(conn->io);
	if (len == 0)
		return -1;
	data = io_data(conn->io);
	io_drop(conn->io, 1);

	log_debug("%08x eof %d", conn->id, (int)*data);

	if (*data != '\0') {
		lpr_close(conn);
		return -1;
	}

	conn->state = STATE_READ_COMMAND;
	lpr_ack(conn, LPR_ACK);
	return 0;
}

static int
lpr_parsejobfilter(struct lpr_conn *conn, struct lp_jobfilter *jf, int argc,
    char **argv)
{
	const char *errstr;
	char *arg;
	int i, jobnum;

	memset(jf, 0, sizeof(*jf));

	for (i = 0; i < argc; i++) {
		arg = argv[i];
		if (isdigit((unsigned int)arg[0])) {
			if (jf->njob == LP_MAXREQUESTS) {
				lpr_reply(conn, "Too many requests");
				return -1;
			}
			errstr = NULL;
			jobnum = strtonum(arg, 0, INT_MAX, &errstr);
			if (errstr) {
				lpr_reply(conn, "Invalid job number");
				return -1;
			}
			jf->jobs[jf->njob++] = jobnum;
		}
		else {
			if (jf->nuser == LP_MAXUSERS) {
				lpr_reply(conn, "Too many users");
				return -1;
			}
			jf->users[jf->nuser++] = arg;
		}
	}

	return 0;
}

static void
lpr_free(struct lpr_conn *conn)
{
	if ((conn->flags & F_WAITADDRINFO) == 0)
		free(conn);
}

static void
lpr_close(struct lpr_conn *conn)
{
	uint32_t connid = conn->id;

	SPLAY_REMOVE(lpr_conn_tree, &conns, conn);

	if (conn->recvjob)
		m_compose(p_engine, IMSG_LPR_RECVJOB_ROLLBACK, conn->id, 0, -1,
		    NULL, 0);

	io_free(conn->io);
	free(conn->cmd);
	if (conn->ofp)
		fclose(conn->ofp);
	if (conn->ifd != -1)
		close(conn->ifd);
	if (conn->ai)
		freeaddrinfo(conn->ai);
	if (conn->iofwd)
		io_free(conn->iofwd);

	conn->flags |= F_ZOMBIE;
	lpr_free(conn);

	frontend_conn_closed(connid);
}

static void
lpr_ack(struct lpr_conn *conn, char c)
{
	if (c == 0)
		log_debug("%08x ack", conn->id);
	else
		log_debug("%08x nack %d", conn->id, (int)c);

	io_write(conn->io, &c, 1);
}

static void
lpr_reply(struct lpr_conn *conn, const char *s)
{
	log_debug("%08x reply: %s", conn->id, s);

	io_printf(conn->io, "%s\n", s);
}

/*
 * Stream reponse file to the client.
 */
static void
lpr_stream(struct lpr_conn *conn)
{
	char buf[BUFSIZ];
	ssize_t r;

	for (;;) {
		if (io_queued(conn->io) > 65536)
			return;

		r = read(conn->ifd, buf, sizeof(buf));
		if (r == -1) {
			if (errno == EINTR)
				continue;
			log_warn("%s: read", __func__);
			break;
		}

		if (r == 0) {
			log_debug("%08x stream done", conn->id);
			break;
		}
		log_debug("%08x stream %zu bytes", conn->id, r);

		if (io_write(conn->io, buf, r) == -1) {
			log_warn("%s: io_write", __func__);
			break;
		}
	}

	close(conn->ifd);
	conn->ifd = -1;

	if (conn->cmd)
		lpr_forward(conn);

	else if (io_queued(conn->io) == 0)
		lpr_close(conn);
}

/*
 * Forward request to the remote printer.
 */
static void
lpr_forward(struct lpr_conn *conn)
{
	/*
	 * Do not start forwarding the command if the address is not resolved
	 * or if the local response is still being sent to the client.
	 */
	if (!conn->ai_done || conn->ifd == -1)
		return;

	if (conn->ai == NULL) {
		if (io_queued(conn->io) == 0)
			lpr_close(conn);
		return;
	}

	log_debug("%08x forward start", conn->id);

	conn->iofwd = io_new();
	if (conn->iofwd == NULL) {
		log_warn("%s: io_new", __func__);
		if (io_queued(conn->io) == 0)
			lpr_close(conn);
		return;
	}
	io_set_callback(conn->iofwd, lpr_iofwd_dispatch, conn);
	io_set_timeout(conn->io, SERVER_TIMEOUT);
	io_connect(conn->iofwd, conn->ai);
	conn->ai = NULL;
}

static void
lpr_iofwd_dispatch(struct io *io, int evt, void *arg)
{
	struct lpr_conn *conn = arg;

	switch (evt) {
	case IO_CONNECTED:
		log_debug("%08x forward connected", conn->id);
		/* Send the request. */
		io_print(io, conn->cmd);
		io_print(io, "\n");
		io_set_write(io);
		return;

	case IO_DATAIN:
		/* Relay. */
		io_write(conn->io, io_data(io), io_datalen(io));
		io_drop(io, io_datalen(io));
		return;

	case IO_LOWAT:
		/* Read response. */
		io_set_read(io);
		return;

	case IO_CLOSED:
		break;

	case IO_DISCONNECTED:
		log_debug("%08x forward disconnected", conn->id);
		break;

	case IO_TIMEOUT:
		log_debug("%08x forward timeout", conn->id);
		break;

	case IO_ERROR:
		log_debug("%08x forward io-error", conn->id);
		break;

	default:
		fatalx("%s: unexpected event %d", __func__, evt);
	}

	log_debug("%08x forward done", conn->id);

	io_free(io);
	free(conn->cmd);
	conn->cmd = NULL;
	conn->iofwd = NULL;
	if (io_queued(conn->io) == 0)
		lpr_close(conn);
}

static int
lpr_conn_cmp(struct lpr_conn *a, struct lpr_conn *b)
{
	if (a->id < b->id)
		return -1;
	if (a->id > b->id)
		return 1;
	return 0;
}

SPLAY_GENERATE(lpr_conn_tree, lpr_conn, entry, lpr_conn_cmp);