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