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