PATCH 3/3: Add ztcp -s to shutdown() a session fd

[Mikael Magnusson <mikachu@xxxxxxxxx>]

This lets you close tcp connections without losing data more reliably,
in addition to being able to signal you are done sending data.

---

I'm far from a networking expert, so feel free to point out any mistakes
in my understanding of how this works.

Updating zshtcpsys is left as an exercise for someone else. I think zftp
is fine because it handles draining all the data on the protocol level.

 Doc/Zsh/mod_tcp.yo | 30 ++++++++++++++++++++++++++---
 Src/Modules/tcp.c  | 48 +++++++++++++++++++++++++++++++++++++++++++---
 Src/Modules/tcp.h  |  5 +++--
 3 files changed, 75 insertions(+), 8 deletions(-)

diff --git a/Doc/Zsh/mod_tcp.yo b/Doc/Zsh/mod_tcp.yo
index d19e8cc927..ea55c80ae8 100644
--- a/Doc/Zsh/mod_tcp.yo
+++ b/Doc/Zsh/mod_tcp.yo
@@ -7,7 +7,7 @@ startitem()
 findex(ztcp)
 cindex(TCP)
 cindex(sockets, TCP)
-item(tt(ztcp) [ tt(-acflLtv) ] [ tt(-d) var(fd) ] [ var(args) ])(
+item(tt(ztcp) [ tt(-acflLstv) ] [ tt(-d) var(fd) ] [ var(args) ])(
 tt(ztcp) is implemented as a builtin to allow full use of shell
 command line editing, file I/O, and job control mechanisms.
 
@@ -24,7 +24,7 @@ startitem()
 item(File descriptor)(
 The file descriptor in use for the connection.  For normal inbound (tt(I))
 and outbound (tt(O)) connections this may be read and written by the usual
-shell mechanisms.  However, it should only be close with `tt(ztcp -c)'.
+shell mechanisms.  However, it should only be closed with `tt(ztcp -c)'.
 )
 item(Connection type)(
 A letter indicating how the session was created:
@@ -42,6 +42,8 @@ An inbound connection accepted with `tt(ztcp -a)'.
 item(tt(O))(
 An outbound connection created with `tt(ztcp) var(host) var(...)'.
 )
+item(tt(S))(
+A connection that has been shut down with `tt(ztcp -s)'.
 enditem()
 
 )
@@ -121,6 +123,24 @@ enditem()
 subsect(Closing Connections)
 cindex(sockets, closing TCP)
 
+startitem()
+item(tt(ztcp) tt(-s) var(fd))(
+tt(ztcp -s) will shut down the socket associated
+with var(fd). This means a graceful termination
+of the connection is initiated, flushing unsent
+data (due to Nagle's algorithm, for example). To
+complete it, you should also drain any remaining
+data on the fd by reading it. Afterwards you should
+use tt(ztcp -c) to close the fd. If you close it
+without first using tt(ztcp -s) and draining the
+data, the kernel may reset the connection, causing
+data you sent to not be received by the other end.
+
+Note that this is only useful for inbound and
+outbound connections, not for listening sockets.
+)
+enditem()
+
 startitem()
 xitem(tt(ztcp) tt(-cf) [ tt(-v) ] [ var(fd) ])
 item(tt(ztcp) tt(-c) [ tt(-v) ] [ var(fd) ])(
@@ -129,6 +149,8 @@ with var(fd).  The socket will be removed from the
 session table.  If var(fd) is not specified,
 tt(ztcp) will close everything in the session table.
 
+See also the previous tt(ztcp -s) entry.
+
 Normally, sockets registered by zftp (see
 sectref(The zsh/zftp Module)(zshmodules))
 cannot be closed this way.  In order
@@ -166,6 +188,8 @@ prints `tt(This is a message)'.
 
 To tidy up, on tt(host1):
 example(ztcp -c $listenfd
+ztcp -s $fd
 ztcp -c $fd)
 and on tt(host2)
-example(ztcp -c $fd)
+example(ztcp -s $fd
+ztcp -c $fd)
diff --git a/Src/Modules/tcp.c b/Src/Modules/tcp.c
index 0bbce5d49f..4366f3c026 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)
@@ -594,6 +633,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 +647,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 +738,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 = {
diff --git a/Src/Modules/tcp.h b/Src/Modules/tcp.h
index f69e246e0e..a9ea7df2d1 100644
--- a/Src/Modules/tcp.h
+++ b/Src/Modules/tcp.h
@@ -81,8 +81,9 @@ union tcp_sockaddr {
 
 typedef struct tcp_session *Tcp_session;
 
-#define ZTCP_LISTEN  1
-#define ZTCP_INBOUND 2
+#define ZTCP_LISTEN   1
+#define ZTCP_INBOUND  2
+#define ZTCP_SHUTDOWN 4
 #define ZTCP_ZFTP    16
 
 struct tcp_session {
-- 
2.38.1