Luna/libc/src/env.cpp

#include <errno.h>
#include <luna/ScopeGuard.h>
#include <luna/Vector.h>
#include <stdlib.h>

extern "C" char** environ;
Vector<char*> g_dynamic_env = {};
bool env_is_dynamic { false };

/* If an environment variable matching key is found, stores its value in value and returns its index. Otherwise,
 * returns -1. */
static isize _findenv(const char* key, char** value)
{
    char** env = environ;
    size_t len = strlen(key);

    if (!env) return -1;

    while (*env)
    {
        // Retrieve an environment variable from environ.
        char* var = *(env++);

        // Check for an equals sign, else we just skip this one.
        char* delim = strchr(var, '=');
        if (!delim) continue;

        // Get the length of this variable's key (the part before the equals sign)
        size_t key_length = strcspn(var, "=");
        if (len != key_length) continue;

        // If the keys match, we found it! Return the value (which is just after the equals sign)
        if (!memcmp(key, var, key_length))
        {
            if (value) *value = delim + 1;
            return (env - environ) - 1;
        }
    }

    return -1;
}

/* 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)
    {
        guard.deactivate();
        env_is_dynamic = true;
        environ = g_dynamic_env.data();
        check(!environ);
        return {};
    }

    while (*env)
    {
        char* ptr = strdup(*(env++));
        if (!ptr) return err(errno);

        auto ptr_guard = make_scope_guard([=] { free(ptr); });
        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();

    return {};
}

static int _move_env()
{
    auto rc = _try_move_env();
    if (rc.has_error())
    {
        errno = rc.error();
        return -1;
    }
    return 0;
}

static void _update_env()
{
    environ = g_dynamic_env.data();
}

extern "C"
{
    char** environ = nullptr;

    int clearenv()
    {
        if (env_is_dynamic)
        {
            for (auto element : g_dynamic_env)
            {
                if (element) free(element);
            }
        }

        env_is_dynamic = false;
        environ = nullptr;

        return 0;
    }

    int unsetenv(const char* key)
    {
        if (!key || *key == 0 || strchr(key, '='))
        {
            errno = EINVAL;
            return -1;
        }

        auto index = _findenv(key, nullptr);
        if (index < 0) return 0;

        if (!env_is_dynamic)
        {
            if (_move_env() < 0) return -1;
        }

        g_dynamic_env.remove_at(index);
        _update_env();

        return 0;
    }

    int setenv(const char* key, const char* value, int overwrite)
    {
        if (!key || *key == 0 || strchr(key, '='))
        {
            errno = EINVAL;
            return -1;
        }

        auto index = _findenv(key, nullptr);
        if (index >= 0 && !overwrite) return 0;

        if (!env_is_dynamic)
        {
            if (_move_env() < 0) return -1;
        }

        usize len = strlen(key) + strlen(value) + 2;

        char* str = (char*)calloc(len, 1);
        if (!str) return -1;
        auto guard = make_scope_guard([=] { free(str); });

        strncpy(str, key, strlen(key) + 1);
        strncat(str, "=", 2);
        strncat(str, value, strlen(value) + 1);

        if (index >= 0)
        {
            free(environ[index]);
            environ[index] = str;
            guard.deactivate();
            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;

        auto rc = g_dynamic_env.try_append(nullptr);
        if (rc.has_error())
        {
            errno = rc.error();
            return -1;
        }
        guard.deactivate();
        _update_env();

        environ[index] = str;

        return 0;
    }

    char* getenv(const char* key)
    {
        char* result;
        if (_findenv(key, &result) < 0) return nullptr;
        return result;
    }
}