Add support for output redirection

I had hoped that we could provide access to a subprocess.Popen()
instance directly so users could do whatever they want with the .stdout
and .stderr pipes. However, that's not always possible (e.g. the chroot
backend can't return the Popen object it creates to the caller, because
it's in a different process).

The current approach isn't groundbreaking but it is quite simple.

3 files changed, 96 insertions, 42 deletions

diff --git a/sandboxlib/__init__.py b/sandboxlib/__init__.py
index 3aeaf1c..d620777 100644
--- a/sandboxlib/__init__.py
+++ b/sandboxlib/__init__.py
@@ -23,6 +23,7 @@ docstrings that describe the different parameters.
 
 
 import logging
+import os
 import platform
 import shutil
 import subprocess
@@ -58,8 +59,20 @@ def maximum_possible_isolation():
     raise NotImplementedError()
 
 
-def run_sandbox(rootfs_path, command, cwd=None, extra_env=None,
-                network='undefined'):
+
+# Special value for 'stderr' and 'stdout' parameters to indicate 'capture
+# and return the data'.
+CAPTURE = subprocess.PIPE
+
+# Special value for 'stderr' parameter to indicate 'forward to stdout'.
+STDOUT = subprocess.STDOUT
+
+
+def run_sandbox(command, cwd=None, extra_env=None,
+                filesystem_root='/', filesystem_writable_paths='all',
+                mounts='undefined', extra_mounts=None,
+                network='undefined',
+                stderr=CAPTURE, stdout=CAPTURE):
     '''Run 'command' in a sandboxed environment.
 
     Parameters:
@@ -89,29 +102,28 @@ def run_sandbox(rootfs_path, command, cwd=None, extra_env=None,
             no attempt is made to either prevent or provide networking
             inside the sandbox. Backends may support 'isolated' and/or other
             values as well.
+      - stdout: whether to capture stdout, or redirect stdout to a file handle.
+            If set to sandboxlib.CAPTURE, the function will return the stdout
+            data, if not, it will return None for that. If stdout=None, the
+            data will be discarded -- it will NOT inherit the parent process's
+            stdout, unlike with subprocess.Popen(). Set 'stdout=sys.stdout' if
+            you want that.
+      - stderr: same as stdout
+
+    Returns:
+      a tuple of (exit code, stdout output, stderr output).
 
     '''
     raise NotImplementedError()
 
 
-def Popen(command, stderr=None, stdout=None, **sandbox_config):
-    '''Start a subprocess in a sandbox and return straight away.
+def run_sandbox_with_redirection(command, **sandbox_config):
+    '''Start a subprocess in a sandbox, redirecting stderr and/or stdout.
 
-    This function aims to function like subprocess.Popen(), but with the
-    subprocess running inside a sandbox. It returns a subprocess.Popen
-    instance as soon as 'command' starts executing.
+    The sandbox_config arguments are the same as the run_command() function.
 
-    The 'stderr' and 'stdout' parameters accept None, a file-like object, a
-    file descriptor (integer), or subprocess.PIPE. The only difference from
-    the subprocess.Popen() function is that 'None' means 'ignore all output'
-    rather than 'inherit parent's stdout': if you want to forward output from
-    the subprocess to stdout, you must pass `stdout=sys.stdout`.
-
-    The sandbox_config arguments are the same as the run_command() function. In
-    most cases you should use run_command() instead of this, but there are
-    certain cases where Popen() could be useful. The run_command() function
-    buffers all data from stdout and stderr of the subprocess in memory, which
-    is impractical if there is a huge amount of data.
+    This returns just the exit code, because if stdout or stderr are redirected
+    those values will be None in any case.
 
     '''
     raise NotImplementedError()
@@ -205,26 +217,51 @@ def validate_extra_mounts(extra_mounts):
     return new_extra_mounts
 
 
-def _run_command(argv, cwd=None, env=None, preexec_fn=None):
+def _run_command(argv, stdout, stderr, cwd=None, env=None):
     '''Wrapper around subprocess.Popen() with common settings.
 
-    This function blocks until the subprocesses has terminated. It then
-    returns a tuple of (exit code, stdout output, stderr output).
+    This function blocks until the subprocess has terminated.
+
+    Unlike the subprocess.Popen() function, if stdout or stderr are None then
+    output is discarded.
+
+    It then returns a tuple of (exit code, stdout output, stderr output).
+    If stdout was not equal to subprocess.PIPE, stdout will be None. Same for
+    stderr.
 
     '''
-    process = subprocess.Popen(
-        argv,
-        # The default is to share file descriptors from the parent process
-        # to the subprocess, which is rarely good for sandboxing.
-        close_fds=True,
-        cwd=cwd,
-        env=env,
-        preexec_fn=preexec_fn,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE
-    )
-    process.wait()
-    return process.returncode, process.stdout.read(), process.stderr.read()
+    if stdout is None or stderr is None:
+        dev_null = open(os.devnull, 'w')
+        stdout = stdout or dev_null
+        stderr = stderr or dev_null
+    else:
+        dev_null = None
+
+    try:
+        process = subprocess.Popen(
+            argv,
+            # The default is to share file descriptors from the parent process
+            # to the subprocess, which is rarely good for sandboxing.
+            close_fds=True,
+            cwd=cwd,
+            env=env,
+            stdout=stdout,
+            stderr=stderr,
+        )
+
+        # The 'out' variable will be None unless subprocess.PIPE was passed as
+        # 'stdout' to subprocess.Popen(). Same for 'err' and 'stderr'. If
+        # subprocess.PIPE wasn't passed for either it'd be safe to use .wait()
+        # instead of .communicate(), but if they were then we must use
+        # .communicate() to avoid blocking the subprocess if one of the pipes
+        # becomes full. It's safe to use .communicate() in all cases.
+
+        out, err = process.communicate()
+    finally:
+        if dev_null is not None:
+            dev_null.close()
+
+    return process.returncode, out, err
 
 
 # Executors
diff --git a/sandboxlib/chroot.py b/sandboxlib/chroot.py
index 530e667..97391de 100644
--- a/sandboxlib/chroot.py
+++ b/sandboxlib/chroot.py
@@ -87,7 +87,7 @@ def mount(source, path, mount_type, mount_options):
     # should do that instead.
     argv = [
         'mount', '-t', mount_type, '-o', mount_options, source, path]
-    exit, out, err = sandboxlib._run_command(argv)
+    exit, out, err = sandboxlib._run_command(argv, stdout=None, stderr=None)
 
     if exit != 0:
         raise RuntimeError(
@@ -97,7 +97,7 @@ def mount(source, path, mount_type, mount_options):
 
 def unmount(path):
     argv = ['umount', path]
-    exit, out, err = sandboxlib._run_command(argv)
+    exit, out, err = sandboxlib._run_command(argv, stdout=None, stderr=None)
 
     if exit != 0:
         warnings.warn("%s failed: %s" % (
@@ -127,7 +127,8 @@ def mount_all(rootfs_path, mount_info_list):
             unmount(mountpoint)
 
 
-def run_command_in_chroot(pipe, extra_mounts, chroot_path, command, cwd, env):
+def run_command_in_chroot(pipe, stdout, stderr, extra_mounts, chroot_path,
+                          command, cwd, env):
     # This function should be run in a multiprocessing.Process() subprocess,
     # because it calls os.chroot(). There's no 'unchroot()' function! After
     # chrooting, it calls sandboxlib._run_command(), which uses the
@@ -157,7 +158,8 @@ def run_command_in_chroot(pipe, extra_mounts, chroot_path, command, cwd, env):
                 raise RuntimeError(
                     "Unable to set current working directory: %s" % e)
 
-        exit, out, err = sandboxlib._run_command(command, env=env)
+        exit, out, err = sandboxlib._run_command(
+            command, stdout, stderr, env=env)
         pipe.send([exit, out, err])
         result = 0
     except Exception as e:
@@ -169,7 +171,8 @@ def run_command_in_chroot(pipe, extra_mounts, chroot_path, command, cwd, env):
 def run_sandbox(command, cwd=None, extra_env=None,
                 filesystem_root='/', filesystem_writable_paths='all',
                 mounts='undefined', extra_mounts=None,
-                network='undefined'):
+                network='undefined',
+                stdout=sandboxlib.CAPTURE, stderr=sandboxlib.CAPTURE):
     if type(command) == str:
         command = [command]
 
@@ -186,7 +189,8 @@ def run_sandbox(command, cwd=None, extra_env=None,
     with mount_all(filesystem_root, extra_mounts):
         process = multiprocessing.Process(
             target=run_command_in_chroot,
-            args=(pipe_child, extra_mounts, filesystem_root, command, cwd, env))
+            args=(pipe_child, stdout, stderr, extra_mounts, filesystem_root,
+                  command, cwd, env))
         process.start()
         process.join()
 
@@ -198,3 +202,9 @@ def run_sandbox(command, cwd=None, extra_env=None,
         # will be within the _run_command_in_chroot() function somewhere.
         exception = pipe_parent.recv()
         raise exception
+
+
+def run_sandbox_with_redirection(command, **sandbox_config):
+    exit, out, err = run_sandbox(command, **sandbox_config)
+    # out and err will be None
+    return exit
diff --git a/sandboxlib/linux_user_chroot.py b/sandboxlib/linux_user_chroot.py
index c0715c7..4244d99 100644
--- a/sandboxlib/linux_user_chroot.py
+++ b/sandboxlib/linux_user_chroot.py
@@ -262,7 +262,8 @@ def process_writable_paths(fs_root, writable_paths):
 def run_sandbox(command, cwd=None, extra_env=None,
                 filesystem_root='/', filesystem_writable_paths='all',
                 mounts='undefined', extra_mounts=None,
-                network='undefined'):
+                network='undefined',
+                stdout=sandboxlib.CAPTURE, stderr=sandboxlib.CAPTURE):
     if type(command) == str:
         command = [command]
 
@@ -284,5 +285,11 @@ def run_sandbox(command, cwd=None, extra_env=None,
     env = sandboxlib.environment_vars(extra_env)
 
     argv = (unshare_command + linux_user_chroot_command + command)
-    exit, out, err = sandboxlib._run_command(argv, env=env)
+    exit, out, err = sandboxlib._run_command(argv, stdout, stderr, env=env)
     return exit, out, err
+
+
+def run_sandbox_with_redirection(command, **sandbox_config):
+    exit, out, err = run_sandbox(command, **sandbox_config)
+    # out and err will be None
+    return exit