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",
  22         "SSH_CLIENT",
  23         "SSH_CONNECTION",
  24         "SSH_TTY",
  25         "XDG_SESSION_COOKIE",
  26         NULL
  27 };
  28
  29
  30 /* The cached environment. */
  31
  32 struct env {
  33         pid_t pid;
  34         char *vars[sizeof(grabenv) / sizeof(*grabenv)];
  35         char buf_[4096];
  36 };
  37
  38 static struct env env0;
  39
  40
  41 /* Parse environment in /proc/.../environ and save it in *env. */
  42
  43 int loadenv(pid_t pid, struct env *env)
  44 {
  45         env->pid = pid;
  46
  47         /* XXX: All the code below is utterly stupid since we should've just
  48          * mmap()'d the file and parse it in place. */
  49
  50         char envname[128];
  51         snprintf(envname, 128, "/proc/%d/environ", pid);
  52         int fd = open(envname, O_RDONLY);
  53         if (fd < 0) return -1;
  54
  55         char *envp = env->buf_;
  56
  57         char buf[8192], *bufp;
  58         size_t bl = 0;
  59
  60         /* What do we do in this read() iteration; we keep enough context
  61          * necessary in the buffer from the previous iteration; if the
  62          * variable is not interesting or is too long for our taste, we
  63          * just skip to next one. */
  64         enum { VARNAME_LOAD, VALUE_LOAD, VAR_SKIP } state = VARNAME_LOAD;
  65         int var;
  66
  67         while (1) {
  68                 int l;
  69 next_data:
  70                 l = read(fd, buf + bl, sizeof(buf) - bl);
  71                 if (l < 0) {
  72                         close(fd);
  73                         return -1;
  74                 }
  75                 bl += l;
  76                 if (bl <= 0) break;
  77                 bufp = buf;
  78
  79                 while (bl > 0) {
  80                         char *tok;
  81                         switch (state) {
  82                                 case VARNAME_LOAD:
  83                                         tok = memchr(bufp, '=', bl);
  84                                         if (!tok) {
  85                                                 if (buf == bufp) {
  86                                                         state = VAR_SKIP;
  87                                                         bl = 0;
  88                                                 } else {
  89                                                         memmove(buf, bufp, bl);
  90                                                 }
  91                                                 goto next_data;
  92                                         }
  93                                         *tok++ = 0;
  94                                         for (var = 0; grabenv[var]; var++)
  95                                                 if (!strcmp(bufp, grabenv[var]))
  96                                                         break;
  97                                         bl -= tok - bufp; bufp = tok;
  98                                         if (grabenv[var]) {
  99                                                 state = VALUE_LOAD;
 100                                         } else {
 101                                                 state = VAR_SKIP;
 102                                         }
 103                                         break;
 104                                 case VAR_SKIP:
 105                                         tok = memchr(bufp, 0, bl);
 106                                         if (!tok) {
 107                                                 bl = 0;
 108                                                 goto next_data;
 109                                         }
 110                                         tok++;
 111                                         bl -= tok - bufp; bufp = tok;
 112                                         state = VARNAME_LOAD;
 113                                         break;
 114                                 case VALUE_LOAD:
 115                                         tok = memchr(bufp, 0, bl);
 116                                         if (!tok) {
 117                                                 if (buf == bufp) {
 118                                                         state = VAR_SKIP;
 119                                                         bl = 0;
 120                                                 } else {
 121                                                         memmove(buf, bufp, bl);
 122                                                 }
 123                                                 goto next_data;
 124                                         }
 125                                         tok++;
 126                                         int vallen = tok - bufp;
 127                                         if (sizeof(env->buf_) - (envp - env->buf_) < vallen) {
 128                                                 /* Environment space full. */
 129                                                 close(fd);
 130                                                 return -1;
 131                                         }
 132                                         memcpy(envp, bufp, vallen);
 133                                         // printf("[%d] `%s'\n", var, envp);
 134                                         env->vars[var] = envp; envp += vallen;
 135                                         bl -= tok - bufp; bufp = tok;
 136                                         state = VARNAME_LOAD;
 137                                         break;
 138                         }
 139                 }
 140         }
 141
 142         close(fd);
 143         return 0;
 144 }
 145
 146
 147 /* The getenv() wrapper itself! */
 148
 149 __attribute__((visibility("default"))) char *
 150 getenv(const char *env)
 151 {
 152         static char *(*up_getenv)(const char *);
 153         if (__builtin_expect(!up_getenv, 0))
 154                 up_getenv = dlsym(RTLD_NEXT, "getenv");
 155
 156         /* The fast way out, not to hog non-screen processes. */
 157         static int in_screen = -1;
 158         if (__builtin_expect(in_screen < 0, 0)) {
 159                 char *term = up_getenv("TERM");
 160                 in_screen = term && !strcmp(term, "screen");
 161         }
 162         if (__builtin_expect(!in_screen, 1))
 163                 return up_getenv(env);
 164
 165         /* TERM=screen - so probably we will get client pid? */
 166         pid_t pid = get_client_pid(getpid());
 167         if (__builtin_expect(pid <= 0, 0))
 168                 return up_getenv(env);
 169
 170         /* Do we look for a hijacked variable? */
 171         int var;
 172         for (var = 0; grabenv[var]; var++)
 173                 if (__builtin_expect(!strcmp(env, grabenv[var]), 0))
 174                         break;
 175         if (__builtin_expect(!grabenv[var], 1))
 176                 return up_getenv(env);
 177
 178         /* Should we reload the environment? */
 179         if (env0.pid != pid) {
 180                 memset(&env0, 0, sizeof(env0));
 181                 if (loadenv(pid, &env0) < 0)
 182                         return up_getenv(env);
 183         }
 184
 185         return env0.vars[var];
 186 }
 187
 188 #if 0
 189 void putsn(char *str)
 190 {
 191         puts(str ? str : "(null)");
 192 }
 193
 194 int main(void)
 195 {
 196         putsn(getenv("USER"));
 197         putsn(getenv("DISPLAY"));
 198         putsn(getenv("TERM"));
 199         putsn(getenv("XDG_SESSION_COOKIE"));
 200         return 0;
 201 }
 202 #endif