From 5a1adcb2a6708ae5c4ee0da08dd85aadb739d4e3 Mon Sep 17 00:00:00 2001
From: apio <blobs.trading@gmail.com>
Date: Fri, 11 Aug 2023 17:59:04 +0200
Subject: [PATCH] libc: Add putenv

---
 libc/include/stdlib.h  |  3 ++
 libc/src/env.cpp       | 75 ++++++++++++++++++++++++++++++++----------
 tests/libc/TestEnv.cpp | 30 +++++++++++++++++
 3 files changed, 90 insertions(+), 18 deletions(-)

diff --git a/libc/include/stdlib.h b/libc/include/stdlib.h
index 9fa130f0..b06e1861 100644
--- a/libc/include/stdlib.h
+++ b/libc/include/stdlib.h
@@ -128,6 +128,9 @@ extern "C"
     /* Clear all environment variables. */
     int clearenv(void);
 
+    /* Add a new variable to the environment. */
+    int putenv(char* string);
+
     /* Sort an array of arbitrary elements using a comparison function. */
     void qsort(void* base, size_t nmemb, size_t size, int (*compar)(const void*, const void*));
 
diff --git a/libc/src/env.cpp b/libc/src/env.cpp
index 25812a96..b7c96d35 100644
--- a/libc/src/env.cpp
+++ b/libc/src/env.cpp
@@ -1,11 +1,13 @@
 #include <bits/errno-return.h>
 #include <errno.h>
+#include <luna/HashTable.h>
 #include <luna/ScopeGuard.h>
 #include <luna/Vector.h>
 #include <stdlib.h>
 
 extern "C" char** environ;
 Vector<char*> g_dynamic_env = {};
+HashTable<char*> g_dynamic_vars = {};
 bool env_is_dynamic { false };
 
 /* If an environment variable matching key is found, stores its value in value and returns its index. Otherwise,
@@ -41,22 +43,24 @@ static isize _findenv(const char* key, char** value)
     return -1;
 }
 
+static void free_if_needed(char* element)
+{
+    if (g_dynamic_vars.try_find(element))
+    {
+        g_dynamic_vars.try_remove(element);
+        free(element);
+    }
+}
+
 /* Move environ to a heap-allocated array. */
 static Result<void> _try_move_env()
 {
     char** env = environ;
     g_dynamic_env.clear();
-    auto guard = make_scope_guard([] {
-        for (auto element : g_dynamic_env)
-        {
-            if (element) free(element);
-        }
-    });
 
     if (!env)
     {
         TRY(g_dynamic_env.try_append(nullptr));
-        guard.deactivate();
         env_is_dynamic = true;
         environ = g_dynamic_env.data();
         return {};
@@ -64,18 +68,12 @@ static Result<void> _try_move_env()
 
     while (*env)
     {
-        char* ptr = strdup(*(env++));
-        if (!ptr) return err(errno);
-
-        auto ptr_guard = make_scope_guard([=] { free(ptr); });
+        char* ptr = *(env++);
         TRY(g_dynamic_env.try_append(ptr));
-        ptr_guard.deactivate();
     }
 
     TRY(g_dynamic_env.try_append(nullptr));
 
-    guard.deactivate();
-
     env_is_dynamic = true;
     environ = g_dynamic_env.data();
 
@@ -105,7 +103,7 @@ static void _check_dynamic_env()
     {
         for (auto element : g_dynamic_env)
         {
-            if (element) free(element);
+            if (element) free_if_needed(element);
         }
     }
 }
@@ -122,7 +120,7 @@ extern "C"
         {
             for (auto element : g_dynamic_env)
             {
-                if (element) free(element);
+                if (element) free_if_needed(element);
             }
 
             g_dynamic_env.clear();
@@ -152,9 +150,11 @@ extern "C"
             if (_move_env() < 0) return -1;
         }
 
-        g_dynamic_env.remove_at(index);
+        char* p = g_dynamic_env.remove_at(index);
         _update_env();
 
+        free_if_needed(p);
+
         return 0;
     }
 
@@ -188,12 +188,14 @@ extern "C"
 
         if (index >= 0)
         {
-            free(environ[index]);
+            free_if_needed(environ[index]);
             environ[index] = str;
             guard.deactivate();
             return 0;
         }
 
+        TRY_OR_SET_ERRNO(g_dynamic_vars.try_set(str), int, -1);
+
         // Add a new NULL at the end of the array and replace the previous one with our string.
         index = g_dynamic_env.size() - 1;
 
@@ -206,6 +208,43 @@ extern "C"
         return 0;
     }
 
+    int putenv(char* string)
+    {
+        char* p = strchr(string, '=');
+        if (!p) return unsetenv(string);
+        size_t key_len = p - string;
+
+        char* key = strndup(string, key_len);
+        auto guard = make_scope_guard([key] { free(key); });
+
+        auto index = _findenv(key, nullptr);
+
+        _check_dynamic_env();
+
+        if (!env_is_dynamic)
+        {
+            if (_move_env() < 0) return -1;
+        }
+
+        if (index >= 0)
+        {
+            free_if_needed(environ[index]);
+            environ[index] = string;
+            return 0;
+        }
+
+        // Add a new NULL at the end of the array and replace the previous one with our string.
+        index = g_dynamic_env.size() - 1;
+
+        TRY_OR_SET_ERRNO(g_dynamic_env.try_append(nullptr), int, -1);
+        guard.deactivate();
+        _update_env();
+
+        environ[index] = string;
+
+        return 0;
+    }
+
     char* getenv(const char* key)
     {
         char* result;
diff --git a/tests/libc/TestEnv.cpp b/tests/libc/TestEnv.cpp
index b7f076a5..6652b677 100644
--- a/tests/libc/TestEnv.cpp
+++ b/tests/libc/TestEnv.cpp
@@ -72,6 +72,35 @@ TestResult test_setenv_before_and_after_manual_environ_modification()
     test_success;
 }
 
+TestResult test_putenv()
+{
+    clearenv();
+
+    char lit[] = "NAME=VALUE";
+
+    validate(putenv(lit) == 0);
+
+    char* value = getenv("NAME");
+    validate(value && !strcmp(value, "VALUE"));
+
+    lit[0] = 'F';
+
+    value = getenv("NAME");
+    validate(!value);
+
+    value = getenv("FAME");
+    validate(value && !strcmp(value, "VALUE"));
+
+    validate(unsetenv("FAME") == 0);
+
+    value = getenv("FAME");
+    validate(!value);
+
+    validate(!strcmp(lit, "FAME=VALUE"));
+
+    test_success;
+}
+
 Result<void> test_main()
 {
     test_prelude;
@@ -80,6 +109,7 @@ Result<void> test_main()
     run_test(test_setenv_dont_overwrite);
     run_test(test_unsetenv);
     run_test(test_setenv_before_and_after_manual_environ_modification);
+    run_test(test_putenv);
 
     return {};
 }