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
38static 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
56extern const char *__progname;
57
58static const char *
59getprogname(void)
60{
61	return (__progname);
62}
63
64static void
65setprogname(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
79static inline size_t
80spt_min(size_t a, size_t b)
81{
82	return ((a < b) ? a : b);
83}
84
85static int
86spt_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
137static int
138spt_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
157void
158zfs_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
219void
220zfs_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