Move ClearEnv and IterateEnv to util

commit_hash:c42da6c8b3d8e0c63f1cf4457c245bbd97025cc2

3 files changed, 154 insertions, 4 deletions

diff --git a/util/system/env.cpp b/util/system/env.cpp
index 1c1fc0cea85..24bd7707025 100644
--- a/util/system/env.cpp
+++ b/util/system/env.cpp
@@ -5,11 +5,19 @@
 #include <util/generic/yexception.h>
 
 #ifdef _win_
+    #include <util/generic/scope.h>
     #include <util/generic/vector.h>
+    #include <util/system/yassert.h>
     #include "winint.h"
 #else
+    #ifndef _linux_
+        #include <util/generic/vector.h> // ClearEnv impl
+    #endif
+
     #include <cerrno>
     #include <cstdlib>
+
+extern char** environ;
 #endif
 
 /**
@@ -84,3 +92,73 @@ void UnsetEnv(const TString& key) {
 #endif
     Y_ENSURE_EX(ok || notFound, TSystemError() << "failed to unset environment variable " << key.Quote());
 }
+
+void IterateEnv(const std::function<void(TStringBuf, TStringBuf)>& f, bool ignoreMalformedStrings) {
+#ifdef _win_
+    // Env block format:
+    // Var1=Value1\0
+    // Var2=Value2\0
+    // ...
+    // VarN=ValueN\0\0
+    // Or "\0" if empty
+
+    auto envBlock = GetEnvironmentStringsA();
+    if (!envBlock) {
+        ythrow TSystemError() << "failed to get environment variables";
+    }
+    Y_DEFER {
+        bool ok = FreeEnvironmentStringsA(envBlock);
+        Y_ABORT_UNLESS(ok, "failed to free env block"); // ¯\_(ツ)_/¯
+    };
+    const char* charEnv = envBlock;
+    while (*charEnv) {
+        TStringBuf varStr(charEnv);
+        TStringBuf name, value;
+        if (varStr.TrySplit('=', name, value)) { // Not optimal (we could scan for null byte and '=' in one pass), but less bugprone
+            f(name, value);
+        } else {
+            Y_ENSURE(ignoreMalformedStrings, "Environment contains a string without '=': \"" << varStr << '"');
+        }
+        charEnv = varStr.data() + varStr.size() + 1;
+    }
+#else
+    if (!environ) { // may be null after clearenv
+        return;
+    }
+    for (char** var = environ; *var; ++var) {
+        TStringBuf varStr(*var);
+        TStringBuf name, value;
+        if (varStr.TrySplit('=', name, value)) {
+            f(name, value);
+        } else {
+            Y_ENSURE(ignoreMalformedStrings, "Environment contains a string without '=': \"" << varStr << '"');
+        }
+    }
+#endif
+}
+
+void ClearEnv() {
+#if defined(_win_)
+    wchar_t emptyEnv[] = {L'\0'};
+    // At least on Wine, SetEnvironmentStringsA expects a multi-byte string that is internally converted to UTF-16.
+    // So it's easier to use SetEnvironmentStringsW directly.
+    bool ok = SetEnvironmentStringsW(emptyEnv);
+    Y_ENSURE_EX(ok, TSystemError() << "failed to clear environment");
+#elif defined(_linux_)
+    bool ok = !clearenv();
+    Y_ENSURE_EX(ok, TSystemError() << "failed to clear environment");
+#else
+    // Darwin or other unix platform that may not have clearenv
+    TVector<TString> keys;
+    IterateEnv(
+        [&keys](TStringBuf name, TStringBuf) {
+            keys.emplace_back(name);
+        },
+        true // Ignore malformed strings - they don't seem to be accessible anyway.
+             // IterateEnv won't expose these strings, and GetEnv/SetEnv/UnsetEnv ignore them (at least in glibc).
+    );
+    for (const auto& key : keys) {
+        UnsetEnv(key); // duplicate UnsetEnv is ok.
+    }
+#endif
+}
diff --git a/util/system/env.h b/util/system/env.h
index 6458e918b83..5aa4fbe400f 100644
--- a/util/system/env.h
+++ b/util/system/env.h
@@ -2,6 +2,12 @@
 
 #include <util/generic/fwd.h>
 #include <util/generic/string.h>
+#include <util/generic/strbuf.h>
+
+#include <functional>
+
+// NOTE: On Windows, functions from this header may have unexpected behavior when used with codepages that don't support all Unicode characters.
+//       For more info, see https://blog.orange.tw/posts/2025-01-worstfit-unveiling-hidden-transformers-in-windows-ansi/
 
 /**
  * Search the environment list provided by the host environment for associated variable.
@@ -14,7 +20,7 @@
  *
  * @note        Use it only in pair with `SetEnv` as there may be inconsistency in their behaviour
  *              otherwise.
- * @note        Calls to `GetEnv` and environment modifying functions (`SetEnv` or `UnsetEnv`) from different threads must be synchronized.
+ * @note        Calls to `GetEnv` and environment modifying functions (like `SetEnv`) from different threads must be synchronized.
  * @see         SetEnv
  */
 TString GetEnv(const TString& key, const TString& def = TString());
@@ -28,7 +34,7 @@ TString GetEnv(const TString& key, const TString& def = TString());
  *              variable or empty optional value if such variable is missing.
  *
  * @note        Use it only in pair with `SetEnv` as there may be inconsistency in their behaviour otherwise.
- * @note        Calls to `TryGetEnv` and environment modifying functions (`SetEnv` or `UnsetEnv`) from different threads must be synchronized.
+ * @note        Calls to `TryGetEnv` and environment modifying functions (like `SetEnv`) from different threads must be synchronized.
  * @see         SetEnv
  */
 TMaybe<TString> TryGetEnv(const TString& key);
@@ -43,7 +49,7 @@ TMaybe<TString> TryGetEnv(const TString& key);
  *
  * @note        Use it only in pair with `GetEnv` as there may be inconsistency in their behaviour
  *              otherwise.
- * @note        Calls to `SetEnv` and `GetEnv`, `TryGetEnv`, `UnsetEnv` from different threads must be synchronized.
+ * @note        Calls to `SetEnv` and all other env-related functions from different threads must be synchronized.
  * @see         GetEnv
  */
 void SetEnv(const TString& key, const TString& value);
@@ -57,7 +63,34 @@ void SetEnv(const TString& key, const TString& value);
  *
  * @note                    If key does not exist in the environment, then the environment is unchanged,
  *                          and the function returns normally.
- * @note                    Calls to `UnsetEnv` and `GetEnv`, `TryGetEnv`, `SetEnv` from different threads must be synchronized.
+ * @note                    Calls to `UnsetEnv` and and all other env-related functions from different threads must be synchronized.
  * @see                     GetEnv
  */
 void UnsetEnv(const TString& key);
+
+/**
+ * Clear the host environment.
+ *
+ * @throws TSystemError     On error
+ *
+ * @note                    Calls to `ClearEnv` and all other env-related functions from different threads must be synchronized.
+ */
+void ClearEnv();
+
+/**
+ * Call the provided function `f` with name and value as arguments for each environment variable.
+ *
+ * @param f                         The function to call
+ * @param ignoreMalformedStrings    If false, any malformed string in host environment (one that doesn't contain a '=') will cause an exception.
+ *                                  If true, such strings will be skipped.
+ *
+ * @throws TSystemError         On error
+ * @throws yexception           On encountering a malformed string (see `ignoreMalformedStrings`)
+ *
+ * @note                        The function `f` may be called with the same variable name multiple times.
+ *                              Lifetime for contents of `name` and `value` is not guaranteed to last beyond the scope of the invocation of `f`.
+ *
+ * @note                        Calls to `IterateEnv` and environment modifying functions (like `SetEnv`) from different threads must be synchronized.
+ *                              Also, `f` must not call any environment modifying functions.
+ */
+void IterateEnv(const std::function<void(TStringBuf /*name*/, TStringBuf /*value*/)>& f, bool ignoreMalformedStrings = false);
diff --git a/util/system/env_ut.cpp b/util/system/env_ut.cpp
index acc799d556a..fff32b9abd6 100644
--- a/util/system/env_ut.cpp
+++ b/util/system/env_ut.cpp
@@ -56,4 +56,43 @@ Y_UNIT_TEST_SUITE(EnvTest) {
         UNIT_ASSERT_EXCEPTION(SetEnv("", "value"), yexception);
         UNIT_ASSERT_EXCEPTION(SetEnv("A=B", "C=D"), yexception);
     }
+
+    Y_UNIT_TEST(ClearEnv) {
+        TString key1 = "util_GETENV_TestVar1";
+        TString key2 = "util_GETENV_TestVar2";
+        SetEnv(key1, "some value");
+        SetEnv(key2, "some other value");
+        ClearEnv();
+        UNIT_ASSERT(TryGetEnv(key1).Empty());
+        UNIT_ASSERT(TryGetEnv(key2).Empty());
+    }
+
+    Y_UNIT_TEST(IterateEnv) {
+        ClearEnv();
+        IterateEnv([](TStringBuf name, TStringBuf value) {
+            UNIT_FAIL(TString::Join("Got unexpected env variable ", name, " with value ", value));
+        });
+        TString key1 = "util_GETENV_TestVar1";
+        TString key2 = "util_GETENV_TestVar2";
+        TString value1 = "some value";
+        TString value2 = "some other value";
+        SetEnv(key1, value1);
+        SetEnv(key2, value2);
+
+        bool foundVar1 = false;
+        bool foundVar2 = false;
+        IterateEnv([&](TStringBuf name, TStringBuf value) {
+            if (name == key1) {
+                foundVar1 = true;
+                UNIT_ASSERT_EQUAL(value, value1);
+            } else if (name == key2) {
+                foundVar2 = true;
+                UNIT_ASSERT_EQUAL(value, value2);
+            } else {
+                UNIT_FAIL(TString::Join("Got unexpected env variable ", name, " with value ", value));
+            }
+        });
+        UNIT_ASSERT(foundVar1);
+        UNIT_ASSERT(foundVar2);
+    }
 } // Y_UNIT_TEST_SUITE(EnvTest)