// Copyright 2015 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "components/gcm_driver/crypto/gcm_message_cryptographer.h"

#include <stddef.h>
#include <stdint.h>

#include <algorithm>
#include <sstream>

#include "base/logging.h"
#include "base/numerics/safe_math.h"
#include "base/strings/string_util.h"
#include "base/sys_byteorder.h"
#include "crypto/hkdf.h"
#include "third_party/boringssl/src/include/openssl/aead.h"

namespace gcm {
namespace {

// Size, in bytes, of the nonce for a record. This must be at least the size
// of a uint64_t, which is used to indicate the record sequence number.
const uint64_t kNonceSize = 12;

// The default record size as defined by httpbis-encryption-encoding-06.
const size_t kDefaultRecordSize = 4096;

// Key size, in bytes, of a valid AEAD_AES_128_GCM key.
const size_t kContentEncryptionKeySize = 16;

// The BoringSSL functions used to seal (encrypt) and open (decrypt) a payload
// follow the same prototype, declared as follows.
using EVP_AEAD_CTX_TransformFunction =
    int(const EVP_AEAD_CTX *ctx, uint8_t *out, size_t *out_len,
        size_t max_out_len, const uint8_t *nonce, size_t nonce_len,
        const uint8_t *in, size_t in_len, const uint8_t *ad, size_t ad_len);

// Creates the info parameter for an HKDF value for the given |content_encoding|
// in accordance with draft-thomson-http-encryption.
//
// cek_info = "Content-Encoding: aesgcm" || 0x00 || context
// nonce_info = "Content-Encoding: nonce" || 0x00 || context
//
// context = "P-256" || 0x00 ||
//           length(recipient_public) || recipient_public ||
//           length(sender_public) || sender_public
//
// The length of the public keys must be written as a two octet unsigned integer
// in network byte order (big endian).
std::string InfoForContentEncoding(
    const char* content_encoding,
    const base::StringPiece& recipient_public_key,
    const base::StringPiece& sender_public_key) {
  DCHECK_EQ(recipient_public_key.size(), 65u);
  DCHECK_EQ(sender_public_key.size(), 65u);

  std::stringstream info_stream;
  info_stream << "Content-Encoding: " << content_encoding << '\x00';
  info_stream << "P-256" << '\x00';

  uint16_t local_len =
      base::HostToNet16(static_cast<uint16_t>(recipient_public_key.size()));
  info_stream.write(reinterpret_cast<char*>(&local_len), sizeof(local_len));
  info_stream << recipient_public_key;

  uint16_t peer_len =
      base::HostToNet16(static_cast<uint16_t>(sender_public_key.size()));
  info_stream.write(reinterpret_cast<char*>(&peer_len), sizeof(peer_len));
  info_stream << sender_public_key;

  return info_stream.str();
}

}  // namespace

const size_t GCMMessageCryptographer::kAuthenticationTagBytes = 16;
const size_t GCMMessageCryptographer::kSaltSize = 16;

GCMMessageCryptographer::GCMMessageCryptographer(
    const base::StringPiece& recipient_public_key,
    const base::StringPiece& sender_public_key,
    const std::string& auth_secret)
    : content_encryption_key_info_(
          InfoForContentEncoding("aesgcm", recipient_public_key,
                                 sender_public_key)),
      nonce_info_(
          InfoForContentEncoding("nonce", recipient_public_key,
                                 sender_public_key)),
      auth_secret_(auth_secret) {
}

GCMMessageCryptographer::~GCMMessageCryptographer() {}

bool GCMMessageCryptographer::Encrypt(const base::StringPiece& plaintext,
                                      const base::StringPiece& ikm,
                                      const base::StringPiece& salt,
                                      size_t* record_size,
                                      std::string* ciphertext) const {
  DCHECK(ciphertext);
  DCHECK(record_size);

  if (salt.size() != kSaltSize)
    return false;

  std::string prk = DerivePseudoRandomKey(ikm);

  std::string content_encryption_key = DeriveContentEncryptionKey(prk, salt);
  std::string nonce = DeriveNonce(prk, salt);

  // Prior to the plaintext, draft-thomson-http-encryption has a two-byte
  // padding length followed by zero to 65535 bytes of padding. There is no need
  // for payloads created by Chrome to be padded so the padding length is set to
  // zero.
  std::string record;
  record.reserve(sizeof(uint16_t) + plaintext.size());
  record.append(sizeof(uint16_t), '\0');

  plaintext.AppendToString(&record);

  std::string encrypted_record;
  if (!EncryptDecryptRecordInternal(ENCRYPT, record, content_encryption_key,
                                    nonce, &encrypted_record)) {
    return false;
  }

  // The advertised record size must be at least one more than the padded
  // plaintext to ensure only one record.
  *record_size = std::max(kDefaultRecordSize, record.size() + 1);

  ciphertext->swap(encrypted_record);
  return true;
}

bool GCMMessageCryptographer::Decrypt(const base::StringPiece& ciphertext,
                                      const base::StringPiece& ikm,
                                      const base::StringPiece& salt,
                                      size_t record_size,
                                      std::string* plaintext) const {
  DCHECK(plaintext);

  if (salt.size() != kSaltSize || record_size <= 1)
    return false;

  // The |ciphertext| must be at least of size kAuthenticationTagBytes plus
  // len(uint16) to hold the message's padding length, which is the case when an
  // empty message with a zero padding length has been received. Per
  // https://tools.ietf.org/html/draft-thomson-http-encryption-02#section-3, the
  // |record_size| parameter must be large enough to use only one record.
  if (ciphertext.size() < sizeof(uint16_t) + kAuthenticationTagBytes ||
      ciphertext.size() > record_size + kAuthenticationTagBytes) {
    return false;
  }

  std::string prk = DerivePseudoRandomKey(ikm);

  std::string content_encryption_key = DeriveContentEncryptionKey(prk, salt);
  std::string nonce = DeriveNonce(prk, salt);

  std::string decrypted_record_string;
  if (!EncryptDecryptRecordInternal(DECRYPT, ciphertext, content_encryption_key,
                                    nonce, &decrypted_record_string)) {
    return false;
  }

  DCHECK(!decrypted_record_string.empty());

  base::StringPiece decrypted_record(decrypted_record_string);

  // Records must be at least two octets in size (to hold the padding). Records
  // that are smaller, i.e. a single octet, are invalid.
  if (decrypted_record.size() < sizeof(uint16_t))
    return false;

  // Records contain a two-byte, big-endian padding length followed by zero to
  // 65535 bytes of padding. Padding bytes must be zero but, since AES-GCM
  // authenticates the plaintext, checking and removing padding need not be done
  // in constant-time.
  uint16_t padding_length = (static_cast<uint8_t>(decrypted_record[0]) << 8) |
                            static_cast<uint8_t>(decrypted_record[1]);
  decrypted_record.remove_prefix(sizeof(uint16_t));

  if (padding_length > decrypted_record.size()) {
    return false;
  }

  for (size_t i = 0; i < padding_length; ++i) {
    if (decrypted_record[i] != 0)
      return false;
  }

  decrypted_record.remove_prefix(padding_length);
  decrypted_record.CopyToString(plaintext);
  return true;
}

bool GCMMessageCryptographer::EncryptDecryptRecordInternal(
    Mode mode,
    const base::StringPiece& input,
    const base::StringPiece& key,
    const base::StringPiece& nonce,
    std::string* output) const {
  DCHECK(output);

  const EVP_AEAD* aead = EVP_aead_aes_128_gcm();

  EVP_AEAD_CTX context;
  if (!EVP_AEAD_CTX_init(&context, aead,
                         reinterpret_cast<const uint8_t*>(key.data()),
                         key.size(), EVP_AEAD_DEFAULT_TAG_LENGTH, nullptr)) {
    return false;
  }

  base::CheckedNumeric<size_t> maximum_output_length(input.size());
  if (mode == ENCRYPT)
    maximum_output_length += kAuthenticationTagBytes;

  // WriteInto requires the buffer to finish with a NULL-byte.
  maximum_output_length += 1;

  size_t output_length = 0;
  uint8_t* raw_output = reinterpret_cast<uint8_t*>(
      base::WriteInto(output, maximum_output_length.ValueOrDie()));

  EVP_AEAD_CTX_TransformFunction* transform_function =
      mode == ENCRYPT ? EVP_AEAD_CTX_seal : EVP_AEAD_CTX_open;

  if (!transform_function(
          &context, raw_output, &output_length, output->size(),
          reinterpret_cast<const uint8_t*>(nonce.data()), nonce.size(),
          reinterpret_cast<const uint8_t*>(input.data()), input.size(),
          nullptr, 0)) {
    EVP_AEAD_CTX_cleanup(&context);
    return false;
  }

  EVP_AEAD_CTX_cleanup(&context);

  base::CheckedNumeric<size_t> expected_output_length(input.size());
  if (mode == ENCRYPT)
    expected_output_length += kAuthenticationTagBytes;
  else
    expected_output_length -= kAuthenticationTagBytes;

  DCHECK_EQ(expected_output_length.ValueOrDie(), output_length);

  output->resize(output_length);
  return true;
}

std::string GCMMessageCryptographer::DerivePseudoRandomKey(
    const base::StringPiece& ikm) const {
  if (allow_empty_auth_secret_for_tests_ && auth_secret_.empty())
    return ikm.as_string();

  CHECK(!auth_secret_.empty());

  std::stringstream info_stream;
  info_stream << "Content-Encoding: auth" << '\x00';

  crypto::HKDF hkdf(ikm, auth_secret_,
                    info_stream.str(),
                    32, /* key_bytes_to_generate */
                    0,  /* iv_bytes_to_generate */
                    0   /* subkey_secret_bytes_to_generate */);

  return hkdf.client_write_key().as_string();
}

std::string GCMMessageCryptographer::DeriveContentEncryptionKey(
    const base::StringPiece& prk,
    const base::StringPiece& salt) const {
  crypto::HKDF hkdf(prk, salt,
                    content_encryption_key_info_,
                    kContentEncryptionKeySize,
                    0,  /* iv_bytes_to_generate */
                    0   /* subkey_secret_bytes_to_generate */);

  return hkdf.client_write_key().as_string();
}

std::string GCMMessageCryptographer::DeriveNonce(
    const base::StringPiece& prk,
    const base::StringPiece& salt) const {
  crypto::HKDF hkdf(prk, salt,
                    nonce_info_,
                    kNonceSize,
                    0,  /* iv_bytes_to_generate */
                    0   /* subkey_secret_bytes_to_generate */);

  // draft-thomson-http-encryption defines that the result should be XOR'ed with
  // the record's sequence number, however, Web Push encryption is limited to a
  // single record per draft-ietf-webpush-encryption.

  return hkdf.client_write_key().as_string();
}

}  // namespace gcm