screenenv.c

   1 /* The customized getenv() function itself. */
   2
   3 /* We keep a cached environment for the latest client, with only
   4  * the interesting variables picked out. */
   5 /* FIXME: We should keep the old clients' environments as well,
   6  * in case the app uses stale getenv() pointers! */
   7
   8 #define _GNU_SOURCE 1
   9 #include <stdio.h>
  10 #include <stdlib.h>
  11 #include <string.h>
  12 #include <unistd.h>
  13 #include <fcntl.h>
  14 #include <dlfcn.h>
  15
  16 #include "screenenv.h"
  17
  18 /* Just add more if you need. */
  19 const static char *grabenv[] = {
  20         "DISPLAY",
  21         // "SSH_AUTH_SOCK", - enable if you use ssh agent forwarding rigoriously
  22         "SSH_CLIENT",
  23         "SSH_CONNECTION",
  24         "SSH_TTY",
  25         "XDG_SESSION_COOKIE",
  26         NULL
  27 };
  28
  29
  30 /* getenv() upcall setup. */
  31
  32 static char *(*up_getenv)(const char *);
  33 static int in_screen = 0;
  34
  35 static void __attribute__((constructor))
  36 up_getenv_setup(void)
  37 {
  38         /* XXX: We do this here instead of lazily on first call, since
  39          * there is a possibility of dlsym()-malloc() race when called
  40          * within a getenv() call. */
  41         up_getenv = dlsym(RTLD_NEXT, "getenv");
  42
  43         char *term = up_getenv("TERM");
  44         in_screen = term && !strcmp(term, "screen");
  45 }
  46
  47
  48 /* The cached environment. */
  49
  50 struct env {
  51         pid_t pid;
  52         char *vars[sizeof(grabenv) / sizeof(*grabenv)];
  53         char buf_[4096];
  54 };
  55
  56 static struct env env0;
  57
  58
  59 /* Parse environment in /proc/.../environ and save it in *env. */
  60
  61 int loadenv(pid_t pid, struct env *env)
  62 {
  63         env->pid = pid;
  64
  65         /* XXX: All the code below is utterly stupid since we should've just
  66          * mmap()'d the file and parse it in place. */
  67
  68         char envname[128];
  69         snprintf(envname, 128, "/proc/%d/environ", pid);
  70         int fd = open(envname, O_RDONLY);
  71         if (fd < 0) return -1;
  72
  73         char *envp = env->buf_;
  74
  75         char buf[8192], *bufp;
  76         size_t bl = 0;
  77
  78         /* What do we do in this read() iteration; we keep enough context
  79          * necessary in the buffer from the previous iteration; if the
  80          * variable is not interesting or is too long for our taste, we
  81          * just skip to next one. */
  82         enum { VARNAME_LOAD, VALUE_LOAD, VAR_SKIP } state = VARNAME_LOAD;
  83         int var;
  84
  85         while (1) {
  86                 int l;
  87 next_data:
  88                 l = read(fd, buf + bl, sizeof(buf) - bl);
  89                 if (l < 0) {
  90                         close(fd);
  91                         return -1;
  92                 }
  93                 bl += l;
  94                 if (bl <= 0) break;
  95                 bufp = buf;
  96
  97                 while (bl > 0) {
  98                         char *tok;
  99                         switch (state) {
 100                                 case VARNAME_LOAD:
 101                                         tok = memchr(bufp, '=', bl);
 102                                         if (!tok) {
 103                                                 if (buf == bufp) {
 104                                                         state = VAR_SKIP;
 105                                                         bl = 0;
 106                                                 } else {
 107                                                         memmove(buf, bufp, bl);
 108                                                 }
 109                                                 goto next_data;
 110                                         }
 111                                         *tok++ = 0;
 112                                         for (var = 0; grabenv[var]; var++)
 113                                                 if (!strcmp(bufp, grabenv[var]))
 114                                                         break;
 115                                         bl -= tok - bufp; bufp = tok;
 116                                         if (grabenv[var]) {
 117                                                 state = VALUE_LOAD;
 118                                         } else {
 119                                                 state = VAR_SKIP;
 120                                         }
 121                                         break;
 122                                 case VAR_SKIP:
 123                                         tok = memchr(bufp, 0, bl);
 124                                         if (!tok) {
 125                                                 bl = 0;
 126                                                 goto next_data;
 127                                         }
 128                                         tok++;
 129                                         bl -= tok - bufp; bufp = tok;
 130                                         state = VARNAME_LOAD;
 131                                         break;
 132                                 case VALUE_LOAD:
 133                                         tok = memchr(bufp, 0, bl);
 134                                         if (!tok) {
 135                                                 if (buf == bufp) {
 136                                                         state = VAR_SKIP;
 137                                                         bl = 0;
 138                                                 } else {
 139                                                         memmove(buf, bufp, bl);
 140                                                 }
 141                                                 goto next_data;
 142                                         }
 143                                         tok++;
 144                                         int vallen = tok - bufp;
 145                                         if (sizeof(env->buf_) - (envp - env->buf_) < vallen) {
 146                                                 /* Environment space full. */
 147                                                 close(fd);
 148                                                 return -1;
 149                                         }
 150                                         memcpy(envp, bufp, vallen);
 151                                         // printf("[%d] `%s'\n", var, envp);
 152                                         env->vars[var] = envp; envp += vallen;
 153                                         bl -= tok - bufp; bufp = tok;
 154                                         state = VARNAME_LOAD;
 155                                         break;
 156                         }
 157                 }
 158         }
 159
 160         close(fd);
 161         return 0;
 162 }
 163
 164
 165 /* The getenv() wrapper itself! */
 166
 167 __attribute__((visibility("default"))) char *
 168 getenv(const char *env)
 169 {
 170         /* We may be called by a crazy
 171          *      up_getenv_setup() -> dlsym() -> calloc() -> getenv()
 172          * sequence. */
 173         if (__builtin_expect(up_getenv == NULL, 0))
 174                 return ""; /* XXX */
 175
 176         /* The fast way out, not to hog non-screen processes. */
 177         if (__builtin_expect(!in_screen, 1))
 178                 return up_getenv(env);
 179
 180         /* TERM=screen - so probably we will get client pid? */
 181         pid_t pid = get_client_pid(getpid());
 182         if (__builtin_expect(pid <= 0, 0))
 183                 return up_getenv(env);
 184
 185         /* Do we look for a hijacked variable? */
 186         int var;
 187         for (var = 0; grabenv[var]; var++)
 188                 if (__builtin_expect(!strcmp(env, grabenv[var]), 0))
 189                         break;
 190         if (__builtin_expect(!grabenv[var], 1))
 191                 return up_getenv(env);
 192
 193         /* Should we reload the environment? */
 194         if (env0.pid != pid) {
 195                 memset(&env0, 0, sizeof(env0));
 196                 if (loadenv(pid, &env0) < 0)
 197                         return up_getenv(env);
 198         }
 199
 200         return env0.vars[var];
 201 }
 202
 203 #if 0
 204 void putsn(char *str)
 205 {
 206         puts(str ? str : "(null)");
 207 }
 208
 209 int main(void)
 210 {
 211         putsn(getenv("USER"));
 212         putsn(getenv("DISPLAY"));
 213         putsn(getenv("TERM"));
 214         putsn(getenv("XDG_SESSION_COOKIE"));
 215         return 0;
 216 }
 217 #endif