/*-
 * Copyright (c) 2006 Sean Farley <sean-freebsd@FreeBSD.org>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer,
 *    without modification, immediately at the beginning of the file.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
 * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
#include <stdlib.h>
#include <stddef.h>
#include <string.h>
#include <stdbool.h>
#include <sys/types.h>
#include <err.h>
#if defined(DMALLOC_DEBUG)
#include <dmalloc.h>
#endif


/* Standard environ that is exposed to all code. */
extern char **environ;
static int environSize = 0;

/*
 * Array of environment variables built from environ.  Each element records:
 *	name:		Pointer to name=value string
 *	value:		Pointer to value within same string as name
 *	value size:	Size (not length) of space for value not counting nul
 *			character
 *	active state:	true/false value to signify whether variable is active.
 *			Useful since multiple variable with the same name can
 *			co-exist.  At most, one variable can be active at any
 *			time.
 */
static struct envVars {
	int32_t valueSize;
	char *name;
	char *value;
	bool active;
} *envVars = NULL;
static int envActive = 0;
static int envVarsSize = 0;


/*
 * Using the environment, rebuild the environ array for use by other C library
 * calls that depend upon it.
 */
static int
__rebuild_environ(int newEnvSize)
{
	int envNdx;
	int environNdx;

	/* Resize environ. */
	if (newEnvSize > environSize) {
		environSize = newEnvSize;
		environ = realloc(environ, sizeof (*environ) * (environSize +
			    1));
		if (environ == NULL)
			return (-1);
	}
	envActive = newEnvSize;

	/* Assign active variables to environ. */
	for (envNdx = envVarsSize - 1, environNdx = 0; envNdx >= 0; envNdx--)
		if (envVars[envNdx].active)
			environ[environNdx++] = envVars[envNdx].name;
	environ[environNdx] = NULL;

	return (0);
}


/*
 * Using environ, build an environment for use by standard C library calls.
 */
int
__build_env(void)
{
	char **env;
	int envNdx;

	/* Count environment variables. */
	for (env = environ, envVarsSize = 0; *env != NULL; env++)
		envVarsSize++;

	/* Create new environment. */
	envVars = malloc(sizeof (*envVars) * envVarsSize);
	if (envVars == NULL)
		return (-1);

	/* Copy environ values and keep track of them. */
	for (envNdx = 0; envNdx < envVarsSize; envNdx++) {
		envVars[envNdx].name =
		    strdup(environ[envVarsSize - envNdx - 1]);
		if (envVars[envNdx].name == NULL)
			return (-1);
		envVars[envNdx].value = strstr(envVars[envNdx].name, "=");
		if (envVars[envNdx].value != NULL) {
			envVars[envNdx].value++;
			envVars[envNdx].valueSize =
			    strlen(envVars[envNdx].value);
		}
		else
			errx(EXIT_FAILURE, "environ corrupt");
		envVars[envNdx].active = true;
	}

	/* Create a new environ. */
	environ = NULL;
	return (__rebuild_environ(envVarsSize));
}


/*
 * Using environment, returns pointer to value associated with name, if any,
 * else NULL.  Explicitly removes '=' in argument name in the search for the
 * variable.  If the onlyActive flag is set to true, only variables that are
 * active are returned else all are.
 */
static char *
__findenv(const char **name, int *nameLen, int *envNdx, bool onlyActive)
{
	if (nameLen == 0 || envNdx == NULL)
		return (NULL);

	/*
	 * Remove '=' from end of name for compatibility with other environment
	 * conventions.
	 */
	if ((*name)[*nameLen - 1] == '=')
		(*nameLen)--;
	if (*nameLen == 0)
		return (NULL);

	/* Find environment variable. */
	if (*envNdx == -1)
		*envNdx = envVarsSize - 1;
	for (; *envNdx >= 0; (*envNdx)--)
		if ((!onlyActive || envVars[*envNdx].active) &&
		    strncmp(envVars[*envNdx].name, *name, *nameLen) == 0 &&
		    (envVars[*envNdx].name)[*nameLen] == '=')
			return (envVars[*envNdx].value);

	return (NULL);
}


/*
 * Using environ, returns pointer to value associated with name, if any, else
 * NULL.  Explicitly removes '=' in argument name in the search for the
 * variable.
 */
static char *
__findenv_environ(const char *name)
{
	int envNdx;
	int nameLen;

	if (name == NULL || environ == NULL)
		return (NULL);

	/* Find variable within environ. */
	nameLen = strlen(name);
	for (envNdx = 0; environ[envNdx] != NULL; envNdx++)
		if (strncmp(environ[envNdx], name, nameLen) == 0 &&
		    (environ[envNdx])[nameLen] == '=') {
			return (&(environ[envNdx][nameLen + strlen("=")]));
		}

	return (NULL);
}


/*
 * Returns the value of a variable or NULL if none are found.  Special handling
 * is provided for a MALLOC_OPTIONS variable to prevent recursion since
 * malloc(), during its initialization, calls getenv().
 */
char *
getenv(const char *name)
{
	int envNdx;
	int nameLen;

	if (name == NULL || (nameLen = strlen(name)) == 0)
		return (NULL);

	/*
	 * Initialize environment.  To prevent recursion, handle malloc()
	 * initialization calling getenv().
	 */
	if (envVars == NULL) {
#ifndef DMALLOC_DEBUG
		if (strcmp(name, "MALLOC_OPTIONS") == 0)
#else
		if (strcmp(name, "DMALLOC_OPTIONS") == 0)
#endif
			return (__findenv_environ(name));
		else if (__build_env() == -1)
			return (NULL);
	}

	/* Find environment variable. */
	envNdx = -1;
	return (__findenv(&name, &nameLen, &envNdx, true));
}


/*
 * Set the value of a variable.  Older settings are labeled as inactive.  If an
 * older setting has enough room to store the new value, it will be reused.  No
 * previous variables are ever freed here to avoid causing a segmentation fault
 * in a user's code.
 */
int
setenv(const char *name, const char *value, int overwrite)
{
	char *env;
	char *savedEnv;
	int envNdx;
	int nameLen;
	int newEnvSize;
	int valueLen;

	/* Initialize environment. */
	if (envVars == NULL && __build_env() == -1)
		return (-1);

	/*
	 * Remove '=' from beginning of value for compatibility with other
	 * environment conventions.
	 */
	if (value[0] == '=')
		value++;

	/* Find existing environment variable large enough to use. */
	newEnvSize = envActive;
	envNdx = -1;
	nameLen = strlen(name);
	valueLen = strlen(value);
	savedEnv = NULL;
	while ((env = __findenv(&name, &nameLen, &envNdx, false))) {
		/* Deactivate entry. */
		if (envVars[envNdx].active) {
			if (!overwrite)
				return (0);
			envVars[envNdx].active = false;
			newEnvSize--;
		}

		/* Entry is large enough to use. */
		if (envVars[envNdx].valueSize >= valueLen) {
			if (overwrite)
				break;
			else if (!savedEnv)
				savedEnv = env;
		}

		envNdx--;
	}
	if (nameLen == 0)
		return (-1);
	if (savedEnv)
		env = savedEnv;

	/* Create new variable if none were found of sufficient size. */
	if (env == NULL) {
		/* Enlarge environment. */
		envNdx = envVarsSize;
		envVarsSize++;
		envVars = realloc(envVars, sizeof (*envVars) * envVarsSize);
		if (envVars == NULL)
			return (-1);

		/* Create environment entry. */
		envVars[envNdx].name = malloc(nameLen + sizeof ("=") +
		    valueLen);
		if (envVars[envNdx].name == NULL)
			return (-1);
		envVars[envNdx].valueSize = valueLen;
	}

	/* Save variable name/value pair. */
	env = stpcpy(envVars[envNdx].name, name);
	if ((envVars[envNdx].name)[nameLen] != '=')
		env = stpcpy(env, "=");
	stpcpy(env, value);
	envVars[envNdx].value = env;
	envVars[envNdx].active = true;
	newEnvSize++;

	return (__rebuild_environ(newEnvSize));
}


/*
 * Unset all variables with the same name by flagging them as inactive.  No
 * variable is ever freed.
 */
void
unsetenv(const char *name)
{
	int envNdx;
	int nameLen;
	int newEnvSize;

	/* Initialize environment. */
	if (envVars == NULL && __build_env() == -1)
		return;

	/* Unset all instances of specified variable. */
	nameLen = strlen(name);
	newEnvSize = envActive;
	envNdx = -1;
	for (envNdx = envVarsSize - 1;
	    __findenv(&name, &nameLen, &envNdx, true); envNdx--) {
		envVars[envNdx].active = false;
		newEnvSize--;
	}

	if (newEnvSize != envActive)
		__rebuild_environ(newEnvSize);

	return;
}


/*
 * Deallocate the environment built from environ as well as environ then set
 * both to NULL.  Eases debugging of memory leaks.
 */
void
__clean_env(void)
{
	int envNdx;

	if (envVars == NULL)
		return;

	/* Deallocate environment and environ. */
	for (envNdx = 0; envNdx < envVarsSize; envNdx++)
		free(envVars[envNdx].name);
	free(envVars);
	free(environ);
	envVars = NULL;
	environ = NULL;

	return;
}