1 /*
2 ** Zabbix
3 ** Copyright (C) 2001-2021 Zabbix SIA
4 **
5 ** This program is free software; you can redistribute it and/or modify
6 ** it under the terms of the GNU General Public License as published by
7 ** the Free Software Foundation; either version 2 of the License, or
8 ** (at your option) any later version.
9 **
10 ** This program is distributed in the hope that it will be useful,
11 ** but WITHOUT ANY WARRANTY; without even the implied warranty of
12 ** MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13 ** GNU General Public License for more details.
14 **
15 ** You should have received a copy of the GNU General Public License
16 ** along with this program; if not, write to the Free Software
17 ** Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
18 **/
19
20 /*
21 ** Ideas from PostgreSQL implementation (src/backend/utils/misc/ps_status.c)
22 ** were used in development of this file. Thanks to PostgreSQL developers!
23 **/
24
25 #include "common.h"
26 #include "setproctitle.h"
27
28 #if defined(PS_DARWIN_ARGV)
29 #include <crt_externs.h>
30 #endif
31
32 #if defined(PS_OVERWRITE_ARGV)
33 /* external environment we got on startup */
34 extern char **environ;
35 static int argc_ext_copied_first = 0, argc_ext_copied_last = 0, environ_ext_copied = 0;
36 static char **environ_ext = NULL;
37
38 /* internal copy of argv[] and environment variables */
39 static char **argv_int = NULL, **environ_int = NULL;
40 static const char *empty_str = "";
41
42 /* ps display buffer */
43 static char *ps_buf = NULL;
44 static size_t ps_buf_size = 0, prev_msg_size = 0;
45 #elif defined(PS_PSTAT_ARGV)
46 #define PS_BUF_SIZE 512
47 static char ps_buf[PS_BUF_SIZE], *p_msg = NULL;
48 static size_t ps_buf_size = PS_BUF_SIZE, ps_buf_size_msg = PS_BUF_SIZE;
49 #endif
50
51 /******************************************************************************
52 * *
53 * Function: setproctitle_save_env *
54 * *
55 * Purpose: prepare for changing process commandline to display status *
56 * messages with "ps" command on platforms which do not support *
57 * setproctitle(). Depending on platform: *
58 * - make a copy of argc, argv[] and environment variables to *
59 * enable overwriting original argv[]. *
60 * - prepare a buffer with common part of status message. *
61 * *
62 * Comments: call this function soon after main process start, before using *
63 * argv[] and environment variables. *
64 * *
65 ******************************************************************************/
66 #if defined(PS_OVERWRITE_ARGV)
setproctitle_save_env(int argc,char ** argv)67 char **setproctitle_save_env(int argc, char **argv)
68 {
69 int i;
70 char *arg_next = NULL;
71
72 if (NULL == argv || 0 == argc)
73 return argv;
74
75 /* measure a size of continuous argv[] area and make a copy */
76
77 argv_int = (char **)zbx_malloc(argv_int, ((unsigned int)argc + 1) * sizeof(char *));
78
79 #if defined(PS_APPEND_ARGV)
80 argc_ext_copied_first = argc - 1;
81 #else
82 argc_ext_copied_first = 0;
83 #endif
84 for (i = 0; i < argc_ext_copied_first; i++)
85 argv_int[i] = argv[i];
86
87 for (i = argc_ext_copied_first, arg_next = argv[argc_ext_copied_first]; arg_next == argv[i]; i++)
88 {
89 arg_next = argv[i] + strlen(argv[i]) + 1;
90 argv_int[i] = zbx_strdup(NULL, argv[i]);
91
92 /* argv[argc_ext_copied_first] will be used to display status messages. The rest of arguments can be */
93 /* overwritten and their argv[] pointers will point to wrong strings. */
94 if (argc_ext_copied_first < i)
95 argv[i] = (char *)empty_str;
96 }
97
98 argc_ext_copied_last = i - 1;
99
100 for (; i < argc; i++)
101 argv_int[i] = argv[i];
102
103 argv_int[argc] = NULL; /* C standard: "argv[argc] shall be a null pointer" */
104
105 if (argc_ext_copied_last == argc - 1)
106 {
107 int envc = 0;
108
109 while (NULL != environ[envc])
110 envc++;
111
112 /* measure a size of continuous environment area and make a copy */
113
114 environ_int = (char **)zbx_malloc(environ_int, ((unsigned int)envc + 1) * sizeof(char *));
115
116 for (i = 0; arg_next == environ[i]; i++)
117 {
118 arg_next = environ[i] + strlen(environ[i]) + 1;
119 environ_int[i] = zbx_strdup(NULL, environ[i]);
120
121 /* environment variables can be overwritten by status messages in argv[0] */
122 /* and environ[] pointers will point to wrong strings */
123 environ[i] = (char *)empty_str;
124 }
125
126 environ_ext_copied = i;
127
128 for (; i < envc; i++)
129 environ_int[i] = environ[i];
130
131 environ_int[envc] = NULL;
132 }
133
134 ps_buf_size = (size_t)(arg_next - argv[argc_ext_copied_first]);
135 ps_buf = argv[argc_ext_copied_first];
136
137 #if defined(PS_CONCAT_ARGV)
138 {
139 char *p = ps_buf;
140 size_t size = ps_buf_size, len;
141
142 for (i = argc_ext_copied_first + 1; i < argc; i++)
143 {
144 len = strlen(argv_int[i - 1]);
145 p += len;
146 size -= len;
147 if (2 >= size)
148 break;
149 zbx_strlcpy(p++, " ", size--);
150 zbx_strlcpy(p, argv_int[i], size);
151 }
152 }
153 #endif
154
155 #if defined(PS_DARWIN_ARGV)
156 *_NSGetArgv() = argv_int;
157 #endif
158 environ_ext = environ;
159 environ = environ_int; /* switch environment to internal copy */
160
161 return argv_int;
162 }
163 #elif defined(PS_PSTAT_ARGV)
setproctitle_save_env(int argc,char ** argv)164 char **setproctitle_save_env(int argc, char **argv)
165 {
166 size_t len0;
167
168 len0 = strlen(argv[0]);
169
170 if (len0 + 2 < ps_buf_size) /* is there space for ": " ? */
171 {
172 zbx_strlcpy(ps_buf, argv[0], ps_buf_size);
173 zbx_strlcpy(ps_buf + len0, ": ", (size_t)3);
174 p_msg = ps_buf + len0 + 2;
175 ps_buf_size_msg = ps_buf_size - len0 - 2; /* space after "argv[0]: " for status message */
176 }
177 return argv;
178 }
179 #endif /* defined(PS_PSTAT_ARGV) */
180
181 /******************************************************************************
182 * *
183 * Function: setproctitle_set_status *
184 * *
185 * Purpose: set a process command line displayed by "ps" command. *
186 * *
187 * Comments: call this function when a process starts some interesting task. *
188 * Program name argv[0] will be displayed "as-is" followed by ": " *
189 * and a status message. *
190 * *
191 ******************************************************************************/
setproctitle_set_status(const char * status)192 void setproctitle_set_status(const char *status)
193 {
194 #if defined(PS_OVERWRITE_ARGV)
195 static int initialized = 0;
196
197 if (1 == initialized)
198 {
199 size_t msg_size;
200
201 msg_size = zbx_strlcpy(ps_buf, status, ps_buf_size);
202
203 if (prev_msg_size > msg_size)
204 memset(ps_buf + msg_size + 1, '\0', ps_buf_size - msg_size - 1);
205
206 prev_msg_size = msg_size;
207 }
208 else if (NULL != ps_buf)
209 {
210 size_t start_pos;
211
212 /* Initialization has not been moved to setproctitle_save_env() because setproctitle_save_env() */
213 /* is called from the main process and we do not change its command line. */
214 /* argv[] changing takes place only in child processes. */
215
216 #if defined(PS_CONCAT_ARGV)
217 start_pos = strlen(argv_int[0]);
218 #else
219 start_pos = strlen(ps_buf);
220 #endif
221 if (start_pos + 2 < ps_buf_size) /* is there space for ": " ? */
222 {
223 zbx_strlcpy(ps_buf + start_pos, ": ", (size_t)3);
224 ps_buf += start_pos + 2;
225 ps_buf_size -= start_pos + 2; /* space after "argv[copy_first]: " for status message */
226
227 memset(ps_buf, '\0', ps_buf_size);
228 prev_msg_size = zbx_strlcpy(ps_buf, status, ps_buf_size);
229
230 initialized = 1;
231 }
232 }
233 #elif defined(PS_PSTAT_ARGV)
234 if (NULL != p_msg)
235 {
236 union pstun pst;
237
238 zbx_strlcpy(p_msg, status, ps_buf_size_msg);
239 pst.pst_command = ps_buf;
240 pstat(PSTAT_SETCMD, pst, strlen(ps_buf), 0, 0);
241 }
242 #endif
243 }
244
245 /******************************************************************************
246 * *
247 * Function: setproctitle_free_env *
248 * *
249 * Purpose: release memory allocated in setproctitle_save_env(). *
250 * *
251 * Comments: call this function when process terminates and argv[] and *
252 * environment variables are not used anymore. *
253 * *
254 ******************************************************************************/
255 #if defined(PS_OVERWRITE_ARGV)
setproctitle_free_env(void)256 void setproctitle_free_env(void)
257 {
258 int i;
259
260 /* restore the original environment variable to safely free our internally allocated environ array */
261 if (environ == environ_int)
262 environ = environ_ext;
263
264 for (i = argc_ext_copied_first; i <= argc_ext_copied_last; i++)
265 zbx_free(argv_int[i]);
266
267 for (i = 0; i <= environ_ext_copied; i++)
268 zbx_free(environ_int[i]);
269
270 zbx_free(argv_int);
271 zbx_free(environ_int);
272 }
273 #endif /* PS_OVERWRITE_ARGV */
274