From e57d43de33c288a4b6bff5a121982fb3fad8d24f Mon Sep 17 00:00:00 2001
From: lohhiiccc <96543753+lohhiiccc@users.noreply.github.com>
Date: Tue, 27 Jan 2026 21:30:49 +0100
Subject: [PATCH] feat: add max_packets parameter to make `icmp_process`
 non-blocking

---
 includes/icmp.h           |  3 ++-
 src/recv/api/process.c    | 10 +++++++---
 tests/recv/test_process.c | 10 +++++-----
 3 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/includes/icmp.h b/includes/icmp.h
index e745b5b..a193dcf 100644
--- a/includes/icmp.h
+++ b/includes/icmp.h
@@ -41,7 +41,8 @@ int icmp_send_echo(icmp_handle_t *h, struct in_addr dest, uint16_t id,
 					uint16_t seq, uint8_t ttl);
 
 /* Receive function */
-int icmp_process(icmp_handle_t *h, icmp_callback_t cb, void *userdata);
+int icmp_process(icmp_handle_t *h, icmp_callback_t cb, void *userdata,
+				size_t max_packets);
 
 /* Error handling */
 const char *icmp_strerror(const icmp_handle_t *h);
diff --git a/src/recv/api/process.c b/src/recv/api/process.c
index b4a8422..50a15ea 100644
--- a/src/recv/api/process.c
+++ b/src/recv/api/process.c
@@ -6,15 +6,18 @@
 #define MAX_PACKET_SIZE 1500
 
 int
-icmp_process(icmp_handle_t *h, icmp_callback_t cb, void *userdata)
+icmp_process(icmp_handle_t *h, icmp_callback_t cb, void *userdata,
+			size_t max_packets)
 {
 	uint8_t buffer[MAX_PACKET_SIZE];
 	struct sockaddr_in from;
 	ssize_t recv_len;
+	int processed;
 
 	if (0 == recv_validate_params(h, cb))
 		return -1;
-	while (1)
+	processed = 0;
+	while (max_packets == 0 || (size_t)processed < max_packets)
 	{
 		recv_len = recv_receive_packet(h->fd, buffer, sizeof(buffer), &from);
 		if (0 == recv_len)
@@ -22,7 +25,8 @@ icmp_process(icmp_handle_t *h, icmp_callback_t cb, void *userdata)
 		if (recv_len < 0)
 			return recv_handle_receive_error(h);
 		recv_process_single_packet(buffer, recv_len, cb, userdata);
+		processed++;
 	}
 	icmp_clear_error(h);
-	return 0;
+	return processed;
 }
diff --git a/tests/recv/test_process.c b/tests/recv/test_process.c
index 11430b2..e5189c6 100644
--- a/tests/recv/test_process.c
+++ b/tests/recv/test_process.c
@@ -26,7 +26,7 @@ Test(process, null_handle)
 {
 	int ret;
 
-	ret = icmp_process(NULL, dummy_callback, NULL);
+	ret = icmp_process(NULL, dummy_callback, NULL, 0);
 	cr_assert_eq(ret, -1);
 }
 
@@ -39,7 +39,7 @@ Test(process, null_callback)
 		cr_skip("Test requires CAP_NET_RAW or root privileges");
 	h = icmp_create();
 	cr_assert_not_null(h);
-	ret = icmp_process(h, NULL, NULL);
+	ret = icmp_process(h, NULL, NULL, 0);
 	cr_assert_eq(ret, -1);
 	cr_assert_str_eq(icmp_strerror(h), "Callback cannot be NULL");
 	icmp_destroy(h);
@@ -54,7 +54,7 @@ Test(process, no_packets_available)
 		cr_skip("Test requires CAP_NET_RAW or root privileges");
 	h = icmp_create();
 	cr_assert_not_null(h);
-	ret = icmp_process(h, dummy_callback, NULL);
+	ret = icmp_process(h, dummy_callback, NULL, 0);
 	cr_assert_eq(ret, 0);
 	icmp_destroy(h);
 }
@@ -75,8 +75,8 @@ Test(process, echo_loopback, .disabled = true)
 	cr_assert_eq(ret, 0);
 	sleep(1);
 	count = 0;
-	ret = icmp_process(h, counter_callback, &count);
-	cr_assert_eq(ret, 0);
+	ret = icmp_process(h, counter_callback, &count, 0);
+	cr_assert_geq(ret, 0);
 	cr_assert_geq(count, 1, "Should receive at least one packet");
 	icmp_destroy(h);
 }