tftp: Fix handling of port name or number

MFC after:	1 week
Reviewed by:	markj
Differential Revision:	https://reviews.freebsd.org/D57105

(cherry picked from commit 47e4571df15804fe2b010a3c82b22691d80ee680)

2 files changed, 26 insertions, 21 deletions

diff --git a/usr.bin/tftp/main.c b/usr.bin/tftp/main.c
index 3db01f53dd78..10343bff412e 100644
--- a/usr.bin/tftp/main.c
+++ b/usr.bin/tftp/main.c
@@ -344,6 +344,8 @@ setpeer0(char *host, const char *lport)
 	}
 
 	freeaddrinfo(res0);
+	free(port);
+	port = strdup(lport);
 }
 
 static void
@@ -363,10 +365,9 @@ setpeer(int argc, char *argv[])
 		printf("usage: %s [host [port]]\n", argv[0]);
 		return;
 	}
-	if (argc == 3) {
-		port = argv[2];
+	if (argc == 3)
 		setpeer0(argv[1], argv[2]);
-	} else
+	else
 		setpeer0(argv[1], NULL);
 }
 
diff --git a/usr.bin/tftp/tftp.c b/usr.bin/tftp/tftp.c
index 9d727fe8bc2c..d1923864a06d 100644
--- a/usr.bin/tftp/tftp.c
+++ b/usr.bin/tftp/tftp.c
@@ -70,8 +70,9 @@ int
 xmitfile(int peer, char *port, int fd, char *name, char *mode)
 {
 	struct tftphdr *rp;
+	struct servent *se;
 	int n, i, ret = 0;
-	uint16_t block;
+	uint16_t block, portn;
 	struct sockaddr_storage serv;	/* valid server port number */
 	char recvbuffer[MAXPKTSIZE];
 	struct tftp_stats tftp_stats;
@@ -81,14 +82,15 @@ xmitfile(int peer, char *port, int fd, char *name, char *mode)
 	memset(&serv, 0, sizeof(serv));
 	rp = (struct tftphdr *)recvbuffer;
 
-	if (port == NULL) {
-		struct servent *se;
-		se = getservbyname("tftp", "udp");
-		assert(se != NULL);
-		((struct sockaddr_in *)&peer_sock)->sin_port = se->s_port;
-	} else
-		((struct sockaddr_in *)&peer_sock)->sin_port =
-		    htons(atoi(port));
+	if ((se = getservbyname(port ? port : "tftp", "udp")) != NULL)
+		portn = se->s_port;
+	else
+		portn = htons(atoi(port));
+	if (portn == 0) {
+		printf("Invalid port '%s'.\n", port ? port : "tftp");
+		return -1;
+	}
+	((struct sockaddr_in *)&peer_sock)->sin_port = portn;
 
 	for (i = 0; i < 12; i++) {
 		struct sockaddr_storage from;
@@ -172,7 +174,8 @@ int
 recvfile(int peer, char *port, int fd, char *name, char *mode)
 {
 	struct tftphdr *rp;
-	uint16_t block;
+	struct servent *se;
+	uint16_t block, portn;
 	char recvbuffer[MAXPKTSIZE];
 	int n, i, ret = 0;
 	struct tftp_stats tftp_stats;
@@ -181,14 +184,15 @@ recvfile(int peer, char *port, int fd, char *name, char *mode)
 
 	rp = (struct tftphdr *)recvbuffer;
 
-	if (port == NULL) {
-		struct servent *se;
-		se = getservbyname("tftp", "udp");
-		assert(se != NULL);
-		((struct sockaddr_in *)&peer_sock)->sin_port = se->s_port;
-	} else
-		((struct sockaddr_in *)&peer_sock)->sin_port =
-		    htons(atoi(port));
+	if ((se = getservbyname(port ? port : "tftp", "udp")) != NULL)
+		portn = se->s_port;
+	else
+		portn = htons(atoi(port));
+	if (portn == 0) {
+		printf("Invalid port '%s'.\n", port ? port : "tftp");
+		return (-1);
+	}
+	((struct sockaddr_in *)&peer_sock)->sin_port = portn;
 
 	for (i = 0; i < 12; i++) {
 		struct sockaddr_storage from;