Add Random Pool

fw/Core/Hitcon/Logic/RandomPool.cc

@@ -0,0 +1,98 @@
+#include <Logic/RandomPool.h>
+#include <Service/Sched/PeriodicTask.h>
+#include <Service/Sched/Scheduler.h>
+
+using hitcon::service::sched::task_callback_t;
+
+namespace hitcon {
+
+FastRandomPool g_fast_random_pool;
+SecureRandomPool g_secure_random_pool;
+
+SecureRandomPool::SecureRandomPool()
+    : init_finished(false), seed_count(0),
+      routine_task(950, (task_callback_t)&SecureRandomPool::Routine, this, 20),
+      seed_queue_head(0), seed_queue_tail(0), random_queue_head(0),
+      random_queue_tail(0) {}
+
+void SecureRandomPool::Init() {
+  sha3_Init256(&keccak_context);
+  hitcon::service::sched::scheduler.Queue(&routine_task, nullptr);
+  hitcon::service::sched::scheduler.EnablePeriodic(&routine_task);
+  init_finished = true;
+}
+
+bool SecureRandomPool::Seed(uint64_t seed_val) {
+  if (seed_queue_head == (seed_queue_tail + 1) % kSeedQueueSize) {
+    // Queue is full
+    return false;
+  }
+  seed_queue[seed_queue_tail] = seed_val;
+  seed_queue_tail = (seed_queue_tail + 1) % kSeedQueueSize;
+  return true;
+}
+
+bool SecureRandomPool::GetRandom(uint64_t* res) {
+  if (!init_finished || seed_count < kMinSeedCountBeforeReady ||
+      random_queue_head == random_queue_tail) {
+    // Not ready or queue is empty
+    return false;
+  }
+  *res = random_queue[random_queue_head];
+  random_queue_head = (random_queue_head + 1) % kRandomQueueSize;
+  return true;
+}
+
+void SecureRandomPool::Routine(void* unused) {
+  if (TrySeed()) {
+    return;
+  }
+  if (TryRandom()) {
+    return;
+  }
+}
+
+bool SecureRandomPool::TrySeed() {
+  if (seed_queue_head == seed_queue_tail) {
+    // Queue is empty
+    return false;
+  }
+  uint64_t seed_val = seed_queue[seed_queue_head];
+  seed_queue_head = (seed_queue_head + 1) % kSeedQueueSize;
+  for (size_t i = 0; i < sizeof(uint64_t); i++) {
+    keccak_context.u.sb[i] ^= static_cast<uint8_t>(seed_val >> (8 * i));
+  }
+  keccakf(keccak_context.u.s);
+  seed_count++;
+  return true;
+}
+
+bool SecureRandomPool::TryRandom() {
+  if (seed_count < kMinSeedCountBeforeReady ||
+      random_queue_head == (random_queue_tail + 1) % kRandomQueueSize) {
+    // Not ready or queue is full
+    return false;
+  }
+  uint64_t random_val = 0;
+  for (size_t i = 0; i < sizeof(uint64_t); i++) {
+    random_val |= static_cast<uint64_t>(keccak_context.u.sb[i]) << (8 * i);
+  }
+  keccakf(keccak_context.u.s);
+  random_queue[random_queue_tail] = random_val;
+  random_queue_tail = (random_queue_tail + 1) % kRandomQueueSize;
+  return true;
+}
+
+FastRandomPool::FastRandomPool() : prng(0) {}
+
+void FastRandomPool::Init() {
+  // Left blank as it was not specified what should be done here,
+  // it might be used to set the seed to a certain value or initialize other
+  // resources.
+}
+
+uint32_t FastRandomPool::GetRandom() { return prng.GetRandom(); }
+
+void FastRandomPool::Seed(uint64_t seed) { prng.MixState(seed); }
+
+}  // namespace hitcon

fw/Core/Hitcon/Logic/RandomPool.h

@@ -0,0 +1,110 @@
+#ifndef LOGIC_RANDOM_POOL_DOT_H_
+#define LOGIC_RANDOM_POOL_DOT_H_
+
+#include <Logic/keccak.h>
+#include <Logic/pcg32.h>
+#include <Service/Sched/Scheduler.h>
+#include <stddef.h>
+#include <stdint.h>
+
+namespace hitcon {
+
+class SecureRandomPool;
+class FastRandomPool;
+
+// SecureRandomPool maintains a pool of entropy in a keccak state, and everytime
+// randomness is requested it'll run the keccakf() and output the bytes. This
+// should be used for anything that needs cryptography randomness and doesn't
+// mind being slow.
+class SecureRandomPool {
+ public:
+  SecureRandomPool();
+
+  void Init();
+
+  // Add entropy to the secure random pool.
+  // Returns true if successful, otherwise the pool is busy (or not ready yet)
+  // and the caller should retry later on. Generally this adds the seed to
+  // seeding queue and Routine() will take care of the rest.
+  bool Seed(uint64_t seed_val);
+
+  // Pull entropy from the secure random pool.
+  // Return true if successful and the randomness will be in res. Otherwise the
+  // pool is busy (or not ready yet) and the caller should retry later on.
+  // Generally this involves pulling from the random queue and Routine() is in
+  // charge of filling the queue.
+  bool GetRandom(uint64_t* res);
+
+  // Will be called routinely by the scheduler, and will try to empty the seed
+  // queue first then fill the random queue. Each invocation is limited to 1
+  // keccakf() run due to scheduling.
+  void Routine(void* unused);
+
+ private:
+  // Try to seed once into the state if there's seed in the queue.
+  // Seeding involves xor'ing the corresponding bytes in the state then run
+  // keccakf(). Return true if one such operation has been done and that signals
+  // to the caller that we can't run another keccakf() in the same Routine()
+  // run.
+  bool TrySeed();
+
+  // Try to pull one random out if we've enough entropy considering seed_count
+  // and there's room in the random queue. This involves getting the
+  // corresponding bytes in the state. After getting the bytes, keccakf() is
+  // invoked. Return true if one such operation has been done and that signals
+  // to the caller that we can't run another keccakf() in the same Routine()
+  // run.
+  bool TryRandom();
+
+  // Set to true after Init.
+  bool init_finished;
+
+  // The internal keccak state.
+  sha3_context keccak_context;
+
+  // Number of times we've been seeded.
+  int seed_count;
+
+  // Scheduler task for running Routine(), runs every 20ms with low priority
+  // (950).
+  hitcon::service::sched::PeriodicTask routine_task;
+
+  // A circular queue for holding the to be seeded values.
+  static constexpr size_t kSeedQueueSize = 8;
+  uint64_t seed_queue[kSeedQueueSize];
+  size_t seed_queue_head;
+  size_t seed_queue_tail;
+
+  // A circular queue for holding the newly minted random values.
+  static constexpr size_t kRandomQueueSize = 8;
+  uint64_t random_queue[kRandomQueueSize];
+  size_t random_queue_head;
+  size_t random_queue_tail;
+
+  // Must be seeded this amount of times before GetRandom() will be ready.
+  static constexpr int kMinSeedCountBeforeReady = 64;
+};
+
+// Fast random pool uses the PCG32 for faster but non-secure random generation.
+class FastRandomPool {
+ public:
+  FastRandomPool();
+
+  void Init();
+
+  // Never blocks, simply retrieves the random.
+  uint32_t GetRandom();
+
+  // Add to the PRNG state, aka MixState().
+  void Seed(uint64_t seed);
+
+ private:
+  PCG32 prng;
+};
+
+extern FastRandomPool g_fast_random_pool;
+extern SecureRandomPool g_secure_random_pool;
+
+}  // namespace hitcon
+
+#endif  // LOGIC_RANDOM_POOL_DOT_H_