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 /*
86  * For discussion on the portability of the various methods, see
87  * https://lists.freebsd.org/pipermail/freebsd-stable/2008-June/043136.html
88  */
89 static int
90 spt_clearenv(void)
91 {
92 	char **tmp;
93 
94 	tmp = malloc(sizeof (*tmp));
95 	if (tmp == NULL)
96 		return (errno);
97 
98 	tmp[0] = NULL;
99 	environ = tmp;
100 
101 	return (0);
102 }
103 
104 static int
105 spt_copyenv(int envc, char *envp[])
106 {
107 	char **envcopy;
108 	char *eq;
109 	int envsize;
110 	int i, error;
111 
112 	if (environ != envp)
113 		return (0);
114 
115 	/*
116 	 * Make a copy of the old environ array of pointers, in case
117 	 * clearenv() or setenv() is implemented to free the internal
118 	 * environ array, because we will need to access the old environ
119 	 * contents to make the new copy.
120 	 */
121 	envsize = (envc + 1) * sizeof (char *);
122 	envcopy = malloc(envsize);
123 	if (envcopy == NULL)
124 		return (errno);
125 	memcpy(envcopy, envp, envsize);
126 
127 	error = spt_clearenv();
128 	if (error) {
129 		environ = envp;
130 		free(envcopy);
131 		return (error);
132 	}
133 
134 	for (i = 0; envcopy[i]; i++) {
135 		eq = strchr(envcopy[i], '=');
136 		if (eq == NULL)
137 			continue;
138 
139 		*eq = '\0';
140 		if (setenv(envcopy[i], eq + 1, 1) < 0)
141 			error = errno;
142 		*eq = '=';
143 
144 		if (error) {
145 			environ = envp;
146 			free(envcopy);
147 			return (error);
148 		}
149 	}
150 
151 	/*
152 	 * Dispose of the shallow copy, now that we've finished transfering
153 	 * the old environment.
154 	 */
155 	free(envcopy);
156 
157 	return (0);
158 }
159 
160 static int
161 spt_copyargs(int argc, char *argv[])
162 {
163 	char *tmp;
164 	int i;
165 
166 	for (i = 1; i < argc || (i >= argc && argv[i]); i++) {
167 		if (argv[i] == NULL)
168 			continue;
169 
170 		tmp = strdup(argv[i]);
171 		if (tmp == NULL)
172 			return (errno);
173 
174 		argv[i] = tmp;
175 	}
176 
177 	return (0);
178 }
179 
180 void
181 zfs_setproctitle_init(int argc, char *argv[], char *envp[])
182 {
183 	char *base, *end, *nul, *tmp;
184 	int i, envc, error;
185 
186 	/* Try to make sure we got called with main() arguments. */
187 	if (argc < 0)
188 		return;
189 
190 	base = argv[0];
191 	if (base == NULL)
192 		return;
193 
194 	nul = base + strlen(base);
195 	end = nul + 1;
196 
197 	for (i = 0; i < argc || (i >= argc && argv[i]); i++) {
198 		if (argv[i] == NULL || argv[i] != end)
199 			continue;
200 
201 		end = argv[i] + strlen(argv[i]) + 1;
202 	}
203 
204 	for (i = 0; envp[i]; i++) {
205 		if (envp[i] != end)
206 			continue;
207 
208 		end = envp[i] + strlen(envp[i]) + 1;
209 	}
210 	envc = i;
211 
212 	SPT.arg0 = strdup(argv[0]);
213 	if (SPT.arg0 == NULL) {
214 		SPT.error = errno;
215 		return;
216 	}
217 
218 	tmp = strdup(getprogname());
219 	if (tmp == NULL) {
220 		SPT.error = errno;
221 		return;
222 	}
223 	setprogname(tmp);
224 
225 	error = spt_copyenv(envc, envp);
226 	if (error) {
227 		SPT.error = error;
228 		return;
229 	}
230 
231 	error = spt_copyargs(argc, argv);
232 	if (error) {
233 		SPT.error = error;
234 		return;
235 	}
236 
237 	SPT.nul  = nul;
238 	SPT.base = base;
239 	SPT.end  = end;
240 }
241 
242 void
243 zfs_setproctitle(const char *fmt, ...)
244 {
245 	/* Use buffer in case argv[0] is passed. */
246 	char buf[SPT_MAXTITLE + 1];
247 	va_list ap;
248 	char *nul;
249 	int len;
250 	if (SPT.base == NULL) {
251 		if (!SPT.warned) {
252 			warnx("setproctitle not initialized, please"
253 			    "call zfs_setproctitle_init()");
254 			SPT.warned = B_TRUE;
255 		}
256 		return;
257 	}
258 
259 	if (fmt) {
260 		if (fmt[0] == '-') {
261 			/* Skip program name prefix. */
262 			fmt++;
263 			len = 0;
264 		} else {
265 			/* Print program name heading for grep. */
266 			snprintf(buf, sizeof (buf), "%s: ", getprogname());
267 			len = strlen(buf);
268 		}
269 
270 		va_start(ap, fmt);
271 		len += vsnprintf(buf + len, sizeof (buf) - len, fmt, ap);
272 		va_end(ap);
273 	} else {
274 		len = snprintf(buf, sizeof (buf), "%s", SPT.arg0);
275 	}
276 
277 	if (len <= 0) {
278 		SPT.error = errno;
279 		return;
280 	}
281 
282 	if (!SPT.reset) {
283 		memset(SPT.base, 0, SPT.end - SPT.base);
284 		SPT.reset = B_TRUE;
285 	} else {
286 		memset(SPT.base, 0, spt_min(sizeof (buf), SPT.end - SPT.base));
287 	}
288 
289 	len = spt_min(len, spt_min(sizeof (buf), SPT.end - SPT.base) - 1);
290 	memcpy(SPT.base, buf, len);
291 	nul = SPT.base + len;
292 
293 	if (nul < SPT.nul) {
294 		*SPT.nul = '.';
295 	} else if (nul == SPT.nul && nul + 1 < SPT.end) {
296 		*SPT.nul = ' ';
297 		*++nul = '\0';
298 	}
299 }
300