ydb/util/random/entropy.cpp

#include "fast.h"
#include "random.h"
#include "entropy.h"
#include "mersenne.h"
#include "shuffle.h"
#include "init_atfork.h"

#include <util/stream/output.h>
#include <util/stream/mem.h>
#include <util/stream/zlib.h>
#include <util/stream/buffer.h>

#include <util/system/fs.h>
#include <util/system/info.h>
#include <util/system/spinlock.h>
#include <util/system/thread.h>
#include <util/system/execpath.h>
#include <util/system/datetime.h>
#include <util/system/hostname.h>
#include <util/system/getpid.h>
#include <util/system/mem_info.h>
#include <util/system/rusage.h>
#include <util/system/cpu_id.h>
#include <util/system/unaligned_mem.h>

#include <util/generic/buffer.h>
#include <util/generic/singleton.h>

#include <util/digest/murmur.h>
#include <util/digest/city.h>

#include <util/ysaveload.h>

namespace {
    inline void Permute(char* buf, size_t len, ui32 seed) noexcept {
        Shuffle(buf, buf + len, TReallyFastRng32(seed));
    }

    struct THostEntropy: public TBuffer {
        inline THostEntropy() {
            {
                TBufferOutput buf(*this);
                TZLibCompress out(&buf);

                Save(&out, GetPID());
                Save(&out, GetCycleCount());
                Save(&out, MicroSeconds());
                Save(&out, TThread::CurrentThreadId());
                Save(&out, NSystemInfo::CachedNumberOfCpus());
                Save(&out, NSystemInfo::TotalMemorySize());
                Save(&out, HostName());

                try {
                    Save(&out, GetExecPath());
                } catch (...) {
                    // workaround - sometimes fails on FreeBSD
                }

                Save(&out, (size_t)Data());
                Save(&out, (size_t)&buf);

                {
                    double la[3];

                    NSystemInfo::LoadAverage(la, Y_ARRAY_SIZE(la));

                    out.Write(la, sizeof(la));
                }

                {
                    auto mi = NMemInfo::GetMemInfo();

                    out.Write(&mi, sizeof(mi));
                }

                {
                    auto ru = TRusage::Get();

                    out.Write(&ru, sizeof(ru));
                }

                {
                    ui32 store[12];

                    out << TStringBuf(CpuBrand(store));
                }

                out << NFs::CurrentWorkingDirectory();

                out.Finish();
            }

            {
                TMemoryOutput out(Data(), Size());

                // replace zlib header with hash
                Save(&out, CityHash64(Data(), Size()));
            }

            Permute(Data(), Size(), MurmurHash<ui32>(Data(), Size()));
        }
    };

    // not thread-safe
    class TMersenneInput: public IInputStream {
        using TKey = ui64;
        using TRnd = TMersenne<TKey>;

    public:
        inline explicit TMersenneInput(const TBuffer& rnd)
            : Rnd_((const TKey*)rnd.Data(), rnd.Size() / sizeof(TKey))
        {
        }

        ~TMersenneInput() override = default;

        size_t DoRead(void* buf, size_t len) override {
            size_t toRead = len;

            while (toRead) {
                const TKey next = Rnd_.GenRand();
                const size_t toCopy = Min(toRead, sizeof(next));

                memcpy(buf, &next, toCopy);

                buf = (char*)buf + toCopy;
                toRead -= toCopy;
            }

            return len;
        }

    private:
        TRnd Rnd_;
    };

    class TEntropyPoolStream: public IInputStream {
    public:
        inline explicit TEntropyPoolStream(const TBuffer& buffer)
            : Mi_(buffer)
            , Bi_(&Mi_, 8192)
        {
        }

        size_t DoRead(void* buf, size_t len) override {
            auto guard = Guard(Mutex_);

            return Bi_.Read(buf, len);
        }

    private:
        TAdaptiveLock Mutex_;
        TMersenneInput Mi_;
        TBufferedInput Bi_;
    };

    struct TSeedStream: public IInputStream {
        size_t DoRead(void* inbuf, size_t len) override {
            char* buf = (char*)inbuf;

#define DO_STEP(type)                                    \
    while (len >= sizeof(type)) {                        \
        WriteUnaligned<type>(buf, RandomNumber<type>()); \
        buf += sizeof(type);                             \
        len -= sizeof(type);                             \
    }

            DO_STEP(ui64);
            DO_STEP(ui32);
            DO_STEP(ui16);
            DO_STEP(ui8);

#undef DO_STEP

            return buf - (char*)inbuf;
        }
    };

    struct TDefaultTraits {
        THolder<TEntropyPoolStream> EP;
        TSeedStream SS;

        inline TDefaultTraits() {
            Reset();
        }

        inline IInputStream& EntropyPool() noexcept {
            return *EP;
        }

        inline IInputStream& Seed() noexcept {
            return SS;
        }

        inline void Reset() noexcept {
            EP.Reset(new TEntropyPoolStream(THostEntropy()));
        }

        static inline TDefaultTraits& Instance() {
            auto res = SingletonWithPriority<TDefaultTraits, 0>();

            RNGInitAtForkHandlers();

            return *res;
        }
    };

    using TRandomTraits = TDefaultTraits;
}

IInputStream& EntropyPool() {
    return TRandomTraits::Instance().EntropyPool();
}

IInputStream& Seed() {
    return TRandomTraits::Instance().Seed();
}

void ResetEntropyPool() {
    TRandomTraits::Instance().Reset();
}