lua-luxio 13-1 / usr / share / lua / 5.1 / luxio / subprocess.lua

-- Light Unix I/O for Lua
-- Copyright 2012 Daniel Silverstone <dsilvers@digital-scurf.org>
--
-- Distributed under the same terms as Lua itself (MIT).
--
-- A wrapper for managing sub processes
-------------------------------------------------------------------------------
--
-- sp = require "luxio.subprocess"
--
--
-- proc = sp.spawn {
--     "sleep", "10",
--     exe = "/usr/bin/sleep"
-- }
--
-- or
--
-- proc = sp.spawn_simple {
--     "sleep", "10",
--     exe = "/usr/bin/sleep"
-- }
--
-- spawn_simple wrappers any returned FDs in sio handles
--
--
-- The spawn functions take an argv, an optional exe to use instead of
-- the first argument (which is argv[0])
--
-- Also, you can optionally pass in an (incomplete) set of FDs for the standard
-- IO streams (stdin, stdout, stderr) for the child process.  If they're
-- sp.PIPE then a pipe will be created for you and the child's stream FDs will
-- be the other end of it.
--
-- For example:
--
-- proc = sp.spawn_simple {
--     "ls", "/",
--     stdout = sp.PIPE
-- }
--
-- listing = proc.stdout:read("*a")
-- proc.stdout:close()
--
-- If you want the child (during pre-exec) to close additional FDs then
-- you can specify them in a table called close_in_child
--
-- If you want to change directory before exec in the child, you can pass
-- cwd="/some/path".
--
-- To run some code before forking, pass prefork=somefunction and to
-- run code before exec (in the child) pass preexec=somefunction.
--
-- If the prefork function raises an error then the spawn will be stopped and
-- the error propagated.  If the preexec function raises an error then the exec
-- will be stopped and the child will exit with status code 1.
--
-- spawn_simple has one additional tweak.  If 'stdin' is provided as
-- a string then it will automatically write that string to the stdin
-- of the child (and close the FD afterwards) for you automatically.
--
-- You can provide an env table which modifies the environment of
-- the child.  key=value pairs are set in the child environment and
-- key=false pairs cause an unsetenv allowing you to remove environment
-- variables from the child if you need to.
--
-- For example:
--
-- proc = sp.spawn_simple {
--     "tar", "czf", "foo.tar.gz", "foo",
--     env = { GZIP = "-9" }
-- }
--
-- Then, once you are done and wait the result of the process execution,
-- you run:
--
-- how, why = proc:wait(noblock)
--     noblock: If true, wait won't block until proc exits.
--
-- how will be a string (or -1 on error, with why set to the errno)
--
-- If the wait() didn't error, then how will be one of:
--         running, exit, signal, unknown
-- In the first case, the process is still running.  In the second it has
-- exited and the 'why' will be the exit code.  In the third case, the process
-- was terminated by a signal and the signal in question is in why.
-- If wait() cannot determine why a process exited, then unknown is returned
-- and the why is the full status value from the waitpid() system call.
--

local l = require 'luxio'
local sio = require 'luxio.simple'

local fork = l.fork
local waitpid = l.waitpid
local exec = l.exec
local execp = l.execp
local pipe = l.pipe
local close = l.close
local dup = l.dup
local dup2 = l.dup2
local fcntl = l.fcntl
local bclear = l.bit.bclear
local wrap_fd = sio.wrap_fd
local chdir = sio.chdir
local stderr_fh = sio.stderr
local setenv = l.setenv
local unsetenv = l.unsetenv
local _exit = l._exit
local _TRACEBACK = debug.traceback

local F_SETFD = l.F_SETFD
local F_GETFD = l.F_GETFD

local WIFEXITED = l.WIFEXITED
local WEXITSTATUS = l.WEXITSTATUS
local WIFSIGNALED = l.WIFSIGNALED
local WTERMSIG = l.WTERMSIG
local WNOHANG = l.WNOHANG

local assert = assert
local unpack = unpack
local pairs = pairs
local error = error
local tostring = tostring

l = nil
sio = nil

local _PIPE = {}

local function _wait(proc, nohang)
   local pid, status = waitpid(proc.pid, nohang and WNOHANG or 0)
   if pid == -1 then
      return "error", status
   end
   if pid == 0 then
      return "running", 0
   end

   if WIFEXITED(status) ~= 0 then
      return "exit", WEXITSTATUS(status)
   end

   if WIFSIGNALED(status) ~= 0 then
      return "signal", WTERMSIG(status)
   end

   return "unknown", status
end

local function _spawn(t)
   -- Stage one, normalise our process
   local proc = {args={}}
   for k, v in pairs(t) do
      proc.args[k] = v
   end

   if proc.args[0] == nil then
      if proc.args.exe then 
	 proc.args[0] = proc.args.exe
      else
	 proc.args[0] = proc.args[1]
	 if proc.args[0]:match("^/") then
	    proc.args.exe = proc.args[0]
	 end
      end
   end

   -- Choose which exec function to use, based on whether or not
   -- we have an explicit exe.
   local exec_fn = proc.args.exe and exec or execp
   -- And prepare a local copy of the exe based on the args.
   local exe = proc.args.exe or proc.args[0]

   assert(exe, "Unable to continue, nothing to run")

   -- Simple verification of the FDs.
   local function check(v)
      if proc.args[v] then
	 -- Something to do for handle
	 if proc.args[v] == _PIPE then
	    -- For now, nothing, later we'll handle this
	 elseif type(proc.args[v]) == "table" then
	    -- Simple handle, extract the FD
	    proc.args[v] = proc.args[v].fd
	 elseif type(proc.args[v]) == "number" then
	    -- FD, nothing to do for now
	 else
	    assert(false, v .. " isn't a PIPE, Simple Handle or FD")
	 end
      end
   end

   check "stdin"
   check "stdout"
   check "stderr"
   
   

   -- Beyond here, if we get errors, we *MUST* clean up FDs
   local parent_close = {}
   local child_close = {}

   local function _assert(m,v)
      if not m then
	 for fd in pairs(parent_close) do
	    close(fd)
	 end
	 for fd in pairs(child_close) do
	    close(fd)
	 end
	 error(v)
      end
   end

   local function do_pipe(v)
      if proc.args[v] == _PIPE then
	 local fds = {}
	 local res, errno = pipe(fds)
	 _assert(res == 0, "Unable to pipe()")
	 proc.args[v] = fds[2]
	 proc[v] = fds[1]
	 parent_close[proc.args[v]] = true
	 child_close[proc[v]] = true
      end
   end
   do_pipe "stdin"
   do_pipe "stdout"
   do_pipe "stderr"

   if proc.stdin then
      -- Swap stdin over.
      proc.stdin, proc.args.stdin = proc.args.stdin, proc.stdin
      parent_close[proc.stdin] = nil
      parent_close[proc.args.stdin] = true
      child_close[proc.stdin] = true
      child_close[proc.args.stdin] = nil
   end

   proc.args.stdin = proc.args.stdin or 0
   proc.args.stdout = proc.args.stdout or 1
   proc.args.stderr = proc.args.stderr or 2

   -- Finally, let's fork

   if proc.args.prefork then
      local ok, msg = xpcall(proc.args.prefork, function(...)
	 for fd in pairs(parent_close) do
	    close(fd)
	 end
	 for fd in pairs(child_close) do
	    close(fd)
	 end
	 return _TRACEBACK(...)
      end)
      if not ok then
	 error(msg)
      end
   end

   local pid, errno = fork()

   _assert(pid > -1, "Unable to fork()")

   if pid == 0 then
      local function check(m)
	 if m == -1 then
	    _exit(1)
	 end
	 return m
      end

      -- Set up the FDs.
      if proc.args.stdin == -1 then
	 close(0)
      elseif proc.args.stdin ~= 0 then
	 check(dup2(proc.args.stdin, 0))
      end
      if proc.args.stdout == -1 then
	 close(1)
      elseif proc.args.stdout ~= 1 then
	 if proc.args.stderr == 1 then
	    proc.args.old_stdout = check(dup(1))
	 end
	 check(dup2(proc.args.stdout, 1))
      end
      if proc.args.stderr == -1 then
	 close(2)
      elseif proc.args.stderr ~= 2 then
	 if proc.args.stderr == 1 and proc.args.old_stdout then
	    check(dup2(proc.args.old_stdout, 2))
	 else
	    check(dup2(proc.args.stderr, 2))
	 end
      end
      if proc.args.old_stdout then
	 close(proc.args.old_stdout)
      end
      -- FDs 0, 1, 2 initialised, ensure they won't close on exec
      fcntl(0, F_SETFD, bclear(fcntl(0, F_GETFD), 1))
      fcntl(1, F_SETFD, bclear(fcntl(1, F_GETFD), 1))
      fcntl(2, F_SETFD, bclear(fcntl(2, F_GETFD), 1))
      
      -- Close any FDs left which we were asked to
      for _, fd in ipairs(proc.args.close_in_child or {}) do
	 close(fd)
      end

      -- Finally, close the parent halves of any pipes
      for fd in pairs(child_close) do
	 close(fd)
      end

      -- perform environment twiddling
      if proc.args.env then
	 for k, v in pairs(proc.args.env) do
	    if not v then
	       unsetenv(tostring(k))
	    else
	       check(setenv(tostring(k), tostring(v)))
	    end
	 end
      end

      -- Change directory if needed
      if proc.args.cwd then
	 local cwd = tostring(proc.args.cwd)
	 local ret, err = chdir(cwd)
	 if ret ~= 0 then
	    stderr_fh:write("chdir(" .. cwd .. "): " .. err .. "\n")
	    _exit(1)
	 end
      end

      if proc.args.preexec then
	 local ok, msg = xpcall(proc.args.preexec, _TRACEBACK)
	 if not ok then
	    stderr_fh:write(msg .. "\n")
	    _exit(1)
	 end
      end

      -- Run the child process
      exec_fn(exe, unpack(proc.args))
      
      -- Exec failed, exit without doing lua cleanups
      _exit(1)
   end

   -- The parent
   for fd in pairs(parent_close) do
      close(fd)
   end

   proc.pid = pid

   proc.wait = _wait

   return proc
end

local function _spawn_simple(t)
   local t_copy = {}
   for k, v in pairs(t) do
      t_copy[k] = v
   end

   local piping = nil

   if type(t_copy.stdin) == "string" then
      piping = t_copy.stdin
      t_copy.stdin = _PIPE
   end

   local proc = _spawn(t_copy)

   local function doit(v)
      if proc[v] then
	 proc[v] = wrap_fd(proc[v])
      end
   end
   doit "stdin"
   doit "stdout"
   doit "stderr"

   if piping then
      proc.stdin:write(piping)
      proc.stdin:close()
      proc.stdin = nil
   end

   return proc
end

return {
   spawn = _spawn,
   spawn_simple = _spawn_simple,
   PIPE = _PIPE,
}