/*
 * Copyright (c) 2017 Antoine Kaufmann <toni@famkaufmann.info>
 *
 * 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/un.h>

#include <inttypes.h>
#include <stdlib.h>
#include <string.h>

#include "smtpd.h"
#include "log.h"

/*
 * The PROXYv2 protocol is described here:
 * http://www.haproxy.org/download/1.8/doc/proxy-protocol.txt
 */

#define PROXY_CLOCAL 0x0
#define PROXY_CPROXY 0x1
#define PROXY_AF_UNSPEC 0x0
#define PROXY_AF_INET 0x1
#define PROXY_AF_INET6 0x2
#define PROXY_AF_UNIX 0x3
#define PROXY_TF_UNSPEC 0x0
#define PROXY_TF_STREAM 0x1
#define PROXY_TF_DGRAM 0x2

#define PROXY_SESSION_TIMEOUT	300

static const uint8_t pv2_signature[] = {
	0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D,
	0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A
};

struct proxy_hdr_v2 {
	uint8_t sig[12];
	uint8_t ver_cmd;
	uint8_t fam;
	uint16_t len;
} __attribute__((packed));


struct proxy_addr_ipv4 {
	uint32_t src_addr;
	uint32_t dst_addr;
	uint16_t src_port;
	uint16_t dst_port;
} __attribute__((packed));

struct proxy_addr_ipv6 {
	uint8_t src_addr[16];
	uint8_t dst_addr[16];
	uint16_t src_port;
	uint16_t dst_port;
} __attribute__((packed));

struct proxy_addr_unix {
	uint8_t src_addr[108];
	uint8_t dst_addr[108];
} __attribute__((packed));

union proxy_addr {
	struct proxy_addr_ipv4 ipv4;
	struct proxy_addr_ipv6 ipv6;
	struct proxy_addr_unix un;
} __attribute__((packed));

struct proxy_session {
	struct listener	*l;
	struct io	*io;

	uint64_t	id;
	int		fd;
	uint16_t	header_len;
	uint16_t	header_total;
	uint16_t	addr_len;
	uint16_t	addr_total;

	struct sockaddr_storage	ss;
	struct proxy_hdr_v2	hdr;
	union proxy_addr	addr;

	void (*cb_accepted)(struct listener *, int,
	    const struct sockaddr_storage *, struct io *);
	void (*cb_dropped)(struct listener *, int,
	    const struct sockaddr_storage *);
};

static void proxy_io(struct io *, int, void *);
static void proxy_error(struct proxy_session *, const char *, const char *);
static int proxy_header_validate(struct proxy_session *);
static int proxy_translate_ss(struct proxy_session *);

int
proxy_session(struct listener *listener, int sock,
    const struct sockaddr_storage *ss,
    void (*accepted)(struct listener *, int,
	const struct sockaddr_storage *, struct io *),
    void (*dropped)(struct listener *, int,
	const struct sockaddr_storage *));

int
proxy_session(struct listener *listener, int sock,
    const struct sockaddr_storage *ss,
    void (*accepted)(struct listener *, int,
	const struct sockaddr_storage *, struct io *),
    void (*dropped)(struct listener *, int,
	const struct sockaddr_storage *))
{
	struct proxy_session *s;

	if ((s = calloc(1, sizeof(*s))) == NULL)
		return (-1);

	if ((s->io = io_new()) == NULL) {
		free(s);
		return (-1);
	}

	s->id = generate_uid();
	s->l = listener;
	s->fd = sock;
	s->header_len = 0;
	s->addr_len = 0;
	s->ss = *ss;
	s->cb_accepted = accepted;
	s->cb_dropped = dropped;

	io_set_callback(s->io, proxy_io, s);
	io_set_fd(s->io, sock);
	io_set_timeout(s->io, PROXY_SESSION_TIMEOUT * 1000);
	io_set_read(s->io);

	log_info("%016"PRIx64" smtp event=proxy address=%s",
	    s->id, ss_to_text(&s->ss));

	return 0;
}

static void
proxy_io(struct io *io, int evt, void *arg)
{
	struct proxy_session	*s = arg;
	struct proxy_hdr_v2	*h = &s->hdr;
	uint8_t *buf;
	size_t len, off;

	switch (evt) {

	case IO_DATAIN:
		buf = io_data(io);
		len = io_datalen(io);

		if (s->header_len < sizeof(s->hdr)) {
			/* header is incomplete */
			off = sizeof(s->hdr) - s->header_len;
			off = (len < off ? len : off);
			memcpy((uint8_t *) &s->hdr + s->header_len, buf, off);

			s->header_len += off;
			buf += off;
			len -= off;
			io_drop(s->io, off);

			if (s->header_len < sizeof(s->hdr)) {
				/* header is still not complete */
				return;
			}

			if (proxy_header_validate(s) != 0)
				return;
		}

		if (s->addr_len < s->addr_total) {
			/* address is incomplete */
			off = s->addr_total - s->addr_len;
			off = (len < off ? len : off);
			memcpy((uint8_t *) &s->addr + s->addr_len, buf, off);

			s->header_len += off;
			s->addr_len += off;
			buf += off;
			len -= off;
			io_drop(s->io, off);

			if (s->addr_len < s->addr_total) {
				/* address is still not complete */
				return;
			}
		}

		if (s->header_len < s->header_total) {
			/* additional parameters not complete */
			/* these are ignored for now, but we still need to drop
			 * the bytes from the buffer */
			off = s->header_total - s->header_len;
			off = (len < off ? len : off);

			s->header_len += off;
			io_drop(s->io, off);

			if (s->header_len < s->header_total)
				/* not complete yet */
				return;
		}

		switch(h->ver_cmd & 0xF) {
		case PROXY_CLOCAL:
			/* local address, no need to modify ss */
			break;

		case PROXY_CPROXY:
			if (proxy_translate_ss(s) != 0)
				return;
			break;

		default:
			proxy_error(s, "protocol error", "unknown command");
			return;
		}

		s->cb_accepted(s->l, s->fd, &s->ss, s->io);
		/* we passed off s->io, so it does not need to be freed here */
		free(s);
		break;

	case IO_TIMEOUT:
		proxy_error(s, "timeout", NULL);
		break;

	case IO_DISCONNECTED:
		proxy_error(s, "disconnected", NULL);
		break;

	case IO_ERROR:
		proxy_error(s, "IO error", io_error(io));
		break;

	default:
		fatalx("proxy_io()");
	}

}

static void
proxy_error(struct proxy_session *s, const char *reason, const char *extra)
{
	if (extra)
		log_info("proxy %p event=closed address=%s reason=\"%s (%s)\"",
		    s, ss_to_text(&s->ss), reason, extra);
	else
		log_info("proxy %p event=closed address=%s reason=\"%s\"",
		    s, ss_to_text(&s->ss), reason);

	s->cb_dropped(s->l, s->fd, &s->ss);
	io_free(s->io);
	free(s);
}

static int
proxy_header_validate(struct proxy_session *s)
{
	struct proxy_hdr_v2 *h = &s->hdr;

	if (memcmp(h->sig, pv2_signature,
		sizeof(pv2_signature)) != 0) {
		proxy_error(s, "protocol error", "invalid signature");
		return (-1);
	}

	if ((h->ver_cmd >> 4) != 2) {
		proxy_error(s, "protocol error", "invalid version");
		return (-1);
	}

	switch (h->fam) {
	case (PROXY_AF_UNSPEC << 4 | PROXY_TF_UNSPEC):
		s->addr_total = 0;
		break;

	case (PROXY_AF_INET << 4 | PROXY_TF_STREAM):
		s->addr_total = sizeof(s->addr.ipv4);
		break;

	case (PROXY_AF_INET6 << 4 | PROXY_TF_STREAM):
		s->addr_total = sizeof(s->addr.ipv6);
		break;

	case (PROXY_AF_UNIX << 4 | PROXY_TF_STREAM):
		s->addr_total = sizeof(s->addr.un);
		break;

	default:
		proxy_error(s, "protocol error", "unsupported address family");
		return (-1);
	}

	s->header_total = ntohs(h->len);
	if (s->header_total > UINT16_MAX - sizeof(struct proxy_hdr_v2)) {
		proxy_error(s, "protocol error", "header too long");
		return (-1);
	}
	s->header_total += sizeof(struct proxy_hdr_v2);

	if (s->header_total < sizeof(struct proxy_hdr_v2) + s->addr_total) {
		proxy_error(s, "protocol error", "address info too short");
		return (-1);
	}

	return 0;
}

static int
proxy_translate_ss(struct proxy_session *s)
{
	struct sockaddr_in *sin = (struct sockaddr_in *) &s->ss;
	struct sockaddr_in6 *sin6 = (struct sockaddr_in6 *) &s->ss;
	struct sockaddr_un *sun = (struct sockaddr_un *) &s->ss;
	size_t sun_len;

	switch (s->hdr.fam) {
	case (PROXY_AF_UNSPEC << 4 | PROXY_TF_UNSPEC):
		/* unspec: only supported for local */
		proxy_error(s, "address translation", "UNSPEC family not "
		    "supported for PROXYing");
		return (-1);

	case (PROXY_AF_INET << 4 | PROXY_TF_STREAM):
		memset(&s->ss, 0, sizeof(s->ss));
		sin->sin_family = AF_INET;
		sin->sin_port = s->addr.ipv4.src_port;
		sin->sin_addr.s_addr = s->addr.ipv4.src_addr;
		break;

	case (PROXY_AF_INET6 << 4 | PROXY_TF_STREAM):
		memset(&s->ss, 0, sizeof(s->ss));
		sin6->sin6_family = AF_INET6;
		sin6->sin6_port = s->addr.ipv6.src_port;
		memcpy(sin6->sin6_addr.s6_addr, s->addr.ipv6.src_addr,
		    sizeof(s->addr.ipv6.src_addr));
		break;

	case (PROXY_AF_UNIX << 4 | PROXY_TF_STREAM):
		memset(&s->ss, 0, sizeof(s->ss));
		sun_len = strnlen(s->addr.un.src_addr,
		    sizeof(s->addr.un.src_addr));
		if (sun_len > sizeof(sun->sun_path)) {
			proxy_error(s, "address translation", "Unix socket path"
			    " longer than supported");
			return (-1);
		}
		sun->sun_family = AF_UNIX;
		memcpy(sun->sun_path, s->addr.un.src_addr, sun_len);
		break;

	default:
		fatalx("proxy_translate_ss()");
	}

	return 0;
}