stapsh: Add OK/ERROR replies to file and run commands

* runtime/staprun/stapsh.c (reply): Print a reply (and dbug it too).
  (do_hello): Send uname info using reply().
  (do_file, do_run): Add OK reply, or ERROR from former stderr chatter.
  (qpdecode): Convert stderr chatter to dbug.
* remote.cxx (stapsh::get_reply): Read a one-line reply from the remote.
  (stapsh::send_file): Look for an OK reply after sending the file.
  (stapsh::start): Look for an OK reply after telling it to run.
  (stapsh::set_child_fds): Use get_reply to read the uname back.

diff --git a/remote.cxx b/remote.cxx
index f2fb4fc63440531367900b347d7156f29bdb622f..fdbbbf89ccd7a2c2cc20fd5c8854bb27ba0f2fc5 100644 (file)
--- a/remote.cxx
+++ b/remote.cxx
@@ -160,6 +160,14 @@ class stapsh : public remote {
             }
       }
 
+    string get_reply()
+      {
+        char reply[4096];
+        if (!fgets(reply, sizeof(reply), OUT))
+          reply[0] = '\0';
+        return reply;
+      }
+
     int send_command(const string& cmd)
       {
         if (!IN)
@@ -209,6 +217,23 @@ class stapsh : public remote {
           rc = fflush(IN);
 
         fclose(f);
+
+        if (!rc)
+          {
+            string reply = get_reply();
+            if (reply != "OK\n")
+              {
+                rc = 1;
+                if (s->verbose > 1)
+                  {
+                    if (reply.empty())
+                      clog << "stapsh file ERROR: no reply" << endl;
+                    else
+                      clog << "stapsh file " << reply;
+                  }
+              }
+          }
+
         return rc;
       }
 
@@ -269,6 +294,22 @@ class stapsh : public remote {
 
         int rc = send_command(run.str());
 
+        if (!rc)
+          {
+            string reply = get_reply();
+            if (reply != "OK\n")
+              {
+                rc = 1;
+                if (s->verbose > 1)
+                  {
+                    if (reply.empty())
+                      clog << "stapsh run ERROR: no reply" << endl;
+                    else
+                      clog << "stapsh run " << reply;
+                  }
+              }
+          }
+
         if (!rc)
           {
             long flags = fcntl(fdout, F_GETFL) | O_NONBLOCK;
@@ -307,8 +348,8 @@ class stapsh : public remote {
         if (send_command("stap " VERSION "\n"))
           throw runtime_error("error sending hello to stapsh!");
 
-        char reply[1024];
-        if (!fgets(reply, sizeof(reply), OUT))
+        string reply = get_reply();
+        if (reply.empty())
           throw runtime_error("error receiving hello from stapsh!");
 
         // stapsh VERSION MACHINE RELEASE

diff --git a/runtime/staprun/stapsh.c b/runtime/staprun/stapsh.c
index dcbeeed50e9b47e32dcfd38f1fde940da14202e8..9d8bb2d5f2a65d24d2b661a8c0b5b046fbd89df5 100644 (file)
--- a/runtime/staprun/stapsh.c
+++ b/runtime/staprun/stapsh.c
@@ -20,13 +20,13 @@
 //
 //   command: file SIZE NAME
 //            DATA
-//     reply: (none)
+//     reply: OK / error message
 //      desc: Create a file of SIZE bytes, called NAME.  The NAME is a basename
 //            only, and limited to roughly "[a-z0-9][a-z0-9._]*".  The DATA is
 //            read as raw bytes following the command's newline.
 //
 //   command: run ARG1 ARG2 ...
-//     reply: (none)
+//     reply: OK / error message
 //      desc: Start staprun with the given quoted-printable arguments.  When
 //            the child exits, stapsh will clean up and then exit with the same
 //            return code.  Note that whitespace has significance in stapsh
@@ -207,7 +207,7 @@ qpdecode(char* s)
           s += 3;
         else if (!s[1] || !s[2])
           {
-            fprintf(stderr, "truncated quoted-printable escape\n");
+            dbug(2, "truncated quoted-printable escape \"%s\"\n", s);
             return 1;
           }
         else
@@ -217,7 +217,7 @@ qpdecode(char* s)
             unsigned char c = strtol(hex, &end, 16);
             if (errno || end != hex + 2)
               {
-                fprintf(stderr, "invalid quoted-printable escape =%s\n", hex);
+                dbug(2, "invalid quoted-printable escape \"=%s\"\n", hex);
                 return 1;
               }
             *o++ = c;
@@ -229,6 +229,22 @@ qpdecode(char* s)
 }
 
 
+// Send a reply back to the client on stdout
+static int __attribute__ ((format (printf, 1, 2)))
+reply(const char* format, ...)
+{
+  va_list args, dbug_args;
+  va_start (args, format);
+  va_copy (dbug_args, args);
+  vdbug (1, format, dbug_args);
+  int ret = vprintf (format, args);
+  fflush (stdout);
+  va_end (dbug_args);
+  va_end (args);
+  return ret;
+}
+
+
 static int
 do_hello()
 {
@@ -241,9 +257,7 @@ do_hello()
   if (uname(&uts))
     return 1;
 
-  printf("stapsh %s %s %s\n", VERSION, uts.machine, uts.release);
-  fflush(stdout);
-
+  reply ("stapsh %s %s %s\n", VERSION, uts.machine, uts.release);
   return 0;
 }
 
@@ -259,31 +273,19 @@ do_file()
   if (arg)
     size = atoi(arg);
   if (size <= 0 || size > STAPSH_MAX_FILE_SIZE)
-    {
-      fprintf(stderr, "bad file size %d\n", size);
-      return 1;
-    }
+    return reply ("ERROR: bad file size %d\n", size);
 
   const char* name = strtok(NULL, STAPSH_TOK_DELIM);
   if (!name)
-    {
-      fprintf(stderr, "missing file name\n");
-      return 1;
-    }
+    return reply ("ERROR: missing file name\n");
   for (arg = name; *arg; ++arg)
     if (!isalnum(*arg) &&
         !(arg > name && (*arg == '.' || *arg == '_')))
-      {
-        fprintf(stderr, "bad character in file name '%c'\n", *arg);
-        return 1;
-      }
+      return reply ("ERROR: bad character '%c' in file name\n", *arg);
 
   FILE* f = fopen(name, "w");
   if (!f)
-    {
-      fprintf(stderr, "can't open file for writing\n");
-      return 1;
-    }
+    return reply ("ERROR: can't open file \"%s\" for writing\n", name);
   while (size > 0 && ret == 0)
     {
       char buf[1024];
@@ -291,30 +293,30 @@ do_file()
       if ((size_t)size < sizeof(buf))
        r = size;
       r = fread(buf, 1, r, stdin);
-      if (!r)
+      if (!r && feof(stdin))
+        ret = reply ("ERROR: reached EOF while reading file data\n");
+      else if (!r)
+        ret = reply ("ERROR: unable to read file data\n");
+      else
         {
-          fprintf(stderr, "error reading file data\n");
-          ret = 1;
-          break;
-        }
-      size -= r;
+          size -= r;
 
-      const char* bufp = buf;
-      while (bufp < buf + r)
-        {
-          size_t w = (buf + r) - bufp;
-          w = fwrite(bufp, 1, w, f);
-          if (!w)
+          const char* bufp = buf;
+          while (bufp < buf + r && ret == 0)
             {
-              fprintf(stderr, "error writing file data\n");
-              ret = 1;
-              break;
+              size_t w = (buf + r) - bufp;
+              w = fwrite(bufp, 1, w, f);
+              if (!w)
+                ret = reply ("ERROR: unable to write file data\n");
+              else
+                bufp += w;
             }
-          bufp += w;
         }
     }
   fclose(f);
 
+  if (ret == 0)
+    reply ("OK\n");
   return ret;
 }
 
@@ -332,23 +334,20 @@ do_run()
   while ((arg = strtok(NULL, STAPSH_TOK_DELIM)))
     {
       if (nargs + 1 > STAPSH_MAX_ARGS)
-        {
-          fprintf(stderr, "too many args\n");
-          return 1;
-        }
+        return reply ("ERROR: too many arguments\n");
       if (qpdecode(arg) != 0)
-        return 1;
+        return reply ("ERROR: invalid encoding in argument \"%s\"\n", arg);
       args[nargs++] = arg;
     }
 
   int ret = 0;
   posix_spawn_file_actions_t fa;
   if (posix_spawn_file_actions_init(&fa) != 0)
-    return 1;
+    return reply ("ERROR: can't initialize posix_spawn actions\n");
 
   // no stdin for staprun
   if (posix_spawn_file_actions_addopen(&fa, 0, "/dev/null", O_RDONLY, 0) != 0)
-    ret = 1;
+    ret = reply ("ERROR: can't set posix_spawn actions\n");
   else
     {
       pid_t pid;
@@ -356,10 +355,13 @@ do_run()
       if (ret == 0)
         staprun_pid = pid;
       else
-        fprintf(stderr, "error launching staprun\n");
+        reply("ERROR: can't launch staprun\n");
     }
 
   posix_spawn_file_actions_destroy(&fa);
+
+  if (ret == 0)
+    reply ("OK\n");
   return ret;
 }