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