Merge pull request #978 '[RDY] implement system() with pipes'

6 files changed, 329 insertions, 79 deletions

diff --git a/src/nvim/eval.c b/src/nvim/eval.c
index 069715ffee..fffeedfd53 100644
--- a/src/nvim/eval.c
+++ b/src/nvim/eval.c
@@ -14077,76 +14077,58 @@ static void f_synstack(typval_T *argvars, typval_T *rettv)
   }
 }
 
-/*
- * "system()" function
- */
+/// f_system - the VimL system() function
 static void f_system(typval_T *argvars, typval_T *rettv)
 {
-  char_u      *res = NULL;
-  char_u      *p;
-  char_u      *infile = NULL;
+  rettv->v_type = VAR_STRING;
+  rettv->vval.v_string = NULL;
+
+  if (check_restricted() || check_secure()) {
+    return;
+  }
+
+  // get input to the shell command (if any), and its length
   char_u buf[NUMBUFLEN];
-  int err = FALSE;
-  FILE        *fd;
+  const char *input = (argvars[1].v_type != VAR_UNKNOWN)
+      ? (char *) get_tv_string_buf_chk(&argvars[1], buf): NULL;
+  size_t input_len = input ? strlen(input) : 0;
 
-  if (check_restricted() || check_secure())
-    goto done;
+  // get shell command to execute
+  const char *cmd = (char *) get_tv_string(&argvars[0]);
 
-  if (argvars[1].v_type != VAR_UNKNOWN) {
-    /*
-     * Write the string to a temp file, to be used for input of the shell
-     * command.
-     */
-    if ((infile = vim_tempname()) == NULL) {
-      EMSG(_(e_notmp));
-      goto done;
-    }
+  // execute the command
+  size_t nread = 0;
+  char *res = NULL;
+  int status = os_system(cmd, input, input_len, &res, &nread);
 
-    fd = mch_fopen((char *)infile, WRITEBIN);
-    if (fd == NULL) {
-      EMSG2(_(e_notopen), infile);
-      goto done;
-    }
-    p = get_tv_string_buf_chk(&argvars[1], buf);
-    if (p == NULL) {
-      fclose(fd);
-      goto done;                /* type error; errmsg already given */
-    }
-    if (fwrite(p, STRLEN(p), 1, fd) != 1)
-      err = TRUE;
-    if (fclose(fd) != 0)
-      err = TRUE;
-    if (err) {
-      EMSG(_("E677: Error writing temp file"));
-      goto done;
+  set_vim_var_nr(VV_SHELL_ERROR, (long) status);
+
+#if defined(USE_CR)
+  // translate <CR> into <NL>
+  if (res != NULL) {
+    for (char *s = res; *s; ++s) {
+      if (*s == CAR) {
+        *s = NL;
+      }
     }
   }
-
-  res = get_cmd_output(get_tv_string(&argvars[0]), infile,
-      kShellOptSilent | kShellOptCooked);
-
-#ifdef USE_CRNL
-  /* translate <CR><NL> into <NL> */
+#elif defined(USE_CRNL)
+  // translate <CR><NL> into <NL>
   if (res != NULL) {
-    char_u  *s, *d;
-
-    d = res;
-    for (s = res; *s; ++s) {
-      if (s[0] == CAR && s[1] == NL)
+    char *d = res;
+    for (char *s = res; *s; ++s) {
+      if (s[0] == CAR && s[1] == NL) {
         ++s;
+      }
+
       *d++ = *s;
     }
+
     *d = NUL;
   }
 #endif
 
-done:
-  if (infile != NULL) {
-    os_remove((char *)infile);
-    free(infile);
-  }
-  rettv->v_type = VAR_STRING;
-  rettv->vval.v_string = res;
+  rettv->vval.v_string = (char_u *) res;
 }
 
 /*
diff --git a/src/nvim/os/job.c b/src/nvim/os/job.c
index 5c45c3727e..01a529d533 100644
--- a/src/nvim/os/job.c
+++ b/src/nvim/os/job.c
@@ -14,17 +14,20 @@
 #include "nvim/os/event_defs.h"
 #include "nvim/os/time.h"
 #include "nvim/os/shell.h"
+#include "nvim/os/signal.h"
 #include "nvim/vim.h"
 #include "nvim/memory.h"
 #include "nvim/term.h"
 
 #define EXIT_TIMEOUT 25
 #define MAX_RUNNING_JOBS 100
-#define JOB_BUFFER_SIZE 1024
+#define JOB_BUFFER_SIZE 0xFFFF
 
 struct job {
   // Job id the index in the job table plus one.
   int id;
+  // Exit status code of the job process
+  int64_t status;
   // Number of polls after a SIGTERM that will trigger a SIGKILL
   int exit_timeout;
   // exit_cb may be called while there's still pending data from stdout/stderr.
@@ -163,6 +166,7 @@ Job *job_start(char **argv,
   // Initialize
   job->id = i + 1;
   *status = job->id;
+  job->status = -1;
   job->pending_refs = 3;
   job->pending_closes = 4;
   job->data = data;
@@ -257,6 +261,101 @@ void job_stop(Job *job)
   job->stopped = true;
 }
 
+/// job_wait - synchronously wait for a job to finish
+///
+/// @param job The job instance
+/// @param ms Number of milliseconds to wait, 0 for not waiting, -1 for
+///        waiting until the job quits.
+/// @return returns the status code of the exited job. -1 if the job is
+///         still running and the `timeout` has expired. Note that this is
+///         indistinguishable from the process returning -1 by itself. Which
+///         is possible on some OS.
+int job_wait(Job *job, int ms) FUNC_ATTR_NONNULL_ALL
+{
+  // switch to cooked so `got_int` will be set if the user interrupts
+  int old_mode = cur_tmode;
+  settmode(TMODE_COOK);
+
+  EventSource sources[] = {job_event_source(job), signal_event_source(), NULL};
+
+  // keep track of the elapsed time if ms > 0
+  uint64_t before = (ms > 0) ? os_hrtime() : 0;
+
+  while (1) {
+    // check if the job has exited (and the status is available).
+    if (job->pending_refs == 0) {
+      break;
+    }
+
+    event_poll(ms, sources);
+
+    // we'll assume that a user frantically hitting interrupt doesn't like
+    // the current job. Signal that it has to be killed.
+    if (got_int) {
+      job_stop(job);
+    }
+
+    if (ms == 0) {
+      break;
+    }
+
+    // check if the poll timed out, if not, decrease the ms to wait for the
+    // next run
+    if (ms > 0) {
+      uint64_t now = os_hrtime();
+      ms -= (int) ((now - before) / 1000000);
+      before = now;
+
+      // if the time elapsed is greater than the `ms` wait time, break
+      if (ms <= 0) {
+        break;
+      }
+    }
+  }
+
+  settmode(old_mode);
+
+  // return -1 for a timeout, the job status otherwise
+  return (job->pending_refs) ? -1 : (int) job->status;
+}
+
+/// Close the pipe used to write to the job.
+///
+/// This can be used for example to indicate to the job process that no more
+/// input is coming, and that it should shut down cleanly.
+///
+/// It has no effect when the input pipe doesn't exist or was already
+/// closed.
+///
+/// @param job The job instance
+void job_close_in(Job *job) FUNC_ATTR_NONNULL_ALL
+{
+  if (!job->in) {
+    return;
+  }
+
+  // let other functions in the job module know that the in pipe is no more
+  wstream_free(job->in);
+  job->in = NULL;
+
+  uv_close((uv_handle_t *)&job->proc_stdin, close_cb);
+}
+
+/// All writes that complete after calling this function will be reported
+/// to `cb`.
+///
+/// Use this function to be notified about the status of an in-flight write.
+///
+/// @see {wstream_set_write_cb}
+///
+/// @param job The job instance
+/// @param cb The function that will be called on write completion or
+///        failure. It will be called with the job as the `data` argument.
+void job_write_cb(Job *job, wstream_cb cb) FUNC_ATTR_NONNULL_ALL
+{
+  wstream_set_write_cb(job->in, cb, job);
+}
+
 /// Writes data to the job's stdin. This is a non-blocking operation, it
 /// returns when the write request was sent.
 ///
@@ -329,7 +428,9 @@ static bool is_alive(Job *job)
 static void free_job(Job *job)
 {
   uv_close((uv_handle_t *)&job->proc_stdout, close_cb);
-  uv_close((uv_handle_t *)&job->proc_stdin, close_cb);
+  if (job->in) {
+    uv_close((uv_handle_t *)&job->proc_stdin, close_cb);
+  }
   uv_close((uv_handle_t *)&job->proc_stderr, close_cb);
   uv_close((uv_handle_t *)&job->proc, close_cb);
 }
@@ -377,6 +478,7 @@ static void exit_cb(uv_process_t *proc, int64_t status, int term_signal)
 {
   Job *job = handle_get_job((uv_handle_t *)proc);
 
+  job->status = status;
   if (--job->pending_refs == 0) {
     emit_exit_event(job);
   }
@@ -401,7 +503,9 @@ static void close_cb(uv_handle_t *handle)
     // closed by libuv
     rstream_free(job->out);
     rstream_free(job->err);
-    wstream_free(job->in);
+    if (job->in) {
+      wstream_free(job->in);
+    }
 
     // Free data memory of process and pipe handles, that was allocated
     // by handle_set_job in job_start.
diff --git a/src/nvim/os/rstream.c b/src/nvim/os/rstream.c
index d7ab5b8a64..b3a5196351 100644
--- a/src/nvim/os/rstream.c
+++ b/src/nvim/os/rstream.c
@@ -26,7 +26,7 @@ struct rstream {
   uv_file fd;
   rstream_cb cb;
   size_t buffer_size, rpos, wpos, fpos;
-  bool reading, free_handle;
+  bool free_handle;
   EventSource source_override;
 };
 
@@ -150,7 +150,6 @@ void rstream_start(RStream *rstream)
   if (rstream->file_type == UV_FILE) {
     uv_idle_start(rstream->fread_idle, fread_idle_cb);
   } else {
-    rstream->reading = false;
     uv_read_start(rstream->stream, alloc_cb, read_cb);
   }
 }
@@ -236,16 +235,8 @@ static void alloc_cb(uv_handle_t *handle, size_t suggested, uv_buf_t *buf)
 {
   RStream *rstream = handle_get_rstream(handle);
 
-  if (rstream->reading) {
-    buf->len = 0;
-    return;
-  }
-
   buf->len = rstream->buffer_size - rstream->wpos;
   buf->base = rstream->buffer + rstream->wpos;
-
-  // Avoid `alloc_cb`, `alloc_cb` sequences on windows
-  rstream->reading = true;
 }
 
 // Callback invoked by libuv after it copies the data into the buffer provided
@@ -287,7 +278,6 @@ static void read_cb(uv_stream_t *stream, ssize_t cnt, const uv_buf_t *buf)
          rstream_event_source(rstream));
   }
 
-  rstream->reading = false;
   emit_read_event(rstream, false);
 }
 
diff --git a/src/nvim/os/shell.c b/src/nvim/os/shell.c
index 68cd2ad527..02dc396d78 100644
--- a/src/nvim/os/shell.c
+++ b/src/nvim/os/shell.c
@@ -5,6 +5,10 @@
 #include <uv.h>
 
 #include "nvim/ascii.h"
+#include "nvim/lib/kvec.h"
+#include "nvim/log.h"
+#include "nvim/os/job.h"
+#include "nvim/os/rstream.h"
 #include "nvim/os/shell.h"
 #include "nvim/os/signal.h"
 #include "nvim/types.h"
@@ -31,6 +35,11 @@ typedef struct {
   garray_T ga;
 } ProcessData;
 
+typedef struct {
+  char *data;
+  size_t cap;
+  size_t len;
+} dyn_buffer_t;
 
 #ifdef INCLUDE_GENERATED_DECLARATIONS
 # include "os/shell.c.generated.h"
@@ -47,16 +56,13 @@ typedef struct {
 /// @param extra_shell_opt Extra argument to the shell. If NULL it is ignored
 /// @return A newly allocated argument vector. It must be freed with
 ///         `shell_free_argv` when no longer needed.
-char ** shell_build_argv(char_u *cmd, char_u *extra_shell_opt)
+char **shell_build_argv(const char_u *cmd, const char_u *extra_shell_opt)
 {
-  int i;
-  char **rv;
   int argc = tokenize(p_sh, NULL) + tokenize(p_shcf, NULL);
-
-  rv = (char **)xmalloc((unsigned)((argc + 4) * sizeof(char *)));
+  char **rv = xmalloc((unsigned)((argc + 4) * sizeof(char *)));
 
   // Split 'shell'
-  i = tokenize(p_sh, rv);
+  int i = tokenize(p_sh, rv);
 
   if (extra_shell_opt != NULL) {
     // Push a copy of `extra_shell_opt`
@@ -237,6 +243,128 @@ int os_call_shell(char_u *cmd, ShellOpts opts, char_u *extra_shell_arg)
   return proc_cleanup_exit(&pdata, &proc_opts, opts);
 }
 
+/// os_system - synchronously execute a command in the shell
+///
+/// example:
+///   char *output = NULL;
+///   size_t nread = 0;
+///   int status = os_sytem("ls -la", NULL, 0, &output, &nread);
+///
+/// @param cmd The full commandline to be passed to the shell
+/// @param input The input to the shell (NULL for no input), passed to the
+///              stdin of the resulting process.
+/// @param len The length of the input buffer (not used if `input` == NULL)
+/// @param[out] output A pointer to to a location where the output will be
+///                    allocated and stored. Will point to NULL if the shell
+///                    command did not output anything. NOTE: it's not
+///                    allowed to pass NULL yet
+/// @param[out] nread the number of bytes in the returned buffer (if the
+///             returned buffer is not NULL)
+/// @return the return code of the process, -1 if the process couldn't be
+///         started properly
+int os_system(const char *cmd,
+              const char *input,
+              size_t len,
+              char **output,
+              size_t *nread) FUNC_ATTR_NONNULL_ARG(1, 4)
+{
+  // the output buffer
+  dyn_buffer_t buf;
+  memset(&buf, 0, sizeof(buf));
+
+  char **argv = shell_build_argv((char_u *) cmd, NULL);
+
+  int i;
+  Job *job = job_start(argv,
+                       &buf,
+                       system_data_cb,
+                       system_data_cb,
+                       NULL,
+                       false,
+                       0,
+                       &i);
+
+  if (i <= 0) {
+    // couldn't even start the job
+    ELOG("Couldn't start job, error code: '%d'", i);
+    return -1;
+  }
+
+  // write the input, if any
+  if (input) {
+    WBuffer *input_buffer = wstream_new_buffer((char *) input, len, 1, NULL);
+
+    // we want to be notified when the write completes
+    job_write_cb(job, system_write_cb);
+
+    if (!job_write(job, input_buffer)) {
+      // couldn't write, stop the job and tell the user about it
+      job_stop(job);
+      return -1;
+    }
+  } else {
+    // close the input stream, let the process know that no input is coming
+    job_close_in(job);
+  }
+
+  int status = job_wait(job, -1);
+
+  // prepare the out parameters if requested
+  if (buf.len == 0) {
+    // no data received from the process, return NULL
+    *output = NULL;
+    free(buf.data);
+  } else {
+    // NUL-terminate to make the output directly usable as a C string
+    buf.data[buf.len] = NUL;
+    *output = buf.data;
+  }
+
+  if (nread) {
+    *nread = buf.len;
+  }
+
+  return status;
+}
+
+/// dyn_buf_ensure - ensures at least `desired` bytes in buffer
+///
+/// TODO(aktau): fold with kvec/garray
+static void dyn_buf_ensure(dyn_buffer_t *buf, size_t desired)
+{
+  if (buf->cap >= desired) {
+    return;
+  }
+
+  buf->cap = desired;
+  kv_roundup32(buf->cap);
+  buf->data = xrealloc(buf->data, buf->cap);
+}
+
+static void system_data_cb(RStream *rstream, void *data, bool eof)
+{
+  Job *job = data;
+  dyn_buffer_t *buf = job_data(job);
+
+  size_t nread = rstream_available(rstream);
+
+  dyn_buf_ensure(buf, buf->len + nread + 1);
+  rstream_read(rstream, buf->data + buf->len, nread);
+
+  buf->len += nread;
+}
+
+static void system_write_cb(WStream *wstream,
+                            void *data,
+                            size_t pending,
+                            int status)
+{
+  if (pending == 0) {
+    Job *job = data;
+    job_close_in(job);
+  }
+}
+
 /// Parses a command string into a sequence of words, taking quotes into
 /// consideration.
 ///
@@ -244,10 +372,10 @@ int os_call_shell(char_u *cmd, ShellOpts opts, char_u *extra_shell_arg)
 /// @param argv The vector that will be filled with copies of the parsed
 ///        words. It can be NULL if the caller only needs to count words.
 /// @return The number of words parsed.
-static int tokenize(char_u *str, char **argv)
+static int tokenize(const char_u *str, char **argv)
 {
   int argc = 0, len;
-  char_u *p = str;
+  char_u *p = (char_u *) str;
 
   while (*p != NUL) {
     len = word_length(p);
@@ -271,9 +399,9 @@ static int tokenize(char_u *str, char **argv)
 ///
 /// @param str A pointer to the first character of the word
 /// @return The offset from `str` at which the word ends.
-static int word_length(char_u *str)
+static int word_length(const char_u *str)
 {
-  char_u *p = str;
+  const char_u *p = str;
   bool inquote = false;
   int length = 0;
 
diff --git a/src/nvim/os/wstream.c b/src/nvim/os/wstream.c
index 0978d33a10..194bf757e4 100644
--- a/src/nvim/os/wstream.c
+++ b/src/nvim/os/wstream.c
@@ -21,6 +21,9 @@ struct wstream {
   // Number of pending requests
   size_t pending_reqs;
   bool freed;
+  // (optional) Write callback and data
+  wstream_cb cb;
+  void *data;
 };
 
 struct wbuffer {
@@ -57,6 +60,7 @@ WStream * wstream_new(size_t maxmem)
   rv->curmem = 0;
   rv->pending_reqs = 0;
   rv->freed = false;
+  rv->cb = NULL;
 
   return rv;
 }
@@ -83,6 +87,25 @@ void wstream_set_stream(WStream *wstream, uv_stream_t *stream)
   wstream->stream = stream;
 }
 
+/// Sets a callback that will be called on completion of a write request,
+/// indicating failure/success.
+///
+/// This affects all requests currently in-flight as well. Overwrites any
+/// possible earlier callback.
+///
+/// @note This callback will not fire if the write request couldn't even be
+///       queued properly (i.e.: when `wstream_write() returns an error`).
+///
+/// @param wstream The `WStream` instance
+/// @param cb The callback
+/// @param data User-provided data that will be passed to `cb`
+void wstream_set_write_cb(WStream *wstream, wstream_cb cb, void *data)
+  FUNC_ATTR_NONNULL_ARG(1)
+{
+  wstream->cb = cb;
+  wstream->data = data;
+}
+
 /// Queues data for writing to the backing file descriptor of a `WStream`
 /// instance. This will fail if the write would cause the WStream use more
 /// memory than specified by `maxmem`.
@@ -162,6 +185,14 @@ static void write_cb(uv_write_t *req, int status)
   release_wbuffer(data->buffer);
 
   data->wstream->pending_reqs--;
+
+  if (data->wstream->cb) {
+    data->wstream->cb(data->wstream,
+                      data->wstream->data,
+                      data->wstream->pending_reqs,
+                      status);
+  }
+
   if (data->wstream->freed && data->wstream->pending_reqs == 0) {
     // Last pending write, free the wstream;
     free(data->wstream);
@@ -173,7 +204,10 @@ static void write_cb(uv_write_t *req, int status)
 static void release_wbuffer(WBuffer *buffer)
 {
   if (!--buffer->refcount) {
-    buffer->cb(buffer->data);
+    if (buffer->cb) {
+      buffer->cb(buffer->data);
+    }
+
     free(buffer);
   }
 }
diff --git a/src/nvim/os/wstream_defs.h b/src/nvim/os/wstream_defs.h
index 1bf61ffce1..e42481f283 100644
--- a/src/nvim/os/wstream_defs.h
+++ b/src/nvim/os/wstream_defs.h
@@ -5,5 +5,17 @@ typedef struct wbuffer WBuffer;
 typedef struct wstream WStream;
 typedef void (*wbuffer_data_finalizer)(void *data);
 
+/// Type of function called when the WStream has information about a write
+/// request.
+///
+/// @param wstream The `WStream` instance
+/// @param data User-defined data
+/// @param pending The number of write requests that are still pending
+/// @param status 0 on success, anything else indicates failure
+typedef void (*wstream_cb)(WStream *wstream,
+                           void *data,
+                           size_t pending,
+                           int status);
+
 #endif  // NVIM_OS_WSTREAM_DEFS_H