/*************************************************************************
 *
 * Copyright (C) 2018-2025 Ruilin Peng (Nick) <pymumu@gmail.com>.
 *
 * smartdns 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 3 of the License, or
 * (at your option) any later version.
 *
 * smartdns 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, see <http://www.gnu.org/licenses/>.
 */
#define _GNU_SOURCE

#include "server_tls.h"
#include "connection.h"
#include "dns_server.h"
#include "server_https.h"
#include "server_socket.h"
#include "server_tcp.h"
#include "server_http2.h"

#include "smartdns/http2.h"

#include <errno.h>
#include <netinet/tcp.h>
#include <openssl/err.h>
#include <openssl/ssl.h>
#include <pthread.h>
#include <string.h>
#include <sys/epoll.h>

static int alpn_select_cb(SSL *ssl, const unsigned char **out, unsigned char *outlen, const unsigned char *in,
						  unsigned int inlen, void *arg)
{
	struct dns_bind_ip *bind_ip = (struct dns_bind_ip *)arg;
	const char *alpn = bind_ip->alpn;
	if (alpn[0] == '\0') {
		alpn = "h2,http/1.1";
	}

	/* Parse server ALPN list */
	char alpn_copy[256];
	safe_strncpy(alpn_copy, alpn, sizeof(alpn_copy));
	char *saveptr = NULL;
	char *proto = strtok_r(alpn_copy, ",", &saveptr);

	while (proto) {
		unsigned int proto_len = strlen(proto);
		for (unsigned int i = 0; i < inlen;) {
			unsigned int len = in[i++];
			if (i + len > inlen) {
				break;
			}

			if (len == proto_len && memcmp(&in[i], proto, len) == 0) {
				*out = &in[i];
				*outlen = len;
				return SSL_TLSEXT_ERR_OK;
			}
			i += len;
		}
		proto = strtok_r(NULL, ",", &saveptr);
	}

	/* No match found */
	tlog(TLOG_DEBUG, "ALPN negotiation failed: no matching protocol found");
	return SSL_TLSEXT_ERR_NOACK;
}

static ssize_t _ssl_read(struct dns_server_conn_tls_client *conn, void *buff, int num)
{
	ssize_t ret = 0;
	pthread_mutex_lock(&conn->ssl_lock);
	if (conn == NULL || buff == NULL) {
		pthread_mutex_unlock(&conn->ssl_lock);
		return SSL_ERROR_SYSCALL;
	}

	ret = SSL_read(conn->ssl, buff, num);
	pthread_mutex_unlock(&conn->ssl_lock);
	return ret;
}

static ssize_t _ssl_write(struct dns_server_conn_tls_client *conn, const void *buff, int num)
{
	ssize_t ret = 0;
	pthread_mutex_lock(&conn->ssl_lock);
	if (conn == NULL || buff == NULL || conn->ssl == NULL) {
		pthread_mutex_unlock(&conn->ssl_lock);
		return SSL_ERROR_SYSCALL;
	}

	ret = SSL_write(conn->ssl, buff, num);
	pthread_mutex_unlock(&conn->ssl_lock);
	return ret;
}

static int _ssl_get_error(struct dns_server_conn_tls_client *conn, int ret)
{
	int err = 0;
	pthread_mutex_lock(&conn->ssl_lock);
	if (conn == NULL || conn->ssl == NULL) {
		pthread_mutex_unlock(&conn->ssl_lock);
		return SSL_ERROR_SYSCALL;
	}

	err = SSL_get_error(conn->ssl, ret);
	pthread_mutex_unlock(&conn->ssl_lock);
	return err;
}

static int _ssl_do_accept(struct dns_server_conn_tls_client *conn)
{
	int err = 0;
	pthread_mutex_lock(&conn->ssl_lock);
	if (conn == NULL || conn->ssl == NULL) {
		pthread_mutex_unlock(&conn->ssl_lock);
		return SSL_ERROR_SYSCALL;
	}

	err = SSL_accept(conn->ssl);
	pthread_mutex_unlock(&conn->ssl_lock);
	return err;
}

int _dns_server_socket_ssl_send(struct dns_server_conn_tls_client *tls_client, const void *buf, int num)
{
	int ret = 0;
	int ssl_ret = 0;
	unsigned long ssl_err = 0;

	if (tls_client->ssl == NULL) {
		errno = EINVAL;
		return -1;
	}

	if (num < 0) {
		errno = EINVAL;
		return -1;
	}

	ret = _ssl_write(tls_client, buf, num);
	if (ret > 0) {
		return ret;
	}

	ssl_ret = _ssl_get_error(tls_client, ret);
	switch (ssl_ret) {
	case SSL_ERROR_NONE:
	case SSL_ERROR_ZERO_RETURN:
		return 0;
		break;
	case SSL_ERROR_WANT_READ:
		errno = EAGAIN;
		ret = -SSL_ERROR_WANT_READ;
		break;
	case SSL_ERROR_WANT_WRITE:
		errno = EAGAIN;
		ret = -SSL_ERROR_WANT_WRITE;
		break;
	case SSL_ERROR_SSL:
		ssl_err = ERR_get_error();
		int ssl_reason = ERR_GET_REASON(ssl_err);
		if (ssl_reason == SSL_R_UNINITIALIZED || ssl_reason == SSL_R_PROTOCOL_IS_SHUTDOWN ||
			ssl_reason == SSL_R_BAD_LENGTH || ssl_reason == SSL_R_SHUTDOWN_WHILE_IN_INIT ||
			ssl_reason == SSL_R_BAD_WRITE_RETRY) {
			errno = EAGAIN;
			return -1;
		}

		tlog(TLOG_ERROR, "SSL write fail error no:  %s(%d)\n", ERR_reason_error_string(ssl_err), ssl_reason);
		errno = EFAULT;
		ret = -1;
		break;
	case SSL_ERROR_SYSCALL:
		if (errno == 0) {
			tlog(TLOG_DEBUG, "SSL connection closed");
		} else {
			tlog(TLOG_DEBUG, "SSL syscall failed, %s", strerror(errno));
		}
		return ret;
	default:
		errno = EFAULT;
		ret = -1;
		break;
	}

	return ret;
}

int _dns_server_socket_ssl_recv(struct dns_server_conn_tls_client *tls_client, void *buf, int num)
{
	ssize_t ret = 0;
	int ssl_ret = 0;
	unsigned long ssl_err = 0;

	if (tls_client->ssl == NULL) {
		errno = EFAULT;
		return -1;
	}

	ret = _ssl_read(tls_client, buf, num);
	if (ret > 0) {
		return ret;
	}

	ssl_ret = _ssl_get_error(tls_client, ret);
	switch (ssl_ret) {
	case SSL_ERROR_NONE:
	case SSL_ERROR_ZERO_RETURN:
		return 0;
		break;
	case SSL_ERROR_WANT_READ:
		errno = EAGAIN;
		ret = -SSL_ERROR_WANT_READ;
		break;
	case SSL_ERROR_WANT_WRITE:
		errno = EAGAIN;
		ret = -SSL_ERROR_WANT_WRITE;
		break;
	case SSL_ERROR_SSL:
		ssl_err = ERR_get_error();
		int ssl_reason = ERR_GET_REASON(ssl_err);
		if (ssl_reason == SSL_R_UNINITIALIZED) {
			errno = EAGAIN;
			return -1;
		}

		if (ssl_reason == SSL_R_SHUTDOWN_WHILE_IN_INIT || ssl_reason == SSL_R_PROTOCOL_IS_SHUTDOWN) {
			return 0;
		}

#ifdef SSL_R_UNEXPECTED_EOF_WHILE_READING
		if (ssl_reason == SSL_R_UNEXPECTED_EOF_WHILE_READING) {
			return 0;
		}
#endif

		tlog(TLOG_DEBUG, "SSL read fail error no: %s(%lx), reason: %d\n", ERR_reason_error_string(ssl_err), ssl_err,
			 ssl_reason);
		errno = EFAULT;
		ret = -1;
		break;
	case SSL_ERROR_SYSCALL:
		if (errno == 0) {
			return 0;
		}

		ret = -1;
		return ret;
	default:
		errno = EFAULT;
		ret = -1;
		break;
	}

	return ret;
}

int _dns_server_ssl_poll_event(struct dns_server_conn_tls_client *tls_client, int ssl_ret)
{
	struct epoll_event fd_event;

	memset(&fd_event, 0, sizeof(fd_event));

	if (ssl_ret == SSL_ERROR_WANT_READ) {
		fd_event.events = EPOLLIN;
	} else if (ssl_ret == SSL_ERROR_WANT_WRITE) {
		fd_event.events = EPOLLOUT | EPOLLIN;
	} else {
		goto errout;
	}

	fd_event.data.ptr = tls_client;
	if (epoll_ctl(server.epoll_fd, EPOLL_CTL_MOD, tls_client->tcp.head.fd, &fd_event) != 0) {
		if (errno == ENOENT) {
			/* fd not found, ignore */
			return 0;
		}
		tlog(TLOG_ERROR, "epoll ctl failed, %s", strerror(errno));
		goto errout;
	}

	return 0;

errout:
	return -1;
}

int _dns_server_tls_accept(struct dns_server_conn_tls_server *tls_server, struct epoll_event *event, unsigned long now)
{
	struct sockaddr_storage addr;
	struct dns_server_conn_tls_client *tls_client = NULL;
	DNS_CONN_TYPE conn_type;
	socklen_t addr_len = sizeof(addr);
	int fd = -1;
	SSL *ssl = NULL;

	fd = accept4(tls_server->head.fd, (struct sockaddr *)&addr, &addr_len, SOCK_NONBLOCK | SOCK_CLOEXEC);
	if (fd < 0) {
		tlog(TLOG_ERROR, "accept failed, %s", strerror(errno));
		return -1;
	}

	if (tls_server->head.type == DNS_CONN_TYPE_TLS_SERVER) {
		conn_type = DNS_CONN_TYPE_TLS_CLIENT;
	} else if (tls_server->head.type == DNS_CONN_TYPE_HTTPS_SERVER) {
		conn_type = DNS_CONN_TYPE_HTTPS_CLIENT;
	} else {
		tlog(TLOG_ERROR, "invalid http server type.");
		goto errout;
	}

	tls_client = zalloc(1, sizeof(*tls_client));
	if (tls_client == NULL) {
		tlog(TLOG_ERROR, "malloc for tls_client failed.");
		goto errout;
	}
	_dns_server_conn_head_init(&tls_client->tcp.head, fd, conn_type);
	tls_client->tcp.head.server_flags = tls_server->head.server_flags;
	tls_client->tcp.head.dns_group = tls_server->head.dns_group;
	tls_client->tcp.head.ipset_nftset_rule = tls_server->head.ipset_nftset_rule;
	tls_client->tcp.conn_idle_timeout = dns_conf.tcp_idle_time;

	atomic_set(&tls_client->tcp.head.refcnt, 0);
	memcpy(&tls_client->tcp.addr, &addr, addr_len);
	tls_client->tcp.addr_len = addr_len;
	tls_client->tcp.localaddr_len = sizeof(struct sockaddr_storage);
	if (_dns_server_epoll_ctl(&tls_client->tcp.head, EPOLL_CTL_ADD, EPOLLIN) != 0) {
		tlog(TLOG_ERROR, "epoll ctl failed.");
		return -1;
	}

	if (getsocket_inet(tls_client->tcp.head.fd, (struct sockaddr *)&tls_client->tcp.localaddr,
					   &tls_client->tcp.localaddr_len) != 0) {
		tlog(TLOG_ERROR, "get local addr failed, %s", strerror(errno));
		goto errout;
	}

	ssl = SSL_new(tls_server->ssl_ctx);
	if (ssl == NULL) {
		tlog(TLOG_ERROR, "SSL_new failed.");
		goto errout;
	}

	if (SSL_set_fd(ssl, fd) != 1) {
		tlog(TLOG_ERROR, "SSL_set_fd failed.");
		goto errout;
	}

	tls_client->ssl = ssl;
	tls_client->tcp.status = DNS_SERVER_CLIENT_STATUS_CONNECTING;
	pthread_mutex_init(&tls_client->ssl_lock, NULL);
	_dns_server_client_touch(&tls_client->tcp.head);

	pthread_mutex_lock(&server.conn_list_lock);
	list_add(&tls_client->tcp.head.list, &server.conn_list);
	pthread_mutex_unlock(&server.conn_list_lock);

	_dns_server_conn_get(&tls_client->tcp.head);

	set_sock_keepalive(fd, 30, 3, 5);

	return 0;
errout:
	if (fd > 0) {
		close(fd);
	}

	if (ssl) {
		SSL_free(ssl);
	}

	if (tls_client) {
		free(tls_client);
	}
	return -1;
}

int _dns_server_process_tls(struct dns_server_conn_tls_client *tls_client, struct epoll_event *event, unsigned long now)
{
	int ret = 0;
	int ssl_ret = 0;
	struct epoll_event fd_event;

	if (tls_client->tcp.status == DNS_SERVER_CLIENT_STATUS_CONNECTING) {
		/* do SSL hand shake */
		ret = _ssl_do_accept(tls_client);
		if (ret <= 0) {
			memset(&fd_event, 0, sizeof(fd_event));
			ssl_ret = _ssl_get_error(tls_client, ret);
			if (_dns_server_ssl_poll_event(tls_client, ssl_ret) == 0) {
				return 0;
			}

			if (ssl_ret != SSL_ERROR_SYSCALL) {
				unsigned long ssl_err = ERR_get_error();
				int ssl_reason = ERR_GET_REASON(ssl_err);
				char name[DNS_MAX_CNAME_LEN];
				tlog(TLOG_DEBUG, "Handshake with %s failed, error no: %s(%d, %d, %d)\n",
					 get_host_by_addr(name, sizeof(name), (struct sockaddr *)&tls_client->tcp.addr),
					 ERR_reason_error_string(ssl_err), ret, ssl_ret, ssl_reason);
				ret = 0;
			}

			goto errout;
		}

		tls_client->tcp.status = DNS_SERVER_CLIENT_STATUS_CONNECTED;

		/* Get negotiated ALPN */
		const unsigned char *alpn_data = NULL;
		unsigned int alpn_len = 0;
		SSL_get0_alpn_selected(tls_client->ssl, &alpn_data, &alpn_len);
		if (alpn_data && alpn_len > 0 && alpn_len < sizeof(tls_client->alpn_selected)) {
			memcpy(tls_client->alpn_selected, alpn_data, alpn_len);
			tls_client->alpn_selected[alpn_len] = '\0';
		} else {
			safe_strncpy(tls_client->alpn_selected, "http/1.1", sizeof(tls_client->alpn_selected));
		}

		memset(&fd_event, 0, sizeof(fd_event));
		fd_event.events = EPOLLIN | EPOLLOUT;
		fd_event.data.ptr = tls_client;
		if (epoll_ctl(server.epoll_fd, EPOLL_CTL_MOD, tls_client->tcp.head.fd, &fd_event) != 0) {
			tlog(TLOG_ERROR, "epoll ctl failed, %s", strerror(errno));
			goto errout;
		}
	}

	/* if HTTP/2 was negotiated */
	if (strcmp(tls_client->alpn_selected, "h2") == 0) {
		ret = _dns_server_process_http2(tls_client, event, now);
		if (ret != 0) {
			goto errout;
		}
		return ret;
	}

	return _dns_server_process_tcp((struct dns_server_conn_tcp_client *)tls_client, event, now);
errout:
	_dns_server_client_close(&tls_client->tcp.head);
	return ret;
}

static int _dns_server_socket_tls_ssl_pass_callback(char *buf, int size, int rwflag, void *userdata)
{
	struct dns_bind_ip *bind_ip = userdata;
	if (bind_ip->ssl_cert_key_pass == NULL || bind_ip->ssl_cert_key_pass[0] == '\0') {
		return 0;
	}
	safe_strncpy(buf, bind_ip->ssl_cert_key_pass, size);
	return strlen(buf);
}

int _dns_server_socket_tls(struct dns_bind_ip *bind_ip, DNS_CONN_TYPE conn_type)
{
	const char *host_ip = NULL;
	const char *ssl_cert_file = NULL;
	const char *ssl_cert_key_file = NULL;

	struct dns_server_conn_tls_server *conn = NULL;
	int fd = -1;
	const SSL_METHOD *method = NULL;
	SSL_CTX *ssl_ctx = NULL;
	const int on = 1;

	host_ip = bind_ip->ip;
	ssl_cert_file = bind_ip->ssl_cert_file;
	ssl_cert_key_file = bind_ip->ssl_cert_key_file;

	if (ssl_cert_file == NULL || ssl_cert_key_file == NULL) {
		tlog(TLOG_WARN, "no cert or cert key file");
		goto errout;
	}

	if (ssl_cert_file[0] == '\0' || ssl_cert_key_file[0] == '\0') {
		tlog(TLOG_WARN, "no cert or cert key file");
		goto errout;
	}

	fd = _dns_create_socket(host_ip, SOCK_STREAM);
	if (fd <= 0) {
		goto errout;
	}

	setsockopt(fd, SOL_TCP, TCP_FASTOPEN, &on, sizeof(on));

#if (OPENSSL_VERSION_NUMBER >= 0x10100000L)
	method = TLS_server_method();
	if (method == NULL) {
		goto errout;
	}
#else
	method = SSLv23_server_method();
#endif

	ssl_ctx = SSL_CTX_new(method);
	if (ssl_ctx == NULL) {
		goto errout;
	}

	SSL_CTX_set_session_cache_mode(ssl_ctx,
								   SSL_SESS_CACHE_BOTH | SSL_SESS_CACHE_NO_INTERNAL | SSL_SESS_CACHE_NO_AUTO_CLEAR);
	SSL_CTX_set_default_passwd_cb(ssl_ctx, _dns_server_socket_tls_ssl_pass_callback);
	SSL_CTX_set_default_passwd_cb_userdata(ssl_ctx, bind_ip);

	/* Set the key and cert */
	if (ssl_cert_file[0] != '\0' && SSL_CTX_use_certificate_chain_file(ssl_ctx, ssl_cert_file) <= 0) {
		tlog(TLOG_ERROR, "load cert %s failed, %s", ssl_cert_file, ERR_error_string(ERR_get_error(), NULL));
		goto errout;
	}

	if (ssl_cert_key_file[0] != '\0' &&
		SSL_CTX_use_PrivateKey_file(ssl_ctx, ssl_cert_key_file, SSL_FILETYPE_PEM) <= 0) {
		tlog(TLOG_ERROR, "load cert key %s failed, %s", ssl_cert_key_file, ERR_error_string(ERR_get_error(), NULL));
		goto errout;
	}

	/* Set ALPN */
	SSL_CTX_set_alpn_select_cb(ssl_ctx, alpn_select_cb, bind_ip);

	conn = zalloc(1, sizeof(struct dns_server_conn_tls_server));
	if (conn == NULL) {
		goto errout;
	}
	_dns_server_conn_head_init(&conn->head, fd, conn_type);
	conn->ssl_ctx = ssl_ctx;
	_dns_server_set_flags(&conn->head, bind_ip);
	_dns_server_conn_get(&conn->head);

	return 0;
errout:
	if (ssl_ctx) {
		SSL_CTX_free(ssl_ctx);
		ssl_ctx = NULL;
	}

	if (conn) {
		free(conn);
		conn = NULL;
	}

	if (fd > 0) {
		close(fd);
	}
	return -1;
}
