From a0c1624125665e55f42171d43d772aa8f2ad81ce Mon Sep 17 00:00:00 2001
From: "Joseph C. Lehner" <joseph.c.lehner@gmail.com>
Date: Mon, 23 Oct 2023 16:07:49 +0200
Subject: [PATCH] Don't try to drain late NMRP packets in one go

---
 nmrp.c  | 33 ++++++++++++++-------------------
 nmrpd.h |  7 ++++---
 tftp.c  | 10 +++-------
 3 files changed, 21 insertions(+), 29 deletions(-)

diff --git a/nmrp.c b/nmrp.c
index 9424983..95416ea 100644
--- a/nmrp.c
+++ b/nmrp.c
@@ -348,34 +348,30 @@ static void sigh(int sig)
 	g_interrupted = 1;
 }
 
-static void nmrp_drain(void* arg)
+void nmrp_discard(struct ethsock *sock)
 {
 	// between nmrpflash sending the TFTP WRQ packet, and the router
 	// responding with ACK(0)/OACK, some devices send extraneous
 	// CONF_REQ and/or TFTP_UL_REQ packets.
 	//
-	// we drain the NMRP receive buffer here, otherwise it might seem
-	// as if these packets arrived *after* the TFTP upload.
+	// the TFTP code thus calls this function in each loop iteration,
+	// to discard any late NMRP packets.
+	//
+	// without this it might seem  as if these packets arrived after
+	// the TFTP upload completed successfuly, confusing the NMRP code.
 
-	struct ethsock* sock = (struct ethsock*)arg;
 	unsigned timeout = ethsock_get_timeout(sock);
-	ethsock_set_timeout(sock, 0);
-	long long beg = millis();
+	// don't set this to 0, as this would cause pkt_recv to block!
+	ethsock_set_timeout(sock, 1);
 
 	struct nmrp_pkt rx;
-	int i = 0;
 
-	while (pkt_recv(sock, &rx) == 0) {
+	if (pkt_recv(sock, &rx) == 0) {
 		if (rx.msg.code != NMRP_C_CONF_REQ && rx.msg.code != NMRP_C_TFTP_UL_REQ) {
-			if (verbosity > 1) {
-				printf("Drained unexpected packet type %s\n", msg_code_str(rx.msg.code));
-			}
+			printf("Discarding unexpected %s packet\n", msg_code_str(rx.msg.code));
+		} else if (verbosity > 1) {
+			printf("Discarding %s packet\n", msg_code_str(rx.msg.code));
 		}
-		++i;
-	}
-
-	if (verbosity > 1) {
-		printf("Drained %d packet(s) from rx buffer in %lld ms\n", i, millis() - beg);
 	}
 
 	ethsock_set_timeout(sock, timeout);
@@ -470,6 +466,8 @@ int nmrp_do(struct nmrpd_args *args)
 		return 1;
 	}
 
+	args->sock = sock;
+
 	sigh_orig = signal(SIGINT, sigh);
 
 	if (ethsock_is_unplugged(sock)) {
@@ -597,9 +595,6 @@ int nmrp_do(struct nmrpd_args *args)
 	ulreqs = 0;
 	ka_reqs = 0;
 
-	args->on_tftp_ack0_callback = &nmrp_drain;
-	args->on_tftp_ack0_arg = sock;
-
 	while (!g_interrupted) {
 		if (expect != NMRP_C_NONE && rx.msg.code != expect) {
 			fprintf(stderr, "Received %s while waiting for %s!\n",
diff --git a/nmrpd.h b/nmrpd.h
index f647c09..15584f2 100644
--- a/nmrpd.h
+++ b/nmrpd.h
@@ -96,6 +96,8 @@ enum nmrp_op {
 	NMRP_SET_REGION = 2,
 };
 
+struct ethsock;
+
 struct nmrpd_args {
 	unsigned rx_timeout;
 	unsigned ul_timeout;
@@ -113,8 +115,7 @@ struct nmrpd_args {
 	const char *region;
 	off_t offset;
 
-	void (*on_tftp_ack0_callback)(void*);
-	void* on_tftp_ack0_arg;
+	struct ethsock *sock;
 };
 
 const char *leafname(const char *path);
@@ -122,6 +123,7 @@ ssize_t tftp_put(struct nmrpd_args *args);
 bool tftp_is_valid_filename(const char *filename);
 
 int nmrp_do(struct nmrpd_args *args);
+void nmrp_discard(struct ethsock *sock);
 
 int select_fd(int fd, unsigned timeout);
 const char *mac_to_str(uint8_t *mac);
@@ -135,7 +137,6 @@ void sock_perror(const char *msg);
 
 extern int verbosity;
 
-struct ethsock;
 struct ethsock_arp_undo;
 struct ethsock_ip_undo;
 
diff --git a/tftp.c b/tftp.c
index bf18f08..007912d 100644
--- a/tftp.c
+++ b/tftp.c
@@ -328,7 +328,7 @@ ssize_t tftp_put(struct nmrpd_args *args)
 	char rx[2048], tx[2048];
 	const char *file_remote = args->file_remote;
 	char *val, *end;
-	bool rollover, tftp_ack0_callback_called;
+	bool rollover;
 	const unsigned rx_timeout = MAX(args->rx_timeout / (args->blind ? 50 : 5), 2000);
 	const unsigned max_timeouts = args->blind ? 3 : 5;
 #ifndef NMRPFLASH_WINDOWS
@@ -414,7 +414,6 @@ ssize_t tftp_put(struct nmrpd_args *args)
 	rollover = false;
 	/* Not really, but this way the loop sends our WRQ before receiving */
 	timeouts = 1;
-	tftp_ack0_callback_called = false;
 
 	pkt_mkwrq(tx, file_remote, TFTP_BLKSIZE);
 
@@ -440,11 +439,6 @@ ssize_t tftp_put(struct nmrpd_args *args)
 					}
 				}
 			}
-
-			if (!ackblock && args->on_tftp_ack0_callback && !tftp_ack0_callback_called) {
-				args->on_tftp_ack0_callback(args->on_tftp_ack0_arg);
-				tftp_ack0_callback_called = true;
-			}
 		}
 
 		if (timeouts || ackblock == block) {
@@ -496,6 +490,8 @@ ssize_t tftp_put(struct nmrpd_args *args)
 		}
 
 		ret = tftp_recvfrom(sock, rx, &port, rx_timeout, blksize + 4);
+		nmrp_discard(args->sock);
+
 		if (ret < 0) {
 			goto cleanup;
 		} else if (!ret) {