/*
 * Copyright (C) 2025 OpenSIPS Solutions
 *
 * This file is part of opensips, a free SIP server.
 *
 * opensips is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * opensips 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 General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 *
 */

#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include "proxy_protocol.h"
#include "../socket_info.h"
#include "tcp_conn_defs.h"
#include "tcp_common.h"
#include "../ut.h"
#include "../tsend.h"

#define PROXY_PROTOCOL_V1_HDR "PROXY "
#define PROXY_PROTOCOL_V1_HDR_LEN (sizeof(PROXY_PROTOCOL_V1_HDR) - 1)
#define PROXY_PROTOCOL_V1_UNKN "PROXY UNKNOWN\r\n"
#define PROXY_PROTOCOL_V1_UNKN_LEN (sizeof(PROXY_PROTOCOL_V1_UNKN) - 1)
#define PROXY_PROTOCOL_TCP4 "TCP4 "
#define PROXY_PROTOCOL_TCP4_LEN (sizeof(PROXY_PROTOCOL_TCP4) - 1)
#define PROXY_PROTOCOL_TCP6 "TCP6 "
#define PROXY_PROTOCOL_TCP6_LEN (sizeof(PROXY_PROTOCOL_TCP6) - 1)
#define PROXY_PROTOCOL_UNKN "UNKNOWN"
#define PROXY_PROTOCOL_UNKN_LEN (sizeof(PROXY_PROTOCOL_UNKN) - 1)
#define PROXY_PROTOCOL_END "\r\n"
#define PROXY_PROTOCOL_END_LEN (sizeof(PROXY_PROTOCOL_END) - 1)
#define PROXY_PROTOCOL_SEP ' '
/*
 * these are defined in the specs:
 * https://www.haproxy.org/download/1.8/doc/proxy-protocol.txt
 */
#define PROXY_PROTOCOL_TCP4_MAX 56
#define PROXY_PROTOCOL_TCP4_PAYLOAD_MAX \
	(PROXY_PROTOCOL_TCP4_MAX - PROXY_PROTOCOL_V1_HDR_LEN - PROXY_PROTOCOL_TCP4_LEN)
#define PROXY_PROTOCOL_TCP6_MAX 104
#define PROXY_PROTOCOL_TCP6_PAYLOAD_MAX \
	(PROXY_PROTOCOL_TCP6_MAX - PROXY_PROTOCOL_V1_HDR_LEN - PROXY_PROTOCOL_TCP6_LEN)
#define PROXY_PROTOCOL_UNKN_MAX PROXY_PROTOCOL_BUF_MAX
#define PROXY_PROTOCOL_UNKN_PAYLOAD_MAX \
	(PROXY_PROTOCOL_UNKN_MAX - PROXY_PROTOCOL_V1_HDR_LEN - PROXY_PROTOCOL_UNKN_LEN)

enum proxy_tok_match {
	PP_TOK_MISMATCH = -1,
	PP_TOK_PARTIAL = 0,
	PP_TOK_MATCH = 1,
};

static inline enum proxy_tok_match match_proxy_token(const str *buf,
		const char *tok, int tok_len)
{
	int cmp_len;

	cmp_len = (buf->len < tok_len) ? buf->len : tok_len;
	if (strncasecmp(buf->s, tok, cmp_len) != 0)
		return PP_TOK_MISMATCH;
	if (buf->len < tok_len)
		return PP_TOK_PARTIAL;
	return PP_TOK_MATCH;
}

int is_net_proxy_protocol(char *buf, int size)
{
	if (size < PROXY_PROTOCOL_V1_HDR_LEN)
		return -1; /* don't know yet */
	if (memcmp(buf, PROXY_PROTOCOL_V1_HDR, PROXY_PROTOCOL_V1_HDR_LEN))
		return 0;
	/* TODO: support v2 as well? */
	return 1;
}

static int build_proxy_protocol_v1_hdr(const struct ip_addr *src_ip,
		unsigned short src_port, const struct ip_addr *dst_ip,
		unsigned short dst_port, char *buf, int size)
{
	const char *src_s, *dst_s, *proto;
	int len;

	if (!src_ip || !dst_ip || !buf || size <= 0)
		return -1;

	if (src_ip->af != dst_ip->af ||
			(src_ip->af != AF_INET && src_ip->af != AF_INET6)) {
		if (size < PROXY_PROTOCOL_V1_UNKN_LEN + 1)
			return -1;
		memcpy(buf, PROXY_PROTOCOL_V1_UNKN, PROXY_PROTOCOL_V1_UNKN_LEN + 1);
		return PROXY_PROTOCOL_V1_UNKN_LEN;
	}

	src_s = ip_addr2a((struct ip_addr *)src_ip);
	dst_s = ip_addr2a((struct ip_addr *)dst_ip);
	proto = (src_ip->af == AF_INET) ? "TCP4" : "TCP6";
	len = snprintf(buf, size, "PROXY %s %s %s %hu %hu\r\n",
			proto, src_s, dst_s, src_port, dst_port);
	if (len <= 0 || len >= size)
		return -1;

	return len;
}

int build_outbound_proxy_protocol_v1_hdr(const struct receive_info *ri,
		const struct ip_addr *fallback_src_ip,
		unsigned short fallback_src_port,
		const struct ip_addr *fallback_dst_ip,
		unsigned short fallback_dst_port,
		char *buf, int size)
{
	const struct ip_addr *src_ip = fallback_src_ip;
	const struct ip_addr *dst_ip = fallback_dst_ip;
	unsigned short src_port = fallback_src_port;
	unsigned short dst_port = fallback_dst_port;

	if (ri) {
		if (ri->real_ep.flags == PP_OK) {
			src_ip = &ri->real_ep.src_ip;
			src_port = ri->real_ep.src_port;
			dst_ip = &ri->real_ep.dst_ip;
			dst_port = ri->real_ep.dst_port;
		} else {
			src_ip = &ri->src_ip;
			src_port = ri->src_port;
			dst_ip = &ri->dst_ip;
			dst_port = ri->dst_port;
		}
	}

	return build_proxy_protocol_v1_hdr(src_ip, src_port, dst_ip, dst_port,
			buf, size);
}

static inline int should_send_stream_proxy_protocol(const struct tcp_connection *c)
{
	return c && c->rcv.bind_address &&
		(c->flags & F_CONN_ACCEPTED) == 0 &&
		(c->rcv.bind_address->flags & SI_PROXY_OUT) &&
		(c->flags & F_CONN_PROXY_OUT_SENT) == 0;
}

int send_stream_proxy_protocol_v1(struct tcp_connection *c, int fd,
		int write_timeout, int lock, const struct receive_info *ri,
		const char *proto_name)
{
	char pp_hdr[PROXY_PROTOCOL_BUF_MAX];
	int pp_len, rc, dbg_len;
	const char *action;

	if (!should_send_stream_proxy_protocol(c))
		return 0;

	if (fd < 0 && ((c->flags & F_CONN_INIT) == 0 || c->fd < 0))
		return 0;

	pp_len = build_outbound_proxy_protocol_v1_hdr(ri,
			&c->rcv.dst_ip, c->rcv.dst_port,
			&c->rcv.src_ip, c->rcv.src_port,
			pp_hdr, sizeof(pp_hdr));
	if (pp_len < 0) {
		LM_ERR("failed to build outbound PROXY header\n");
		return -1;
	}

	dbg_len = pp_len;
	if (dbg_len >= 2 && pp_hdr[dbg_len - 2] == '\r' &&
			pp_hdr[dbg_len - 1] == '\n')
		dbg_len -= 2;

	action = (fd < 0) ? "queueing" : "sending";
	LM_DBG("%s outbound PROXY header on %s conn %u: %.*s\n",
			action, proto_name ? proto_name : "stream", c->id,
			dbg_len, pp_hdr);

	if (fd < 0) {
		rc = tcp_async_add_chunk(c, pp_hdr, pp_len, lock);
	} else {
		if (lock)
			lock_get(&c->write_lock);
		rc = tsend_stream(fd, pp_hdr, pp_len, write_timeout);
		if (lock)
			lock_release(&c->write_lock);
		if (rc != pp_len) {
			LM_ERR("failed to send outbound PROXY header on %s\n",
					proto_name ? proto_name : "stream");
			return -1;
		}
		tcp_conn_reset_lifetime(c);
		rc = 0;
	}

	if (rc < 0) {
		LM_ERR("failed to send outbound PROXY header on %s\n",
				proto_name ? proto_name : "stream");
		return -1;
	}

	c->flags |= F_CONN_PROXY_OUT_SENT;
	return 0;
}

char *parse_net_proxy_protocol(char *buf, int size, struct proxy_protocol *proxy)
{
	int af, max;
	char *p, *end;
	str _buf, tmp;
	struct ip_addr *ip;
	unsigned int port;
	enum proxy_tok_match m;

	if (size < PROXY_PROTOCOL_V1_HDR_LEN ||
			memcmp(buf, PROXY_PROTOCOL_V1_HDR, PROXY_PROTOCOL_V1_HDR_LEN) != 0)
		goto error;

	_buf.s = buf + PROXY_PROTOCOL_V1_HDR_LEN;
	_buf.len = size - PROXY_PROTOCOL_V1_HDR_LEN;

	m = match_proxy_token(&_buf, PROXY_PROTOCOL_TCP4, PROXY_PROTOCOL_TCP4_LEN);
	if (m == PP_TOK_MATCH) {
		af = AF_INET;
		_buf.s += PROXY_PROTOCOL_TCP4_LEN;
		_buf.len -= PROXY_PROTOCOL_TCP4_LEN;
		if (_buf.len > PROXY_PROTOCOL_TCP4_PAYLOAD_MAX)
			_buf.len = PROXY_PROTOCOL_TCP4_PAYLOAD_MAX;
		max = PROXY_PROTOCOL_TCP4_MAX;
	} else if (m == PP_TOK_PARTIAL) {
		return NULL;
	} else {
		m = match_proxy_token(&_buf, PROXY_PROTOCOL_TCP6, PROXY_PROTOCOL_TCP6_LEN);
		if (m == PP_TOK_MATCH) {
			af = AF_INET6;
			_buf.s += PROXY_PROTOCOL_TCP6_LEN;
			_buf.len -= PROXY_PROTOCOL_TCP6_LEN;
			if (_buf.len > PROXY_PROTOCOL_TCP6_PAYLOAD_MAX)
				_buf.len = PROXY_PROTOCOL_TCP6_PAYLOAD_MAX;
			max = PROXY_PROTOCOL_TCP6_MAX;
		} else if (m == PP_TOK_PARTIAL) {
			return NULL;
		} else {
			m = match_proxy_token(&_buf, PROXY_PROTOCOL_UNKN, PROXY_PROTOCOL_UNKN_LEN);
			if (m == PP_TOK_MATCH) {
				_buf.s += PROXY_PROTOCOL_UNKN_LEN;
				_buf.len -= PROXY_PROTOCOL_UNKN_LEN;
				if (_buf.len == 0)
					return NULL;
				if (_buf.s[0] != PROXY_PROTOCOL_SEP &&
						_buf.s[0] != PROXY_PROTOCOL_END[0]) {
					LM_DBG("bad separator after proxy_protocol UNKNOWN: 0x%x\n",
							(unsigned char)_buf.s[0]);
					goto error;
				}
				if (_buf.len > PROXY_PROTOCOL_UNKN_PAYLOAD_MAX)
					_buf.len = PROXY_PROTOCOL_UNKN_PAYLOAD_MAX;
				af = AF_UNSPEC;
				max = PROXY_PROTOCOL_UNKN_MAX;
			} else if (m == PP_TOK_PARTIAL) {
				return NULL;
			} else {
				LM_DBG("unknown proxy_protocol proto [%.*s]\n", _buf.len, _buf.s);
				goto error;
			}
		}
	}

	/* search for the end within the maximum payload */
	end = str_strstr(&_buf, _str(PROXY_PROTOCOL_END));
	if (!end) {
		if (size < max) /* we still don't have enough data to find the end */
			return NULL;
		LM_DBG("could not find proxy_protocol end in [%.*s]\n",
				_buf.len, _buf.s);
		goto error;
	}
	if (af == AF_UNSPEC) {
		/* we have an unknown protocol - stop it here */
		proxy->flags = PP_UNKNOWN;
		return end + PROXY_PROTOCOL_END_LEN;
	}
	_buf.len = end - _buf.s;
	/* src ip */
	p = q_memchr(_buf.s, PROXY_PROTOCOL_SEP, _buf.len);
	if (!p)
		goto error;
	tmp.s = _buf.s;
	tmp.len = p - _buf.s;
	ip = (af == AF_INET?str2ip(&tmp):str2ip6(&tmp));
	if (!ip) {
		LM_DBG("could not parse proxy_protocol src ip %.*s for af %d\n",
				tmp.len, tmp.s, af);
		goto error;
	}
	memcpy(&proxy->src_ip, ip, sizeof *ip);
	_buf.len -= tmp.len + 1;
	_buf.s = p + 1;

	/* dst ip */
	p = q_memchr(_buf.s, PROXY_PROTOCOL_SEP, _buf.len);
	if (!p)
		goto error;
	tmp.s = _buf.s;
	tmp.len = p - _buf.s;
	ip = (af == AF_INET?str2ip(&tmp):str2ip6(&tmp));
	if (!ip) {
		LM_DBG("could not parse proxy_protocol dst ip %.*s for af %d\n",
				tmp.len, tmp.s, af);
		goto error;
	}
	memcpy(&proxy->dst_ip, ip, sizeof *ip);
	_buf.len -= tmp.len + 1;
	_buf.s = p + 1;

	/* src port */
	p = q_memchr(_buf.s, PROXY_PROTOCOL_SEP, _buf.len);
	if (!p)
		goto error;
	tmp.s = _buf.s;
	tmp.len = p - _buf.s;
	if (str2int(&tmp, &port) < 0) {
		LM_DBG("could not parse proxy_protocol src port %.*s\n", tmp.len, tmp.s);
		goto error;
	}
	if (port > 65535) {
		LM_DBG("bad proxy_protocol src port %d\n", port);
		goto error;
	}
	proxy->src_port = port;
	_buf.len -= tmp.len + 1;
	_buf.s = p + 1;

	/* dst port */
	/* remaining should be just the dst port */
	if (str2int(&_buf, &port) < 0) {
		LM_DBG("could not parse proxy_protocol dst port %.*s\n", _buf.len, _buf.s);
		goto error;
	}
	if (port > 65535) {
		LM_DBG("bad proxy_protocol dst port %d\n", port);
		goto error;
	}
	proxy->dst_port = port;
	proxy->flags = PP_OK;
	return end + PROXY_PROTOCOL_END_LEN;
error:
	proxy->flags = PP_ERROR;
	return NULL;
}

static int tcp_peek(struct tcp_connection *c, char *buf, int len)
{
	int bytes_read;
	int fd;

	fd=c->fd;
again:
	bytes_read=recv(fd, buf, len, MSG_PEEK);

	if(bytes_read==-1){
		if (errno == EWOULDBLOCK || errno == EAGAIN){
			return 0; /* nothing has been read */
		} else if (errno == EINTR) {
			goto again;
		} else if (errno == ECONNRESET) {
			c->state=S_CONN_EOF;
			LM_DBG("CONN RESET on %p, FD %d\n", c, fd);
			return -1;
		} else {
			LM_ERR("error reading: %s\n",strerror(errno));
			return -1;
		}
	}else if (bytes_read==0){
		c->state=S_CONN_EOF;
		LM_DBG("EOF on %p, FD %d\n", c, fd);
		return -1;
	}
#ifdef EXTRA_DEBUG
	LM_DBG("peek %d bytes:\n%.*s\n", bytes_read, bytes_read, buf);
#endif
	return bytes_read;
}

static int tcp_drain(struct tcp_connection *c, char *buf, int len)
{
	int bytes_read, total = 0;
	int retries = 0;
	int fd = c->fd;
again:
	bytes_read = read(fd, buf + total, len - total);
	if (bytes_read < 0) {
		if (errno == EINTR) {
			goto again;
		} else if (errno == EAGAIN || errno == EWOULDBLOCK) {
			if (++retries > 100) {
				LM_ERR("failed consuming proxy_protocol header due to repeated EAGAIN\n");
				return -1;
			}
			goto again;
		} else {
			LM_ERR("error consuming proxy_protocol header: %s\n", strerror(errno));
			return -1;
		}
	} else if (bytes_read == 0) {
		c->state = S_CONN_EOF;
		LM_DBG("EOF on %p, FD %d while consuming proxy_protocol header\n", c, fd);
		return -1;
	}
	total += bytes_read;
	retries = 0;
	if (total < len)
		goto again;

	return total;
}

int check_tcp_proxy_protocol(struct tcp_connection *c)
{
	static char pp_buf[PROXY_PROTOCOL_BUF_MAX];
	char *p;
	int len;

	if (c->flags & F_CONN_DATA_READY)
		return 1;

	if (!c->rcv.bind_address ||
			(c->rcv.bind_address->flags & SI_PROXY_IN) == 0) {
		c->flags |= F_CONN_DATA_READY;
		return 1;
	}

	len = tcp_peek(c, pp_buf, PROXY_PROTOCOL_BUF_MAX);
	if (len < 0)
		return -1;

	switch (is_net_proxy_protocol(pp_buf, len)) {
		case -1:
			return 0;
		case 0:
			c->flags |= F_CONN_DATA_READY;
			return 1;
		default:
			/* parse proxy_protocol fields */
			break;
	}

	p = parse_net_proxy_protocol(pp_buf, len, &c->rcv.real_ep);
	if (!p) {
		if (c->rcv.real_ep.flags == PP_ERROR) {
			LM_ERR("could not parse proxy_protocol header\n");
			return -1;
		}
		return 0;
	}
	if (c->rcv.real_ep.flags == PP_OK) {
		LM_DBG("message proxy_protocol %s:%hu -> %s:%hu\n",
				ip_addr2a(&c->rcv.real_ep.src_ip), c->rcv.real_ep.src_port,
				ip_addr2a(&c->rcv.real_ep.dst_ip), c->rcv.real_ep.dst_port);
	} else {
		LM_DBG("message proxy_protocol UNKNOWN\n");
	}

	if (tcp_drain(c, pp_buf, p - pp_buf) < 0) {
		LM_ERR("could not consume PROXY header\n");
		return -1;
	}

	c->flags |= F_CONN_DATA_READY;
	return 1;
}

int check_udp_proxy_protocol(char **buf, int *size, struct receive_info *ri)
{
	char *p, *msg;
	int len;

	if (!buf || !*buf || !size || !ri)
		return -1;

	ri->real_ep.flags = PP_INIT;

	if (!ri->bind_address || (ri->bind_address->flags & SI_PROXY_IN) == 0)
		return 1;

	msg = *buf;
	len = *size;

	switch (is_net_proxy_protocol(msg, len)) {
	case -1:
	case 0:
		return 1;
	default:
		break;
	}

	p = parse_net_proxy_protocol(msg, len, &ri->real_ep);
	if (!p) {
		if (ri->real_ep.flags == PP_ERROR) {
			LM_ERR("could not parse proxy_protocol header\n");
		} else {
			LM_ERR("incomplete proxy_protocol header in UDP packet\n");
		}
		return -1;
	}

	if (ri->real_ep.flags == PP_OK) {
		LM_DBG("message proxy_protocol %s:%hu -> %s:%hu\n",
				ip_addr2a(&ri->real_ep.src_ip), ri->real_ep.src_port,
				ip_addr2a(&ri->real_ep.dst_ip), ri->real_ep.dst_port);
	} else {
		LM_DBG("message proxy_protocol UNKNOWN\n");
	}

	*size -= p - msg;
	*buf = p;
	return 1;
}
