1 files changed, 64 insertions, 9 deletions

diff --git a/Src/Modules/tcp.c b/Src/Modules/tcp.c
index 0bbce5d49..36c3dca37 100644
--- a/Src/Modules/tcp.c
+++ b/Src/Modules/tcp.c
@@ -279,6 +279,29 @@ zts_byfd(int fd)
     return NULL;
 }
 
+/**/
+mod_export int
+tcp_shutdown(Tcp_session sess)
+{
+    if (sess && sess->fd != -1)
+    {
+	if (sess->flags & ZTCP_LISTEN) {
+	    zwarn("can't shutdown a listening socket");
+	    return 1;
+	}
+	if (sess->flags & ZTCP_SHUTDOWN) {
+	    zwarn("session is already shutdown");
+	    return 1;
+	}
+	int err = shutdown(sess->fd, SHUT_WR);
+	if (err)
+	    zwarn("shutdown failed: %e", errno);
+	sess->flags |= ZTCP_SHUTDOWN;
+	return 0;
+    }
+    return -1;
+}
+
 static void
 tcp_cleanup(void)
 {
@@ -366,6 +389,22 @@ bin_ztcp(char *nam, char **args, Options ops, UNUSED(int func))
 	}
     }
 
+    if (OPT_ISSET(ops,'s')) {
+	targetfd = atoi(args[0]);
+	if(!targetfd) {
+	    zwarnnam(nam, "%s is an invalid argument to -s", args[0]);
+	    return 1;
+	}
+
+	sess = zts_byfd(targetfd);
+	if (sess) {
+	    tcp_shutdown(sess);
+	    return 0;
+	} else {
+	    zwarnnam(nam, "fd %s not found in tcp table", args[0]);
+	    return 1;
+	}
+    }
 
     if (OPT_ISSET(ops,'c')) {
 	if (!args[0]) {
@@ -373,12 +412,12 @@ bin_ztcp(char *nam, char **args, Options ops, UNUSED(int func))
 	}
 	else {
 	    targetfd = atoi(args[0]);
-	    sess = zts_byfd(targetfd);
 	    if(!targetfd) {
 		zwarnnam(nam, "%s is an invalid argument to -c", args[0]);
 		return 1;
 	    }
 
+	    sess = zts_byfd(targetfd);
 	    if (sess)
 	    {
 		if ((sess->flags & ZTCP_ZFTP) && !force)
@@ -450,18 +489,26 @@ bin_ztcp(char *nam, char **args, Options ops, UNUSED(int func))
 
 	if (targetfd) {
 	    sess->fd = redup(sess->fd, targetfd);
+	    if (sess->fd < 0) {
+		zerrnam(nam, "could not duplicate socket fd to %d: %e", targetfd, errno);
+		tcp_close(sess);
+		return 1;
+	    }
 	}
 	else {
-	    /* move the fd since no one will want to read from it */
+	    targetfd = sess->fd; /* just stashing here for error message */
 	    sess->fd = movefd(sess->fd);
 	}
 
 	if (sess->fd == -1) {
-	    zwarnnam(nam, "cannot duplicate fd %d: %e", sess->fd, errno);
+	    zwarnnam(nam, "cannot duplicate fd %d: %e", targetfd, errno);
 	    tcp_close(sess);
 	    return 1;
 	}
 
+	/* disallow to be closed explicitly, after movefd set it to INTERNAL */
+	fdtable[sess->fd] = FDT_MODULE;
+
 	setiparam_no_convert("REPLY", (zlong)sess->fd);
 
 	if (verbose)
@@ -549,19 +596,24 @@ bin_ztcp(char *nam, char **args, Options ops, UNUSED(int func))
 	    return 1;
 	}
 
-	/* redup expects fd is already registered */
-	addmodulefd(rfd, FDT_MODULE);
-
 	if (targetfd) {
 	    sess->fd = redup(rfd, targetfd);
 	    if (sess->fd < 0) {
 		zerrnam(nam, "could not duplicate socket fd to %d: %e", targetfd, errno);
+		tcp_close(sess);
 		return 1;
 	    }
 	}
 	else {
-	    sess->fd = rfd;
+	    sess->fd = movefd(rfd);
+	}
+	if (sess->fd < 0) {
+	    zerrnam(nam, "could not move socket fd %d: %e", rfd, errno);
+	    tcp_close(sess);
+	    return 1;
 	}
+	/* disallow to be closed explicitly */
+	fdtable[sess->fd] = FDT_MODULE;
 
 	setiparam_no_convert("REPLY", (zlong)sess->fd);
 
@@ -594,6 +646,8 @@ bin_ztcp(char *nam, char **args, Options ops, UNUSED(int func))
 			    schar = 'Z';
 			else if (sess->flags & ZTCP_LISTEN)
 			    schar = 'L';
+			else if (sess->flags & ZTCP_SHUTDOWN)
+			    schar = 'S';
 			else if (sess->flags & ZTCP_INBOUND)
 			    schar = 'I';
 			else
@@ -606,7 +660,8 @@ bin_ztcp(char *nam, char **args, Options ops, UNUSED(int func))
 			printf("%s:%d %s %s:%d is on fd %d%s\n",
 			       localname, ntohs(sess->sock.in.sin_port),
 			       ((sess->flags & ZTCP_LISTEN) ? "-<" :
-				((sess->flags & ZTCP_INBOUND) ? "<-" : "->")),
+				((sess->flags & ZTCP_SHUTDOWN) ? "<>" :
+				 ((sess->flags & ZTCP_INBOUND) ? "<-" : "->"))),
 			       remotename, ntohs(sess->peer.in.sin_port),
 			       sess->fd,
 			       (sess->flags & ZTCP_ZFTP) ? " ZFTP" : "");
@@ -696,7 +751,7 @@ bin_ztcp(char *nam, char **args, Options ops, UNUSED(int func))
 }
 
 static struct builtin bintab[] = {
-    BUILTIN("ztcp", 0, bin_ztcp, 0, 3, 0, "acd:flLtv", NULL),
+    BUILTIN("ztcp", 0, bin_ztcp, 0, 3, 0, "acd:flLstv", NULL),
 };
 
 static struct features module_features = {