/*
 * Copyright 2018 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef FCP_SECAGG_SHARED_AES_CTR_PRNG_H_
#define FCP_SECAGG_SHARED_AES_CTR_PRNG_H_

#include <cstdint>
#include <string>

#include "fcp/secagg/shared/aes_key.h"
#include "fcp/secagg/shared/prng.h"
#include "openssl/evp.h"

namespace fcp {
namespace secagg {

// A cryptographically strong Deterministic Pseudorandom Number Generator based
// on AES-CTR in OpenSSL. The seed must be supplied by the user.
//
// This code is used for the very specific purpose of generating *reproduceable*
// numbers which appear pseudorandom to any party not in possession of the input
// seed. DO NOT use AesCtrPrng in any situation where real randomness would be
// useful, because it uses no entropy at all and will always produce the same
// numbers if given the same seed.
//
// A single instance of AesCtrPrng can generate up to 2^36 bytes of
// pseudorandom output. If more than 2^36 bytes of output are needed, multiple
// instances of AesCtrPrng with different seeds should be used.
//
// This class is not thread-safe.

class AesCtrPrng : public SecureBatchPrng {
 public:
  // This class should only be instantiated via AesCtrPrngFactory.
  friend class AesCtrPrngFactory;

  ~AesCtrPrng() override;

  // Returns a new pseudorandom number of the specified size, generating new
  // pseudorandom bytes as needed.
  uint8_t Rand8() override;
  uint64_t Rand64() override;

  // Fills the provided buffer with pseudorandom bytes. Returns the number of
  // bytes that has been generated, which can be smaller than the requested
  // buffer_size if it exceeds the maximum buffer size returned by
  // GetMaxBufferSize().
  int RandBuffer(uint8_t* buffer, int buffer_size) override;

  // Get the maximum size of a buffer that can be filled by RandBuffer() in a
  // single call.
  size_t GetMaxBufferSize() const override { return kCacheSize; }

 private:
  static constexpr size_t kIvSize = 16;  // IV size, in bytes

  // Constructs the PRNG with the given seed, and an IV of all zeroes.
  // This is ONLY secure if the seed is never used more than once.
  explicit AesCtrPrng(const AesKey& seed);

  // Number of AES blocks in the cache.
  // The number of blocks is optimized to make kCacheSize to be a multiple
  // of any possible number of bytes in a SecAgg output (i.e. 1 to 8).
  static constexpr size_t kBatchSize = 3 * 5 * 7;

  // Block size, in bytes
  static constexpr size_t kBlockSize = 16;

  // Size of our cache, in bytes. We cache blocks to save leftover bytes.
  static constexpr int kCacheSize = kBatchSize * kBlockSize;

  // For security, we don't want to generate more than 2^32-1 blocks.
  static constexpr size_t kMaxBlocks = 0xFFFFFFFF;

  // Fills the selected cache with deterministic pseudorandomly generated bytes.
  // After this, the associated next_byte_pos counter must be set to 0.
  void GenerateBytes(uint8_t* cache, int cache_size);

  // Cache used by both Rand8() and Rand64()
  uint8_t cache_[kCacheSize];
  size_t next_byte_pos_;

  // This is used to generate bytes.
  static constexpr uint8_t kAllZeroes[kCacheSize] = {0};

  EVP_CIPHER_CTX* ctx_;
  size_t blocks_generated_;
};

}  // namespace secagg
}  // namespace fcp

#endif  // FCP_SECAGG_SHARED_AES_CTR_PRNG_H_