attach: implement remaining options of lxc_attach_set_environment

src/lxc/attach.c

@@ -254,23 +254,72 @@ int lxc_attach_drop_privs(struct lxc_proc_context_info *ctx)
 
 int lxc_attach_set_environment(enum lxc_attach_env_policy_t policy, char** extra_env, char** extra_keep)
 {
-	/* TODO: implement extra_env, extra_keep
-	 * Rationale:
-	 *  - extra_env is an array of strings of the form
-	 *    "VAR=VALUE", which are to be set (after clearing or not,
-	 *    depending on the value of the policy variable)
-	 *  - extra_keep is an array of strings of the form
-	 *    "VAR", which are extra environment variables to be kept
-	 *    around after clearing (if that is done, otherwise, the
-	 *    remain anyway)
-	 */
-	(void) extra_env;
-	(void) extra_keep;
-
 	if (policy == LXC_ATTACH_CLEAR_ENV) {
+		char **extra_keep_store = NULL;
+		char *path_env;
+		size_t n;
+		int path_kept = 0;
+
+		if (extra_keep) {
+			size_t count, i;
+
+			for (count = 0; extra_keep[count]; count++);
+
+			extra_keep_store = calloc(count, sizeof(char *));
+			if (!extra_keep_store) {
+				SYSERROR("failed to allocate memory for storing current "
+				         "environment variable values that will be kept");
+				return -1;
+			}
+			for (i = 0; i < count; i++) {
+				char *v = getenv(extra_keep[i]);
+				if (v) {
+					extra_keep_store[i] = strdup(v);
+					if (!extra_keep_store[i]) {
+						SYSERROR("failed to allocate memory for storing current "
+						         "environment variable values that will be kept");
+						while (i > 0)
+							free(extra_keep_store[--i]);
+						free(extra_keep_store);
+						return -1;
+					}
+					if (strcmp(extra_keep[i], "PATH") == 0)
+						path_kept = 1;
+				}
+				/* calloc sets entire array to zero, so we don't
+				 * need an else */
+			}
+		}
+
 		if (clearenv()) {
 			SYSERROR("failed to clear environment");
-			/* don't error out though */
+			return -1;
+		}
+
+		if (extra_keep_store) {
+			size_t i;
+			for (i = 0; extra_keep[i]; i++) {
+				if (extra_keep_store[i])
+					setenv(extra_keep[i], extra_keep_store[i], 1);
+				free(extra_keep_store[i]);
+			}
+			free(extra_keep_store);
+		}
+
+		/* always set a default path; shells and execlp tend
+		 * to be fine without it, but there is a disturbing
+		 * number of C programs out there that just assume
+		 * that getenv("PATH") is never NULL and then die a
+		 * painful segfault death. */
+		if (!path_kept) {
+			n = confstr(_CS_PATH, NULL, 0);
+			path_env = malloc(n);
+			if (path_env) {
+				confstr(_CS_PATH, path_env, n);
+				setenv("PATH", path_env, 1);
+				free(path_env);
+			}
+			/* don't error out, this is just an extra service */
 		}
 	}
 
@@ -279,6 +328,24 @@ int lxc_attach_set_environment(enum lxc_attach_env_policy_t policy, char** extra
 		return -1;
 	}
 
+	/* set extra environment variables */
+	if (extra_env) {
+		for (; *extra_env; extra_env++) {
+			/* duplicate the string, just to be on
+			 * the safe side, because putenv does not
+			 * do it for us */
+			char *p = strdup(*extra_env);
+			/* we just assume the user knows what they
+			 * are doing, so we don't do any checks */
+			if (!p) {
+				SYSERROR("failed to allocate memory for additional environment "
+				         "variables");
+				return -1;
+			}
+			putenv(p);
+		}
+	}
+
 	return 0;
 }
 

src/lxc/lxc_attach.c

@@ -24,6 +24,7 @@
 #define _GNU_SOURCE
 #include <sys/wait.h>
 #include <sys/types.h>
+#include <stdlib.h>
 
 #include "attach.h"
 #include "arguments.h"
@@ -44,6 +45,8 @@ static const struct option my_longopts[] = {
 	/* TODO: decide upon short option names */
 	{"clear-env", no_argument, 0, 500},
 	{"keep-env", no_argument, 0, 501},
+	{"keep-var", required_argument, 0, 502},
+	{"set-var", required_argument, 0, 'v'},
 	LXC_COMMON_OPTIONS
 };
 
@@ -52,6 +55,32 @@ static signed long new_personality = -1;
 static int namespace_flags = -1;
 static int remount_sys_proc = 0;
 static lxc_attach_env_policy_t env_policy = LXC_ATTACH_KEEP_ENV;
+static char **extra_env = NULL;
+static ssize_t extra_env_size = 0;
+static char **extra_keep = NULL;
+static ssize_t extra_keep_size = 0;
+
+static int add_to_simple_array(char ***array, ssize_t *capacity, char *value)
+{
+	ssize_t count = 0;
+
+	if (*array)
+		for (; (*array)[count]; count++);
+
+	/* we have to reallocate */
+	if (count >= *capacity - 1) {
+		ssize_t new_capacity = ((count + 1) / 32 + 1) * 32;
+		char **new_array = realloc((void*)*array, sizeof(char *) * new_capacity);
+		if (!new_array)
+			return -1;
+		memset(&new_array[count], 0, sizeof(char*)*(new_capacity - count));
+		*array = new_array;
+		*capacity = new_capacity;
+	}
+
+	(*array)[count] = value;
+	return 0;
+}
 
 static int my_parser(struct lxc_arguments* args, int c, char* arg)
 {
@@ -81,6 +110,20 @@ static int my_parser(struct lxc_arguments* args, int c, char* arg)
         case 501: /* keep-env */
                 env_policy = LXC_ATTACH_KEEP_ENV;
                 break;
+	case 502: /* keep-var */
+		ret = add_to_simple_array(&extra_keep, &extra_keep_size, arg);
+		if (ret < 0) {
+			lxc_error(args, "memory allocation error");
+			return -1;
+		}
+		break;
+	case 'v':
+		ret = add_to_simple_array(&extra_env, &extra_env_size, arg);
+		if (ret < 0) {
+			lxc_error(args, "memory allocation error");
+			return -1;
+		}
+		break;
 	}
 
 	return 0;
@@ -113,14 +156,18 @@ Options :\n\
                     mount namespace when using -s in order to properly\n\
                     reflect the correct namespace context. See the\n\
                     lxc-attach(1) manual page for details.\n\
-      --clear-env\n\
-                    Clear all environment variables before attaching.\n\
+      --clear-env   Clear all environment variables before attaching.\n\
                     The attached shell/program will start with only\n\
                     container=lxc set.\n\
-      --keep-env\n\
-                    Keep all current enivornment variables. This\n\
+      --keep-env    Keep all current enivornment variables. This\n\
                     is the current default behaviour, but is likely to\n\
-                    change in the future.\n",
+                    change in the future.\n\
+  -v, --set-var     Set an additional variable that is seen by the\n\
+                    attached program in the container. May be specified\n\
+                    multiple times.\n\
+      --keep-var    Keep an additional environment variable. Only\n\
+                    applicable if --clear-env is specified. May be used\n\
+                    multiple times.\n",
 	.options  = my_longopts,
 	.parser   = my_parser,
 	.checker  = NULL,
@@ -153,6 +200,8 @@ int main(int argc, char *argv[])
 	attach_options.namespaces = namespace_flags;
 	attach_options.personality = new_personality;
 	attach_options.env_policy = env_policy;
+	attach_options.extra_env_vars = extra_env;
+	attach_options.extra_keep_env = extra_keep;
 
 	if (my_args.argc) {
 		command.program = my_args.argv[0];