1 /* This file is part of GNU Pies.
2 Copyright (C) 2007-2020 Sergey Poznyakoff
3
4 GNU Pies is free software; you can redistribute it and/or modify
5 it under the terms of the GNU General Public License as published by
6 the Free Software Foundation; either version 3, or (at your option)
7 any later version.
8
9 GNU Pies is distributed in the hope that it will be useful,
10 but WITHOUT ANY WARRANTY; without even the implied warranty of
11 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12 GNU General Public License for more details.
13
14 You should have received a copy of the GNU General Public License
15 along with GNU Pies. If not, see <http://www.gnu.org/licenses/>. */
16
17 #ifdef HAVE_CONFIG_H
18 # include <config.h>
19 #endif
20 #include <sys/types.h>
21 #include <unistd.h>
22 #include <stdlib.h>
23 #include <stdio.h>
24 #include <stdarg.h>
25 #include <string.h>
26 #if MF_PROCTITLE_TYPE == MF_PROCTITLE_PSTAT
27 # include <sys/pstat.h>
28 #elif MF_PROCTITLE_TYPE == MF_PROCTITLE_PSSTRINGS
29 # include <sys/proc.h>
30 # include <vm/pmap.h>
31 # include <machine/pmap.h>
32 # include <machine/vmparam.h>
33 # include <sys/exec.h>
34 #endif
35
36 extern char **environ;
37
38 /* Move the environment to prepare more space for argv */
39 static int
move_env(char * env[])40 move_env (char *env[])
41 {
42 size_t size;
43 int i;
44 char **p;
45
46 size = 0;
47 for (i = 0; env[i]; i++)
48 size += strlen (env[i]) + 1;
49
50 p = calloc (i + 1, sizeof (*p));
51 if (!p)
52 return 1;
53
54 for (i = 0; env[i]; i++)
55 if ((p[i] = strdup(env[i])) == NULL) {
56 int j;
57 /* Free allocated memory and return */
58 for (j = 0; j < i; j++)
59 free(p[i]);
60 free(p);
61 return 1;
62 }
63 p[i] = NULL;
64 environ = p;
65 return 0;
66 }
67
68 static int orig_argc;
69 static char **orig_argv;
70 static char *orig_argv_end;
71 static char *proctitle_buffer;
72 #ifdef HAVE___PROGNAME
73 extern char *__progname;
74 extern char *__progname_full;
75 #else
76 static char *__progname;
77 #endif
78
79 void
mf_proctitle_init(int argc,char * argv[],char * env[])80 mf_proctitle_init (int argc, char *argv[], char *env[])
81 {
82 int i;
83
84 move_env (env);
85
86 orig_argc = argc;
87 orig_argv = argv;
88 orig_argv_end = argv[0] + strlen (argv[0]);
89 __progname = strrchr (argv[0], '/');
90 if (__progname)
91 __progname++;
92 else
93 __progname = argv[0];
94 __progname = strdup (__progname);
95 #ifdef HAVE___PROGNAME
96 __progname_full = strdup(argv[0]);
97 #endif
98
99 for (i = 0; i < orig_argc; i++) {
100 if (orig_argv_end + 1 == argv[i])
101 orig_argv_end = argv[i] + strlen(argv[i]);
102 }
103
104 for (i = 0; env[i]; i++) {
105 if ((orig_argv_end + 1) == env[i])
106 orig_argv_end = env[i] + strlen(env[i]);
107 }
108 }
109
110 static void
mf_proctitle_flush(void)111 mf_proctitle_flush (void)
112 {
113 #if MF_PROCTITLE_TYPE == MF_PROCTITLE_SETPROCTITLE
114 setproctitle ("%s", proctitle_buffer);
115 #elif MF_PROCTITLE_TYPE == MF_PROCTITLE_REPLACE_ARGV
116 orig_argv[0] = proctitle_buffer;
117 for (i = 1; i < orig_argc; i++) {
118 orig_argv[i] = "";
119 }
120 #elif MF_PROCTITLE_TYPE == MF_PROCTITLE_REWRITE_ARGV
121 size_t argv_size = orig_argv_end - orig_argv[0] - 2;
122 size_t len = strlen (proctitle_buffer);
123 memset (orig_argv[0], 0, argv_size);
124 if (len > argv_size)
125 len = argv_size;
126 memcpy (orig_argv[0], proctitle_buffer, len);
127 orig_argv[0][len] = 0;
128 #elif MF_PROCTITLE_TYPE == MF_PROCTITLE_PSTAT
129 union pstun pst;
130 pst.pst_command = proc_title_buf;
131 pstat(PSTAT_SETCMD, pst, strlen (proctitle_buffer), 0, 0);
132 #elif MF_PROCTITLE_TYPE == MF_PROCTITLE_PSSTRINGS
133 PS_STRINGS->ps_nargvstr = 1;
134 PS_STRINGS->ps_argvstr = proctitle_buffer;
135 #endif
136 }
137
138 void
mf_proctitle_format(const char * fmt,...)139 mf_proctitle_format (const char *fmt, ...)
140 {
141 va_list ap;
142 char *tmp;
143 int rc;
144 if (!orig_argc)
145 return;
146 va_start (ap, fmt);
147 rc = vasprintf (&tmp, fmt, ap);
148 va_end (ap);
149 if (rc > 0) {
150 free (proctitle_buffer);
151 #if __FreeBSD__ >= 4
152 /* On FreeBSD the process name is prepended automatically */
153 proctitle_buffer = tmp;
154 #else
155 /* Otherwise we need to do that manually */
156 if (asprintf (&proctitle_buffer, "%s: %s", __progname, tmp) > 0)
157 free (tmp);
158 else
159 proctitle_buffer = tmp;
160 #endif
161 mf_proctitle_flush ();
162 }
163 }
164