Fix Winsock recvfrom() timeouts

ethsock.c

@@ -29,12 +29,11 @@ struct ethsock
 {
 	pcap_t *pcap;
 #ifndef NMRPFLASH_WINDOWS
-	struct timeval timeout;
 	int fd;
 #else
-	DWORD timeout;
 	HANDLE handle;
 #endif
+	unsigned timeout;
 	uint8_t hwaddr[6];
 };
 
@@ -314,31 +313,34 @@ cleanup_malloc:
 	return NULL;
 }
 
+int select_fd(int fd, unsigned timeout)
+{
+	struct timeval tv;
+	int status;
+	fd_set fds;
+
+	FD_ZERO(&fds);
+	FD_SET(fd, &fds);
+
+	tv.tv_sec = timeout / 1000;
+	tv.tv_usec = 1000 * (timeout % 1000);
+
+	status = select(fd + 1, &fds, NULL, NULL, &tv);
+	if (status < 0) {
+		sock_perror("select");
+	}
+
+	return status;
+}
+
 ssize_t ethsock_recv(struct ethsock *sock, void *buf, size_t len)
 {
 	struct pcap_pkthdr* hdr;
 	const u_char *capbuf;
 	int status;
-#ifndef NMRPFLASH_WINDOWS
+#ifdef NMRPFLASH_WINDOWS
-	fd_set fds;
-#else
 	DWORD ret;
-#endif
 
-#ifndef NMRPFLASH_WINDOWS
-	if (sock->timeout.tv_sec || sock->timeout.tv_usec) {
-		FD_ZERO(&fds);
-		FD_SET(sock->fd, &fds);
-
-		status = select(sock->fd + 1, &fds, NULL, NULL, &sock->timeout);
-		if (status == -1) {
-			perror("select");
-			return -1;
-		} else if (status == 0) {
-			return 0;
-		}
-	}
-#else
 	if (sock->timeout) {
 		ret = WaitForSingleObject(sock->handle, sock->timeout);
 		if (ret == WAIT_TIMEOUT) {
@@ -348,6 +350,15 @@ ssize_t ethsock_recv(struct ethsock *sock, void *buf, size_t len)
 			return -1;
 		}
 	}
+#else
+	if (sock->timeout) {
+		status = select_fd(sock->fd, sock->timeout);
+		if (status < 0) {
+			return -1;
+		} else if (status == 0) {
+			return 0;
+		}
+	}
 #endif
 
 	status = pcap_next_ex(sock->pcap, &hdr, &capbuf);
@@ -392,14 +403,9 @@ int ethsock_close(struct ethsock *sock)
 	return 0;
 }
 
-int ethsock_set_timeout(struct ethsock *sock, unsigned msec)
+inline int ethsock_set_timeout(struct ethsock *sock, unsigned msec)
 {
-#ifndef NMRPFLASH_WINDOWS
-	sock->timeout.tv_sec = msec / 1000;
-	sock->timeout.tv_usec = (msec % 1000) * 1000;
-#else
 	sock->timeout = msec;
-#endif
 	return 0;
 }
 

nmrpd.h

@@ -68,6 +68,14 @@ struct nmrpd_args {
 
 int tftp_put(struct nmrpd_args *args);
 int nmrp_do(struct nmrpd_args *args);
+int select_fd(int fd, unsigned timeout);
+
+#ifdef NMRPFLASH_WINDOWS
+void win_perror2(const char *msg, DWORD err);
+void sock_perror(const char *msg);
+#else
+#define sock_perror(x) perror(x)
+#endif
 
 extern int verbosity;
 

tftp.c

@@ -111,18 +111,22 @@ static inline void pkt_print(char *pkt, FILE *fp)
 	}
 }
 
-static ssize_t tftp_recvfrom(int sock, char *pkt, struct sockaddr_in *src)
+static ssize_t tftp_recvfrom(int sock, char *pkt, struct sockaddr_in *src,
+		unsigned timeout)
 {
 	ssize_t len;
 
-	len = recvfrom(sock, pkt, TFTP_PKT_SIZE, 0, NULL, NULL);
+	len = select_fd(sock, timeout);
 	if (len < 0) {
-		if (errno != EAGAIN) {
-			perror("recvfrom");
 		return -1;
+	} else if (!len) {
+		return 0;
 	}
 
-		return -2;
+	len = recvfrom(sock, pkt, TFTP_PKT_SIZE, 0, NULL, NULL);
+	if (len < 0) {
+		sock_perror("recvfrom");
+		return -1;
 	}
 
 	uint16_t opcode = pkt_num(pkt);
@@ -176,54 +180,46 @@ static ssize_t tftp_sendto(int sock, char *pkt, size_t len,
 
 	sent = sendto(sock, pkt, len, 0, (struct sockaddr*)dst, sizeof(*dst));
 	if (sent < 0) {
-		perror("sendto");
+		sock_perror("sendto");
 	}
 
 	return sent;
 }
 
-static int sock_set_rx_timeout(int fd, unsigned msec)
+#ifdef NMRPFLASH_WINDOWS
+void sock_perror(const char *msg)
 {
-	struct timeval tv;
+	win_perror2(msg, WSAGetLastError());
-
-	if (msec) {
-		tv.tv_usec = (msec % 1000) * 1000;
-		tv.tv_sec = msec / 1000;
-		if (setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, (char*)&tv, sizeof(tv)) < 0) {
-			perror("setsockopt(SO_RCVTIMEO)");
-			return 1;
 }
+#else
+inline void sock_perror(const char *msg)
+{
+	perror(msg);
 }
-
+#endif
-	return 0;
-}
 
 int tftp_put(struct nmrpd_args *args)
 {
 	struct sockaddr_in addr;
 	uint16_t block;
 	ssize_t len;
-	int fd, sock, err, timeout, last_len;
+	int fd, sock, ret, timeout, last_len;
 	char rx[TFTP_PKT_SIZE], tx[TFTP_PKT_SIZE];
 
 	sock = -1;
+	ret = -1;
 
 	fd = open(args->filename, O_RDONLY);
 	if (fd < 0) {
 		perror("open");
-		err = fd;
+		ret = fd;
 		goto cleanup;
 	}
 
 	sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
 	if (sock < 0) {
-		perror("socket");
+		sock_perror("socket");
-		err = sock;
+		ret = sock;
-		goto cleanup;
-	}
-
-	err = sock_set_rx_timeout(sock, args->rx_timeout);
-	if (err) {
 		goto cleanup;
 	}
 
@@ -252,7 +248,7 @@ int tftp_put(struct nmrpd_args *args)
 				len = read(fd, tx + 4, 512);
 				if (len < 0) {
 					perror("read");
-					err = len;
+					ret = len;
 					goto cleanup;
 				} else if (!len) {
 					if (last_len != 512) {
@@ -263,8 +259,8 @@ int tftp_put(struct nmrpd_args *args)
 				last_len = len;
 			}
 
-			err = tftp_sendto(sock, tx, len, &addr);
+			ret = tftp_sendto(sock, tx, len, &addr);
-			if (err < 0) {
+			if (ret < 0) {
 				goto cleanup;
 			}
 		} else if (pkt_num(rx) != ACK) {
@@ -273,22 +269,22 @@ int tftp_put(struct nmrpd_args *args)
 			fprintf(stderr, "!\n");
 		}
 
-		err = tftp_recvfrom(sock, rx, &addr);
+		ret = tftp_recvfrom(sock, rx, &addr, args->rx_timeout);
-		if (err < 0) {
+		if (ret < 0) {
-			if (err == -2) {
+			goto cleanup;
+		} else if (!ret) {
 			if (++timeout < 5) {
 				continue;
 			}
 			fprintf(stderr, "Timeout while waiting for ACK(%d).\n", block);
-			}
 			goto cleanup;
 		} else {
 			timeout = 0;
-			err = 0;
+			ret = 0;
 		}
 	} while(1);
 
-	err = 0;
+	ret = 0;
 
 cleanup:
 	if (fd >= 0) {
@@ -305,5 +301,5 @@ cleanup:
 #endif
 	}
 
-	return err;
+	return ret;
 }