folly/Random.cpp

   1 /*
   2  * Copyright 2015 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #include <folly/Random.h>
  18
  19 #include <atomic>
  20 #include <unistd.h>
  21 #include <sys/time.h>
  22 #include <mutex>
  23 #include <random>
  24 #include <array>
  25
  26 #include <glog/logging.h>
  27 #include <folly/File.h>
  28 #include <folly/FileUtil.h>
  29
  30 #ifdef _MSC_VER
  31 # include <wincrypt.h>
  32 #endif
  33
  34 namespace folly {
  35
  36 namespace {
  37
  38 void readRandomDevice(void* data, size_t size) {
  39 #ifdef _MSC_VER
  40   static std::once_flag flag;
  41   static HCRYPTPROV cryptoProv;
  42   std::call_once(flag, [&] {
  43     PCHECK(CryptAcquireContext(&cryptoProv, nullptr, nullptr,
  44                                PROV_RSA_FULL, 0));
  45   });
  46   CHECK(size <= std::numeric_limits<DWORD>::max());
  47   PCHECK(CryptGenRandom(cryptoProv, (DWORD)size, (BYTE*)data));
  48 #else
  49   // Keep the random device open for the duration of the program.
  50   static int randomFd = ::open("/dev/urandom", O_RDONLY);
  51   PCHECK(randomFd >= 0);
  52   auto bytesRead = readFull(randomFd, data, size);
  53   PCHECK(bytesRead >= 0 && size_t(bytesRead) == size);
  54 #endif
  55 }
  56
  57 class BufferedRandomDevice {
  58  public:
  59   static constexpr size_t kDefaultBufferSize = 128;
  60
  61   explicit BufferedRandomDevice(size_t bufferSize = kDefaultBufferSize);
  62
  63   void get(void* data, size_t size) {
  64     if (LIKELY(size <= remaining())) {
  65       memcpy(data, ptr_, size);
  66       ptr_ += size;
  67     } else {
  68       getSlow(static_cast<unsigned char*>(data), size);
  69     }
  70   }
  71
  72  private:
  73   void getSlow(unsigned char* data, size_t size);
  74
  75   inline size_t remaining() const {
  76     return buffer_.get() + bufferSize_ - ptr_;
  77   }
  78
  79   const size_t bufferSize_;
  80   std::unique_ptr<unsigned char[]> buffer_;
  81   unsigned char* ptr_;
  82 };
  83
  84 BufferedRandomDevice::BufferedRandomDevice(size_t bufferSize)
  85   : bufferSize_(bufferSize),
  86     buffer_(new unsigned char[bufferSize]),
  87     ptr_(buffer_.get() + bufferSize) {  // refill on first use
  88 }
  89
  90 void BufferedRandomDevice::getSlow(unsigned char* data, size_t size) {
  91   DCHECK_GT(size, remaining());
  92   if (size >= bufferSize_) {
  93     // Just read directly.
  94     readRandomDevice(data, size);
  95     return;
  96   }
  97
  98   size_t copied = remaining();
  99   memcpy(data, ptr_, copied);
 100   data += copied;
 101   size -= copied;
 102
 103   // refill
 104   readRandomDevice(buffer_.get(), bufferSize_);
 105   ptr_ = buffer_.get();
 106
 107   memcpy(data, ptr_, size);
 108   ptr_ += size;
 109 }
 110
 111
 112 }  // namespace
 113
 114 void Random::secureRandom(void* data, size_t size) {
 115   static ThreadLocal<BufferedRandomDevice> bufferedRandomDevice;
 116   bufferedRandomDevice->get(data, size);
 117 }
 118
 119 ThreadLocalPRNG::ThreadLocalPRNG() {
 120   static folly::ThreadLocal<ThreadLocalPRNG::LocalInstancePRNG> localInstance;
 121   local_ = localInstance.get();
 122 }
 123
 124 class ThreadLocalPRNG::LocalInstancePRNG {
 125  public:
 126   LocalInstancePRNG() : rng(Random::create()) { }
 127
 128   Random::DefaultGenerator rng;
 129 };
 130
 131 uint32_t ThreadLocalPRNG::getImpl(LocalInstancePRNG* local) {
 132   return local->rng();
 133 }
 134
 135 }