carb/extras/EnvironmentVariable.h

File members: carb/extras/EnvironmentVariable.h

// Copyright (c) 2018-2024, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto. Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
//
//
#pragma once

#include "../Defines.h"
#include "../cpp/Optional.h"

#include <cstring>
#include <string>
#include <utility>

#if CARB_PLATFORM_LINUX
#    include <cstdlib>
#endif

#if CARB_PLATFORM_WINDOWS
#    include "../CarbWindows.h"
#    include "../../omni/extras/ScratchBuffer.h"
#    include "Unicode.h"
#    include <vector>
#endif

namespace carb
{
namespace extras
{

class EnvironmentVariable final
{
    CARB_PREVENT_COPY(EnvironmentVariable);

    template <typename T>
    using optional = ::carb::cpp::optional<T>;

public:
    EnvironmentVariable() = delete;

    EnvironmentVariable(std::string name)
        : m_name(std::move(name)), m_restore(false), m_restoreValue(carb::cpp::nullopt)
    {
        CARB_ASSERT(!m_name.empty());
    }

    EnvironmentVariable(std::string name, const optional<const std::string>& value)
        : m_name(std::move(name)), m_restore(false), m_restoreValue(carb::cpp::nullopt)
    {
        CARB_ASSERT(!m_name.empty());

        // attempt to get and store current value
        std::string currentValue;
        if (getValue(m_name.c_str(), currentValue))
        {
            m_restoreValue = currentValue;
        }

        // attempt to set to new value
        if (setValue(m_name.c_str(), value ? value->c_str() : nullptr))
        {
            m_restore = true;
        }
    }

    ~EnvironmentVariable()
    {
        if (m_restore)
        {
            bool ret = setValue(m_name.c_str(), m_restoreValue ? m_restoreValue->c_str() : nullptr);
            CARB_ASSERT(ret);
            CARB_UNUSED(ret);
        }
    }

    EnvironmentVariable(EnvironmentVariable&& other)
        : m_name(std::move(other.m_name)),
          m_restore(std::exchange(other.m_restore, false)),
          m_restoreValue(std::move(other.m_restoreValue))
    {
    }

    EnvironmentVariable& operator=(EnvironmentVariable&& other)
    {
        m_name = std::move(other.m_name);
        m_restore = std::exchange(other.m_restore, false);
        m_restoreValue = std::move(other.m_restoreValue);
        return *this;
    }

    const std::string& getName() const noexcept
    {
        return m_name;
    }

    optional<const std::string> getValue() const noexcept
    {
        optional<const std::string> result;
        std::string value;
        if (getValue(m_name.c_str(), value))
        {
            result = value;
        }
        return result;
    }

    static bool setValue(const char* name, const char* value)
    {
        bool result;

        // Set the new value
#if CARB_PLATFORM_WINDOWS
        std::wstring nameWide = convertUtf8ToWide(name);
        if (value)
        {
            std::wstring valueWide = convertUtf8ToWide(value);
            result = (SetEnvironmentVariableW(nameWide.c_str(), valueWide.c_str()) != 0);
        }
        else
        {
            result = (SetEnvironmentVariableW(nameWide.c_str(), nullptr) != 0);
        }
#else
        if (value)
        {
            result = (setenv(name, value, /*overwrite=*/1) == 0);
        }
        else
        {
            result = (unsetenv(name) == 0);
        }
#endif
        return result;
    }

    static bool getValue(const char* name, std::string& value)
    {
#if CARB_PLATFORM_WINDOWS
        std::wstring nameWide = convertUtf8ToWide(name);

        omni::extras::ScratchBuffer<wchar_t> buffer;
        while (true)
        {
            DWORD count = GetEnvironmentVariableW(nameWide.c_str(), buffer.data(), static_cast<DWORD>(buffer.size()));
            if (count == 0)
            {
                if (GetLastError() == 0)
                {
                    // no error but no count means we got an empty env var
                    value = "";
                    return true;
                }
                else
                {
                    // the error case here means the env var does not exist
                    return false;
                }
            }
            else if (count < buffer.size())
            {
                // success case
                value = convertWideToUtf8(buffer.data());
                return true;
            }
            else
            {
                // env var is too large for the current buffer -- grow it and try again
                buffer.resize(count + 1U); // include null termination
                continue;
            }
        }
#else
        const char* buffer = getenv(name);
        if (buffer == nullptr)
        {
            return false;
        }
        value = buffer; // Copy string
#endif
        return true;
    }

private:
    std::string m_name;
    bool m_restore;
    optional<const std::string> m_restoreValue;
};

} // namespace extras
} // namespace carb