/*****************************************************************************\
 *  src/common/env.c - add an environment variable to environment vector
 *****************************************************************************
 *  Copyright (C) 2002-2007 The Regents of the University of California.
 *  Copyright (C) 2008-2009 Lawrence Livermore National Security.
 *  Produced at Lawrence Livermore National Laboratory (cf, DISCLAIMER).
 *  Written by Mark Grondona <mgrondona@llnl.gov>, Danny Auble <da@llnl.gov>.
 *  CODE-OCEC-09-009. All rights reserved.
 *
 *  This file is part of Slurm, a resource management program.
 *  For details, see <https://slurm.schedmd.com/>.
 *  Please also read the included file: DISCLAIMER.
 *
 *  Slurm is free software; you can redistribute it and/or modify it under
 *  the terms of the GNU General Public License as published by the Free
 *  Software Foundation; either version 2 of the License, or (at your option)
 *  any later version.
 *
 *  In addition, as a special exception, the copyright holders give permission
 *  to link the code of portions of this program with the OpenSSL library under
 *  certain conditions as described in each individual source file, and
 *  distribute linked combinations including the two. You must obey the GNU
 *  General Public License in all respects for all of the code used other than
 *  OpenSSL. If you modify file(s) with this exception, you may extend this
 *  exception to your version of the file(s), but you are not obligated to do
 *  so. If you do not wish to do so, delete this exception statement from your
 *  version.  If you delete this exception statement from all source files in
 *  the program, then also delete it here.
 *
 *  Slurm is distributed in the hope that it will be useful, but WITHOUT ANY
 *  WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 *  FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
 *  details.
 *
 *  You should have received a copy of the GNU General Public License along
 *  with Slurm; if not, write to the Free Software Foundation, Inc.,
 *  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301  USA.
\*****************************************************************************/

#include "config.h"

#include <fcntl.h>
#include <poll.h>
#include <signal.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/param.h>		/* MAXPATHLEN */
#include <unistd.h>

#include "slurm/slurm.h"
#include "src/common/cpu_frequency.h"
#include "src/common/log.h"
#include "src/common/env.h"
#include "src/common/fd.h"
#include "src/common/node_select.h"
#include "src/common/macros.h"
#include "src/common/proc_args.h"
#include "src/common/read_config.h"
#include "src/common/slurm_opt.h"
#include "src/common/slurm_protocol_api.h"
#include "src/common/slurm_protocol_defs.h"
#include "src/common/slurm_step_layout.h"
#include "src/common/slurmdb_defs.h"
#include "src/common/strlcpy.h"
#include "src/common/xassert.h"
#include "src/common/xmalloc.h"
#include "src/common/xstring.h"

/*
 * Define slurm-specific aliases for use by plugins, see slurm_xlator.h
 * for details.
 */
strong_alias(setenvf,			slurm_setenvpf);
strong_alias(unsetenvp,			slurm_unsetenvp);
strong_alias(getenvp,			slurm_getenvp);
strong_alias(env_array_create,		slurm_env_array_create);
strong_alias(env_array_merge,		slurm_env_array_merge);
strong_alias(env_array_copy,		slurm_env_array_copy);
strong_alias(env_array_free,		slurm_env_array_free);
strong_alias(env_array_append,		slurm_env_array_append);
strong_alias(env_array_append_fmt,	slurm_env_array_append_fmt);
strong_alias(env_array_overwrite,	slurm_env_array_overwrite);
strong_alias(env_array_overwrite_fmt,	slurm_env_array_overwrite_fmt);
strong_alias(env_array_overwrite_het_fmt, slurm_env_array_overwrite_het_fmt);
strong_alias(env_unset_environment,	slurm_env_unset_environment);

#define ENV_BUFSIZE (256 * 1024)
#define MAX_ENV_STRLEN (32 * 4096)	/* Needed for CPU_BIND and MEM_BIND on
					 * SGI systems with huge CPU counts */

/*
 *  Return pointer to `name' entry in environment if found, or
 *   pointer to the last entry (i.e. NULL) if `name' is not
 *   currently set in `env'
 *
 */
static char **
_find_name_in_env(char **env, const char *name)
{
	char **ep;

	ep = env;
	while (*ep != NULL) {
		size_t cnt = 0;

		while ( ((*ep)[cnt] == name[cnt])
		      && ( name[cnt] != '\0')
		      && ((*ep)[cnt] != '\0')    )
			++cnt;

		if (name[cnt] == '\0' && (*ep)[cnt] == '=') {
			break;
		} else
			++ep;
	}

	return (ep);
}

/*
 *  Extend memory allocation for env by 1 entry. Make last entry == NULL.
 *   return pointer to last env entry;
 */
static char **
_extend_env(char ***envp)
{
	char **ep;
	size_t newcnt = (xsize (*envp) / sizeof (char *)) + 1;

	*envp = xrealloc (*envp, newcnt * sizeof (char *));

	(*envp)[newcnt - 1] = NULL;
	ep = &((*envp)[newcnt - 2]);

	/*
	 *  Find last non-NULL entry
	 */
	while (*ep == NULL)
		--ep;

	return (++ep);
}

/* return true if the environment variables should not be set for
 *	srun's --get-user-env option */
static bool _discard_env(char *name, char *value)
{
	if ((xstrcmp(name, "DISPLAY")     == 0) ||
	    (xstrcmp(name, "ENVIRONMENT") == 0) ||
	    (xstrcmp(name, "HOSTNAME")    == 0))
		return true;

	return false;
}

/*
 * Return the number of elements in the environment `env'
 */
int
envcount (char **env)
{
	int envc = 0;
	while (env && env[envc])
		envc++;
	return (envc);
}

/*
 * _setenvfs() (stolen from pdsh)
 *
 * Set a variable in the callers environment.  Args are printf style.
 * XXX Space is allocated on the heap and will never be reclaimed.
 * Example: setenvfs("RMS_RANK=%d", rank);
 */
int
setenvfs(const char *fmt, ...)
{
	va_list ap;
	char *buf, *bufcpy, *loc;
	int rc, size;

	buf = xmalloc(ENV_BUFSIZE);
	va_start(ap, fmt);
	vsnprintf(buf, ENV_BUFSIZE, fmt, ap);
	va_end(ap);

	size = strlen(buf);
	bufcpy = xstrdup(buf);
	xfree(buf);

	if (size >= MAX_ENV_STRLEN) {
		if ((loc = strchr(bufcpy, '=')))
			loc[0] = '\0';
		error("environment variable %s is too long", bufcpy);
		xfree(bufcpy);
		rc = ENOMEM;
	} else {
		rc = putenv(bufcpy);
	}

	return rc;
}

int setenvf(char ***envp, const char *name, const char *fmt, ...)
{
	char *value;
	va_list ap;
	int size, rc;

	if (!name || name[0] == '\0')
		return EINVAL;

	value = xmalloc(ENV_BUFSIZE);
	va_start(ap, fmt);
	vsnprintf(value, ENV_BUFSIZE, fmt, ap);
	va_end(ap);

	size = strlen(name) + strlen(value) + 2;
	if (size >= MAX_ENV_STRLEN) {
		error("environment variable %s is too long", name);
		return ENOMEM;
	}

	if (envp && *envp) {
		if (env_array_overwrite(envp, name, value) == 1)
			rc = 0;
		else
			rc = 1;
	} else {
		rc = setenv(name, value, 1);
	}

	xfree(value);
	return rc;
}

/*
 *  Remove environment variable `name' from "environment"
 *   contained in `env'
 *
 *  [ This was taken almost verbatim from glibc's
 *    unsetenv()  code. ]
 */
void unsetenvp(char **env, const char *name)
{
	char **ep;

	if (env == NULL)
		return;

	ep = env;
	while ((ep = _find_name_in_env (ep, name)) && (*ep != NULL)) {
		char **dp = ep;
		xfree (*ep);
		do
			dp[0] = dp[1];
		while (*dp++);

		/*  Continue loop in case `name' appears again. */
		++ep;
	}
	return;
}

char *getenvp(char **env, const char *name)
{
	size_t len;
	char **ep;

	if (!name || !env || !env[0])
		return (NULL);

	len = strlen(name);
	ep = _find_name_in_env (env, name);

	if (*ep != NULL)
		return (&(*ep)[len+1]);

	return NULL;
}

int setup_env(env_t *env, bool preserve_env)
{
	int rc = SLURM_SUCCESS;
	char *addr, *dist = NULL, *lllp_dist = NULL;
	char addrbuf[INET6_ADDRSTRLEN];

	if (env == NULL)
		return SLURM_ERROR;

	/*
	 * Always force SLURM_CONF into the environment. This ensures the
	 * "configless" operation is working, and prevents the client commands
	 * from falling back to separate RPC requests in case the cache dir
	 * is unresponsive.
	 */
	if (setenvf(&env->env, "SLURM_CONF", "%s", getenv("SLURM_CONF"))) {
		error("Unable to set SLURM_CONF environment variable");
		rc = SLURM_ERROR;
	}
	/*
	 * Similarly, prevent this option from leaking in. SLURM_CONF would
	 * always take precedence, but tidy it up anyways.
	 */
	unsetenvp(env->env, "SLURM_CONF_SERVER");

	if (!preserve_env && env->ntasks) {
		if (setenvf(&env->env, "SLURM_NTASKS", "%d", env->ntasks)) {
			error("Unable to set SLURM_NTASKS environment variable");
			rc = SLURM_ERROR;
		}
		if (setenvf(&env->env, "SLURM_NPROCS", "%d", env->ntasks)) {
			error("Unable to set SLURM_NPROCS environment variable");
			rc = SLURM_ERROR;
		}
	}

	if (env->cpus_per_task &&
	    setenvf(&env->env, "SLURM_CPUS_PER_TASK", "%d",
		    env->cpus_per_task) ) {
		error("Unable to set SLURM_CPUS_PER_TASK");
		rc = SLURM_ERROR;
	}

	if (env->ntasks_per_gpu &&
	    setenvf(&env->env, "SLURM_NTASKS_PER_GPU", "%d",
		    env->ntasks_per_gpu)) {
		error("Unable to set SLURM_NTASKS_PER_GPU");
		rc = SLURM_ERROR;
	}

	if (env->ntasks_per_node
	   && setenvf(&env->env, "SLURM_NTASKS_PER_NODE", "%d",
		      env->ntasks_per_node) ) {
		error("Unable to set SLURM_NTASKS_PER_NODE");
		rc = SLURM_ERROR;
	}

	if (env->ntasks_per_socket
	   && setenvf(&env->env, "SLURM_NTASKS_PER_SOCKET", "%d",
		      env->ntasks_per_socket) ) {
		error("Unable to set SLURM_NTASKS_PER_SOCKET");
		rc = SLURM_ERROR;
	}

	if (env->ntasks_per_core
	   && setenvf(&env->env, "SLURM_NTASKS_PER_CORE", "%d",
		      env->ntasks_per_core) ) {
		error("Unable to set SLURM_NTASKS_PER_CORE");
		rc = SLURM_ERROR;
	}

	if (env->ntasks_per_tres
	    && setenvf(&env->env, "SLURM_NTASKS_PER_TRES", "%d",
		      env->ntasks_per_tres) ) {
		error("Unable to set SLURM_NTASKS_PER_TRES");
		rc = SLURM_ERROR;
	}

	if (env->cpus_on_node
	   && setenvf(&env->env, "SLURM_CPUS_ON_NODE", "%d",
		      env->cpus_on_node) ) {
		error("Unable to set SLURM_CPUS_ON_NODE");
		rc = SLURM_ERROR;
	}

	set_distribution(env->distribution, &dist, &lllp_dist);
	if (dist)
		if (setenvf(&env->env, "SLURM_DISTRIBUTION", "%s", dist)) {
			error("Can't set SLURM_DISTRIBUTION env variable");
			rc = SLURM_ERROR;
		}

	if ((env->distribution & SLURM_DIST_STATE_BASE) == SLURM_DIST_PLANE)
		if (setenvf(&env->env, "SLURM_DIST_PLANESIZE", "%u",
			    env->plane_size)) {
			error("Can't set SLURM_DIST_PLANESIZE env variable");
			rc = SLURM_ERROR;
		}

	if (lllp_dist)
		if (setenvf(&env->env, "SLURM_DIST_LLLP", "%s", lllp_dist)) {
			error("Can't set SLURM_DIST_LLLP env variable");
			rc = SLURM_ERROR;
		}


	if (env->cpu_bind_type && !env->batch_flag &&
	    (env->stepid != SLURM_INTERACTIVE_STEP)) {
		char *str_verbose, *str_bind1 = NULL, *str_bind2 = NULL;
		char *str_bind_list, *str_bind_type = NULL, *str_bind = NULL;

		unsetenvp(env->env, "SLURM_CPU_BIND");
		unsetenvp(env->env, "SLURM_CPU_BIND_LIST");
		unsetenvp(env->env, "SLURM_CPU_BIND_TYPE");
		unsetenvp(env->env, "SLURM_CPU_BIND_VERBOSE");

		if (env->cpu_bind_type & CPU_BIND_VERBOSE)
			str_verbose = "verbose";
		else
			str_verbose = "quiet";

		if (env->cpu_bind_type & CPU_BIND_TO_THREADS) {
			str_bind1 = "threads";
		} else if (env->cpu_bind_type & CPU_BIND_TO_CORES) {
			str_bind1 = "cores";
		} else if (env->cpu_bind_type & CPU_BIND_TO_SOCKETS) {
			str_bind1 = "sockets";
		} else if (env->cpu_bind_type & CPU_BIND_TO_LDOMS) {
			str_bind1 = "ldoms";
		} else if (env->cpu_bind_type & CPU_BIND_TO_BOARDS) {
			str_bind1 = "boards";
		}

		if (env->cpu_bind_type & CPU_BIND_NONE) {
			str_bind2 = "none";
		} else if (env->cpu_bind_type & CPU_BIND_RANK) {
			str_bind2 = "rank";
		} else if (env->cpu_bind_type & CPU_BIND_MAP) {
			str_bind2 = "map_cpu:";
		} else if (env->cpu_bind_type & CPU_BIND_MASK) {
			str_bind2 = "mask_cpu:";
		} else if (env->cpu_bind_type & CPU_BIND_LDRANK) {
			str_bind2 = "rank_ldom";
		} else if (env->cpu_bind_type & CPU_BIND_LDMAP) {
			str_bind2 = "map_ldom:";
		} else if (env->cpu_bind_type & CPU_BIND_LDMASK) {
			str_bind2 = "mask_ldom:";
		}

		if (env->cpu_bind)
			str_bind_list = env->cpu_bind;
		else
			str_bind_list = "";

		/* combine first and second part with a comma if needed */
		if (str_bind1)
			xstrcat(str_bind_type, str_bind1);
		if (str_bind1 && str_bind2)
			xstrcatchar(str_bind_type, ',');
		if (str_bind2)
			xstrcat(str_bind_type, str_bind2);

		xstrcat(str_bind, str_verbose);
		if (str_bind_type) {
			xstrcatchar(str_bind, ',');
			xstrcat(str_bind, str_bind_type);
			xstrcat(str_bind, str_bind_list);
		} else
			str_bind_type = xstrdup("");

		if (setenvf(&env->env, "SLURM_CPU_BIND", "%s", str_bind)) {
			error("Unable to set SLURM_CPU_BIND");
			rc = SLURM_ERROR;
		}
		if (setenvf(&env->env, "SLURM_CPU_BIND_LIST", "%s",
			    str_bind_list)) {
			error("Unable to set SLURM_CPU_BIND_LIST");
			rc = SLURM_ERROR;
		}
		if (setenvf(&env->env, "SLURM_CPU_BIND_TYPE", "%s",
			    str_bind_type)) {
			error("Unable to set SLURM_CPU_BIND_TYPE");
			rc = SLURM_ERROR;
		}
		if (setenvf(&env->env, "SLURM_CPU_BIND_VERBOSE", "%s",
			    str_verbose)) {
			error("Unable to set SLURM_CPU_BIND_VERBOSE");
			rc = SLURM_ERROR;
		}

		xfree(str_bind);
		xfree(str_bind_type);
	}

	if (env->mem_bind_type && (env->stepid != SLURM_INTERACTIVE_STEP)) {
		char *str_verbose, *str_bind_type = NULL, *str_bind_list;
		char *str_prefer = NULL, *str_bind = NULL;
		char *str_bind_sort = NULL;

		if (env->batch_flag) {
			unsetenvp(env->env, "SBATCH_MEM_BIND");
			unsetenvp(env->env, "SBATCH_MEM_BIND_LIST");
			unsetenvp(env->env, "SBATCH_MEM_BIND_PREFER");
			unsetenvp(env->env, "SBATCH_MEM_BIND_TYPE");
			unsetenvp(env->env, "SBATCH_MEM_BIND_VERBOSE");
		} else {
			unsetenvp(env->env, "SLURM_MEM_BIND");
			unsetenvp(env->env, "SLURM_MEM_BIND_LIST");
			unsetenvp(env->env, "SLURM_MEM_BIND_PREFER");
			unsetenvp(env->env, "SLURM_MEM_BIND_SORT");
			unsetenvp(env->env, "SLURM_MEM_BIND_TYPE");
			unsetenvp(env->env, "SLURM_MEM_BIND_VERBOSE");
		}

		if (env->mem_bind_type & MEM_BIND_VERBOSE)
			str_verbose = "verbose";
		else
			str_verbose = "quiet";
		if (env->mem_bind_type & MEM_BIND_PREFER)
			str_prefer = "prefer";
		if (env->mem_bind_type & MEM_BIND_NONE) {
			str_bind_type = "none";
		} else if (env->mem_bind_type & MEM_BIND_RANK) {
			str_bind_type = "rank";
		} else if (env->mem_bind_type & MEM_BIND_MAP) {
			str_bind_type = "map_mem:";
		} else if (env->mem_bind_type & MEM_BIND_MASK) {
			str_bind_type = "mask_mem:";
		} else if (env->mem_bind_type & MEM_BIND_LOCAL) {
			str_bind_type = "local";
		}

		if (env->mem_bind_type & MEM_BIND_SORT)
			str_bind_sort = "sort";

		if (env->mem_bind)
			str_bind_list = env->mem_bind;
		else
			str_bind_list = "";

		xstrcat(str_bind, str_verbose);
		if (str_prefer) {
			xstrcatchar(str_bind, ',');
			xstrcat(str_bind, str_prefer);
		}
		if (str_bind_type) {
			xstrcatchar(str_bind, ',');
			xstrcat(str_bind, str_bind_type);
			xstrcat(str_bind, str_bind_list);
		} else
			str_bind_type = "";

		if (env->batch_flag) {
			if (setenvf(&env->env, "SBATCH_MEM_BIND", "%s", str_bind)) {
				error("Unable to set SBATCH_MEM_BIND");
				rc = SLURM_ERROR;
			}
			if (setenvf(&env->env, "SBATCH_MEM_BIND_LIST", "%s",
				    str_bind_list)) {
				error("Unable to set SBATCH_MEM_BIND_LIST");
				rc = SLURM_ERROR;
			}
			if (str_prefer &&
			    setenvf(&env->env, "SBATCH_MEM_BIND_PREFER", "%s",
				    str_prefer)) {
				error("Unable to set SBATCH_MEM_BIND_PREFER");
				rc = SLURM_ERROR;
			}
			if (str_bind_sort &&
			    setenvf(&env->env, "SBATCH_MEM_BIND_SORT", "%s",
				    str_bind_sort)) {
				error("Unable to set SBATCH_MEM_BIND_SORT");
				rc = SLURM_ERROR;
			}
			if (setenvf(&env->env, "SBATCH_MEM_BIND_TYPE", "%s",
				    str_bind_type)) {
				error("Unable to set SBATCH_MEM_BIND_TYPE");
				rc = SLURM_ERROR;
			}
			if (setenvf(&env->env, "SBATCH_MEM_BIND_VERBOSE", "%s",
				    str_verbose)) {
				error("Unable to set SBATCH_MEM_BIND_VERBOSE");
				rc = SLURM_ERROR;
			}
		} else {
			if (setenvf(&env->env, "SLURM_MEM_BIND", "%s", str_bind)) {
				error("Unable to set SLURM_MEM_BIND");
				rc = SLURM_ERROR;
			}
			if (setenvf(&env->env, "SLURM_MEM_BIND_LIST", "%s",
				    str_bind_list)) {
				error("Unable to set SLURM_MEM_BIND_LIST");
				rc = SLURM_ERROR;
			}
			if (str_prefer &&
			    setenvf(&env->env, "SLURM_MEM_BIND_PREFER", "%s",
				    str_prefer)) {
				error("Unable to set SLURM_MEM_BIND_PREFER");
				rc = SLURM_ERROR;
			}
			if (str_bind_sort &&
			    setenvf(&env->env, "SLURM_MEM_BIND_SORT", "%s",
				    str_bind_sort)) {
				error("Unable to set SLURM_MEM_BIND_SORT");
				rc = SLURM_ERROR;
			}
			if (setenvf(&env->env, "SLURM_MEM_BIND_TYPE", "%s",
				    str_bind_type)) {
				error("Unable to set SLURM_MEM_BIND_TYPE");
				rc = SLURM_ERROR;
			}
			if (setenvf(&env->env, "SLURM_MEM_BIND_VERBOSE", "%s",
				    str_verbose)) {
				error("Unable to set SLURM_MEM_BIND_VERBOSE");
				rc = SLURM_ERROR;
			}
		}

		xfree(str_bind);
	}

	if (cpu_freq_set_env("SLURM_CPU_FREQ_REQ", env->cpu_freq_min,
			env->cpu_freq_max, env->cpu_freq_gov) != SLURM_SUCCESS)
		rc = SLURM_ERROR;

	if (env->overcommit
	    && (setenvf(&env->env, "SLURM_OVERCOMMIT", "%s", "1"))) {
		error("Unable to set SLURM_OVERCOMMIT environment variable");
		rc = SLURM_ERROR;
	}

	if (env->slurmd_debug
	    && setenvf(&env->env, "SLURMD_DEBUG", "%d", env->slurmd_debug)) {
		error("Can't set SLURMD_DEBUG environment variable");
		rc = SLURM_ERROR;
	}

	if (env->labelio
	   && setenvf(&env->env, "SLURM_LABELIO", "1")) {
		error("Unable to set SLURM_LABELIO environment variable");
		rc = SLURM_ERROR;
	}

	if (env->jobid >= 0) {
		if (setenvf(&env->env, "SLURM_JOB_ID", "%d", env->jobid)) {
			error("Unable to set SLURM_JOB_ID environment");
			rc = SLURM_ERROR;
		}
		/* and for backwards compatibility... */
		if (setenvf(&env->env, "SLURM_JOBID", "%d", env->jobid)) {
			error("Unable to set SLURM_JOBID environment");
			rc = SLURM_ERROR;
		}
	}

	if (env->job_name) {
		if (setenvf(&env->env, "SLURM_JOB_NAME", "%s", env->job_name)) {
			error("Unable to set SLURM_JOB_NAME environment");
			rc = SLURM_ERROR;
		}
	}

	/*
	 * These aren't relevant to a system not using Slurm as the
	 * launcher. Since there isn't a flag for that we check for
	 * the flags we do have.
	 */
	if (env->task_pid &&
	    setenvf(&env->env, "SLURM_TASK_PID", "%d",
		       (int)env->task_pid)) {
		error("Unable to set SLURM_TASK_PID environment "
		      "variable");
		rc = SLURM_ERROR;
	}
	if ((env->nodeid >= 0) &&
	    setenvf(&env->env, "SLURM_NODEID", "%d", env->nodeid)) {
		error("Unable to set SLURM_NODEID environment");
		rc = SLURM_ERROR;
	}

	if ((env->procid >= 0) &&
	    setenvf(&env->env, "SLURM_PROCID", "%d", env->procid)) {
		error("Unable to set SLURM_PROCID environment");
		rc = SLURM_ERROR;
	}

	if ((env->localid >= 0) &&
	    setenvf(&env->env, "SLURM_LOCALID", "%d", env->localid)) {
		error("Unable to set SLURM_LOCALID environment");
		rc = SLURM_ERROR;
	}

	if (env->stepid >= 0) {
		if (setenvf(&env->env, "SLURM_STEP_ID", "%d", env->stepid)) {
			error("Unable to set SLURM_STEP_ID environment");
			rc = SLURM_ERROR;
		}
		/* and for backwards compatibility... */
		if (setenvf(&env->env, "SLURM_STEPID", "%d", env->stepid)) {
			error("Unable to set SLURM_STEPID environment");
			rc = SLURM_ERROR;
		}
	}

	if (!preserve_env && env->nhosts
	    && setenvf(&env->env, "SLURM_NNODES", "%d", env->nhosts)) {
		error("Unable to set SLURM_NNODES environment var");
		rc = SLURM_ERROR;
	}

	if (env->nhosts
	    && setenvf(&env->env, "SLURM_JOB_NUM_NODES", "%d", env->nhosts)) {
		error("Unable to set SLURM_JOB_NUM_NODES environment var");
		rc = SLURM_ERROR;
	}

	if (env->nodelist &&
	    setenvf(&env->env, "SLURM_NODELIST", "%s", env->nodelist)) {
		error("Unable to set SLURM_NODELIST environment var.");
		rc = SLURM_ERROR;
	}

	if (env->partition
	    && setenvf(&env->env, "SLURM_JOB_PARTITION", "%s", env->partition)) {
		error("Unable to set SLURM_JOB_PARTITION environment var.");
		rc = SLURM_ERROR;
	}

	if (!preserve_env && env->task_count
	    && setenvf (&env->env,
			"SLURM_TASKS_PER_NODE", "%s", env->task_count)) {
		error ("Can't set SLURM_TASKS_PER_NODE env variable");
		rc = SLURM_ERROR;
	}

	if (!preserve_env && env->threads_per_core &&
	    setenvf(&env->env, "SLURM_THREADS_PER_CORE", "%d",
		    env->threads_per_core)) {
		error("Can't set SLURM_THREADS_PER_CORE env variable");
		rc = SLURM_ERROR;
	}

	if (env->comm_port
	    && setenvf (&env->env, "SLURM_SRUN_COMM_PORT", "%u",
			env->comm_port)) {
		error ("Can't set SLURM_SRUN_COMM_PORT env variable");
		rc = SLURM_ERROR;
	}

	if (env->cli) {
		slurm_get_ip_str(env->cli, addrbuf, INET6_ADDRSTRLEN);
		setenvf(&env->env, "SLURM_LAUNCH_NODE_IPADDR", "%s", addrbuf);
	}

	if (env->sgtids &&
	    setenvf(&env->env, "SLURM_GTIDS", "%s", env->sgtids)) {
		error("Unable to set SLURM_GTIDS environment variable");
		rc = SLURM_ERROR;
	}

	if (env->pty_port
	&&  setenvf(&env->env, "SLURM_PTY_PORT", "%hu", env->pty_port)) {
		error("Can't set SLURM_PTY_PORT env variable");
		rc = SLURM_ERROR;
	}
	if (env->ws_col
	&&  setenvf(&env->env, "SLURM_PTY_WIN_COL", "%hu", env->ws_col)) {
		error("Can't set SLURM_PTY_WIN_COL env variable");
		rc = SLURM_ERROR;
	}
	if (env->ws_row
	&&  setenvf(&env->env, "SLURM_PTY_WIN_ROW", "%hu", env->ws_row)) {
		error("Can't set SLURM_PTY_WIN_ROW env variable");
		rc = SLURM_ERROR;
	}

	if (env->restart_cnt &&
	    setenvf(&env->env, "SLURM_RESTART_COUNT", "%u", env->restart_cnt)) {
		error("Can't set SLURM_RESTART_COUNT env variable");
		rc = SLURM_ERROR;
	}

	if (env->user_name) {
		if (setenvf(&env->env, "SLURM_JOB_UID", "%u",
			    (unsigned int) env->uid)) {
			error("Can't set SLURM_JOB_UID env variable");
			rc = SLURM_ERROR;
		}
		if (setenvf(&env->env, "SLURM_JOB_USER", "%s", env->user_name)){
			error("Can't set SLURM_JOB_USER env variable");
			rc = SLURM_ERROR;
		}
	}

	if (env->account) {
		if (setenvf(&env->env,
			    "SLURM_JOB_ACCOUNT",
			    "%s",
			    env->account)) {
			error("%s: can't set SLURM_JOB_ACCOUNT env variable",
			      __func__);
			rc = SLURM_ERROR;
		}
	}
	if (env->qos) {
		if (setenvf(&env->env,
			    "SLURM_JOB_QOS",
			    "%s",
			    env->qos)) {
			error("%s: can't set SLURM_JOB_QOS env variable",
				__func__);
			rc = SLURM_ERROR;
		}
	}
	if (env->resv_name) {
		if (setenvf(&env->env,
			    "SLURM_JOB_RESERVATION",
			    "%s",
			    env->resv_name)) {
			error("%s: can't set SLURM_JOB_RESERVATION env variable",
				__func__);
			rc = SLURM_ERROR;
		}
	}

	if (slurm_conf.slurmctld_addr)
		addr = slurm_conf.slurmctld_addr;
	else
		addr = slurm_conf.control_addr[0];
	setenvf(&env->env, "SLURM_WORKING_CLUSTER", "%s:%s:%d:%d:%d",
	        slurm_conf.cluster_name, addr, slurm_conf.slurmctld_port,
	        SLURM_PROTOCOL_VERSION, select_get_plugin_id());

	return rc;
}

/**********************************************************************
 * From here on are the new environment variable management functions,
 * used by the "new" commands: salloc, sbatch, and the step launch APIs.
 **********************************************************************/

/*
 * Return a string representation of an array of uint16_t elements.
 * Each value in the array is printed in decimal notation and elements
 * are separated by a comma.  If sequential elements in the array
 * contain the same value, the value is written out just once followed
 * by "(xN)", where "N" is the number of times the value is repeated.
 *
 * Example:
 *   The array "1, 2, 1, 1, 1, 3, 2" becomes the string "1,2,1(x3),3,2"
 *
 * Returns an xmalloc'ed string.  Free with xfree().
 */
extern char *uint16_array_to_str(int array_len, const uint16_t *array)
{
	int i;
	int previous = 0;
	char *sep = ",";  /* seperator */
	char *str = xstrdup("");

	if (array == NULL)
		return str;

	for (i = 0; i < array_len; i++) {
		if ((i+1 < array_len) && (array[i] == array[i+1])) {
			previous++;
			continue;
		}

		if (i == array_len-1) /* last time through loop */
			sep = "";
		if (previous > 0) {
			xstrfmtcat(str, "%u(x%u)%s",
				   array[i], previous+1, sep);
		} else {
			xstrfmtcat(str, "%u%s", array[i], sep);
		}
		previous = 0;
	}

	return str;
}


/*
 * The cpus-per-node representation in Slurm (and perhaps tasks-per-node
 * in the future) is stored in a compressed format comprised of two
 * equal-length arrays, and an integer holding the array length.  In one
 * array an element represents a count (number of cpus, number of tasks,
 * etc.), and the corresponding element in the other array contains the
 * number of times the count is repeated sequentially in the uncompressed
 * something-per-node array.
 *
 * This function returns the string representation of the compressed
 * array.  Free with xfree().
 */
extern char *uint32_compressed_to_str(uint32_t array_len,
				      const uint16_t *array,
				      const uint32_t *array_reps)
{
	int i;
	char *sep = ","; /* seperator */
	char *str = xstrdup("");

	if (!array || !array_reps)
		return str;

	for (i = 0; i < array_len; i++) {
		if (i == array_len-1) /* last time through loop */
			sep = "";
		if (array_reps[i] > 1) {
			xstrfmtcat(str, "%u(x%u)%s",
				   array[i], array_reps[i], sep);
		} else {
			xstrfmtcat(str, "%u%s", array[i], sep);
		}
	}

	return str;
}

/*
 * Set in "dest" the environment variables relevant to a Slurm job
 * allocation, overwriting any environment variables of the same name.
 * If the address pointed to by "dest" is NULL, memory will automatically be
 * xmalloc'ed.  The array is terminated by a NULL pointer, and thus is
 * suitable for use by execle() and other env_array_* functions.
 *
 * Sets the variables:
 *	SLURM_JOB_ID
 *	SLURM_JOB_NAME
 *	SLURM_JOB_NUM_NODES
 *	SLURM_JOB_NODELIST
 *	SLURM_JOB_CPUS_PER_NODE
 *	SLURM_NODE_ALIASES
 *	SLURM_NTASKS_PER_NODE
 *
 * dest OUT - array in which to the set environment variables
 * alloc IN - resource allocation response
 * desc IN - job allocation request
 * het_job_offset IN - component offset into hetjob, -1 if not hetjob
 *
 * Sets OBSOLETE variables (needed for MPI, do not remove):
 *	SLURM_JOBID
 *	SLURM_NNODES
 *	SLURM_NODELIST
 *	SLURM_TASKS_PER_NODE
 */
extern int env_array_for_job(char ***dest,
			     const resource_allocation_response_msg_t *alloc,
			     const job_desc_msg_t *desc, int het_job_offset)
{
	char *tmp = NULL;
	char *dist = NULL, *lllp_dist = NULL;
	char *key, *value;
	slurm_step_layout_t *step_layout = NULL;
	int i, rc = SLURM_SUCCESS;
	slurm_step_layout_req_t step_layout_req;
	uint16_t cpus_per_task_array[1];
	uint32_t cpus_task_reps[1];

	if (!alloc || !desc)
		return SLURM_ERROR;

	memset(&step_layout_req, 0, sizeof(slurm_step_layout_req_t));
	step_layout_req.num_tasks = desc->num_tasks;
	step_layout_req.num_hosts = alloc->node_cnt;
	cpus_per_task_array[0] = desc->cpus_per_task;
	cpus_task_reps[0] = alloc->node_cnt;

	if (het_job_offset < 1) {
		env_array_overwrite_fmt(dest, "SLURM_JOB_ID", "%u",
					alloc->job_id);
	}
	env_array_overwrite_het_fmt(dest, "SLURM_JOB_ID", het_job_offset,
				    "%u", alloc->job_id);
	env_array_overwrite_het_fmt(dest, "SLURM_JOB_NAME", het_job_offset,
				    "%s", desc->name);
	env_array_overwrite_het_fmt(dest, "SLURM_JOB_NUM_NODES", het_job_offset,
				    "%u", step_layout_req.num_hosts);
	env_array_overwrite_het_fmt(dest, "SLURM_JOB_NODELIST", het_job_offset,
				    "%s", alloc->node_list);
	env_array_overwrite_het_fmt(dest, "SLURM_NODE_ALIASES", het_job_offset,
				    "%s", alloc->alias_list);
	env_array_overwrite_het_fmt(dest, "SLURM_JOB_PARTITION", het_job_offset,
				    "%s", alloc->partition);

	set_distribution(desc->task_dist, &dist, &lllp_dist);
	if (dist) {
		env_array_overwrite_het_fmt(dest, "SLURM_DISTRIBUTION",
					    het_job_offset, "%s", dist);
	}
	if ((desc->task_dist & SLURM_DIST_STATE_BASE) == SLURM_DIST_PLANE) {
		env_array_overwrite_het_fmt(dest, "SLURM_DIST_PLANESIZE",
					    het_job_offset, "%u",
					    desc->plane_size);
	}
	if (lllp_dist) {
		env_array_overwrite_het_fmt(dest, "SLURM_DIST_LLLP",
					    het_job_offset, "%s", lllp_dist);
	}
	tmp = uint32_compressed_to_str(alloc->num_cpu_groups,
					alloc->cpus_per_node,
					alloc->cpu_count_reps);
	env_array_overwrite_het_fmt(dest, "SLURM_JOB_CPUS_PER_NODE",
				    het_job_offset, "%s", tmp);
	xfree(tmp);

	if (desc->threads_per_core != NO_VAL16)
		env_array_overwrite_het_fmt(dest, "SLURM_THREADS_PER_CORE",
					    het_job_offset, "%d",
					    desc->threads_per_core);

	if (alloc->pn_min_memory & MEM_PER_CPU) {
		uint64_t tmp_mem = alloc->pn_min_memory & (~MEM_PER_CPU);
		env_array_overwrite_het_fmt(dest, "SLURM_MEM_PER_CPU",
					    het_job_offset, "%"PRIu64"",
					    tmp_mem);
	} else if (alloc->pn_min_memory) {
		uint64_t tmp_mem = alloc->pn_min_memory;
		env_array_overwrite_het_fmt(dest, "SLURM_MEM_PER_NODE",
					    het_job_offset, "%"PRIu64"",
					    tmp_mem);
	}

	/* OBSOLETE, but needed by MPI, do not remove */
	env_array_overwrite_het_fmt(dest, "SLURM_JOBID", het_job_offset, "%u",
				    alloc->job_id);
	env_array_overwrite_het_fmt(dest, "SLURM_NNODES", het_job_offset, "%u",
				    step_layout_req.num_hosts);
	env_array_overwrite_het_fmt(dest, "SLURM_NODELIST", het_job_offset, "%s",
				    alloc->node_list);

	if (step_layout_req.num_tasks == NO_VAL) {
		/* If we know how many tasks we are going to do then
		   we set SLURM_TASKS_PER_NODE */
		int i = 0;
		/* If no tasks were given we can figure it out here
		 * by totalling up the cpus and then dividing by the
		 * number of cpus per task */

		step_layout_req.num_tasks = 0;
		for (i = 0; i < alloc->num_cpu_groups; i++) {
			step_layout_req.num_tasks += alloc->cpu_count_reps[i]
				* alloc->cpus_per_node[i];
		}
		if ((int)desc->cpus_per_task > 1
		   && desc->cpus_per_task != NO_VAL16)
			step_layout_req.num_tasks /= desc->cpus_per_task;
		//num_tasks = desc->min_cpus;
	}

	if ((desc->task_dist & SLURM_DIST_STATE_BASE) == SLURM_DIST_ARBITRARY) {
		step_layout_req.node_list = desc->req_nodes;
		env_array_overwrite_het_fmt(dest, "SLURM_ARBITRARY_NODELIST",
					    het_job_offset, "%s",
					     step_layout_req.node_list);
	} else
		step_layout_req.node_list = alloc->node_list;

	step_layout_req.cpus_per_node = alloc->cpus_per_node;
	step_layout_req.cpu_count_reps = alloc->cpu_count_reps;
	step_layout_req.cpus_per_task = cpus_per_task_array;
	step_layout_req.cpus_task_reps = cpus_task_reps;
	step_layout_req.task_dist = desc->task_dist;
	step_layout_req.plane_size = desc->plane_size;

	if (!(step_layout = slurm_step_layout_create(&step_layout_req)))
		return SLURM_ERROR;

	tmp = uint16_array_to_str(step_layout->node_cnt, step_layout->tasks);
	slurm_step_layout_destroy(step_layout);
	env_array_overwrite_het_fmt(dest, "SLURM_TASKS_PER_NODE",
				    het_job_offset,
				    "%s", tmp);
	xfree(tmp);

	if (alloc->account) {
		env_array_overwrite_het_fmt(dest, "SLURM_JOB_ACCOUNT",
					    het_job_offset, "%s",
					    alloc->account);
	}
	if (alloc->qos) {
		env_array_overwrite_het_fmt(dest, "SLURM_JOB_QOS",
					    het_job_offset,
					    "%s", alloc->qos);
	}
	if (alloc->resv_name) {
		env_array_overwrite_het_fmt(dest, "SLURM_JOB_RESERVATION",
					    het_job_offset, "%s",
					     alloc->resv_name);
	}

	if (alloc->env_size) {	/* Used to set Burst Buffer environment */
		for (i = 0; i < alloc->env_size; i++) {
			tmp = xstrdup(alloc->environment[i]);
			key = tmp;
			value = strchr(tmp, '=');
			if (value) {
				value[0] = '\0';
				value++;
				env_array_overwrite_het_fmt(dest, key,
							    het_job_offset,
							    "%s",
							    value);
			}
			xfree(tmp);
		}
	}

	if (desc->acctg_freq) {
		env_array_overwrite_het_fmt(dest, "SLURM_ACCTG_FREQ",
					    het_job_offset, "%s",
					     desc->acctg_freq);
	};

	if (desc->network) {
		env_array_overwrite_het_fmt(dest, "SLURM_NETWORK",
					    het_job_offset, "%s",
					    desc->network);
	}

	if (desc->overcommit != NO_VAL8) {
		env_array_overwrite_het_fmt(dest, "SLURM_OVERCOMMIT",
					    het_job_offset, "%u",
					     desc->overcommit);
	}

	/* Add default task counts for srun, if not already set */
	if (desc->bitflags & JOB_NTASKS_SET) {
		env_array_overwrite_het_fmt(dest, "SLURM_NTASKS",
					    het_job_offset,
					    "%d", desc->num_tasks);
		/* maintain for old scripts */
		env_array_overwrite_het_fmt(dest, "SLURM_NPROCS",
					    het_job_offset,
					    "%d", desc->num_tasks);
	}
	if (desc->bitflags & JOB_CPUS_SET) {
		env_array_overwrite_het_fmt(dest, "SLURM_CPUS_PER_TASK",
					    het_job_offset, "%d",
					     desc->cpus_per_task);
	}
	if (desc->ntasks_per_node && (desc->ntasks_per_node != NO_VAL16)) {
		env_array_overwrite_het_fmt(dest, "SLURM_NTASKS_PER_NODE",
					    het_job_offset, "%d",
					     desc->ntasks_per_node);
	}

	return rc;
}

/*
 * Set in "dest" the environment variables strings relevant to a Slurm batch
 * job allocation, overwriting any environment variables of the same name.
 * If the address pointed to by "dest" is NULL, memory will automatically be
 * xmalloc'ed.  The array is terminated by a NULL pointer, and thus is
 * suitable for use by execle() and other env_array_* functions.
 *
 * Sets the variables:
 *	SLURM_CLUSTER_NAME
 *	SLURM_JOB_ID
 *	SLURM_JOB_NUM_NODES
 *	SLURM_JOB_NODELIST
 *	SLURM_JOB_CPUS_PER_NODE
 *	SLURM_NODE_ALIASES
 *	ENVIRONMENT=BATCH
 *	HOSTNAME
 *
 * Sets OBSOLETE variables (needed for MPI, do not remove):
 *	SLURM_JOBID
 *	SLURM_NNODES
 *	SLURM_NODELIST
 *	SLURM_NTASKS
 *	SLURM_TASKS_PER_NODE
 */
extern int
env_array_for_batch_job(char ***dest, const batch_job_launch_msg_t *batch,
			const char *node_name)
{
	char *tmp = NULL;
	uint32_t num_cpus = 0;
	int i;
	slurm_step_layout_t *step_layout = NULL;
	uint16_t cpus_per_task;
	uint32_t task_dist;
	slurm_step_layout_req_t step_layout_req;
	uint16_t cpus_per_task_array[1];
	uint32_t cpus_task_reps[1];

	if (!batch)
		return SLURM_ERROR;

	memset(&step_layout_req, 0, sizeof(slurm_step_layout_req_t));
	step_layout_req.num_tasks = batch->ntasks;

	/*
	 * There is no explicit node count in the batch structure,
	 * so we need to calculate the node count.
	 */
	for (i = 0; i < batch->num_cpu_groups; i++) {
		step_layout_req.num_hosts += batch->cpu_count_reps[i];
		num_cpus += batch->cpu_count_reps[i] * batch->cpus_per_node[i];
	}

	env_array_overwrite_fmt(dest, "SLURM_CLUSTER_NAME", "%s",
	                        slurm_conf.cluster_name);

	env_array_overwrite_fmt(dest, "SLURM_JOB_ID", "%u", batch->job_id);
	env_array_overwrite_fmt(dest, "SLURM_JOB_NUM_NODES", "%u",
				step_layout_req.num_hosts);
	if (batch->array_task_id != NO_VAL) {
		env_array_overwrite_fmt(dest, "SLURM_ARRAY_JOB_ID", "%u",
					batch->array_job_id);
		env_array_overwrite_fmt(dest, "SLURM_ARRAY_TASK_ID", "%u",
					batch->array_task_id);
	}
	env_array_overwrite_fmt(dest, "SLURM_JOB_NODELIST", "%s", batch->nodes);
	env_array_overwrite_fmt(dest, "SLURM_JOB_PARTITION", "%s",
				batch->partition);
	env_array_overwrite_fmt(dest, "SLURM_NODE_ALIASES", "%s",
				batch->alias_list);

	tmp = uint32_compressed_to_str(batch->num_cpu_groups,
				       batch->cpus_per_node,
				       batch->cpu_count_reps);
	env_array_overwrite_fmt(dest, "SLURM_JOB_CPUS_PER_NODE", "%s", tmp);
	xfree(tmp);

	env_array_overwrite_fmt(dest, "ENVIRONMENT", "BATCH");
	if (node_name)
		env_array_overwrite_fmt(dest, "HOSTNAME", "%s", node_name);

	/* OBSOLETE, but needed by MPI, do not remove */
	env_array_overwrite_fmt(dest, "SLURM_JOBID", "%u", batch->job_id);
	env_array_overwrite_fmt(dest, "SLURM_NNODES", "%u",
				step_layout_req.num_hosts);
	env_array_overwrite_fmt(dest, "SLURM_NODELIST", "%s", batch->nodes);

	if ((batch->cpus_per_task != 0) &&
	    (batch->cpus_per_task != NO_VAL16))
		cpus_per_task = batch->cpus_per_task;
	else
		cpus_per_task = 1;	/* default value */
	cpus_per_task_array[0] = cpus_per_task;
	cpus_task_reps[0] = step_layout_req.num_hosts;

	/* Only overwrite this if it is set.  They are set in
	 * sbatch directly and could have changed. */
	if (getenvp(*dest, "SLURM_CPUS_PER_TASK"))
		env_array_overwrite_fmt(dest, "SLURM_CPUS_PER_TASK", "%u",
					cpus_per_task);

	if (step_layout_req.num_tasks) {
		env_array_append_fmt(dest, "SLURM_NTASKS", "%u",
				     step_layout_req.num_tasks);
		/* keep around for old scripts */
		env_array_append_fmt(dest, "SLURM_NPROCS", "%u",
				     step_layout_req.num_tasks);
	} else {
		step_layout_req.num_tasks = num_cpus / cpus_per_task;
	}

	if ((step_layout_req.node_list =
	     getenvp(*dest, "SLURM_ARBITRARY_NODELIST"))) {
		task_dist = SLURM_DIST_ARBITRARY;
	} else {
		step_layout_req.node_list = batch->nodes;
		task_dist = SLURM_DIST_BLOCK;
	}

	step_layout_req.cpus_per_node = batch->cpus_per_node;
	step_layout_req.cpu_count_reps = batch->cpu_count_reps;
	step_layout_req.cpus_per_task = cpus_per_task_array;
	step_layout_req.cpus_task_reps = cpus_task_reps;
	step_layout_req.task_dist = task_dist;
	step_layout_req.plane_size = NO_VAL16;

	if (!(step_layout = slurm_step_layout_create(&step_layout_req)))
		return SLURM_ERROR;

	tmp = uint16_array_to_str(step_layout->node_cnt, step_layout->tasks);
	slurm_step_layout_destroy(step_layout);
	env_array_overwrite_fmt(dest, "SLURM_TASKS_PER_NODE", "%s", tmp);
	xfree(tmp);

	if (batch->pn_min_memory & MEM_PER_CPU) {
		uint64_t tmp_mem = batch->pn_min_memory & (~MEM_PER_CPU);
		env_array_overwrite_fmt(dest, "SLURM_MEM_PER_CPU", "%"PRIu64"",
					tmp_mem);
	} else if (batch->pn_min_memory) {
		uint64_t tmp_mem = batch->pn_min_memory;
		env_array_overwrite_fmt(dest, "SLURM_MEM_PER_NODE", "%"PRIu64"",
					tmp_mem);
	}

	/* Set the SLURM_JOB_ACCOUNT,  SLURM_JOB_QOS
	 * and SLURM_JOB_RESERVATION if set by
	 * the controller.
	 */
	if (batch->account) {
		env_array_overwrite_fmt(dest,
					"SLURM_JOB_ACCOUNT",
					"%s",
					batch->account);
	}

	if (batch->qos) {
		env_array_overwrite_fmt(dest,
					"SLURM_JOB_QOS",
					"%s",
					batch->qos);
	}

	if (batch->resv_name) {
		env_array_overwrite_fmt(dest,
					"SLURM_JOB_RESERVATION",
					"%s",
					batch->resv_name);
	}

	return SLURM_SUCCESS;
}

/*
 * Set in "dest" the environment variables relevant to a Slurm job step,
 * overwriting any environment variables of the same name.  If the address
 * pointed to by "dest" is NULL, memory will automatically be xmalloc'ed.
 * The array is terminated by a NULL pointer, and thus is suitable for
 * use by execle() and other env_array_* functions.  If preserve_env is
 * true, the variables SLURM_NNODES, SLURM_NTASKS and SLURM_TASKS_PER_NODE
 * remain unchanged.
 *
 * Sets variables:
 *	SLURM_STEP_ID
 *	SLURM_STEP_NUM_NODES
 *	SLURM_STEP_NUM_TASKS
 *	SLURM_STEP_TASKS_PER_NODE
 *	SLURM_STEP_LAUNCHER_PORT
 *	SLURM_STEP_LAUNCHER_IPADDR
 *	SLURM_STEP_RESV_PORTS
 *      SLURM_STEP_SUB_MP
 *
 * Sets OBSOLETE variables:
 *	SLURM_STEPID
 *      SLURM_NNODES
 *	SLURM_NTASKS
 *	SLURM_NODELIST
 *	SLURM_TASKS_PER_NODE
 *	SLURM_SRUN_COMM_PORT
 *	SLURM_LAUNCH_NODE_IPADDR
 *
 */
extern void
env_array_for_step(char ***dest,
		   const job_step_create_response_msg_t *step,
		   launch_tasks_request_msg_t *launch,
		   uint16_t launcher_port,
		   bool preserve_env)
{
	char *tmp, *tpn;
	uint32_t node_cnt, task_cnt;

	if (!step || !launch)
		return;

	node_cnt = step->step_layout->node_cnt;
	env_array_overwrite_fmt(dest, "SLURM_STEP_ID", "%u", step->job_step_id);

	if (launch->het_job_node_list) {
		tmp = launch->het_job_node_list;
		env_array_overwrite_fmt(dest, "SLURM_NODELIST", "%s", tmp);
		env_array_overwrite_fmt(dest, "SLURM_JOB_NODELIST", "%s", tmp);
	} else {
		tmp = step->step_layout->node_list;
		env_array_append_fmt(dest, "SLURM_JOB_NODELIST", "%s", tmp);
	}
	env_array_overwrite_fmt(dest, "SLURM_STEP_NODELIST", "%s", tmp);

	if (launch->het_job_nnodes && (launch->het_job_nnodes != NO_VAL))
		node_cnt = launch->het_job_nnodes;
	env_array_overwrite_fmt(dest, "SLURM_STEP_NUM_NODES", "%u", node_cnt);

	if (launch->het_job_ntasks && (launch->het_job_ntasks != NO_VAL))
		task_cnt = launch->het_job_ntasks;
	else
		task_cnt = step->step_layout->task_cnt;
	env_array_overwrite_fmt(dest, "SLURM_STEP_NUM_TASKS", "%u", task_cnt);

	if (launch->het_job_task_cnts) {
		tpn = uint16_array_to_str(launch->het_job_nnodes,
					  launch->het_job_task_cnts);
		env_array_overwrite_fmt(dest, "SLURM_TASKS_PER_NODE", "%s",
					tpn);
		env_array_overwrite_fmt(dest, "SLURM_NNODES", "%u",
					launch->het_job_nnodes);
	} else {
		tpn = uint16_array_to_str(step->step_layout->node_cnt,
					  step->step_layout->tasks);
		if (!preserve_env) {
			env_array_overwrite_fmt(dest, "SLURM_TASKS_PER_NODE",
						"%s", tpn);
		}
	}
	env_array_overwrite_fmt(dest, "SLURM_STEP_TASKS_PER_NODE", "%s", tpn);

	env_array_overwrite_fmt(dest, "SLURM_STEP_LAUNCHER_PORT",
				"%hu", launcher_port);
	if (step->resv_ports) {
		env_array_overwrite_fmt(dest, "SLURM_STEP_RESV_PORTS",
					"%s", step->resv_ports);
	}

	/* OBSOLETE, but needed by some MPI implementations, do not remove */
	env_array_overwrite_fmt(dest, "SLURM_STEPID", "%u", step->job_step_id);
	if (!preserve_env) {
		env_array_overwrite_fmt(dest, "SLURM_NNODES", "%u", node_cnt);
		env_array_overwrite_fmt(dest, "SLURM_NTASKS", "%u", task_cnt);
		/* keep around for old scripts */
		env_array_overwrite_fmt(dest, "SLURM_NPROCS",
					"%u", step->step_layout->task_cnt);
	}
	env_array_overwrite_fmt(dest, "SLURM_SRUN_COMM_PORT",
				"%hu", launcher_port);

	xfree(tpn);
}

/*
 * Enviroment variables set elsewhere
 * ----------------------------------
 *
 * Set by slurmstepd:
 *	SLURM_STEP_NODEID
 *	SLURM_STEP_PROCID
 *	SLURM_STEP_LOCALID
 *
 * OBSOLETE set by slurmstepd:
 *	SLURM_NODEID
 *	SLURM_PROCID
 *	SLURM_LOCALID
 */

/***********************************************************************
 * Environment variable array support functions
 ***********************************************************************/

/*
 * Return an empty environment variable array (contains a single
 * pointer to NULL).
 */
char **env_array_create(void)
{
	char **env_array;

	env_array = xmalloc(sizeof(char *));
	env_array[0] = NULL;

	return env_array;
}

static int _env_array_update(char ***array_ptr, const char *name,
			     const char *value, bool over_write)
{
	char **ep = NULL;
	char *str = NULL;

	if (array_ptr == NULL)
		return 0;

	if (*array_ptr == NULL)
		*array_ptr = env_array_create();

	ep = _find_name_in_env(*array_ptr, name);
	if (*ep != NULL) {
		if (!over_write)
			return 0;
		xfree (*ep);
	} else {
		ep = _extend_env(array_ptr);
	}

	xstrfmtcat(str, "%s=%s", name, value);
	*ep = str;

	return 1;
}

/*
 * Append a single environment variable to an environment variable array,
 * if and only if a variable by that name does not already exist in the
 * array.
 *
 * "value_fmt" supports printf-style formatting.
 *
 * Return 1 on success, and 0 on error.
 */
int env_array_append_fmt(char ***array_ptr, const char *name,
			 const char *value_fmt, ...)
{
	int rc;
	char *value;
	va_list ap;

	value = xmalloc(ENV_BUFSIZE);
	va_start(ap, value_fmt);
	vsnprintf (value, ENV_BUFSIZE, value_fmt, ap);
	va_end(ap);
	rc = env_array_append(array_ptr, name, value);
	xfree(value);

	return rc;
}

/*
 * Append a single environment variable to an environment variable array,
 * if and only if a variable by that name does not already exist in the
 * array.
 *
 * Return 1 on success, and 0 on error.
 */
int env_array_append(char ***array_ptr, const char *name,
		     const char *value)
{
	return _env_array_update(array_ptr, name, value, false);
}

/*
 * Append a single environment variable to an environment variable array
 * if a variable by that name does not already exist.  If a variable
 * by the same name is found in the array, it is overwritten with the
 * new value.
 *
 * "value_fmt" supports printf-style formatting.
 *
 * Return 1 on success, and 0 on error.
 */
int env_array_overwrite_fmt(char ***array_ptr, const char *name,
			    const char *value_fmt, ...)
{
	int rc;
	char *value;
	va_list ap;

	value = xmalloc(ENV_BUFSIZE);
	va_start(ap, value_fmt);
	vsnprintf (value, ENV_BUFSIZE, value_fmt, ap);
	va_end(ap);
	rc = env_array_overwrite(array_ptr, name, value);
	xfree(value);

	return rc;
}

/*
 * Append a single environment variable to an environment variable array
 * if a variable by that name does not already exist.  If a variable
 * by the same name is found in the array, it is overwritten with the
 * new value.
 *
 * "value_fmt" supports printf-style formatting.
 *
 * Return 1 on success, and 0 on error.
 */
int env_array_overwrite_het_fmt(char ***array_ptr, const char *name,
				    int het_job_offset,
				    const char *value_fmt, ...)
{
	int rc;
	char *value;
	va_list ap;

	value = xmalloc(ENV_BUFSIZE);
	va_start(ap, value_fmt);
	vsnprintf (value, ENV_BUFSIZE, value_fmt, ap);
	va_end(ap);
	if (het_job_offset != -1) {
		char *het_comp_name = NULL;
		/* Continue support for old hetjob terminology. */
		xstrfmtcat(het_comp_name, "%s_PACK_GROUP_%d", name,
			   het_job_offset);
		rc = env_array_overwrite(array_ptr, het_comp_name, value);
		xfree(het_comp_name);
		xstrfmtcat(het_comp_name, "%s_HET_GROUP_%d", name,
			   het_job_offset);
		rc = env_array_overwrite(array_ptr, het_comp_name, value);
		xfree(het_comp_name);
	} else
		rc = env_array_overwrite(array_ptr, name, value);
	xfree(value);

	return rc;
}

/*
 * Append a single environment variable to an environment variable array
 * if a variable by that name does not already exist.  If a variable
 * by the same name is found in the array, it is overwritten with the
 * new value.
 *
 * Return 1 on success, and 0 on error.
 */
int env_array_overwrite(char ***array_ptr, const char *name,
			const char *value)
{
	return _env_array_update(array_ptr, name, value, true);
}

/*
 * Copy env_array must be freed by env_array_free
 */
char **env_array_copy(const char **array)
{
	char **ptr = NULL;

	env_array_merge(&ptr, array);

	return ptr;
}

/*
 * Free the memory used by an environment variable array.
 */
void env_array_free(char **env_array)
{
	char **ptr;

	if (env_array == NULL)
		return;

	for (ptr = env_array; *ptr != NULL; ptr++) {
		xfree(*ptr);
	}
	xfree(env_array);
}

/*
 * Given an environment variable "name=value" string,
 * copy the name portion into the "name" buffer, and the
 * value portion into the "value" buffer.
 *
 * Return 1 on success, 0 on failure.
 */
static int _env_array_entry_splitter(const char *entry,
				     char *name, int name_len,
				     char *value, int value_len)
{
	char *ptr;
	int len;

	ptr = xstrchr(entry, '=');
	if (ptr == NULL)	/* Bad parsing, no '=' found */
		return 0;
	/*
	 * need to consider the  byte pointed by ptr.
	 * example: entry = 0x0 = "a=b"
	 * ptr = 0x1
	 * len = ptr - entry + 1 = 2 because we need
	 * 2 characters to store 'a\0'
	 */
	len = ptr - entry + 1;
	if (len > name_len)
		return 0;
	strlcpy(name, entry, len);

	ptr++;
	/* account for '\0' here */
	len = strlen(ptr) + 1;
	if (len > value_len)
		return 0;
	strlcpy(value, ptr, len);

	return 1;
}

/*
 * Work similarly to putenv() (from C stdlib), but uses setenv()
 * under the covers.  This avoids having pointers from the global
 * array "environ" into "string".
 *
 * Return 1 on success, 0 on failure.
 */
static int _env_array_putenv(const char *string)
{
	int rc = 0;
	char name[256], *value;

	value = xmalloc(ENV_BUFSIZE);
	if ((_env_array_entry_splitter(string, name, sizeof(name),
				       value, ENV_BUFSIZE)) &&
	    (setenv(name, value, 1) != -1))
		rc = 1;

	xfree(value);
	return rc;
}

/*
 * Set all of the environment variables in a supplied environment
 * variable array.
 */
void env_array_set_environment(char **env_array)
{
	char **ptr;

	if (env_array == NULL)
		return;

	for (ptr = env_array; *ptr != NULL; ptr++) {
		_env_array_putenv(*ptr);
	}
}

/*
 * Unset all of the environment variables in a user's current
 * environment.
 *
 * (Note: becuae the environ array is decrementing with each
 *  unsetenv, only increment the ptr on a failure to unset.)
 */
void env_unset_environment(void)
{
	extern char **environ;
	char **ptr;
	char name[256], *value;

	value = xmalloc(ENV_BUFSIZE);
	for (ptr = (char **)environ; *ptr != NULL; ) {
		if ((_env_array_entry_splitter(*ptr, name, sizeof(name),
					       value, ENV_BUFSIZE)) &&
			(unsetenv(name) != -1))
			;
		else
			ptr++;
	}
	xfree(value);
}

/*
 * Merge all of the environment variables in src_array into the
 * array dest_array.  Any variables already found in dest_array
 * will be overwritten with the value from src_array.
 */
void env_array_merge(char ***dest_array, const char **src_array)
{
	char **ptr;
	char name[256], *value;

	if (src_array == NULL)
		return;

	value = xmalloc(ENV_BUFSIZE);
	for (ptr = (char **)src_array; *ptr != NULL; ptr++) {
		if (_env_array_entry_splitter(*ptr, name, sizeof(name),
					      value, ENV_BUFSIZE))
			env_array_overwrite(dest_array, name, value);
	}
	xfree(value);
}

/*
 * Merge the environment variables in src_array beginning with "SLURM" into the
 * array dest_array.  Any variables already found in dest_array will be
 * overwritten with the value from src_array.
 */
void env_array_merge_slurm(char ***dest_array, const char **src_array)
{
	char **ptr;
	char name[256], *value;

	if (src_array == NULL)
		return;

	value = xmalloc(ENV_BUFSIZE);
	for (ptr = (char **)src_array; *ptr != NULL; ptr++) {
		if (_env_array_entry_splitter(*ptr, name, sizeof(name),
					      value, ENV_BUFSIZE) &&
		    (xstrncmp(name, "SLURM", 5) == 0))
			env_array_overwrite(dest_array, name, value);
	}
	xfree(value);
}

/*
 * Strip out trailing carriage returns and newlines
 */
static void _strip_cr_nl(char *line)
{
	int len = strlen(line);
	char *ptr;

	for (ptr = line+len-1; ptr >= line; ptr--) {
		if (*ptr=='\r' || *ptr=='\n') {
			*ptr = '\0';
		} else {
			return;
		}
	}
}

/* Return the net count of curly brackets in a string
 * '{' adds one and '}' subtracts one (zero means it is balanced).
 * Special case: return -1 if no open brackets are found */
static int _bracket_cnt(char *value)
{
	int count = 0, i;
	for (i=0; value[i]; i++) {
		if (value[i] == '{')
			count++;
		else if (value[i] == '}')
			count--;
	}
	return count;
}

/*
 * Load user environment from a specified file or file descriptor.
 *
 * This will read in a user specified file or fd, that is invoked
 * via the --export-file option in sbatch. The NAME=value entries must
 * be NULL separated to support special characters in the environment
 * definitions.
 *
 * (Note: This is being added to a minor release. For the
 * next major release, it might be a consideration to merge
 * this functionality with that of load_env_cache and update
 * env_cache_builder to use the NULL character.)
 */
char **env_array_from_file(const char *fname)
{
	char *buf = NULL, *ptr = NULL, *eptr = NULL;
	char *value, *p;
	char **env = NULL;
	char name[256];
	int buf_size = BUFSIZ, buf_left;
	int file_size = 0, tmp_size;
	int separator = '\0';
	int fd;

	if (!fname)
		return NULL;

	/*
	 * If file name is a numeric value, then it is assumed to be a
	 * file descriptor.
	 */
	fd = (int)strtol(fname, &p, 10);
	if ((*p != '\0') || (fd < 3) || (fd > sysconf(_SC_OPEN_MAX)) ||
	    (fcntl(fd, F_GETFL) < 0)) {
		fd = open(fname, O_RDONLY);
		if (fd == -1) {
			error("Could not open user environment file %s", fname);
			return NULL;
		}
		verbose("Getting environment variables from %s", fname);
	} else
		verbose("Getting environment variables from fd %d", fd);

	/*
	 * Read in the user's environment data.
	 */
	buf = ptr = xmalloc(buf_size);
	buf_left = buf_size;
	while ((tmp_size = read(fd, ptr, buf_left))) {
		if (tmp_size < 0) {
			if (errno == EINTR)
				continue;
			error("read(environment_file): %m");
			break;
		}
		buf_left  -= tmp_size;
		file_size += tmp_size;
		if (buf_left == 0) {
			buf_size += BUFSIZ;
			xrealloc(buf, buf_size);
		}
		ptr = buf + file_size;
		buf_left = buf_size - file_size;
	}
	close(fd);

	/*
	 * Parse the buffer into individual environment variable names
	 * and build the environment.
	 */
	env   = env_array_create();
	value = xmalloc(ENV_BUFSIZE);
	for (ptr = buf; ; ptr = eptr+1) {
		eptr = strchr(ptr, separator);
		if ((ptr == eptr) || (eptr == NULL))
			break;
		if (_env_array_entry_splitter(ptr, name, sizeof(name),
					      value, ENV_BUFSIZE) &&
		    (!_discard_env(name, value))) {
			/*
			 * Unset the SLURM_SUBMIT_DIR if it is defined so
			 * that this new value does not get overwritten
			 * in the subsequent call to env_array_merge().
			 */
			if (xstrcmp(name, "SLURM_SUBMIT_DIR") == 0)
				unsetenv(name);
			env_array_overwrite(&env, name, value);
		}
	}
	xfree(buf);
	xfree(value);

	return env;
}

/*
 * Load user environment from a cache file located in
 * <state_save_location>/env_username
 */
static char **_load_env_cache(const char *username)
{
	char fname[MAXPATHLEN];
	char *line, name[256], *value;
	char **env = NULL;
	FILE *fp;
	int i;

	i = snprintf(fname, sizeof(fname), "%s/env_cache/%s",
		     slurm_conf.state_save_location, username);
	if (i < 0) {
		error("Environment cache filename overflow");
		return NULL;
	}
	if (!(fp = fopen(fname, "r"))) {
		error("Could not open user environment cache at %s: %m",
			fname);
		return NULL;
	}

	verbose("Getting cached environment variables at %s", fname);
	env = env_array_create();
	line  = xmalloc(ENV_BUFSIZE);
	value = xmalloc(ENV_BUFSIZE);
	while (1) {
		if (!fgets(line, ENV_BUFSIZE, fp))
			break;
		_strip_cr_nl(line);
		if (_env_array_entry_splitter(line, name, sizeof(name),
					      value, ENV_BUFSIZE) &&
		    (!_discard_env(name, value))) {
			if (value[0] == '(') {
				/* This is a bash function.
				 * It may span multiple lines */
				while (_bracket_cnt(value) > 0) {
					if (!fgets(line, ENV_BUFSIZE, fp))
						break;
					_strip_cr_nl(line);
					if ((strlen(value) + strlen(line)) >
					    (ENV_BUFSIZE - 2))
						break;
					strcat(value, "\n");
					strcat(value, line);
				}
			}
			env_array_overwrite(&env, name, value);
		}
	}
	xfree(line);
	xfree(value);

	fclose(fp);
	return env;
}

/*
 * Return an array of strings representing the specified user's default
 * environment variables following a two-prongged approach.
 * 1. Execute (more or less): "/bin/su - <username> -c /usr/bin/env"
 *    Depending upon the user's login scripts, this may take a very
 *    long time to complete or possibly never return
 * 2. Load the user environment from a cache file. This is used
 *    in the event that option 1 times out.  This only happens if no_cache isn't
 *    set.  If it is set then NULL will be returned if the normal load fails.
 *
 * timeout value is in seconds or zero for default (2 secs)
 * mode is 1 for short ("su <user>"), 2 for long ("su - <user>")
 * On error, returns NULL.
 *
 * NOTE: The calling process must have an effective uid of root for
 * this function to succeed.
 */
char **env_array_user_default(const char *username, int timeout, int mode,
			      bool no_cache)
{
	char *line = NULL, *last = NULL, name[MAXPATHLEN], *value, *buffer;
	char **env = NULL;
	char *starttoken = "XXXXSLURMSTARTPARSINGHEREXXXX";
	char *stoptoken  = "XXXXSLURMSTOPPARSINGHEREXXXXX";
	char cmdstr[256], *env_loc = NULL;
	char *stepd_path = NULL;
	int fd1, fd2, fildes[2], found, fval, len, rc, timeleft;
	int buf_read, buf_rem, config_timeout;
	pid_t child;
	struct timeval begin, now;
	struct pollfd ufds;
	struct stat buf;

	if (geteuid() != (uid_t)0) {
		error("SlurmdUser must be root to use --get-user-env");
		return NULL;
	}

	if (!slurm_conf.get_env_timeout)	/* just read directly from cache */
		return _load_env_cache(username);

	if (stat(SUCMD, &buf))
		fatal("Could not locate command: "SUCMD);
	if (stat("/bin/echo", &buf))
		fatal("Could not locate command: /bin/echo");
	stepd_path = slurm_get_stepd_loc();
	if (stat(stepd_path, &buf) == 0) {
		xstrcat(stepd_path, " getenv");
		env_loc = stepd_path;
	} else if (stat("/bin/env", &buf) == 0)
		env_loc = "/bin/env";
	else if (stat("/usr/bin/env", &buf) == 0)
		env_loc = "/usr/bin/env";
	else
		fatal("Could not location command: env");
	snprintf(cmdstr, sizeof(cmdstr),
		 "/bin/echo; /bin/echo; /bin/echo; "
		 "/bin/echo %s; %s; /bin/echo %s",
		 starttoken, env_loc, stoptoken);
	xfree(stepd_path);

	if (pipe(fildes) < 0) {
		fatal("pipe: %m");
		return NULL;
	}

	child = fork();
	if (child == -1) {
		fatal("fork: %m");
		return NULL;
	}
	if (child == 0) {
		char **tmp_env = NULL;
		tmp_env = env_array_create();
		env_array_overwrite(&tmp_env, "ENVIRONMENT", "BATCH");
		setpgid(0, 0);
		close(0);
		if ((fd1 = open("/dev/null", O_RDONLY)) == -1)
			error("%s: open(/dev/null): %m", __func__);
		dup2(fildes[1], 1);
		close(2);
		if ((fd2 = open("/dev/null", O_WRONLY)) == -1)
			error("%s: open(/dev/null): %m", __func__);
		if      (mode == 1)
			execle(SUCMD, "su", username, "-c", cmdstr, NULL, tmp_env);
		else if (mode == 2)
			execle(SUCMD, "su", "-", username, "-c", cmdstr, NULL, tmp_env);
		else {	/* Default system configuration */
#ifdef LOAD_ENV_NO_LOGIN
			execle(SUCMD, "su", username, "-c", cmdstr, NULL, tmp_env);
#else
			execle(SUCMD, "su", "-", username, "-c", cmdstr, NULL, tmp_env);
#endif
		}
		if (fd1 >= 0)	/* Avoid Coverity resource leak notification */
			(void) close(fd1);
		if (fd2 >= 0)	/* Avoid Coverity resource leak notification */
			(void) close(fd2);
		_exit(1);
	}

	close(fildes[1]);
	if ((fval = fcntl(fildes[0], F_GETFL, 0)) < 0)
		error("fcntl(F_GETFL) failed: %m");
	else if (fcntl(fildes[0], F_SETFL, fval | O_NONBLOCK) < 0)
		error("fcntl(F_SETFL) failed: %m");

	gettimeofday(&begin, NULL);
	ufds.fd = fildes[0];
	ufds.events = POLLIN;

	/* Read all of the output from /bin/su into buffer */
	if (timeout == 0)
		timeout = slurm_conf.get_env_timeout;	/* != 0 test above */
	found = 0;
	buf_read = 0;
	buffer = xmalloc(ENV_BUFSIZE);
	while (1) {
		gettimeofday(&now, NULL);
		timeleft = timeout * 1000;
		timeleft -= (now.tv_sec -  begin.tv_sec)  * 1000;
		timeleft -= (now.tv_usec - begin.tv_usec) / 1000;
		if (timeleft <= 0) {
			verbose("timeout waiting for "SUCMD" to complete");
			kill(-child, 9);
			break;
		}
		if ((rc = poll(&ufds, 1, timeleft)) <= 0) {
			if (rc == 0) {
				verbose("timeout waiting for "SUCMD" to complete");
				break;
			}
			if ((errno == EINTR) || (errno == EAGAIN))
				continue;
			error("poll(): %m");
			break;
		}
		if (!(ufds.revents & POLLIN)) {
			if (ufds.revents & POLLHUP) {	/* EOF */
				found = 1;		/* success */
			} else if (ufds.revents & POLLERR) {
				error("POLLERR");
			} else {
				error("poll() revents=%d", ufds.revents);
			}
			break;
		}
		buf_rem = ENV_BUFSIZE - buf_read;
		if (buf_rem == 0) {
			error("buffer overflow loading env vars");
			break;
		}
		rc = read(fildes[0], &buffer[buf_read], buf_rem);
		if (rc > 0)
			buf_read += rc;
		else if (rc == 0) {	/* EOF */
			found = 1;	/* success */
			break;
		} else {		/* error */
			error("read(env pipe): %m");
			break;
		}
	}
	close(fildes[0]);
	for (config_timeout=0; ; config_timeout++) {
		kill(-child, SIGKILL);	/* Typically a no-op */
		if (config_timeout)
			sleep(1);
		if (waitpid(child, &rc, WNOHANG) > 0)
			break;
		if (config_timeout >= 2) {
			/*
			 * Non-killable processes are indicative of file system
			 * problems. The process will remain as a zombie, but
			 * slurmd/salloc will not otherwise be effected.
			 */
			error("Failed to kill program loading user environment");
			break;
		}
	}

	if (!found) {
		error("Failed to load current user environment variables");
		xfree(buffer);
		return no_cache ? _load_env_cache(username) : NULL;
	}

	/* First look for the start token in the output */
	len = strlen(starttoken);
	found = 0;
	line = strtok_r(buffer, "\n", &last);
	while (!found && line) {
		if (!xstrncmp(line, starttoken, len)) {
			found = 1;
			break;
		}
		line = strtok_r(NULL, "\n", &last);
	}
	if (!found) {
		error("Failed to get current user environment variables");
		xfree(buffer);
		return no_cache ? _load_env_cache(username) : NULL;
	}

	/* Process environment variables until we find the stop token */
	len = strlen(stoptoken);
	found = 0;
	env = env_array_create();
	line = strtok_r(NULL, "\n", &last);
	value = xmalloc(ENV_BUFSIZE);
	while (!found && line) {
		if (!xstrncmp(line, stoptoken, len)) {
			found = 1;
			break;
		}
		if (_env_array_entry_splitter(line, name, sizeof(name),
					      value, ENV_BUFSIZE) &&
		    (!_discard_env(name, value))) {
			if (value[0] == '(') {
				/* This is a bash function.
				 * It may span multiple lines */
				while (_bracket_cnt(value) > 0) {
					line = strtok_r(NULL, "\n", &last);
					if (!line)
						break;
					if ((strlen(value) + strlen(line)) >
					    (ENV_BUFSIZE - 2))
						break;
					strcat(value, "\n");
					strcat(value, line);
				}
			}
			env_array_overwrite(&env, name, value);
		}
		line = strtok_r(NULL, "\n", &last);
	}
	xfree(value);
	xfree(buffer);
	if (!found) {
		error("Failed to get all user environment variables");
		env_array_free(env);
		return no_cache ? _load_env_cache(username) : NULL;
	}

	return env;
}

/*
 * Set TRES related env vars. Set here rather than env_array_for_job() since
 * we don't have array of opt values and the raw values are not stored in the
 * job_desc_msg_t structure (only the strings with possibly combined TRES)
 *
 * opt IN - options set by command parsing
 * dest IN/OUT - location to write environment variables
 * het_job_offset IN - component offset into hetjob, -1 if not hetjob
 */
extern void set_env_from_opts(slurm_opt_t *opt, char ***dest,
			      int het_job_offset)
{
	if (opt->cpus_per_gpu) {
		env_array_overwrite_het_fmt(dest, "SLURM_CPUS_PER_GPU",
					    het_job_offset, "%d",
					    opt->cpus_per_gpu);
	}
	if (opt->gpus) {
		env_array_overwrite_het_fmt(dest, "SLURM_GPUS",
					    het_job_offset, "%s",
					    opt->gpus);
	}
	if (opt->gpu_bind) {
		env_array_overwrite_het_fmt(dest, "SLURM_GPU_BIND",
					    het_job_offset, "%s",
					    opt->gpu_bind);
	}
	if (opt->gpu_freq) {
		env_array_overwrite_het_fmt(dest, "SLURM_GPU_FREQ",
					    het_job_offset, "%s",
					    opt->gpu_freq);
	}
	if (opt->gpus_per_node) {
		env_array_overwrite_het_fmt(dest, "SLURM_GPUS_PER_NODE",
					    het_job_offset, "%s",
					    opt->gpus_per_node);
	}
	if (opt->gpus_per_socket) {
		env_array_overwrite_het_fmt(dest, "SLURM_GPUS_PER_SOCKET",
					    het_job_offset, "%s",
					    opt->gpus_per_socket);
	}
	if (opt->gpus_per_task) {
		env_array_overwrite_het_fmt(dest, "SLURM_GPUS_PER_TASK",
					    het_job_offset, "%s",
					    opt->gpus_per_task);
	}
	if (opt->mem_per_gpu != NO_VAL64) {
		env_array_overwrite_het_fmt(dest, "SLURM_MEM_PER_GPU",
					    het_job_offset, "%"PRIu64,
					    opt->mem_per_gpu);
	}
}

extern char *find_quote_token(char *tmp, char *sep, char **last)
{
	char *start;
	int i, quote_single = 0, quote_double = 0;

	xassert(last);
	if (*last)
		start = *last;
	else
		start = tmp;
	if (start[0] == '\0')
		return NULL;
	for (i = 0; ; i++) {
		if (start[i] == '\'') {
			if (quote_single)
				quote_single--;
			else
				quote_single++;
		} else if (start[i] == '\"') {
			if (quote_double)
				quote_double--;
			else
				quote_double++;
		} else if (((start[i] == sep[0]) || (start[i] == '\0')) &&
			   (quote_single == 0) && (quote_double == 0)) {
			if (((start[0] == '\'') && (start[i-1] == '\'')) ||
			    ((start[0] == '\"') && (start[i-1] == '\"'))) {
				start++;
				i -= 2;
			}
			if (start[i] == '\0')
				*last = &start[i];
			else
				*last = &start[i] + 1;
			start[i] = '\0';
			return start;
		} else if (start[i] == '\0') {
			error("Improperly formed environment variable (%s)",
			      start);
			*last = &start[i];
			return start;
		}

	}
}
