diff --git a/src/libcmd/command.cc b/src/libcmd/command.cc index 67fef1909..4f93dc83e 100644 --- a/src/libcmd/command.cc +++ b/src/libcmd/command.cc @@ -1,3 +1,4 @@ +#include #include #include "command.hh" @@ -9,8 +10,7 @@ #include "profiles.hh" #include "repl.hh" #include "strings.hh" - -extern char * * environ __attribute__((weak)); +#include "environment-variables.hh" namespace nix { @@ -290,48 +290,63 @@ MixDefaultProfile::MixDefaultProfile() MixEnvironment::MixEnvironment() : ignoreEnvironment(false) { addFlag({ - .longName = "ignore-environment", + .longName = "ignore-env", .shortName = 'i', - .description = "Clear the entire environment (except those specified with `--keep`).", + .description = "Clear the entire environment, except for those specified with `--keep-env-var`.", .handler = {&ignoreEnvironment, true}, }); addFlag({ - .longName = "keep", + .longName = "keep-env-var", .shortName = 'k', .description = "Keep the environment variable *name*.", .labels = {"name"}, - .handler = {[&](std::string s) { keep.insert(s); }}, + .handler = {[&](std::string s) { keepVars.insert(s); }}, }); addFlag({ - .longName = "unset", + .longName = "unset-env-var", .shortName = 'u', .description = "Unset the environment variable *name*.", .labels = {"name"}, - .handler = {[&](std::string s) { unset.insert(s); }}, + .handler = {[&](std::string s) { unsetVars.insert(s); }}, + }); + + addFlag({ + .longName = "set-env-var", + .shortName = 's', + .description = "Add/override an environment variable *name* with *value*.", + .labels = {"name", "value"}, + .handler = {[&](std::string name, std::string value) { setVars.insert_or_assign(name, value); }}, }); } void MixEnvironment::setEnviron() { - if (ignoreEnvironment) { - if (!unset.empty()) - throw UsageError("--unset does not make sense with --ignore-environment"); + if (ignoreEnvironment && !unsetVars.empty()) + throw UsageError("--unset-env-var does not make sense with --ignore-env"); - for (const auto & var : keep) { - auto val = getenv(var.c_str()); - if (val) stringsEnv.emplace_back(fmt("%s=%s", var.c_str(), val)); - } + if (!ignoreEnvironment && !keepVars.empty()) + throw UsageError("--keep-env-var does not make sense without --ignore-env"); - vectorEnv = stringsToCharPtrs(stringsEnv); - environ = vectorEnv.data(); - } else { - if (!keep.empty()) - throw UsageError("--keep does not make sense without --ignore-environment"); + auto env = getEnv(); - for (const auto & var : unset) - unsetenv(var.c_str()); - } + if (ignoreEnvironment) + std::erase_if(env, [&](const auto & var) { + return !keepVars.contains(var.first); + }); + + if (!unsetVars.empty()) + std::erase_if(env, [&](const auto & var) { + return unsetVars.contains(var.first); + }); + + if (!setVars.empty()) + for (const auto & [name, value] : setVars) + env[name] = value; + + replaceEnv(std::move(env)); + + return; } } diff --git a/src/libcmd/command.hh b/src/libcmd/command.hh index 4a72627ed..b823a1c6b 100644 --- a/src/libcmd/command.hh +++ b/src/libcmd/command.hh @@ -315,17 +315,18 @@ struct MixDefaultProfile : MixProfile struct MixEnvironment : virtual Args { - StringSet keep, unset; - Strings stringsEnv; - std::vector vectorEnv; + StringSet keepVars; + StringSet unsetVars; + std::map setVars; bool ignoreEnvironment; MixEnvironment(); /*** - * Modify global environ based on `ignoreEnvironment`, `keep`, and - * `unset`. It's expected that exec will be called before this class - * goes out of scope, otherwise `environ` will become invalid. + * Modify global environ based on `ignoreEnvironment`, `keep`, + * `unset`, and `added`. It's expected that exec will be called + * before this class goes out of scope, otherwise `environ` will + * become invalid. */ void setEnviron(); }; diff --git a/tests/functional/flakes/develop.sh b/tests/functional/flakes/develop.sh index 60abf3ef2..95b5d8968 100755 --- a/tests/functional/flakes/develop.sh +++ b/tests/functional/flakes/develop.sh @@ -43,9 +43,9 @@ echo "\$ENVVAR" EOF )" = "a" ]] -# Test whether `nix develop --ignore-environment` does _not_ pass through environment variables. +# Test whether `nix develop --ignore-env` does _not_ pass through environment variables. [[ -z "$( - ENVVAR=a nix develop --ignore-environment --no-write-lock-file .#hello <