1 /*****************************************************************************\
2 * src/slurmd/common/setproctitle.c - argv manipulation
3 *****************************************************************************
4 * Copyright (C) 2002-2007 The Regents of the University of California.
5 * Copyright (C) 2008-2009 Lawrence Livermore National Security.
6 * Produced at Lawrence Livermore National Laboratory (cf, DISCLAIMER).
7 * Written by Mark Grondona <mgrondona@llnl.gov>.
8 * CODE-OCEC-09-009. All rights reserved.
9 *
10 * This file is part of Slurm, a resource management program.
11 * For details, see <https://slurm.schedmd.com/>.
12 * Please also read the included file: DISCLAIMER.
13 *
14 * Slurm is free software; you can redistribute it and/or modify it under
15 * the terms of the GNU General Public License as published by the Free
16 * Software Foundation; either version 2 of the License, or (at your option)
17 * any later version.
18 *
19 * In addition, as a special exception, the copyright holders give permission
20 * to link the code of portions of this program with the OpenSSL library under
21 * certain conditions as described in each individual source file, and
22 * distribute linked combinations including the two. You must obey the GNU
23 * General Public License in all respects for all of the code used other than
24 * OpenSSL. If you modify file(s) with this exception, you may extend this
25 * exception to your version of the file(s), but you are not obligated to do
26 * so. If you do not wish to do so, delete this exception statement from your
27 * version. If you delete this exception statement from all source files in
28 * the program, then also delete it here.
29 *
30 * Slurm is distributed in the hope that it will be useful, but WITHOUT ANY
31 * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
32 * FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
33 * details.
34 *
35 * You should have received a copy of the GNU General Public License along
36 * with Slurm; if not, write to the Free Software Foundation, Inc.,
37 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
38 \*****************************************************************************/
39
40 /*
41 * Based on src/backend/utils/misc/pg_status.c from
42 * PostgreSQL Database Management System
43 *
44 * Portions Copyright (c) 1996-2001, The PostgreSQL Global Development Group
45 *
46 * Portions Copyright (c) 1994, The Regents of the University of California
47 *
48 * Permission to use, copy, modify, and distribute this software and its
49 * documentation for any purpose, without fee, and without a written agreement
50 * is hereby granted, provided that the above copyright notice and this
51 * paragraph and the following two paragraphs appear in all copies.
52 *
53 * IN NO EVENT SHALL THE UNIVERSITY OF CALIFORNIA BE LIABLE TO ANY PARTY FOR
54 * DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, INCLUDING
55 * LOST PROFITS, ARISING OUT OF THE USE OF THIS SOFTWARE AND ITS
56 * DOCUMENTATION, EVEN IF THE UNIVERSITY OF CALIFORNIA HAS BEEN ADVISED OF THE
57 * POSSIBILITY OF SUCH DAMAGE.
58 *
59 * THE UNIVERSITY OF CALIFORNIA SPECIFICALLY DISCLAIMS ANY WARRANTIES,
60 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY
61 * AND FITNESS FOR A PARTICULAR PURPOSE. THE SOFTWARE PROVIDED HEREUNDER IS
62 * ON AN "AS IS" BASIS, AND THE UNIVERSITY OF CALIFORNIA HAS NO OBLIGATIONS TO
63 * PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.
64 */
65
66 /*--------------------------------------------------------------------
67 * ps_status.c
68 *
69 * Routines to support changing the ps display of PostgreSQL backends
70 * to contain some useful information. Mechanism differs wildly across
71 * platforms.
72 *
73 * $Header: /var/cvs/openssh/openbsd-compat/setproctitle.c,v 1.5 2003/01/20 02:1
74 *
75 * Copyright 2000 by PostgreSQL Global Development Group
76 * various details abducted from various places
77 *--------------------------------------------------------------------
78 */
79
80 #include "config.h"
81
82 #include <errno.h>
83 #if defined(__NetBSD__)
84 #include <stdio.h>
85 #include <stdlib.h>
86 #include <string.h>
87 #endif
88 #if defined(__DragonFly__)
89 #include <stdio.h>
90 #include <string.h>
91 #include <unistd.h>
92 #include <stdlib.h>
93 #endif
94 #ifndef HAVE_SETPROCTITLE
95 #include <stdlib.h>
96 #include <stdio.h>
97 #include <stdarg.h>
98 #include <string.h>
99 #include <unistd.h>
100 #include "src/common/strlcpy.h"
101 #ifdef HAVE_SYS_PSTAT_H
102 #include <sys/pstat.h> /* for HP-UX */
103 #endif
104 #ifdef HAVE_PS_STRINGS
105 #include <machine/vmparam.h> /* for old BSD */
106 #include <sys/exec.h>
107 #endif
108
109 /*------
110 * Alternative ways of updating ps display:
111 *
112 * SETPROCTITLE_STRATEGY == PS_USE_PSTAT
113 * use the pstat(PSTAT_SETCMD, )
114 * (HPUX)
115 * SETPROCTITLE_STRATEGY == PS_USE_PS_STRINGS
116 * assign PS_STRINGS->ps_argvstr = "string"
117 * (some BSD systems)
118 * SETPROCTITLE_STRATEGY == PS_USE_CHANGE_ARGV
119 * assign argv[0] = "string"
120 * (some other BSD systems)
121 * SETPROCTITLE_STRATEGY == PS_USE_CLOBBER_ARGV
122 * write over the argv and environment area
123 * (most SysV-like systems)
124 * SETPROCTITLE_STRATEGY == PS_USE_NONE
125 * don't update ps display
126 * (This is the default, as it is safest.)
127 */
128
129 #define PS_USE_NONE 0
130 #define PS_USE_PSTAT 1
131 #define PS_USE_PS_STRINGS 2
132 #define PS_USE_CHANGE_ARGV 3
133 #define PS_USE_CLOBBER_ARGV 4
134
135 #ifndef SETPROCTITLE_STRATEGY
136 # define SETPROCTITLE_STRATEGY PS_USE_NONE
137 #endif
138
139 #ifndef SETPROCTITLE_PS_PADDING
140 # define SETPROCTITLE_PS_PADDING ' '
141 #endif
142 #endif /* HAVE_SETPROCTITLE */
143
144 extern char **environ;
145
146 /*
147 * argv clobbering uses existing argv space, all other methods need a buffer
148 */
149 #if SETPROCTITLE_STRATEGY != PS_USE_CLOBBER_ARGV
150 static char ps_buffer[256];
151 static const size_t ps_buffer_size = sizeof(ps_buffer);
152 #else
153 static char *ps_buffer; /* will point to argv area */
154 static size_t ps_buffer_size; /* space determined at run time */
155 static char **new_environ = (char **) NULL;
156 #endif
157
158 /* save the original *argv location here */
159 static int save_argc;
160 static char **save_argv;
161
162 #if HAVE__PROGNAME
163 extern char *__progname;
164 #else
165 static char __progname[64];
166 #endif
167
168 #ifndef HAVE_SETPROCTITLE
169 /*
170 * Call this to update the ps status display to a fixed prefix plus an
171 * indication of what you're currently doing passed in the argument.
172 */
173 void
setproctitle(const char * fmt,...)174 setproctitle(const char *fmt, ...)
175 {
176 #if SETPROCTITLE_STRATEGY == PS_USE_PSTAT
177 union pstun pst;
178 #endif
179 #if SETPROCTITLE_STRATEGY != PS_USE_NONE
180 ssize_t used;
181 va_list ap;
182
183 /* no ps display if you didn't call save_ps_display_args() */
184 if (save_argv == NULL)
185 return;
186 #if SETPROCTITLE_STRATEGY == PS_USE_CLOBBER_ARGV
187 /* If ps_buffer is a pointer, it might still be null */
188 if (ps_buffer == NULL)
189 return;
190 #endif /* PS_USE_CLOBBER_ARGV */
191
192 /*
193 * Overwrite *argv to point at appropriate space, if needed
194 */
195 #if SETPROCTITLE_STRATEGY == PS_USE_CHANGE_ARGV
196 save_argv[0] = ps_buffer;
197 save_argv[1] = NULL;
198 #endif /* PS_USE_CHANGE_ARGV */
199
200 #if SETPROCTITLE_STRATEGY == PS_USE_CLOBBER_ARGV
201 save_argv[1] = NULL;
202 #endif /* PS_USE_CLOBBER_ARGV */
203
204 /*
205 * Make fixed prefix of ps display.
206 */
207
208 va_start(ap, fmt);
209 if (fmt == NULL)
210 snprintf(ps_buffer, ps_buffer_size, "%s", __progname);
211 else {
212 used = snprintf(ps_buffer, ps_buffer_size, "%s: ", __progname);
213 if (used == -1 || used >= ps_buffer_size)
214 used = ps_buffer_size;
215 vsnprintf(ps_buffer + used, ps_buffer_size - used, fmt, ap);
216 }
217 va_end(ap);
218
219 #if SETPROCTITLE_STRATEGY == PS_USE_PSTAT
220 pst.pst_command = ps_buffer;
221 pstat(PSTAT_SETCMD, pst, strlen(ps_buffer), 0, 0);
222 #endif /* PS_USE_PSTAT */
223
224 #if SETPROCTITLE_STRATEGY == PS_USE_PS_STRINGS
225 PS_STRINGS->ps_nargvstr = 1;
226 PS_STRINGS->ps_argvstr = ps_buffer;
227 #endif /* PS_USE_PS_STRINGS */
228
229 #if SETPROCTITLE_STRATEGY == PS_USE_CLOBBER_ARGV
230 /* pad unused memory */
231 used = strlen(ps_buffer);
232 memset(ps_buffer + used, SETPROCTITLE_PS_PADDING,
233 ps_buffer_size - used);
234 #endif /* PS_USE_CLOBBER_ARGV */
235
236 #endif /* PS_USE_NONE */
237 }
238
_init__progname(const char * argv0)239 static void _init__progname (const char *argv0)
240 {
241 #if !HAVE__PROGNAME
242 char *start = strrchr (argv0, '/');
243 strlcpy (__progname, start ? (start + 1) : argv0, sizeof (__progname));
244 return;
245 #endif /* !HAVE__PROGNAME */
246 }
247
248 #endif /* HAVE_SETPROCTITLE */
249
250 /*
251 * Call this early in startup to save the original argc/argv values.
252 *
253 * *argv will not be overwritten by this routine, but may be overwritten
254 * during setproctitle. Also, the physical location of the environment
255 * strings may be moved, so this should be called before any code that
256 * might try to hang onto a getenv() result.
257 */
258 void
init_setproctitle(int argc,char ** argv)259 init_setproctitle(int argc, char **argv)
260 {
261 #if SETPROCTITLE_STRATEGY == PS_USE_CLOBBER_ARGV
262 char *end_of_area = NULL;
263 int i;
264 #endif
265
266 save_argc = argc;
267 save_argv = argv;
268
269 #if defined(__NetBSD__) || defined(__DragonFly__)
270 setprogname (argv[0]);
271 #else
272 _init__progname (argv[0]);
273 #endif
274
275 #if SETPROCTITLE_STRATEGY == PS_USE_CLOBBER_ARGV
276 /*
277 * If we're going to overwrite the argv area, count the available
278 * space. Also move the environment to make additional room.
279 */
280
281 /*
282 * check for contiguous argv strings
283 */
284 for (i = 0; i < argc; i++) {
285 if (i == 0 || end_of_area + 1 == argv[i])
286 end_of_area = argv[i] + strlen(argv[i]);
287 }
288
289 /* probably can't happen? */
290 if (end_of_area == NULL) {
291 ps_buffer = NULL;
292 ps_buffer_size = 0;
293 return;
294 }
295
296 /*
297 * check for contiguous environ strings following argv
298 */
299 for (i = 0; environ[i] != NULL; i++) {
300 if (end_of_area + 1 == environ[i])
301 end_of_area = environ[i] + strlen(environ[i]);
302 }
303
304 ps_buffer = argv[0];
305 ps_buffer_size = end_of_area - argv[0] - 1;
306
307 /*
308 * Duplicate and move the environment out of the way
309 */
310 new_environ = malloc(sizeof(char *) * (i + 1));
311 if (!new_environ) {
312 fprintf(stderr, "ERROR: [%s:%d] %s: %s\n",
313 __FILE__, __LINE__, "init_setproctitle",
314 strerror(errno));
315 abort();
316 }
317 for (i = 0; environ[i] != NULL; i++) {
318 new_environ[i] = strdup(environ[i]);
319 }
320 new_environ[i] = NULL;
321 environ = new_environ;
322 #endif /* PS_USE_CLOBBER_ARGV */
323 }
324
325 /* Free memory allocated by init_setproctitle.
326 * Used to verify that all allocated memory gets freed */
fini_setproctitle(void)327 void fini_setproctitle(void)
328 {
329 #if SETPROCTITLE_STRATEGY == PS_USE_CLOBBER_ARGV
330 int i;
331
332 if (!new_environ)
333 return;
334
335 for (i = 0; new_environ[i] != NULL; i++) {
336 free(new_environ[i]);
337 }
338 free(new_environ);
339 new_environ = (char **) NULL;
340 environ = new_environ;
341 #endif /* PS_USE_CLOBBER_ARGV */
342 }
343