New xsetenv() plumbing (repeatedly set same environment variables without

leaking memory), and mod env command to test it.

5 files changed, 130 insertions, 59 deletions

diff --git a/lib/env.c b/lib/env.c
new file mode 100644
index 00000000..94ea0a4e
--- /dev/null
+++ b/lib/env.c
@@ -0,0 +1,115 @@
+// Can't trust libc not to leak enviornment variable memory, so...
+
+#include "toys.h"
+
+// In libc, populated by start code,used by getenv() and exec() and friends.
+extern char **environ;
+
+// Returns the number of bytes taken by the environment variables. For use
+// when calculating the maximum bytes of environment+argument data that can
+// be passed to exec for find(1) and xargs(1).
+long environ_bytes()
+{
+  long bytes = sizeof(char *);
+  char **ev;
+
+  for (ev = environ; *ev; ev++) bytes += sizeof(char *) + strlen(*ev) + 1;
+
+  return bytes;
+}
+
+// This will clear the inherited environment if called first thing.
+// Use this instead of envc so we keep track of what needs to be freed.
+void xclearenv(void)
+{
+  toys.envc = 0;
+  *environ = 0;
+}
+
+// Frees entries we set earlier. Use with libc getenv but not setenv/putenv.
+// if name has an equals and !val, act like putenv (name=val must be malloced!)
+// if !val unset name. (Name with = and val is an error)
+void xsetenv(char *name, char *val)
+{
+  unsigned i, len, envc;
+  char *new;
+
+  // If we haven't snapshot initial environment state yet, do so now.
+  if (!toys.envc) {
+    while (environ[toys.envc++]);
+    memcpy(new = xmalloc(((toys.envc|0xff)+1)*sizeof(char *)),
+      environ, toys.envc*sizeof(char *));
+    environ = (void *)new;
+  }
+
+  new = strchr(name, '=');
+  if (new) {
+    len = new-name;
+    if (val) error_exit("xsetenv %s to %s", name, val);
+    new = name;
+  } else {
+    len = strlen(name);
+    if (val) new = xmprintf("%s=%s", name, val);
+  }
+
+  envc = toys.envc-1;
+  for (i = 0; environ[i]; i++) {
+    // Drop old entry, freeing as appropriate. Assumes no duplicates.
+    if (!memcmp(name, environ[i], len) && environ[i][len]=='=') {
+      if (i>=envc) free(environ[i]);
+      else {
+        char **delete = environ+i;
+
+        // move old entries down, add at end of old data
+        toys.envc = envc--;
+        for (i=0; new ? i<envc : !!delete[i]; i++) delete[i] = delete[i+1];
+        i = envc;
+      }
+      break;
+    }
+  }
+
+  if (!new) return;
+
+  // resize and null terminate if expanding
+  if (!environ[i]) {
+    len = i+1;
+    if (!(len&255)) environ = xrealloc(environ, len*sizeof(char *));
+    environ[len] = 0;
+  }
+  environ[i] = new;
+}
+
+// reset environment for a user, optionally clearing most of it
+void reset_env(struct passwd *p, int clear)
+{
+  int i;
+
+  if (clear) {
+    char *s, *stuff[] = {"TERM", "DISPLAY", "COLORTERM", "XAUTHORITY"};
+
+    for (i=0; i<ARRAY_LEN(stuff); i++)
+      stuff[i] = (s = getenv(stuff[i])) ? xmprintf("%s=%s", stuff[i], s) : 0;
+    xclearenv();
+    for (i=0; i < ARRAY_LEN(stuff); i++) if (stuff[i]) xsetenv(stuff[i], 0);
+    if (chdir(p->pw_dir)) {
+      perror_msg("chdir %s", p->pw_dir);
+      xchdir("/");
+    }
+  } else {
+    char **ev1, **ev2;
+
+    // remove LD_*, IFS, ENV, and BASH_ENV from environment
+    for (ev1 = ev2 = environ;;) {
+      while (*ev2 && (strstart(ev2, "LD_") || strstart(ev2, "IFS=") ||
+        strstart(ev2, "ENV=") || strstart(ev2, "BASH_ENV="))) ev2++;
+      if (!(*ev1++ = *ev2++)) break;
+    }
+  }
+
+  setenv("PATH", _PATH_DEFPATH, 1);
+  setenv("HOME", p->pw_dir, 1);
+  setenv("SHELL", p->pw_shell, 1);
+  setenv("USER", p->pw_name, 1);
+  setenv("LOGNAME", p->pw_name, 1);
+}
diff --git a/lib/lib.c b/lib/lib.c
index c4d871c2..7b7d84d6 100644
--- a/lib/lib.c
+++ b/lib/lib.c
@@ -1398,19 +1398,6 @@ void do_lines(int fd, char delim, void (*call)(char **pline, long len))
   if (fd) fclose(fp);
 }
 
-// Returns the number of bytes taken by the environment variables. For use
-// when calculating the maximum bytes of environment+argument data that can
-// be passed to exec for find(1) and xargs(1).
-long environ_bytes()
-{
-  long bytes = sizeof(char *);
-  char **ev;
-
-  for (ev = environ; *ev; ev++) bytes += sizeof(char *) + strlen(*ev) + 1;
-
-  return bytes;
-}
-
 // Return unix time in milliseconds
 long long millitime(void)
 {
@@ -1432,40 +1419,6 @@ char *format_iso_time(char *buf, size_t len, struct timespec *ts)
   return buf;
 }
 
-// reset environment for a user, optionally clearing most of it
-void reset_env(struct passwd *p, int clear)
-{
-  int i;
-
-  if (clear) {
-    char *s, *stuff[] = {"TERM", "DISPLAY", "COLORTERM", "XAUTHORITY"};
-
-    for (i=0; i<ARRAY_LEN(stuff); i++)
-      stuff[i] = (s = getenv(stuff[i])) ? xmprintf("%s=%s", stuff[i], s) : 0;
-    clearenv();
-    for (i=0; i < ARRAY_LEN(stuff); i++) if (stuff[i]) putenv(stuff[i]);
-    if (chdir(p->pw_dir)) {
-      perror_msg("chdir %s", p->pw_dir);
-      xchdir("/");
-    }
-  } else {
-    char **ev1, **ev2;
-
-    // remove LD_*, IFS, ENV, and BASH_ENV from environment
-    for (ev1 = ev2 = environ;;) {
-      while (*ev2 && (strstart(ev2, "LD_") || strstart(ev2, "IFS=") ||
-        strstart(ev2, "ENV=") || strstart(ev2, "BASH_ENV="))) ev2++;
-      if (!(*ev1++ = *ev2++)) break;
-    }
-  }
-
-  setenv("PATH", _PATH_DEFPATH, 1);
-  setenv("HOME", p->pw_dir, 1);
-  setenv("SHELL", p->pw_shell, 1);
-  setenv("USER", p->pw_name, 1);
-  setenv("LOGNAME", p->pw_name, 1);
-}
-
 // Syslog with the openlog/closelog, autodetecting daemon status via no tty
 
 void loggit(int priority, char *format, ...)
diff --git a/lib/lib.h b/lib/lib.h
index 6ec2c696..82c4c16e 100644
--- a/lib/lib.h
+++ b/lib/lib.h
@@ -261,7 +261,6 @@ int regexec0(regex_t *preg, char *string, long len, int nmatch,
 char *getusername(uid_t uid);
 char *getgroupname(gid_t gid);
 void do_lines(int fd, char delim, void (*call)(char **pline, long len));
-long environ_bytes();
 long long millitime(void);
 char *format_iso_time(char *buf, size_t len, struct timespec *ts);
 void reset_env(struct passwd *p, int clear);
@@ -272,6 +271,12 @@ void loggit(int priority, char *format, ...);
 #define HR_1000  4 // Use decimal instead of binary units
 int human_readable(char *buf, unsigned long long num, int style);
 
+// env.c
+
+long environ_bytes();
+void xsetenv(char *name, char *val);
+void xclearenv(void);
+
 // linestack.c
 
 struct linestack {
diff --git a/tests/env.test b/tests/env.test
index 286fb36c..3098731d 100755
--- a/tests/env.test
+++ b/tests/env.test
@@ -18,3 +18,4 @@ testcmd "-0" "-i five=six seven=eight env -0 | sort -z" "five=six\0seven=eight\0
 unset WALRUS BANANA LETTERS FILTER
 
 testcmd "early fail" '--oops 2> /dev/null ; echo $?' "125\n" "" ""
+testcmd "why is this allowed" "=BLAH env | grep '^=BLAH\$'" "=BLAH\n" "" ""
diff --git a/toys/posix/env.c b/toys/posix/env.c
index 5c7bb789..d1bb580d 100644
--- a/toys/posix/env.c
+++ b/toys/posix/env.c
@@ -28,11 +28,10 @@ GLOBALS(
   struct arg_list *u;
 );
 
-extern char **environ;
-
 void env_main(void)
 {
   char **ev = toys.optargs;
+  struct arg_list *u;
 
   // If first nonoption argument is "-" treat it as -i
   if (*ev && **ev == '-' && !(*ev)[1]) {
@@ -40,21 +39,19 @@ void env_main(void)
     ev++;
   }
 
-  if (toys.optflags & FLAG_i) clearenv();
-  while (TT.u) {
-    unsetenv(TT.u->arg);
-    TT.u = TT.u->next;
-  }
+  if (FLAG(i)) xclearenv();
+  else for (u = TT.u; u; u = u->next)
+    if (strchr(u->arg, '=')) error_msg("bad -u %s", u->arg);
+    else xsetenv(u->arg, 0);
 
   for (; *ev; ev++) {
-    char *name = *ev, *val = strchr(name, '=');
+    char *val = strchr(*ev, '=');
 
     if (val) {
       *(val++) = 0;
-      setenv(name, val, 1);
+      xsetenv(*ev, val);
     } else xexec(ev);
   }
 
-  if (environ) for (ev = environ; *ev; ev++)
-    xprintf("%s%c", *ev, '\n'*!(toys.optflags&FLAG_0));
+  for (ev = environ; *ev; ev++) xprintf("%s%c", *ev, '\n'*!FLAG(0));
 }