Complete support for dtrace's -x setenv option.

This allows one to override the environment for processes created with
dtrace -c. By default, the environment is inherited.

This support was originally merged from illumos in r249367 but was lost
when the commit was later reverted and then brought back piecemeal.

Reported by:	Samuel Lepetit <slepetit@apple.com>
MFC after:	2 weeks

3 files changed, 74 insertions, 19 deletions

diff --git a/cddl/contrib/opensolaris/lib/libdtrace/common/dt_open.c b/cddl/contrib/opensolaris/lib/libdtrace/common/dt_open.c
index 91304548f094..a7027f5e1ba1 100644
--- a/cddl/contrib/opensolaris/lib/libdtrace/common/dt_open.c
+++ b/cddl/contrib/opensolaris/lib/libdtrace/common/dt_open.c
@@ -1182,7 +1182,7 @@ alloc:
 #endif
 	dtp->dt_provbuckets = _dtrace_strbuckets;
 	dtp->dt_provs = calloc(dtp->dt_provbuckets, sizeof (dt_provider_t *));
-	dt_proc_hash_create(dtp);
+	dt_proc_init(dtp);
 	dtp->dt_vmax = DT_VERS_LATEST;
 	dtp->dt_cpp_path = strdup(_dtrace_defcpp);
 	dtp->dt_cpp_argv = malloc(sizeof (char *));
@@ -1199,12 +1199,13 @@ alloc:
 	(void) uname(&dtp->dt_uts);
 
 	if (dtp->dt_mods == NULL || dtp->dt_provs == NULL ||
-	    dtp->dt_procs == NULL || dtp->dt_ld_path == NULL ||
+	    dtp->dt_procs == NULL || dtp->dt_proc_env == NULL ||
+	    dtp->dt_ld_path == NULL || dtp->dt_cpp_path == NULL ||
 #ifdef __FreeBSD__
 	    dtp->dt_kmods == NULL ||
 	    dtp->dt_objcopy_path == NULL ||
 #endif
-	    dtp->dt_cpp_path == NULL || dtp->dt_cpp_argv == NULL)
+	    dtp->dt_cpp_argv == NULL)
 		return (set_open_errno(dtp, errp, EDT_NOMEM));
 
 	for (i = 0; i < DTRACEOPT_MAX; i++)
@@ -1631,7 +1632,7 @@ dtrace_close(dtrace_hdl_t *dtp)
 	int i;
 
 	if (dtp->dt_procs != NULL)
-		dt_proc_hash_destroy(dtp);
+		dt_proc_fini(dtp);
 
 	while ((pgp = dt_list_next(&dtp->dt_programs)) != NULL)
 		dt_program_destroy(dtp, pgp);
diff --git a/cddl/contrib/opensolaris/lib/libdtrace/common/dt_proc.c b/cddl/contrib/opensolaris/lib/libdtrace/common/dt_proc.c
index 6cdd36c24b36..34bcc15b9ac7 100644
--- a/cddl/contrib/opensolaris/lib/libdtrace/common/dt_proc.c
+++ b/cddl/contrib/opensolaris/lib/libdtrace/common/dt_proc.c
@@ -25,6 +25,10 @@
  */
 
 /*
+ * Copyright (c) 2012 by Delphix. All rights reserved.
+ */
+
+/*
  * DTrace Process Control
  *
  * This file provides a set of routines that permit libdtrace and its clients
@@ -504,7 +508,7 @@ dt_proc_control(void *arg)
 	dt_proc_control_data_t *datap = arg;
 	dtrace_hdl_t *dtp = datap->dpcd_hdl;
 	dt_proc_t *dpr = datap->dpcd_proc;
-	dt_proc_hash_t *dph = dpr->dpr_hdl->dt_procs;
+	dt_proc_hash_t *dph = dtp->dt_procs;
 	struct ps_prochandle *P = dpr->dpr_proc;
 	int pid = dpr->dpr_pid;
 
@@ -965,14 +969,18 @@ dt_proc_create(dtrace_hdl_t *dtp, const char *file, char *const *argv,
 	(void) pthread_cond_init(&dpr->dpr_cv, NULL);
 
 #ifdef illumos
-	if ((dpr->dpr_proc = Pcreate(file, argv, &err, NULL, 0)) == NULL) {
+	dpr->dpr_proc = Pxcreate(file, argv, dtp->dt_proc_env, &err, NULL, 0);
+	if (dpr->dpr_proc == NULL) {
+		return (dt_proc_error(dtp, dpr,
+		    "failed to execute %s: %s\n", file, Pcreate_error(err)));
+	}
 #else
-	if ((err = proc_create(file, argv, NULL, pcf, child_arg,
+	if ((err = proc_create(file, argv, dtp->dt_proc_env, pcf, child_arg,
 	    &dpr->dpr_proc)) != 0) {
-#endif
 		return (dt_proc_error(dtp, dpr,
 		    "failed to execute %s: %s\n", file, Pcreate_error(err)));
 	}
+#endif
 
 	dpr->dpr_hdl = dtp;
 #ifdef illumos
@@ -1141,30 +1149,74 @@ dt_proc_unlock(dtrace_hdl_t *dtp, struct ps_prochandle *P)
 }
 
 void
-dt_proc_hash_create(dtrace_hdl_t *dtp)
+dt_proc_init(dtrace_hdl_t *dtp)
 {
+	extern char **environ;
+	static char *envdef[] = {
+		"LD_NOLAZYLOAD=1",	/* linker lazy loading hides funcs */
+		NULL
+	};
+	char **p;
+	int i;
+
 	if ((dtp->dt_procs = dt_zalloc(dtp, sizeof (dt_proc_hash_t) +
-	    sizeof (dt_proc_t *) * _dtrace_pidbuckets - 1)) != NULL) {
+	    sizeof (dt_proc_t *) * _dtrace_pidbuckets - 1)) == NULL)
+		return;
+
+	(void) pthread_mutex_init(&dtp->dt_procs->dph_lock, NULL);
+	(void) pthread_cond_init(&dtp->dt_procs->dph_cv, NULL);
 
-		(void) pthread_mutex_init(&dtp->dt_procs->dph_lock, NULL);
-		(void) pthread_cond_init(&dtp->dt_procs->dph_cv, NULL);
+	dtp->dt_procs->dph_hashlen = _dtrace_pidbuckets;
+	dtp->dt_procs->dph_lrulim = _dtrace_pidlrulim;
 
-		dtp->dt_procs->dph_hashlen = _dtrace_pidbuckets;
-		dtp->dt_procs->dph_lrulim = _dtrace_pidlrulim;
+	/*
+	 * Count how big our environment needs to be.
+	 */
+	for (i = 1, p = environ; *p != NULL; i++, p++)
+		continue;
+	for (p = envdef; *p != NULL; i++, p++)
+		continue;
+
+	if ((dtp->dt_proc_env = dt_zalloc(dtp, sizeof (char *) * i)) == NULL)
+		return;
+
+	for (i = 0, p = environ; *p != NULL; i++, p++) {
+		if ((dtp->dt_proc_env[i] = strdup(*p)) == NULL)
+			goto err;
+	}
+	for (p = envdef; *p != NULL; i++, p++) {
+		if ((dtp->dt_proc_env[i] = strdup(*p)) == NULL)
+			goto err;
 	}
+
+	return;
+
+err:
+	while (--i != 0) {
+		dt_free(dtp, dtp->dt_proc_env[i]);
+	}
+	dt_free(dtp, dtp->dt_proc_env);
+	dtp->dt_proc_env = NULL;
 }
 
 void
-dt_proc_hash_destroy(dtrace_hdl_t *dtp)
+dt_proc_fini(dtrace_hdl_t *dtp)
 {
 	dt_proc_hash_t *dph = dtp->dt_procs;
 	dt_proc_t *dpr;
+	char **p;
 
 	while ((dpr = dt_list_next(&dph->dph_lrulist)) != NULL)
 		dt_proc_destroy(dtp, dpr->dpr_proc);
 
 	dtp->dt_procs = NULL;
 	dt_free(dtp, dph);
+
+	for (p = dtp->dt_proc_env; *p != NULL; p++)
+		dt_free(dtp, *p);
+
+	dt_free(dtp, dtp->dt_proc_env);
+	dtp->dt_proc_env = NULL;
 }
 
 struct ps_prochandle *
diff --git a/cddl/contrib/opensolaris/lib/libdtrace/common/dt_proc.h b/cddl/contrib/opensolaris/lib/libdtrace/common/dt_proc.h
index d1fc7652385d..beae6f6d5cda 100644
--- a/cddl/contrib/opensolaris/lib/libdtrace/common/dt_proc.h
+++ b/cddl/contrib/opensolaris/lib/libdtrace/common/dt_proc.h
@@ -24,11 +24,13 @@
  * Use is subject to license terms.
  */
 
+/*
+ * Copyright (c) 2012 by Delphix. All rights reserved.
+ */
+
 #ifndef	_DT_PROC_H
 #define	_DT_PROC_H
 
-#pragma ident	"%Z%%M%	%I%	%E% SMI"
-
 #include <libproc.h>
 #include <dtrace.h>
 #include <pthread.h>
@@ -106,8 +108,8 @@ extern void dt_proc_lock(dtrace_hdl_t *, struct ps_prochandle *);
 extern void dt_proc_unlock(dtrace_hdl_t *, struct ps_prochandle *);
 extern dt_proc_t *dt_proc_lookup(dtrace_hdl_t *, struct ps_prochandle *, int);
 
-extern void dt_proc_hash_create(dtrace_hdl_t *);
-extern void dt_proc_hash_destroy(dtrace_hdl_t *);
+extern void dt_proc_init(dtrace_hdl_t *);
+extern void dt_proc_fini(dtrace_hdl_t *);
 
 #ifdef	__cplusplus
 }