From 42e6d574a030853f2075df3f1f847c0bfd120d37 Mon Sep 17 00:00:00 2001
From: Vadim Vetrov <vetrovvd@gmail.com>
Date: Sun, 8 Dec 2024 22:08:57 +0300
Subject: [PATCH] Add connbytes to kernel module

---
 README.md    |  2 ++
 args.c       | 14 ++++++++++++++
 config.h     |  3 +++
 kytunblock.c | 32 ++++++++++++++++++++++++++++++++
 4 files changed, 51 insertions(+)

diff --git a/README.md b/README.md
index 9cf9836..0a5ed54 100644
--- a/README.md
+++ b/README.md
@@ -254,6 +254,8 @@ Available flags:
 
 - `--threads=<threads number>` Specifies the amount of threads you want to be running for your program. This defaults to **1** and shouldn't be edited for normal use. But if you really want multiple queue instances of youtubeUnblock, note that you should change --queue-num to --queue balance. For example, with 4 threads, use `--queue-balance 537:540` on iptables and `queue num 537-540` on nftables.
 
+- `--connbytes-limit=<pkts>` **Kernel module only!** Specify how much packets of connection should be processed by kyoutubeUnblock. Pass 0 if you want for each packet to be processed. This flag may be useful for UDP traffic since unlimited youtubeUnblock may lead to traffic flood and unexpected bans. Defaults to 5. In most cases you don't want to change it.
+
 - `--daemonize` Daemonizes the youtubeUnblock (forks and detaches it from the shell). Terminate the program with `killall youtubeUnblock`. If you want to track the logs of youtubeUnblock in logread or journalctl, use **--syslog** flag.
 
 - `--syslog` Redirects logs to the system log. You can read it with `journalctl` or `logread`.
diff --git a/args.c b/args.c
index 89c0284..61cc43f 100644
--- a/args.c
+++ b/args.c
@@ -258,6 +258,7 @@ enum {
 	OPT_CLS,
 	OPT_HELP,
 	OPT_VERSION,
+	OPT_CONNBYTES_LIMIT,
 };
 
 static struct option long_opt[] = {
@@ -300,6 +301,7 @@ static struct option long_opt[] = {
 	{"syslog",		0, 0, OPT_SYSLOG},
 	{"queue-num",		1, 0, OPT_QUEUE_NUM},
 	{"packet-mark",		1, 0, OPT_PACKET_MARK},
+	{"connbytes-limit",	1, 0, OPT_CONNBYTES_LIMIT},
 	{"fbegin",		0, 0, OPT_START_SECTION},
 	{"fend",		0, 0, OPT_END_SECTION},
 	{"cls",			0, 0, OPT_CLS},
@@ -352,6 +354,7 @@ void print_usage(const char *argv0) {
 	printf("\t--udp-filter-quic={disabled|all}\n");
 	printf("\t--threads=<threads number>\n");
 	printf("\t--packet-mark=<mark>\n");
+	printf("\t--connbytes-limit=<pkts>\n");
 	printf("\t--silent\n");
 	printf("\t--trace\n");
 	printf("\t--no-gso\n");
@@ -463,6 +466,13 @@ int yparse_args(int argc, char *argv[]) {
 
 			rep_config.mark = num;
 			break;
+		case OPT_CONNBYTES_LIMIT:
+			num = parse_numeric_option(optarg);
+			if (errno != 0 || num < 0) {
+				goto invalid_opt;
+			}
+			rep_config.connbytes_limit = num;
+			break;
 		case OPT_START_SECTION:
 			if (section_iter != SECT_ITER_DEFAULT && section_iter != SECT_ITER_OUTSIDE)
 				goto invalid_opt;
@@ -954,6 +964,10 @@ size_t print_config(char *buffer, size_t buffer_size) {
 		print_cnf_buf("--no-gso");
 	}
 #endif
+
+#ifdef KERNEL_SPACE
+	print_cnf_buf("--connbytes-limit=%d", config.connbytes_limit);
+#endif
 	if (!config.use_ipv6) {
 		print_cnf_buf("--no-ipv6");
 	}
diff --git a/config.h b/config.h
index 8adb331..5a1cd7a 100644
--- a/config.h
+++ b/config.h
@@ -100,6 +100,8 @@ struct config_t {
 	int noclose;
 	int syslog;
 
+	int connbytes_limit;
+
 #define VERBOSE_INFO	0
 #define VERBOSE_DEBUG	1
 #define VERBOSE_TRACE	2
@@ -241,6 +243,7 @@ enum {
 	.queue_start_num = DEFAULT_QUEUE_NUM,                   \
 	.mark = DEFAULT_RAWSOCKET_MARK,                         \
 	.use_ipv6 = 1,                                          \
+	.connbytes_limit = 8,                                   \
                                                                 \
 	.verbose = VERBOSE_DEBUG,                               \
 	.use_gso = 1,                                           \
diff --git a/kytunblock.c b/kytunblock.c
index b4c789d..7d079c5 100644
--- a/kytunblock.c
+++ b/kytunblock.c
@@ -16,6 +16,9 @@
 #include <linux/netfilter_ipv4.h>
 #include <linux/netfilter_ipv6.h>
 
+#include <net/netfilter/nf_conntrack.h>
+#include <net/netfilter/nf_conntrack_acct.h>
+
 #include "mangle.h"
 #include "config.h"
 #include "utils.h"
@@ -229,6 +232,32 @@ struct instance_config_t instance_config = {
 	.send_delayed_packet = delay_packet_send,
 };
 
+static int connbytes_pkts(const struct sk_buff *skb) {
+	const struct nf_conn *ct;
+	enum ip_conntrack_info ctinfo;
+	u_int64_t pkts = 0;
+	const struct nf_conn_counter *counters;
+
+	ct = nf_ct_get(skb, &ctinfo);
+	if (!ct)
+		return -1;
+
+#if LINUX_VERSION_CODE >= KERNEL_VERSION(3, 13, 0)
+	const struct nf_conn_acct *acct;
+	acct = nf_conn_acct_find(ct);
+	if (!acct)
+		return -1;
+	counters = acct->counter;
+#else 
+	counters = nf_conn_acct_find(ct);
+	if (!counters)
+		return -1;
+#endif
+
+	pkts = atomic64_read(&counters[IP_CT_DIR_ORIGINAL].packets);
+
+	return pkts;
+}
 
 /* If this is a Red Hat-based kernel (Red Hat, CentOS, Fedora, etc)... */
 #ifdef RHEL_RELEASE_CODE
@@ -308,6 +337,9 @@ static NF_CALLBACK(ykb_nf_hook, skb) {
 	if (skb->len > MAX_PACKET_SIZE)
 		goto accept;
 
+	if (config.connbytes_limit != 0 && connbytes_pkts(skb) > config.connbytes_limit)
+		goto accept;
+
 	ret = skb_linearize(skb);
 	if (ret < 0) {
 		lgerror(ret, "Cannot linearize");