// Copyright 2017 The BoringSSL Authors

//

// 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.

#include <openssl/ctrdrbg.h>

#include <assert.h>

#include <openssl/mem.h>

#include <algorithm>

#include "../../mem_internal.h"

#include "../aes/internal.h"

#include "../service_indicator/internal.h"

#include "internal.h"

using namespace bssl;

// Section references in this file refer to SP 800-90Ar1:

// http://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-90Ar1.pdf

// Also see table 3.

constexpr uint64_t kMaxReseedCount = UINT64_C(1) << 48;

// Implements the BCC function as described in Section 10.3.3.

static void bcc(uint8_t out[AES_BLOCK_SIZE], const AES_KEY *aes_key,

                const uint8_t *data, size_t data_len) {

  // 1. chaining_value = 0^outlen.

  uint8_t *chaining_value = out;

  OPENSSL_memset(chaining_value, 0, AES_BLOCK_SIZE);

  // 2. n = len (data)/outlen.

  BSSL_CHECK(data_len % AES_BLOCK_SIZE == 0);

  const size_t n = data_len / AES_BLOCK_SIZE;

  for (size_t i = 0; i < n; i++) {

    const uint8_t *block = data + (i * AES_BLOCK_SIZE);

    uint8_t input_block[AES_BLOCK_SIZE];

    // 4.1: input_block = chaining_value ⊕ block_i.

    CRYPTO_xor16(input_block, chaining_value, block);

    // 4.2: chaining_value = Block_Encrypt (Key, input_block).

    BCM_aes_encrypt(input_block, chaining_value, aes_key);

  }

  // 5. output_block = chaining_value.

}

// Implements the derivation function as described in Section 10.3.2.

static int block_cipher_df(uint8_t *out, size_t out_len, const uint8_t *input,

                           size_t input_len) {

  // Constants for AES-256

  constexpr size_t kAESKeyLen = 32;

  constexpr size_t kAESOutLen = AES_BLOCK_SIZE;

  constexpr size_t kMaxNumBits = 512;

  if (out_len > kMaxNumBits / 8 || input_len > (1u << 30)) {

    return 0;

  }

  // 4. S = L || N || input_string || 0x80.

  const size_t s_rawlen = sizeof(uint32_t) + sizeof(uint32_t) + input_len + 1;

  // S is padded up to a block size.

  const size_t s_len = (s_rawlen + kAESOutLen - 1) & ~(kAESOutLen - 1);

  uint8_t iv_plus_s[/* space used below */ kAESOutLen + 4 + 4 +

                    CTR_DRBG_MAX_ENTROPY_LEN + CTR_DRBG_NONCE_LEN +

                    CTR_DRBG_SEED_LEN + 1 +

                    /* padding */ 7];

  if (kAESOutLen + s_len > sizeof(iv_plus_s)) {

    return 0;

  }

  OPENSSL_memset(iv_plus_s, 0, sizeof(iv_plus_s));

  uint8_t *s_ptr = iv_plus_s + kAESOutLen;

  // 2. L = len (input_string)/8.

  CRYPTO_store_u32_be(s_ptr, (uint32_t)input_len);

  s_ptr += sizeof(uint32_t);

  // 3. N = number_of_bits_to_return/8.

  CRYPTO_store_u32_be(s_ptr, (uint32_t)out_len);

  s_ptr += sizeof(uint32_t);

  OPENSSL_memcpy(s_ptr, input, input_len);

  s_ptr += input_len;

  *s_ptr = 0x80;

  uint8_t temp[kAESKeyLen + kAESOutLen];

  size_t temp_len = 0;

  // 8. K = leftmost (0x00010203...1D1E1F, keylen).

  static const uint8_t kInitialKey[kAESKeyLen] = {

      0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a,

      0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15,

      0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f};

  AES_KEY aes_key;

  bcm_status status =

      BCM_aes_set_encrypt_key(kInitialKey, 8 * sizeof(kInitialKey), &aes_key);

  BSSL_CHECK(status != bcm_status::failure);

  // 7. i = 0.

  uint32_t i = 0;

  while (temp_len < sizeof(temp)) {

    // 9.1 IV = i || 0^(outlen - len(i)).

    CRYPTO_store_u32_be(iv_plus_s, i);

    // 9.2 temp = temp || BCC (K, (IV || S)).

    bcc(temp + temp_len, &aes_key, iv_plus_s, kAESOutLen + s_len);

    temp_len += kAESOutLen;

    // 9.3 i = i + 1.

    i++;

  }

  // 10. K = leftmost (temp, keylen).

  uint8_t *const k = temp;

  // 11. X = select (temp, keylen+1, keylen+outlen).

  uint8_t *const x = temp + kAESKeyLen;

  // 12. temp = the Null string.

  temp_len = 0;

  // Create an AES key schedule for the final encryption steps.

  status = BCM_aes_set_encrypt_key(k, kAESKeyLen * 8, &aes_key);

  BSSL_CHECK(status != bcm_status::failure);

  // 13. While len (temp) < number_of_bits_to_return, do:

  while (temp_len < out_len) {

    // 13.1 X = Block_Encrypt (K, X).

    BCM_aes_encrypt(x, x, &aes_key);

    // 13.2 temp = temp || X.

    size_t to_copy = std::min(kAESOutLen, out_len - temp_len);

    OPENSSL_memcpy(out + temp_len, x, to_copy);

    temp_len += to_copy;

  }

  return 1;

}

CTR_DRBG_STATE *CTR_DRBG_new(const uint8_t entropy[CTR_DRBG_ENTROPY_LEN],

                             const uint8_t *personalization,

                             size_t personalization_len) {

  CTR_DRBG_STATE *drbg = New<CTR_DRBG_STATE>();

  if (drbg == nullptr ||

      !CTR_DRBG_init(drbg, /*df=*/false, entropy, CTR_DRBG_ENTROPY_LEN,

                     /*nonce=*/nullptr, personalization, personalization_len)) {

    CTR_DRBG_free(drbg);

    return nullptr;

  }

  return drbg;

}

CTR_DRBG_STATE *CTR_DRBG_new_df(const uint8_t *entropy, size_t entropy_len,

                                const uint8_t nonce[CTR_DRBG_NONCE_LEN],

                                const uint8_t *personalization,

                                size_t personalization_len) {

  CTR_DRBG_STATE *drbg = New<CTR_DRBG_STATE>();

  if (drbg == nullptr ||

      !CTR_DRBG_init(drbg, /*df=*/true, entropy, entropy_len, nonce,

                     personalization, personalization_len)) {

    CTR_DRBG_free(drbg);

    return nullptr;

  }

  return drbg;

}

void CTR_DRBG_free(CTR_DRBG_STATE *state) { Delete(state); }

int bssl::CTR_DRBG_init(CTR_DRBG_STATE *drbg, int df, const uint8_t *entropy,

                        size_t entropy_len,

                        const uint8_t nonce[CTR_DRBG_NONCE_LEN],

                        const uint8_t *personalization,

                        size_t personalization_len) {

  // Section 10.2.1.3.1 and 10.2.1.3.2

  if (personalization_len > CTR_DRBG_SEED_LEN ||

      (!df && entropy_len != CTR_DRBG_ENTROPY_LEN) ||

      (df && (entropy_len < CTR_DRBG_MIN_ENTROPY_LEN ||

              entropy_len > CTR_DRBG_MAX_ENTROPY_LEN)) ||  //

      (df != (nonce != nullptr))) {

    return 0;

  }

  uint8_t seed_material[CTR_DRBG_SEED_LEN];

  if (df) {

    uint8_t pre_seed_material[CTR_DRBG_MAX_ENTROPY_LEN + CTR_DRBG_NONCE_LEN +

                              CTR_DRBG_SEED_LEN];

    OPENSSL_memcpy(pre_seed_material, entropy, entropy_len);

    OPENSSL_memcpy(pre_seed_material + entropy_len, nonce, CTR_DRBG_NONCE_LEN);

    OPENSSL_memcpy(pre_seed_material + entropy_len + CTR_DRBG_NONCE_LEN,

                   personalization, personalization_len);

    const size_t pre_seed_material_length =

        entropy_len + CTR_DRBG_NONCE_LEN + personalization_len;

    if (!block_cipher_df(seed_material, sizeof(seed_material),

                         pre_seed_material, pre_seed_material_length)) {

      return 0;

    }

  } else {

    OPENSSL_memcpy(seed_material, entropy, CTR_DRBG_ENTROPY_LEN);

    for (size_t i = 0; i < personalization_len; i++) {

      seed_material[i] ^= personalization[i];

    }

  }

  // Section 10.2.1.2

  // kInitMask is the result of encrypting blocks with big-endian value 1, 2

  // and 3 with the all-zero AES-256 key.

  static const uint8_t kInitMask[CTR_DRBG_SEED_LEN] = {

      0x53, 0x0f, 0x8a, 0xfb, 0xc7, 0x45, 0x36, 0xb9, 0xa9, 0x63, 0xb4, 0xf1,

      0xc4, 0xcb, 0x73, 0x8b, 0xce, 0xa7, 0x40, 0x3d, 0x4d, 0x60, 0x6b, 0x6e,

      0x07, 0x4e, 0xc5, 0xd3, 0xba, 0xf3, 0x9d, 0x18, 0x72, 0x60, 0x03, 0xca,

      0x37, 0xa6, 0x2a, 0x74, 0xd1, 0xa2, 0xf5, 0x8e, 0x75, 0x06, 0x35, 0x8e,

  };

  for (size_t i = 0; i < sizeof(kInitMask); i++) {

    seed_material[i] ^= kInitMask[i];

  }

  drbg->df = df;

  drbg->ctr =

      aes_ctr_set_key(&drbg->ks, nullptr, &drbg->block, seed_material, 32);

  OPENSSL_memcpy(drbg->counter, seed_material + 32, 16);

  drbg->reseed_counter = 1;

  return 1;

}

static_assert(CTR_DRBG_SEED_LEN % AES_BLOCK_SIZE == 0,

              "not a multiple of AES block size");

// ctr_inc adds |n| to the last four bytes of |drbg->counter|, treated as a

// big-endian number.

static void ctr32_add(CTR_DRBG_STATE *drbg, uint32_t n) {

  uint32_t ctr = CRYPTO_load_u32_be(drbg->counter + 12);

  CRYPTO_store_u32_be(drbg->counter + 12, ctr + n);

}

static int ctr_drbg_update(CTR_DRBG_STATE *drbg,

                           const uint8_t data[CTR_DRBG_SEED_LEN]) {

  uint8_t temp[CTR_DRBG_SEED_LEN];

  for (size_t i = 0; i < CTR_DRBG_SEED_LEN; i += AES_BLOCK_SIZE) {

    ctr32_add(drbg, 1);

    drbg->block(drbg->counter, temp + i, &drbg->ks);

  }

  for (size_t i = 0; i < CTR_DRBG_SEED_LEN; i++) {

    temp[i] ^= data[i];

  }

  drbg->ctr = aes_ctr_set_key(&drbg->ks, nullptr, &drbg->block, temp, 32);

  OPENSSL_memcpy(drbg->counter, temp + 32, 16);

  return 1;

}

int CTR_DRBG_reseed(CTR_DRBG_STATE *drbg,

                    const uint8_t entropy[CTR_DRBG_ENTROPY_LEN],

                    const uint8_t *additional_data,

                    size_t additional_data_len) {

  return CTR_DRBG_reseed_ex(drbg, entropy, CTR_DRBG_ENTROPY_LEN,

                            additional_data, additional_data_len);

}

int CTR_DRBG_reseed_ex(CTR_DRBG_STATE *drbg, const uint8_t *entropy,

                       size_t entropy_len, const uint8_t *additional_data,

                       size_t additional_data_len) {

  if (additional_data_len > CTR_DRBG_SEED_LEN ||

      (drbg->df && (entropy_len > CTR_DRBG_MAX_ENTROPY_LEN ||

                    entropy_len < CTR_DRBG_MIN_ENTROPY_LEN)) ||

      (!drbg->df && entropy_len != CTR_DRBG_ENTROPY_LEN)) {

    return 0;

  }

  uint8_t seed_material[CTR_DRBG_SEED_LEN];

  if (drbg->df) {

    // Section 10.2.1.4.2

    uint8_t pre_seed_material[CTR_DRBG_MAX_ENTROPY_LEN + CTR_DRBG_SEED_LEN];

    static_assert(CTR_DRBG_MAX_ENTROPY_LEN <= sizeof(pre_seed_material));

    OPENSSL_memcpy(pre_seed_material, entropy, entropy_len);

    OPENSSL_memcpy(pre_seed_material + entropy_len, additional_data,

                   additional_data_len);

    const size_t pre_seed_material_len = entropy_len + additional_data_len;

    if (!block_cipher_df(seed_material, sizeof(seed_material),

                         pre_seed_material, pre_seed_material_len)) {

      return 0;

    }

  } else {

    // Section 10.2.1.4

    static_assert(CTR_DRBG_ENTROPY_LEN == sizeof(seed_material));

    OPENSSL_memcpy(seed_material, entropy, CTR_DRBG_ENTROPY_LEN);

    if (additional_data_len > 0) {

      for (size_t i = 0; i < additional_data_len; i++) {

        seed_material[i] ^= additional_data[i];

      }

    }

  }

  if (!ctr_drbg_update(drbg, seed_material)) {

    return 0;

  }

  drbg->reseed_counter = 1;

  return 1;

}

int CTR_DRBG_generate(CTR_DRBG_STATE *drbg, uint8_t *out, size_t out_len,

                      const uint8_t *additional_data,

                      size_t additional_data_len) {

  // See 9.3.1

  if (out_len > CTR_DRBG_MAX_GENERATE_LENGTH) {

    return 0;

  }

  // See 10.2.1.5.1

  if (drbg->reseed_counter > kMaxReseedCount) {

    return 0;

  }

  uint8_t processed_additional_data[CTR_DRBG_SEED_LEN];

  OPENSSL_memset(processed_additional_data, 0,

                 sizeof(processed_additional_data));

  if (additional_data_len != 0) {

    if (drbg->df) {

      if (!block_cipher_df(processed_additional_data,

                           sizeof(processed_additional_data), additional_data,

                           additional_data_len)) {

        return 0;

      }

    } else {

      if (additional_data_len > sizeof(processed_additional_data)) {

        return 0;

      }

      OPENSSL_memcpy(processed_additional_data, additional_data,

                     additional_data_len);

    }

    if (!ctr_drbg_update(drbg, processed_additional_data)) {

      return 0;

    }

  }

  // kChunkSize is used to interact better with the cache. Since the AES-CTR

  // code assumes that it's encrypting rather than just writing keystream, the

  // buffer has to be zeroed first. Without chunking, large reads would zero

  // the whole buffer, flushing the L1 cache, and then do another pass (missing

  // the cache every time) to “encrypt” it. The code can avoid this by

  // chunking.

  constexpr size_t kChunkSize = 8 * 1024;

  while (out_len >= AES_BLOCK_SIZE) {

    size_t todo = kChunkSize;

    if (todo > out_len) {

      todo = out_len;

    }

    todo &= ~(AES_BLOCK_SIZE - 1);

    const size_t num_blocks = todo / AES_BLOCK_SIZE;

    OPENSSL_memset(out, 0, todo);

    ctr32_add(drbg, 1);

    drbg->ctr(out, out, num_blocks, &drbg->ks, drbg->counter);

    ctr32_add(drbg, (uint32_t)(num_blocks - 1));

    out += todo;

    out_len -= todo;

  }

  if (out_len > 0) {

    uint8_t block[AES_BLOCK_SIZE];

    ctr32_add(drbg, 1);

    drbg->block(drbg->counter, block, &drbg->ks);

    OPENSSL_memcpy(out, block, out_len);

  }

  if (!ctr_drbg_update(drbg, processed_additional_data)) {

    return 0;

  }

  drbg->reseed_counter++;

  FIPS_service_indicator_update_state();

  return 1;

}

void CTR_DRBG_clear(CTR_DRBG_STATE *drbg) {

  OPENSSL_cleanse(drbg, sizeof(CTR_DRBG_STATE));

}