missing/setproctitle.c

   1 /* Based on setproctitle.c from openssh-5.6p1 */
   2 /* Based on conf.c from UCB sendmail 8.8.8 */
   3
   4 /*
   5  * Copyright 2003 Damien Miller
   6  * Copyright (c) 1983, 1995-1997 Eric P. Allman
   7  * Copyright (c) 1988, 1993
   8  *      The Regents of the University of California.  All rights reserved.
   9  *
  10  * Redistribution and use in source and binary forms, with or without
  11  * modification, are permitted provided that the following conditions
  12  * are met:
  13  * 1. Redistributions of source code must retain the above copyright
  14  *    notice, this list of conditions and the following disclaimer.
  15  * 2. Redistributions in binary form must reproduce the above copyright
  16  *    notice, this list of conditions and the following disclaimer in the
  17  *    documentation and/or other materials provided with the distribution.
  18  * 3. Neither the name of the University nor the names of its contributors
  19  *    may be used to endorse or promote products derived from this software
  20  *    without specific prior written permission.
  21  *
  22  * THIS SOFTWARE IS PROVIDED BY THE REGENTS AND CONTRIBUTORS ``AS IS'' AND
  23  * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  24  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  25  * ARE DISCLAIMED.  IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS BE LIABLE
  26  * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
  27  * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
  28  * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
  29  * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
  30  * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
  31  * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
  32  * SUCH DAMAGE.
  33  */
  34
  35 #include "ruby.h"
  36 #include "ruby/util.h"
  37 #define compat_init_setproctitle ruby_init_setproctitle
  38 RUBY_FUNC_EXPORTED void ruby_init_setproctitle(int argc, char *argv[]);
  39
  40 #ifndef HAVE_SETPROCTITLE
  41
  42 #include <stdarg.h>
  43 #include <stdlib.h>
  44 #ifdef HAVE_UNISTD_H
  45 #include <unistd.h>
  46 #endif
  47 #ifdef HAVE_SYS_PSTAT_H
  48 #include <sys/pstat.h>
  49 #endif
  50 #include <string.h>
  51
  52 #if defined(__APPLE__)
  53 # ifdef HAVE_CRT_EXTERNS_H
  54 #  include <crt_externs.h>
  55 #  undef environ
  56 #  define environ (*_NSGetEnviron())
  57 # else
  58 #  include "crt_externs.h"
  59 # endif
  60 #endif
  61
  62 #define SPT_NONE        0       /* don't use it at all */
  63 #define SPT_PSTAT       1       /* use pstat(PSTAT_SETCMD, ...) */
  64 #define SPT_REUSEARGV   2       /* cover argv with title information */
  65
  66 #ifndef SPT_TYPE
  67 # define SPT_TYPE       SPT_NONE
  68 #endif
  69
  70 #ifndef SPT_PADCHAR
  71 # define SPT_PADCHAR    '\0'
  72 #endif
  73
  74 #if SPT_TYPE == SPT_REUSEARGV
  75 static char *argv_start = NULL;
  76 static size_t argv_env_len = 0;
  77 static size_t argv_len = 0;
  78 static char **argv1_addr = NULL;
  79 #endif
  80
  81 #endif /* HAVE_SETPROCTITLE */
  82
  83 #if defined(SPT_TYPE) && SPT_TYPE == SPT_REUSEARGV
  84 # define ALLOCATE_ENVIRON 1
  85 #else
  86 # define ALLOCATE_ENVIRON 0
  87 #endif
  88
  89 #if ALLOCATE_ENVIRON
  90 /* system_environ is the value of environ before we allocate a custom buffer.
  91  *
  92  * We use this to restore environ in ruby_free_proctitle.
  93  */
  94 static char **system_environ = NULL;
  95 /* orig_environ is the buffer we allocate for environ.
  96  *
  97  * We use this to free this buffer in ruby_free_proctitle. When we add new
  98  * environment variables using setenv, the system may change environ to a
  99  * different buffer and will not free the original buffer, so we need to hold
 100  * onto this so we can free it in ruby_free_proctitle.
 101  *
 102  * We must not free any of the contents because it may change if the system
 103  * updates existing environment variables.
 104  */
 105 static char **orig_environ = NULL;
 106 /* alloc_environ is a copy of orig_environ.
 107  *
 108  * We use this to free all the original string copies that were in orig_environ.
 109  * Since environ could be changed to point to strings allocated by the system
 110  * if environment variables are updated, so we need this to point to the
 111  * original strings.
 112  */
 113 static char **alloc_environ = NULL;
 114 #endif
 115
 116 void
 117 compat_init_setproctitle(int argc, char *argv[])
 118 {
 119 #if ALLOCATE_ENVIRON
 120         extern char **environ;
 121         char *lastargv = NULL;
 122         char *lastenvp = NULL;
 123         char **envp = environ;
 124         int i;
 125
 126         /*
 127          * NB: This assumes that argv has already been copied out of the
 128          * way. This is true for sshd, but may not be true for other
 129          * programs. Beware.
 130          */
 131
 132         if (argc == 0 || argv[0] == NULL)
 133                 return;
 134
 135         /* Fail if we can't allocate room for the new environment */
 136         for (i = 0; envp[i] != NULL; i++);
 137
 138         system_environ = environ;
 139
 140         alloc_environ = xcalloc(i + 1, sizeof(*environ));
 141         orig_environ = environ = xcalloc(i + 1, sizeof(*environ));
 142         if (environ == NULL) {
 143                 environ = envp; /* put it back */
 144                 return;
 145         }
 146
 147         /*
 148          * Find the last argv string or environment variable within
 149          * our process memory area.
 150          */
 151         for (i = 0; i < argc; i++) {
 152                 if (lastargv == NULL || lastargv + 1 == argv[i])
 153                         lastargv = argv[i] + strlen(argv[i]);
 154         }
 155         lastenvp = lastargv;
 156         for (i = 0; envp[i] != NULL; i++) {
 157                 if (lastenvp + 1 == envp[i])
 158                         lastenvp = envp[i] + strlen(envp[i]);
 159         }
 160
 161         /* We keep argv[1], argv[2], etc. at this moment,
 162            because the ps command of AIX refers to them. */
 163         argv1_addr = &argv[1];
 164         argv_start = argv[0];
 165         argv_len = lastargv - argv[0];
 166         argv_env_len = lastenvp - argv[0];
 167
 168         for (i = 0; envp[i] != NULL; i++)
 169                 alloc_environ[i] = environ[i] = ruby_strdup(envp[i]);
 170         alloc_environ[i] = environ[i] = NULL;
 171 #endif /* SPT_REUSEARGV */
 172 }
 173
 174 void
 175 ruby_free_proctitle(void)
 176 {
 177 #if ALLOCATE_ENVIRON
 178         extern char **environ;
 179
 180         if (!orig_environ) return; /* environ is allocated by OS */
 181
 182         for (int i = 0; alloc_environ[i] != NULL; i++) {
 183                 xfree(alloc_environ[i]);
 184         }
 185         xfree(alloc_environ);
 186         xfree(orig_environ);
 187
 188         environ = system_environ;
 189 #endif
 190 }
 191
 192 #ifndef HAVE_SETPROCTITLE
 193
 194 void
 195 setproctitle(const char *fmt, ...)
 196 {
 197 #if SPT_TYPE != SPT_NONE
 198         va_list ap;
 199         char ptitle[1024];
 200         size_t len;
 201         size_t argvlen;
 202 #if SPT_TYPE == SPT_PSTAT
 203         union pstun pst;
 204 #endif
 205
 206 #if SPT_TYPE == SPT_REUSEARGV
 207         if (argv_env_len <= 0)
 208                 return;
 209 #endif
 210
 211         /* fmt must be non-NULL */
 212         va_start(ap, fmt);
 213         vsnprintf(ptitle, sizeof(ptitle), fmt, ap);
 214         va_end(ap);
 215
 216 #if SPT_TYPE == SPT_PSTAT
 217         pst.pst_command = ptitle;
 218         pstat(PSTAT_SETCMD, pst, strlen(ptitle), 0, 0);
 219 #elif SPT_TYPE == SPT_REUSEARGV
 220         len = strlcpy(argv_start, ptitle, argv_env_len);
 221         argvlen = len > argv_len ? argv_env_len : argv_len;
 222         for(; len < argvlen; len++)
 223                 argv_start[len] = SPT_PADCHAR;
 224         /* argv[1], argv[2], etc. are no longer valid. */
 225         *argv1_addr = NULL;
 226 #endif
 227
 228 #endif /* SPT_NONE */
 229 }
 230
 231 #endif /* HAVE_SETPROCTITLE */