From 7070ddfc7483d74087567e20abfcf2942a9a8b4a Mon Sep 17 00:00:00 2001
From: Vadim Vetrov <vetrovvd@gmail.com>
Date: Sat, 26 Oct 2024 01:12:52 +0300
Subject: [PATCH] Userspace: better support for ALLOC_MALLOC

An alternative memory allocation strategy for userspace
---
 config.h         |  2 +-
 types.h          |  3 +++
 youtubeUnblock.c | 40 +++++++++++++++++++++++++++++++++-------
 3 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/config.h b/config.h
index f3d6683..05b7a18 100644
--- a/config.h
+++ b/config.h
@@ -185,7 +185,7 @@ if ((fake_bitmask) & strategy)
 
 // The Maximum Transmission Unit size for rawsocket
 // Larger packets will be fragmented. Applicable for Chrome's kyber.
-#define AVAILABLE_MTU 1500
+#define AVAILABLE_MTU 1400
 
 #define DEFAULT_QUEUE_NUM 537
 
diff --git a/types.h b/types.h
index ab63f61..32f88aa 100644
--- a/types.h
+++ b/types.h
@@ -85,6 +85,9 @@
 
 #endif /* not a KERNEL_SPACE */
 
+/* An alternative memory allocation strategy for userspace app */
+// #define ALLOC_MALLOC
+
 /**
  * Use NETBUF_ALLOC and NETBUF_FREE as an abstraction of memory allocation.
  * Do not use it within expressions, consider these defines as separate statements.
diff --git a/youtubeUnblock.c b/youtubeUnblock.c
index c4106aa..5733d8d 100644
--- a/youtubeUnblock.c
+++ b/youtubeUnblock.c
@@ -259,16 +259,27 @@ static int send_raw_socket(const uint8_t *pkt, uint32_t pktlen) {
 		if (config.verbose)
 			printf("Split packet!\n");
 
-		uint8_t buff1[MNL_SOCKET_BUFFER_SIZE];
+		NETBUF_ALLOC(buff1, MNL_SOCKET_BUFFER_SIZE);
+		if (!NETBUF_CHECK(buff1)) {
+			lgerror("Allocation error", -ENOMEM);
+			return -ENOMEM;
+		}
+
+		NETBUF_ALLOC(buff2, MNL_SOCKET_BUFFER_SIZE);
+		if (!NETBUF_CHECK(buff2)) {
+			lgerror("Allocation error", -ENOMEM);
+			NETBUF_FREE(buff1);
+			return -ENOMEM;
+		}
+
 		uint32_t buff1_size = MNL_SOCKET_BUFFER_SIZE;
-		uint8_t buff2[MNL_SOCKET_BUFFER_SIZE];
 		uint32_t buff2_size = MNL_SOCKET_BUFFER_SIZE;
 
 		if ((ret = tcp_frag(pkt, pktlen, AVAILABLE_MTU-128,
 			buff1, &buff1_size, buff2, &buff2_size)) < 0) {
 
 			errno = -ret;
-			return ret;
+			goto free_buffs;
 		}
 
 		int sent = 0;
@@ -276,16 +287,23 @@ static int send_raw_socket(const uint8_t *pkt, uint32_t pktlen) {
 
 		if (status >= 0) sent += status;
 		else {
-			return status;
+			ret = status;
+			goto free_buffs;
 		}
 
 		status = send_raw_socket(buff2, buff2_size);
 		if (status >= 0) sent += status;
 		else {
-			return status;
+			ret = status;
+			goto free_buffs;
 		}
 
-		return sent;
+		ret = sent;
+
+free_buffs:
+		NETBUF_FREE(buff1)
+		NETBUF_FREE(buff2)
+		return ret;
 	}
 	
 	int ipvx = netproto_version(pkt, pktlen);
@@ -452,7 +470,12 @@ int init_queue(int queue_num) {
 	uint32_t portid = mnl_socket_get_portid(nl);
 
 	struct nlmsghdr *nlh;
-	char buf[BUF_SIZE];
+	NETBUF_ALLOC(bbuf, BUF_SIZE);
+	if (!NETBUF_CHECK(bbuf)) {
+		lgerror("Allocation error", -ENOMEM);
+		goto die_alloc;
+	}
+	char *buf = (char *)bbuf;
 
 	/* Support for kernels versions < 3.8 */
 	// Obsolete and ignored in kernel version 3.8
@@ -547,10 +570,13 @@ int init_queue(int queue_num) {
 	}
 
 
+	NETBUF_FREE(bbuf)
 	close_socket(&nl);
 	return 0;
 
 die:
+	NETBUF_FREE(bbuf)
+die_alloc:
 	close_socket(&nl);
 	return -1;
 }