lib/libzutil/os/linux/zutil_setproctitle.c

   1 /*
   2  * Copyright © 2013 Guillem Jover <guillem@hadrons.org>
   3  *
   4  * Redistribution and use in source and binary forms, with or without
   5  * modification, are permitted provided that the following conditions
   6  * are met:
   7  * 1. Redistributions of source code must retain the above copyright
   8  *    notice, this list of conditions and the following disclaimer.
   9  * 2. Redistributions in binary form must reproduce the above copyright
  10  *    notice, this list of conditions and the following disclaimer in the
  11  *    documentation and/or other materials provided with the distribution.
  12  * 3. The name of the author may not be used to endorse or promote products
  13  *    derived from this software without specific prior written permission.
  14  *
  15  * THIS SOFTWARE IS PROVIDED ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES,
  16  * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY
  17  * AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL
  18  * THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
  19  * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
  20  * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
  21  * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
  22  * WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
  23  * OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
  24  * ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  25  */
  26
  27 #include <errno.h>
  28 #include <stddef.h>
  29 #include <stdarg.h>
  30 #include <stdlib.h>
  31 #include <stdio.h>
  32 #include <err.h>
  33 #include <unistd.h>
  34 #include <string.h>
  35 #include <sys/param.h>
  36 #include <libzutil.h>
  37
  38 static struct {
  39         /* Original value. */
  40         const char *arg0;
  41
  42         /* Title space available. */
  43         char *base, *end;
  44
  45         /* Pointer to original nul character within base. */
  46         char *nul;
  47
  48         boolean_t warned;
  49         boolean_t reset;
  50         int error;
  51 } SPT;
  52
  53 #define LIBBSD_IS_PATHNAME_SEPARATOR(c) ((c) == '/')
  54 #define SPT_MAXTITLE 255
  55
  56 extern const char *__progname;
  57
  58 static const char *
  59 getprogname(void)
  60 {
  61         return (__progname);
  62 }
  63
  64 static void
  65 setprogname(const char *progname)
  66 {
  67         size_t i;
  68
  69         for (i = strlen(progname); i > 0; i--) {
  70                 if (LIBBSD_IS_PATHNAME_SEPARATOR(progname[i - 1])) {
  71                         __progname = progname + i;
  72                         return;
  73                 }
  74         }
  75         __progname = progname;
  76 }
  77
  78
  79 static inline size_t
  80 spt_min(size_t a, size_t b)
  81 {
  82         return ((a < b) ? a : b);
  83 }
  84
  85 static int
  86 spt_copyenv(int envc, char *envp[])
  87 {
  88         char **envcopy;
  89         char *eq;
  90         int envsize;
  91         int i, error = 0;
  92
  93         if (environ != envp)
  94                 return (0);
  95
  96         /*
  97          * Make a copy of the old environ array of pointers, in case
  98          * clearenv() or setenv() is implemented to free the internal
  99          * environ array, because we will need to access the old environ
 100          * contents to make the new copy.
 101          */
 102         envsize = (envc + 1) * sizeof (char *);
 103         envcopy = malloc(envsize);
 104         if (envcopy == NULL)
 105                 return (errno);
 106         memcpy(envcopy, envp, envsize);
 107
 108         environ = NULL;
 109
 110         for (i = 0; envcopy[i]; i++) {
 111                 eq = strchr(envcopy[i], '=');
 112                 if (eq == NULL)
 113                         continue;
 114
 115                 *eq = '\0';
 116                 if (setenv(envcopy[i], eq + 1, 1) < 0)
 117                         error = errno;
 118                 *eq = '=';
 119
 120                 if (error) {
 121                         clearenv();
 122                         environ = envp;
 123                         free(envcopy);
 124                         return (error);
 125                 }
 126         }
 127
 128         /*
 129          * Dispose of the shallow copy, now that we've finished transfering
 130          * the old environment.
 131          */
 132         free(envcopy);
 133
 134         return (0);
 135 }
 136
 137 static int
 138 spt_copyargs(int argc, char *argv[])
 139 {
 140         char *tmp;
 141         int i;
 142
 143         for (i = 1; i < argc || (i >= argc && argv[i]); i++) {
 144                 if (argv[i] == NULL)
 145                         continue;
 146
 147                 tmp = strdup(argv[i]);
 148                 if (tmp == NULL)
 149                         return (errno);
 150
 151                 argv[i] = tmp;
 152         }
 153
 154         return (0);
 155 }
 156
 157 void
 158 zfs_setproctitle_init(int argc, char *argv[], char *envp[])
 159 {
 160         char *base, *end, *nul, *tmp;
 161         int i, envc, error;
 162
 163         /* Try to make sure we got called with main() arguments. */
 164         if (argc < 0)
 165                 return;
 166
 167         base = argv[0];
 168         if (base == NULL)
 169                 return;
 170
 171         nul = base + strlen(base);
 172         end = nul + 1;
 173
 174         for (i = 0; i < argc || (i >= argc && argv[i]); i++) {
 175                 if (argv[i] == NULL || argv[i] != end)
 176                         continue;
 177
 178                 end = argv[i] + strlen(argv[i]) + 1;
 179         }
 180
 181         for (i = 0; envp[i]; i++) {
 182                 if (envp[i] != end)
 183                         continue;
 184
 185                 end = envp[i] + strlen(envp[i]) + 1;
 186         }
 187         envc = i;
 188
 189         SPT.arg0 = strdup(argv[0]);
 190         if (SPT.arg0 == NULL) {
 191                 SPT.error = errno;
 192                 return;
 193         }
 194
 195         tmp = strdup(getprogname());
 196         if (tmp == NULL) {
 197                 SPT.error = errno;
 198                 return;
 199         }
 200         setprogname(tmp);
 201
 202         error = spt_copyenv(envc, envp);
 203         if (error) {
 204                 SPT.error = error;
 205                 return;
 206         }
 207
 208         error = spt_copyargs(argc, argv);
 209         if (error) {
 210                 SPT.error = error;
 211                 return;
 212         }
 213
 214         SPT.nul  = nul;
 215         SPT.base = base;
 216         SPT.end  = end;
 217 }
 218
 219 void
 220 zfs_setproctitle(const char *fmt, ...)
 221 {
 222         /* Use buffer in case argv[0] is passed. */
 223         char buf[SPT_MAXTITLE + 1];
 224         va_list ap;
 225         char *nul;
 226         int len;
 227         if (SPT.base == NULL) {
 228                 if (!SPT.warned) {
 229                         warnx("setproctitle not initialized, please"
 230                             "call zfs_setproctitle_init()");
 231                         SPT.warned = B_TRUE;
 232                 }
 233                 return;
 234         }
 235
 236         if (fmt) {
 237                 if (fmt[0] == '-') {
 238                         /* Skip program name prefix. */
 239                         fmt++;
 240                         len = 0;
 241                 } else {
 242                         /* Print program name heading for grep. */
 243                         snprintf(buf, sizeof (buf), "%s: ", getprogname());
 244                         len = strlen(buf);
 245                 }
 246
 247                 va_start(ap, fmt);
 248                 len += vsnprintf(buf + len, sizeof (buf) - len, fmt, ap);
 249                 va_end(ap);
 250         } else {
 251                 len = snprintf(buf, sizeof (buf), "%s", SPT.arg0);
 252         }
 253
 254         if (len <= 0) {
 255                 SPT.error = errno;
 256                 return;
 257         }
 258
 259         if (!SPT.reset) {
 260                 memset(SPT.base, 0, SPT.end - SPT.base);
 261                 SPT.reset = B_TRUE;
 262         } else {
 263                 memset(SPT.base, 0, spt_min(sizeof (buf), SPT.end - SPT.base));
 264         }
 265
 266         len = spt_min(len, spt_min(sizeof (buf), SPT.end - SPT.base) - 1);
 267         memcpy(SPT.base, buf, len);
 268         nul = SPT.base + len;
 269
 270         if (nul < SPT.nul) {
 271                 *SPT.nul = '.';
 272         } else if (nul == SPT.nul && nul + 1 < SPT.end) {
 273                 *SPT.nul = ' ';
 274                 *++nul = '\0';
 275         }
 276 }