Split do_ps() into get_ps() and show_ps() as a start on implementing --sort.

Change readfileat() to pass back length of read.

1 files changed, 244 insertions, 166 deletions

diff --git a/toys/posix/ps.c b/toys/posix/ps.c
index e42e7ec4..e5fcc5c5 100644
--- a/toys/posix/ps.c
+++ b/toys/posix/ps.c
@@ -35,13 +35,13 @@
  * significant. The array index is used in strawberry->which (consumed
  * in do_ps()) and in the bitmasks enabling default fields in ps_main().
 
-USE_PS(NEWTOY(ps, "P(ppid)*aAdeflno*p(pid)*s*t*u*U*g*G*wZ[!ol][+Ae]", TOYFLAG_USR|TOYFLAG_BIN))
+USE_PS(NEWTOY(ps, "k(sort)P(ppid)*aAdeflno*p(pid)*s*t*u*U*g*G*wZ[!ol][+Ae]", TOYFLAG_USR|TOYFLAG_BIN))
 
 config PS
   bool "ps"
   default y
   help
-    usage: ps [-AadeflnwZ] [-gG GROUP] [-o FIELD] [-p PID] [-t TTY] [-uU USER]
+    usage: ps [-AadeflnwZ] [-gG GROUP,] [-k FIELD,] [-o FIELD,] [-p PID,] [-t TTY,] [-uU USER,]
 
     List processes.
 
@@ -62,6 +62,7 @@ config PS
 
     Output modifiers:
 
+    -k	Sort FIELDs in +increasing or -decreasting order (--sort)
     -n	Show numeric USER and GROUP
     -w	Wide output (don't truncate at terminal width)
 
@@ -75,9 +76,9 @@ config PS
     Available -o FIELDs:
 
       ADDR    Instruction pointer
-      CMD     Command line (from /proc/pid/cmdline, including args)
-      CMDLINE Command line (from /proc/pid/cmdline, no args)
-      COMM    Command name (from /proc/pid/stat, no args)
+      CMD     Command name (original)
+      CMDLINE Command name (current argv[0])
+      COMM    Command line (with arguments)
       ETIME   Elapsed time since process start
       F       Process flags (PF_*) from linux source file include/sched.h
               (in octal rather than hex because posix)
@@ -126,14 +127,15 @@ GLOBALS(
   struct arg_list *p;
   struct arg_list *o;
   struct arg_list *P;
+  struct arg_list *k;
 
   struct ptr_len gg, GG, pp, PP, ss, tt, uu, UU, *parsing;
   unsigned width;
   dev_t tty;
   void *fields;
-  long bits;
-  long long ticks;
+  long long ticks, bits;
   size_t header_len;
+  int kcount;
 )
 
 struct strawberry {
@@ -143,16 +145,7 @@ struct strawberry {
   char forever[];
 };
 
-static time_t get_uptime(void)
-{
-  struct sysinfo si;
-
-  sysinfo(&si);
-
-  return si.uptime;
-}
-
-// Return 1 to display, 0 to skip
+// Return 1 to keep, 0 to discard
 static int match_process(long long *slot)
 {
   struct ptr_len *match[] = {
@@ -181,68 +174,32 @@ static int match_process(long long *slot)
   return 1;
 }
 
-// dirtree callback.
-// toybuf used as: 1024 /proc/$PID/stat, 1024 slot[], 2048 /proc/$PID/cmdline
-static int do_ps(struct dirtree *new)
+// Display process data that get_ps() read from /proc, formatting with TT.fields
+static void show_ps(long long *slot)
 {
+  char state, *s, *buf, *strslot[5]; // name, tty, wchan, attr, cmdline
   struct strawberry *field;
-  long long *slot = (void *)(toybuf+1024), ll;
-  char *name, *s, state;
-  int nlen, i, fd, len, width = TT.width;
-
-  if (!new->parent) return DIRTREE_RECURSE|DIRTREE_SHUTUP;
-  if (!(*slot = atol(new->name))) return 0;
-
-  // name field limited to 256 bytes by VFS, plus 40 fields * max 20 chars:
-  // 1000-ish total, but some forced zero so actually there's headroom.
-  sprintf(toybuf, "%lld/stat", *slot);
-  if (!readfileat(dirtree_parentfd(new), toybuf, toybuf, 1024)) return 0;
-
-  // parse oddball fields (name and state)
-  if (!(s = strchr(toybuf, '('))) return 0;
-  for (name = ++s; *s != ')'; s++) if (!*s) return 0;
-  nlen = s++-name;
-  if (1>sscanf(++s, " %c%n", &state, &i)) return 0;
-
-  // parse numeric fields (PID = 0, skip 2, then 4th field goes in slot[1])
-  for (len = 1; len<100; len++)
-    if (1>sscanf(s += i, " %lld%n", slot+len, &i)) break;
-
-  // save uid, ruid, gid, gid, and rgid int slots 31-34 (we don't use sigcatch
-  // or numeric wchan, and the remaining two are always zero).
-  slot[31] = new->st.st_uid;
-  slot[33] = new->st.st_gid;
-
-  // If RGROUP RUSER STAT RUID RGID
-  // Save ruid in slot[34] and rgid in slot[35], which are otherwise zero,
-  // and vmlck into slot[18] (it_real_value, also always zero).
-  if ((TT.bits & 0x38300000) || TT.GG.len || TT.UU.len) {
-    char *out = toybuf+2048;
-
-    sprintf(out, "%lld/status", *slot);
-    if (!readfileat(dirtree_parentfd(new), out, out, 2048)) *out = 0;
-    s = strstr(out, "\nUid:");
-    slot[32] = s ? atol(s+5) : new->st.st_uid;
-    s = strstr(out, "\nGid:");
-    slot[34] = s ? atol(s+5) : new->st.st_gid;
-    s = strstr(out, "\nVmLck:");
-    if (s) slot[18] = atoll(s+5);
+  long long ll;
+  int i, len, width = TT.width;
+
+  // Carve toybuf up into chunks to extract our incoming data
+  s = (char *)(slot+50);
+  state = *s++;
+  for (i=0; i<5; i++) {
+    strslot[i] = s;
+    s += strlen(s)+1;
   }
+  // 64 byte sscratch space reserved for sprintf to output data to
+  buf = toybuf+sizeof(toybuf)-64;
 
-  // skip processes we don't care about.
-  if (!match_process(slot)) return 0;
-
-  // At this point 512 bytes at toybuf+512 are free (already parsed).
-  // Start of toybuf still has name in it.
-
-  // Loop through fields
+  // Loop through fields to display
   for (field = TT.fields; field; field = field->next) {
-    char *out = toybuf+2048, *scratch = toybuf+512;
+    char *out = buf;
 
     // Default: unsupported (5 "C")
     sprintf(out, "-");
 
-    // PID, PPID, PRI, NI, ADDR, SZ, RSS, PGID, VSS, MAJFL, MINFL
+    // stat#s: PID, PPID, PRI, NI, ADDR, SZ, RSS, PGID, VSS, MAJFL, MINFL
     if (-1!=(i = stridx((char[]){3,4,6,7,8,9,24,19,23,25,30,0}, field->which)))
     {
       char *fmt = "%lld";
@@ -254,7 +211,7 @@ static int do_ps(struct dirtree *new)
       else if (i==6) ll <<= 2;
       else if (i==8) ll >>= 10;
       sprintf(out, fmt, ll);
-    // UID USER RUID RUSER GID GROUP RGID RGROUP
+    // user/group: UID USER RUID RUSER GID GROUP RGID RGROUP
     } else if (-1!=(i = stridx((char[]){2,22,28,21,26,17,29,20}, field->which)))
     {
       int id = slot[31+i/2]; // uid, ruid, gid, rgid
@@ -272,97 +229,43 @@ static int do_ps(struct dirtree *new)
           if (pw) out = pw->pw_name;
         }
       }
+    // strslot: CMD TTY WCHAN LABEL (plus CMDLINE handled elsewhere)
+    } else if (-1!=(i = stridx((char[]){15,12,10,31}, field->which)))
+      out = strslot[i];
  
     // F (also assignment of i used by later tests)
     // Posix doesn't specify what flags should say. Man page says
     // 1 for PF_FORKNOEXEC and 4 for PF_SUPERPRIV from linux/sched.h
-    } else if (!(i = field->which)) sprintf(out, "%llo", (slot[6]>>6)&5);
+    else if (!(i = field->which)) sprintf(out, "%llo", (slot[6]>>6)&5);
     // S STAT
     else if (i==1 || i==27) {
-      sprintf(out, "%c", state);
+      s = out;
+      *s++ = state;
       if (i==27) {
         // TODO l = multithreaded
-        s = out+1;
         if (slot[16]<0) *s++ = '<';
         else if (slot[16]>0) *s++ = 'N';
         if (slot[3]==*slot) *s++ = 's';
         if (slot[18]) *s++ = 'L';
         if (slot[5]==*slot) *s++ = '+';
-        *s = 0;
       } 
-    // WCHAN
-    } else if (i==10) {
-      sprintf(scratch, "%lld/wchan", *slot);
-      readfileat(dirtree_parentfd(new), scratch, out, 2047);
-
-    // LABEL
-    } else if (i==31) {
-      sprintf(scratch, "%lld/attr/current", *slot);
-      readfileat(dirtree_parentfd(new), scratch, out, 2047);
-      chomp(out);
-
+      *s = 0;
     // STIME
     } else if (i==11) {
-      time_t t = time(0) - get_uptime() + slot[19]/sysconf(_SC_CLK_TCK);
+      time_t t = time(0)-slot[46]+slot[19]/sysconf(_SC_CLK_TCK);
 
       // Padding behavior's a bit odd: default field size is just hh:mm.
-      // Increasing stime:size reveals more data at left until full
-      // yyyy-mm-dd hh:mm revealed at :16, then adds :ss at end for :19. But
-      // expanding last field just adds :ss.
-      strftime(scratch, 512, "%F %T", localtime(&t));
-      out = scratch+strlen(scratch)-3-abs(field->len);
-      if (out<scratch) out = scratch;
-
-    // TTY
-    } else if (i==12) {
-      int rdev = slot[4];
-      struct stat st;
-
-      // Call no tty "?" rather than "0:0".
-      if (!rdev) strcpy(out, "?");
-      else {
-        // Can we readlink() our way to a name?
-        for (i=0; i<3; i++) {
-          sprintf(scratch, "%lld/fd/%i", *slot, i);
-          fd = dirtree_parentfd(new);
-          if (!fstatat(fd, scratch, &st, 0) && S_ISCHR(st.st_mode)
-            && st.st_rdev == rdev
-            && 0<(len = readlinkat(fd, scratch, out, 2047)))
-          {
-            out[len] = 0;
-            break;
-          }
-        }
-
-        // Couldn't find it, try all the tty drivers.
-        if (i == 3) {
-          FILE *fp = fopen("/proc/tty/drivers", "r");
-          int tty_major = 0, maj = major(rdev), min = minor(rdev);
-
-          if (fp) {
-            while (fscanf(fp, "%*s %256s %d %*s %*s", out, &tty_major) == 2) {
-              // TODO: we could parse the minor range too.
-              if (tty_major == maj) {
-                sprintf(out + strlen(out), "%d", min);
-                if (!stat(out, &st) && S_ISCHR(st.st_mode) && st.st_rdev==rdev)
-                  break;
-              }
-              tty_major = 0;
-            }
-            fclose(fp);
-          }
-
-          // Really couldn't find it, so just show major:minor.
-          if (!tty_major) sprintf(out, "%d:%d", maj, min);
-        }
-
-        strstart(&out, "/dev/");
-      }
+      // Increasing stime:size reveals more data at left until full,
+      // so move start address so yyyy-mm-dd hh:mm revealed on left at :16,
+      // then add :ss on right for :19.
+      strftime(out, 64, "%F %T", localtime(&t));
+      out = out+strlen(out)-3-abs(field->len);
+      if (out<buf) out = buf;
 
     // TIME ELAPSED
     } else if (i==13 || i==16) {
       int unit = 60*60*24, j = sysconf(_SC_CLK_TCK);
-      time_t seconds = (i==16) ? (get_uptime()*j)-slot[19] : slot[11]+slot[12];
+      time_t seconds = (i==16) ? (slot[46]*j)-slot[19] : slot[11]+slot[12];
 
       seconds /= j;
       for (s = 0, j = 0; j<4; j++) {
@@ -377,35 +280,19 @@ static int do_ps(struct dirtree *new)
       }
 
     // COMM - command line including arguments
-    // Command line limited to 2k displayable. We could dynamically malloc, but
-    // it'd almost never get used, querying length of a proc file is awkward,
-    // fixed buffer is nommu friendly... Wait for somebody to complain. :)
-    // CMDLINE - command line from /proc/pid/cmdline without arguments
+    // CMDLINE - command name from /proc/pid/cmdline (no arguments)
     } else if (i==14 || i==32) {
-      int fd;
-
-      len = 0;
-      sprintf(out, "%lld/cmdline", *slot);
-      fd = openat(dirtree_parentfd(new), out, O_RDONLY);
- 
-      if (fd != -1) {
-        if (0<(len = read(fd, out, 2047))) {
-          if (!out[len-1]) len--;
-          else out[len] = 0;
-          if (i==14) for (i = 0; i<len; i++) if (out[i] < ' ') out[i] = ' ';
-        }
-        close(fd);
+      if (slot[47]<1) {
+        out = toybuf+sizeof(toybuf)-300;
+        sprintf(out, "[%s]", *strslot); // kernel thread, use real name
+      } else {
+        out = strslot[4];
+        if (slot[47]!=INT_MAX) out[slot[47]] = ' '*(i==14);
       }
 
-      if (len<1) sprintf(out, "[%.*s]", nlen, name);
-
-    // CMD - command name (without arguments)
-    } else if (i==15) {
-      sprintf(out, "%.*s", nlen, name);
-
     // %CPU
     } else if (i==18) {
-      ll = (get_uptime()*sysconf(_SC_CLK_TCK)-slot[19]);
+      ll = (slot[46]*sysconf(_SC_CLK_TCK)-slot[19]);
       len = ((slot[11]+slot[12])*1000)/ll;
       sprintf(out, "%d.%d", len/10, len%10);
     }
@@ -423,8 +310,172 @@ static int do_ps(struct dirtree *new)
     if (!width) break;
   }
   xputc('\n');
+}
 
-  return 0;
+// dirtree callback: read data about process to display, store, or discard it.
+// Collects slot[50] plus  stat, name, tty, wchan, attr, cmdline
+static int get_ps(struct dirtree *new)
+{
+  struct {
+    char *name;
+    long long bits;
+  } fetch[] = {{"fd/", 1<<12}, {"wchan", 1<<10}, {"attr/current", 1<<31},
+               {"cmdline", (1<<14)|(1LL<<32)}};
+  long long *slot = (void *)toybuf;
+  char *name, *s, *buf = (char *)(slot+50), *end = 0;
+  int i, j, fd, ksave = DIRTREE_SAVE*!!(toys.optflags&FLAG_k);
+  off_t len;
+
+  // Recurse one level into /proc children, skip non-numeric entries
+  if (!new->parent) return DIRTREE_RECURSE|DIRTREE_SHUTUP|ksave;
+  if (!(*slot = atol(new->name))) return 0;
+  fd = dirtree_parentfd(new);
+
+  len = 2048;
+  sprintf(buf, "%lld/stat", *slot);
+  if (!readfileat(fd, buf, buf, &len)) return 0;
+
+  // parse oddball fields (name and state). Name can have embedded ')' so match
+  // _last_ ')' in stat (although VFS limits filenames to 255 bytes max).
+  // All remaining fields should be numeric.
+  if (!(name = strchr(buf, '('))) return 0;
+  for (s = ++name; *s; s++) if (*s == ')') end = s;
+  if (!end || end-name>255) return 0;
+
+  // Put status in first byte, and copy name after with low chars spaced.
+  if (1>sscanf(s = end, ") %c%n", buf++, &i)) return 0;
+  for (i = 0; i<end-name; i++) if ((buf[i] = name[i]) < ' ') buf[i] = ' ';
+  buf += i;
+  *buf++ = 0;
+  i = 3;
+
+  // parse numeric fields (PID = 0, skip 2, then 4th field goes in slot[1])
+  for (j = 1; j<50; j++) if (1>sscanf(s += i, " %lld%n", slot+j, &i)) break;
+
+  // save uid, ruid, gid, gid, and rgid int slots 31-34 (we don't use sigcatch
+  // or numeric wchan, and the remaining two are always zero), and vmlck into
+  // 18 (which is "obsolete, always 0")
+  slot[31] = new->st.st_uid;
+  slot[33] = new->st.st_gid;
+
+  // If RGROUP RUSER STAT RUID RGID happening, or -G or -U, parse "status"
+  // Save ruid in slot[34] and rgid in slot[35], which are otherwise zero,
+  // and vmlck into slot[18] (it_real_value, also always zero).
+  if ((TT.bits & 0x38300000) || TT.GG.len || TT.UU.len) {
+    len = sizeof(toybuf)-(buf-toybuf);
+    sprintf(buf, "%lld/status", *slot);
+    if (!readfileat(fd, buf, buf, &len)) *buf = 0;
+    s = strstr(buf, "\nUid:");
+    slot[32] = s ? atol(s+5) : new->st.st_uid;
+    s = strstr(buf, "\nGid:");
+    slot[34] = s ? atol(s+5) : new->st.st_gid;
+    s = strstr(buf, "\nVmLck:");
+    if (s) slot[18] = atoll(s+5);
+  }
+
+  // We now know enough to skip processes we don't care about.
+  if (!match_process(slot)) return 0;
+
+  // /proc data is generated as it's read, so for maximum accuracy on slow
+  // systems (or ps | more) we re-fetch uptime as we fetch each /proc line.
+  sysinfo((void *)(toybuf+2048));
+  slot[46] = ((struct sysinfo *)toybuf)->uptime;
+
+  // fetch remaining data while parentfd still available, appending to buf.
+  // (There's well over 3k of toybuf left. We could dynamically malloc, but
+  // it'd almost never get used, querying length of a proc file is awkward,
+  // fixed buffer is nommu friendly... Wait for somebody to complain. :)
+  for (j = 0; j<ARRAY_LEN(fetch); j++) { 
+    if (!(TT.bits&fetch[j].bits)) {
+      *buf++ = 0;
+      continue;
+    }
+
+    // Determine remaining space, reserving minimum of 256 bytes/field and
+    // 64 bytes scratch space at the end (for output conversion later).
+    len = sizeof(toybuf)-(buf-toybuf)-64-256*(ARRAY_LEN(fetch)-j);
+    sprintf(buf, "%lld/%s", *slot, fetch[j].name);
+
+    // If it's not the TTY field, data we want is in a file.
+    // Last length saved in slot[] is command line (which has embedded NULs)
+    if (j) {
+      readfileat(fd, buf, buf, &len);
+
+      // When command has no arguments, don't space over the NUL
+      if (len>0) {
+        int temp = 0;
+
+        if (buf[len-1]=='\n') buf[--len] = 0;
+
+        // Always escape spaces because an executable named esc[0m would be bad.
+        // Escaping low ascii does not affect utf8.
+        for (i=0; i<len; i++) {
+          if (!temp && !buf[i]) temp = i;
+          if (buf[i]<' ') buf[i] = ' ';
+        }
+        if (temp) len = temp; // position of _first_ NUL
+        else len = INT_MAX;
+      } else *buf = len = 0;
+      // Store end of argv[0] so COMM and CMDLINE can differ.
+      slot[47] = len;
+    } else {
+      int rdev = slot[4];
+      struct stat st;
+
+      // Call no tty "?" rather than "0:0".
+      strcpy(buf, "?");
+      if (rdev) {
+        // Can we readlink() our way to a name?
+        for (i = 0; i<3; i++) {
+          sprintf(buf, "%lld/fd/%i", *slot, i);
+          if (!fstatat(fd, buf, &st, 0) && S_ISCHR(st.st_mode)
+            && st.st_rdev == rdev && 0<(len = readlinkat(fd, buf, buf, len)))
+          {
+            buf[len] = 0;
+            break;
+          }
+        }
+
+        // Couldn't find it, try all the tty drivers.
+        if (i == 3) {
+          FILE *fp = fopen("/proc/tty/drivers", "r");
+          int tty_major = 0, maj = major(rdev), min = minor(rdev);
+
+          if (fp) {
+            while (fscanf(fp, "%*s %256s %d %*s %*s", buf, &tty_major) == 2) {
+              // TODO: we could parse the minor range too.
+              if (tty_major == maj) {
+                sprintf(buf+strlen(buf), "%d", min);
+                if (!stat(buf, &st) && S_ISCHR(st.st_mode) && st.st_rdev==rdev)
+                  break;
+              }
+              tty_major = 0;
+            }
+            fclose(fp);
+          }
+
+          // Really couldn't find it, so just show major:minor.
+          if (!tty_major) sprintf(buf, "%d:%d", maj, min);
+        }
+
+        s = buf;
+        if (strstart(&s, "/dev/")) memmove(buf, s, strlen(s)+1);;
+      }
+    }
+    buf += strlen(buf)+1;
+  }
+
+  // If we need to sort the output, add it to the list and return.
+  if (ksave) {
+    s = xmalloc(buf-toybuf);
+    new->extra = (long)s;
+    memcpy(s, toybuf, buf-toybuf);
+    TT.kcount++;
+
+  // Otherwise display it now
+  } else show_ps(slot);
+
+  return ksave;
 }
 
 // Traverse arg_list of csv, calling callback on each value
@@ -597,8 +648,16 @@ static char *parse_rest(char *str, int len)
   return str;
 }
 
+// sort for -k
+static int ksort(void *aa, void *bb)
+{
+  // TODO: sort!
+  return 1;
+}
+
 void ps_main(void)
 {
+  struct dirtree *dt;
   int i;
 
   TT.width = 99999;
@@ -658,7 +717,26 @@ void ps_main(void)
   dlist_terminate(TT.fields);
   printf("%s\n", toybuf);
 
-  dirtree_read("/proc", do_ps);
+  dt = dirtree_read("/proc", get_ps);
+
+  if (toys.optflags&FLAG_k) {
+    struct dirtree **dtsort = xmalloc(TT.kcount*sizeof(struct dirtree *));
+
+    // descend into child list
+    *dtsort = dt;
+    dt = dt->child;
+    free(*dtsort);
+
+    // populate array
+    i = 0;
+    for (i = 0; dt; dt = dt->next) dtsort[i++] = dt;
+    qsort(dtsort, TT.kcount, sizeof(struct dirtree *), (void *)ksort);
+    for (i = 0; i<TT.kcount; i++) {
+      show_ps((void *)dtsort[i]->extra);
+      free((void *)dtsort[i]->extra);
+      free(dtsort[i]);
+    }
+  }
 
   if (CFG_TOYBOX_FREE) {
     free(TT.gg.ptr);