36941: Mark file descripors in ztcp as used.

Allow such file descriptors to be either internal and closed on exec
or external and so managed explicitly by module.

3 files changed, 43 insertions, 5 deletions

diff --git a/Src/Modules/tcp.c b/Src/Modules/tcp.c
index bc1765da1..274f01ffc 100644
--- a/Src/Modules/tcp.c
+++ b/Src/Modules/tcp.c
@@ -236,6 +236,8 @@ tcp_socket(int domain, int type, int protocol, int ztflags)
     if (!sess) return NULL;
 
     sess->fd = socket(domain, type, protocol);
+    /* We'll check failure and tidy up in caller */
+    addmodulefd(sess->fd, FALSE);
     return sess;
 }
 
@@ -298,7 +300,7 @@ tcp_close(Tcp_session sess)
     {  
 	if (sess->fd != -1)
 	{
-	    err = close(sess->fd);
+	    err = zclose(sess->fd);
 	    if (err)
 		zwarn("connection close failed: %e", errno);
 	}
@@ -546,6 +548,9 @@ bin_ztcp(char *nam, char **args, Options ops, UNUSED(int func))
 	    return 1;
 	}
 
+	/* redup expects fd is already registered */
+	addmodulefd(rfd, FALSE);
+
 	if (targetfd) {
 	    sess->fd = redup(rfd, targetfd);
 	    if (sess->fd < 0) {
diff --git a/Src/utils.c b/Src/utils.c
index 61cbe84aa..4c69d75a1 100644
--- a/Src/utils.c
+++ b/Src/utils.c
@@ -1892,6 +1892,32 @@ redup(int x, int y)
 }
 
 /*
+ * Add an fd opened ithin a module.
+ *
+ * If is_internal is FALSE, the fd can be used within the shell for
+ * normal I/O but it will not be closed automatically or by normal shell
+ * syntax; it can only be closed by the module (which should included
+ * zclose() as part of the sequence).
+ *
+ * If is_internal is TRUE the fd is treated like others created by the
+ * shell for internall use; it can be closed and will be closed by the
+ * shell if it exec's or performs an exec with a fork optimised out.
+ *.
+ * Safe if fd is -1 to indicate failure.
+ */
+/**/
+mod_export void
+addmodulefd(int fd, bool is_internal)
+{
+    if (fd >= 0) {
+	check_fd_table(fd);
+	fdtable[fd] = is_internal ? FDT_INTERNAL : FDT_MODULE;
+    }
+}
+
+/**/
+
+/*
  * Indicate that an fd has a file lock; if cloexec is 1 it will be closed
  * on exec.
  * The fd should already be known to fdtable (e.g. by movefd).
diff --git a/Src/zsh.h b/Src/zsh.h
index 15fa5e417..f819249c3 100644
--- a/Src/zsh.h
+++ b/Src/zsh.h
@@ -406,25 +406,32 @@ enum {
  */
 #define FDT_EXTERNAL		2
 /*
+ * Entry visible to other processes but controlled by a module.
+ * The difference from FDT_EXTERNAL is that closing this using
+ * standard fd syntax will fail as there is some tidying up that
+ * needs to be done by the module's own mechanism.
+ */
+#define FDT_MODULE		3
+/*
  * Entry used by output from the XTRACE option.
  */
-#define FDT_XTRACE		3
+#define FDT_XTRACE		4
 /*
  * Entry used for file locking.
  */
-#define FDT_FLOCK		4
+#define FDT_FLOCK		5
 /*
  * As above, but the fd is not marked for closing on exec,
  * so the shell can still exec the last process.
  */
-#define FDT_FLOCK_EXEC		5
+#define FDT_FLOCK_EXEC		6
 #ifdef PATH_DEV_FD
 /*
  * Entry used by a process substition.
  * This marker is not tested internally as we associated the file
  * descriptor with a job for closing.
  */
-#define FDT_PROC_SUBST		6
+#define FDT_PROC_SUBST		7
 #endif
 
 /* Flags for input stack */