Add pass_fds parameter to PtyProcess:spawn()

2 files changed, 45 insertions, 3 deletions

diff --git a/ptyprocess/ptyprocess.py b/ptyprocess/ptyprocess.py
index 29b4e43..78d19fd 100644
--- a/ptyprocess/ptyprocess.py
+++ b/ptyprocess/ptyprocess.py
@@ -178,7 +178,7 @@ class PtyProcess(object):
     @classmethod
     def spawn(
             cls, argv, cwd=None, env=None, echo=True, preexec_fn=None,
-            dimensions=(24, 80)):
+            dimensions=(24, 80), pass_fds=()):
         '''Start the given command in a child process in a pseudo terminal.
 
         This does all the fork/exec type of stuff for a pty, and returns an
@@ -190,6 +190,10 @@ class PtyProcess(object):
 
         Dimensions of the psuedoterminal used for the subprocess can be
         specified as a tuple (rows, cols), or the default (24, 80) will be used.
+
+        By default, all file descriptors except 0, 1 and 2 are closed. This
+        behavior can be overridden with pass_fds, a list of file descriptors to
+        keep open between the parent and the child.
         '''
         # Note that it is difficult for this method to fail.
         # You cannot detect if the child process cannot start.
@@ -255,12 +259,14 @@ class PtyProcess(object):
 
             # Do not allow child to inherit open file descriptors from parent,
             # with the exception of the exec_err_pipe_write of the pipe
+            # and pass_fds.
             # Impose ceiling on max_fd: AIX bugfix for users with unlimited
             # nofiles where resource.RLIMIT_NOFILE is 2^63-1 and os.closerange()
             # occasionally raises out of range error
             max_fd = min(1048576, resource.getrlimit(resource.RLIMIT_NOFILE)[0])
-            os.closerange(3, exec_err_pipe_write)
-            os.closerange(exec_err_pipe_write+1, max_fd)
+            spass_fds = sorted(set(pass_fds) | {exec_err_pipe_write})
+            for pair in zip([2] + spass_fds, spass_fds + [max_fd]):
+                os.closerange(pair[0]+1, pair[1])
 
             if cwd is not None:
                 os.chdir(cwd)
diff --git a/tests/test_spawn.py b/tests/test_spawn.py
index 696e126..d69548d 100755
--- a/tests/test_spawn.py
+++ b/tests/test_spawn.py
@@ -1,6 +1,8 @@
+import fcntl
 import os
 import time
 import select
+import tempfile
 import unittest
 from ptyprocess.ptyprocess import which
 from ptyprocess import PtyProcess, PtyProcessUnicode
@@ -111,3 +113,37 @@ class PtyTestCase(unittest.TestCase):
     @unittest.skipIf(which('bc') is None, "bc(1) not found on this server.")
     def test_interactive_repl_unicode_echo(self):
         self._interactive_repl_unicode(echo=True)
+
+    def test_pass_fds(self):
+        with tempfile.NamedTemporaryFile() as temp_file:
+            temp_file_fd = temp_file.fileno()
+            temp_file_name = temp_file.name
+
+            # Temporary files are CLOEXEC by default
+            fcntl.fcntl(temp_file_fd,
+                        fcntl.F_SETFD,
+                        fcntl.fcntl(temp_file_fd, fcntl.F_GETFD) &
+                        ~fcntl.FD_CLOEXEC)
+
+            # You can write with pass_fds
+            p = PtyProcess.spawn(['bash',
+                                  '-c',
+                                  'printf hello >&{}'.format(temp_file_fd)],
+                                 echo=True,
+                                 pass_fds=(temp_file_fd,))
+            p.wait()
+            assert p.status == 0
+
+            with open(temp_file_name, 'r') as temp_file_r:
+                assert temp_file_r.read() == 'hello'
+
+            # You can't write without pass_fds
+            p = PtyProcess.spawn(['bash',
+                                  '-c',
+                                  'printf bye >&{}'.format(temp_file_fd)],
+                                 echo=True)
+            p.wait()
+            assert p.status != 0
+
+            with open(temp_file_name, 'r') as temp_file_r:
+                assert temp_file_r.read() == 'hello'