feat: add flag add to MixEnvironment

src/libcmd/command.cc

@@ -1,3 +1,4 @@
+#include <algorithm>
 #include <nlohmann/json.hpp>
 
 #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;
 }
 
 }

src/libcmd/command.hh

@@ -315,17 +315,18 @@ struct MixDefaultProfile : MixProfile
 
 struct MixEnvironment : virtual Args {
 
-    StringSet keep, unset;
-    Strings stringsEnv;
-    std::vector<char*> vectorEnv;
+    StringSet keepVars;
+    StringSet unsetVars;
+    std::map<std::string, std::string> 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();
 };