From 75c897cf3d2be5cc156a87ec54c6726e8dc2a926 Mon Sep 17 00:00:00 2001
From: matthew <matthewkenigsberg@gmail.com>
Date: Thu, 7 Nov 2019 17:16:48 -0600
Subject: [PATCH] Factor out code to handle environment in run into
 MixEnvironment

---
 src/nix/command.cc | 44 ++++++++++++++++++++++++++++++++++++++++++++
 src/nix/command.hh | 15 ++++++++++++++-
 2 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/src/nix/command.cc b/src/nix/command.cc
index 1cb4cc92a..aa5c8cf1e 100644
--- a/src/nix/command.cc
+++ b/src/nix/command.cc
@@ -128,4 +128,48 @@ MixDefaultProfile::MixDefaultProfile()
     profile = getDefaultProfile();
 }
 
+MixEnvironment::MixEnvironment() : ignoreEnvironment(false) {
+    mkFlag()
+        .longName("ignore-environment")
+        .shortName('i')
+        .description("clear the entire environment (except those specified with --keep)")
+        .set(&ignoreEnvironment, true);
+
+    mkFlag()
+        .longName("keep")
+        .shortName('k')
+        .description("keep specified environment variable")
+        .arity(1)
+        .labels({"name"})
+        .handler([&](std::vector<std::string> ss) { keep.insert(ss.front()); });
+
+    mkFlag()
+        .longName("unset")
+        .shortName('u')
+        .description("unset specified environment variable")
+        .arity(1)
+        .labels({"name"})
+        .handler([&](std::vector<std::string> ss) { unset.insert(ss.front()); });
+}
+
+void MixEnvironment::setEnviron() {
+    if (ignoreEnvironment) {
+        if (!unset.empty())
+            throw UsageError("--unset does not make sense with --ignore-environment");
+
+        for (const auto & var : keep) {
+            auto val = getenv(var.c_str());
+            if (val) stringEnv.emplace_back(fmt("%s=%s", var.c_str(), val));
+        }
+        vectorEnv = stringsToCharPtrs(stringsEnv);
+        environ = vectorEnv.data();
+    } else {
+        if (!keep.empty())
+            throw UsageError("--keep does not make sense without --ignore-environment");
+
+        for (const auto & var : unset)
+            unsetenv(var.c_str());
+    }
+}
+
 }
diff --git a/src/nix/command.hh b/src/nix/command.hh
index 13f3a0dc9..52fb4e1f5 100644
--- a/src/nix/command.hh
+++ b/src/nix/command.hh
@@ -194,4 +194,17 @@ struct MixDefaultProfile : MixProfile
     MixDefaultProfile();
 };
 
-}
+struct MixEnvironment : virtual Args {
+
+    StringSet keep, unset;
+    Strings stringsEnv;
+    std::vector<char*> vectorEnv;
+    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. */
+    void setEnviron();
+};
+
+}
\ No newline at end of file