/// @file dpf/random.hpp
/// @brief
/// @details
/// @author Ryan Henry <ryan.henry@ucalgary.ca>
/// @copyright Copyright (c) 2019-2024 Ryan Henry and [others](@ref authors)
/// @license Released under a GNU General Public v2.0 (GPLv2) license;
///          see [LICENSE.md](@ref license) for details.

#ifndef LIBDPF_INCLUDE_DPF_RANDOM_HPP__
#define LIBDPF_INCLUDE_DPF_RANDOM_HPP__

#include <bsd/stdlib.h>

#include <cerrno>
#include <cstddef>
#include <cstdio>
#include <cstring>
#include <exception>
#include <fcntl.h>
#include <mutex>
#include <type_traits>
#include <unistd.h>
#include <utility>

#include "hedley/hedley.h"

#include "dpf/secret_share.hpp"

namespace dpf
{

namespace detail
{

/// When set, `uniform_fill` copies from this hook and does not read the
/// system RNG. Used to feed the same beaver coins to dealer `make_dpf` and
/// Doerner–Shelat gen. Null in normal use.
inline thread_local void (*uniform_bytes_hook)(void *, std::size_t) = nullptr;

template <typename T>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
bool fill_from_hook(T & buf) noexcept
{
    if (uniform_bytes_hook == nullptr)
    {
        return false;
    }
    uniform_bytes_hook(&buf, sizeof(buf));
    return true;
}

/// `bool` and `enum : bool` (including `dpf::bit`) have only two valid
/// representations. Filling them with a raw entropy byte is undefined.
template <typename T>
HEDLEY_NO_THROW
constexpr bool is_boolean_representation() noexcept
{
    using U = std::remove_cv_t<T>;
    if constexpr (std::is_same_v<U, bool>)
    {
        return true;
    }
    else if constexpr (std::is_enum_v<U>)
    {
        return std::is_same_v<std::underlying_type_t<U>, bool>;
    }
    else
    {
        return false;
    }
}

#if !defined(LIBDPF_USE_ARC4RANDOM)

/// One unbuffered, exclusively locked read of the entropy device.
/// Buffering would copy unread bytes into a `fork()` child, so parent and
/// child would repeat the same key material. The lock keeps concurrent
/// `fread` calls off the shared `FILE`.
struct entropy_source
{
#if defined(LIBDPF_USE_DEV_RANDOM)
    static constexpr const char * path = "/dev/random";
    static constexpr const char * open_error = "dpf: cannot open /dev/random\n";
#else
    static constexpr const char * path = "/dev/urandom";
    static constexpr const char * open_error = "dpf: cannot open /dev/urandom\n";
#endif

    FILE * fp = nullptr;
    std::mutex mu;

    entropy_source() = default;
    entropy_source(const entropy_source &) = delete;
    entropy_source & operator=(const entropy_source &) = delete;
    entropy_source(entropy_source &&) = delete;
    entropy_source & operator=(entropy_source &&) = delete;

    HEDLEY_NO_THROW
    ~entropy_source() noexcept
    {
        if (fp != nullptr)
        {
            std::fclose(fp);
        }
    }

    void open_unlocked()
    {
        if (fp != nullptr)
        {
            return;
        }
        fp = std::fopen(path, "rb");
        if (fp == nullptr)
        {
            std::fputs(open_error, stderr);
            std::terminate();
        }
        // Before any read. A buffered FILE duplicates entropy across fork().
        if (std::setvbuf(fp, nullptr, _IONBF, 0) != 0)
        {
            std::fclose(fp);
            fp = nullptr;
            std::fputs("dpf: cannot disable entropy buffering\n", stderr);
            std::terminate();
        }
        int fd = ::fileno(fp);
        if (fd >= 0)
        {
            ::fcntl(fd, F_SETFD, FD_CLOEXEC);
        }
    }

    void read(void * dst, std::size_t n)
    {
        std::lock_guard<std::mutex> lock(mu);
        if (fp == nullptr)
        {
            open_unlocked();
        }
        auto * p = static_cast<unsigned char *>(dst);
        while (n > 0)
        {
            std::size_t got = std::fread(p, 1, n, fp);
            if (got == 0)
            {
                if (std::ferror(fp) && errno == EINTR)
                {
                    std::clearerr(fp);
                    continue;
                }
                std::fputs("dpf: entropy read failed\n", stderr);
                std::terminate();
            }
            p += got;
            n -= got;
        }
    }
};

inline entropy_source & entropy()
{
    static entropy_source source;
    return source;
}

#endif  // !LIBDPF_USE_ARC4RANDOM

}  // namespace detail

template <typename T>
HEDLEY_NO_THROW
auto & uniform_fill(T & buf) noexcept  // NOLINT(runtime/references)
{
    static_assert(std::is_trivially_copyable_v<std::remove_cv_t<T>>,
        "uniform_fill requires a trivially copyable type");

    if constexpr (detail::is_boolean_representation<T>())
    {
        unsigned char raw = 0;
        uniform_fill(raw);
        buf = static_cast<T>(static_cast<bool>(raw & 1u));
        return buf;
    }
    else
    {
        if (detail::fill_from_hook(buf)) return buf;
#if defined(LIBDPF_USE_ARC4RANDOM)
        arc4random_buf(&buf, sizeof(buf));
#else
        detail::entropy().read(&buf, sizeof(buf));
#endif
        return buf;
    }
}

template <typename T>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
auto uniform_sample() noexcept
{
    using U = std::remove_cv_t<T>;
    U buf;
    uniform_fill(buf);
    return buf;
}

template <typename T>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
auto additively_share(T secret) noexcept
{
    using T_ = std::remove_cv_t<std::remove_reference_t<T>>;
    T_ tmp = uniform_sample<T_>();
    // Signed subtraction overflows for extreme shares. Subtract in the
    // unsigned width and copy the bits back so the group is mod 2^n.
    T_ other;
    if constexpr (std::is_integral_v<T_> && std::is_signed_v<T_>)
    {
        using U = std::make_unsigned_t<T_>;
        U diff = static_cast<U>(static_cast<T_>(secret)) - static_cast<U>(tmp);
        std::memcpy(&other, &diff, sizeof(other));
    }
    else
    {
        other = static_cast<T_>(static_cast<T_>(secret) - tmp);
    }
    return std::make_pair(
        additive_share<T_, 0>::from_raw(tmp),
        additive_share<T_, 1>::from_raw(other));
}

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_RANDOM_HPP__