diff --git a/src/libcmd/command.cc b/src/libcmd/command.cc index 67fef190920a..044148f68b29 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 { @@ -301,37 +301,52 @@ MixEnvironment::MixEnvironment() : ignoreEnvironment(false) .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-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-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-var does not make sense with --ignore-environment"); - 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 does not make sense without --ignore-environment"); - 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 4a72627ed4db..b823a1c6b057 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(); };