/*
* HKDF
* (C) 2013,2015,2017 Jack Lloyd
* (C) 2016 René Korthaus, Rohde & Schwarz Cybersecurity
* (C) 2024 René Meusel, Rohde & Schwarz Cybersecurity
*
* Botan is released under the Simplified BSD License (see license.txt)
*/

#include <botan/internal/hkdf.h>

#include <botan/internal/buffer_stuffer.h>
#include <botan/internal/concat_util.h>
#include <botan/internal/fmt.h>
#include <botan/internal/loadstor.h>
#include <botan/internal/mem_utils.h>

namespace Botan {

std::unique_ptr<KDF> HKDF::new_object() const {
   return std::make_unique<HKDF>(m_prf->new_object());
}

std::string HKDF::name() const {
   return fmt("HKDF({})", m_prf->name());
}

void HKDF::perform_kdf(std::span<uint8_t> key,
                       std::span<const uint8_t> secret,
                       std::span<const uint8_t> salt,
                       std::span<const uint8_t> label) const {
   const HKDF_Extract extract(m_prf->new_object());
   const HKDF_Expand expand(m_prf->new_object());
   secure_vector<uint8_t> prk(m_prf->output_length());

   extract.derive_key(prk, secret, salt, {});
   expand.derive_key(key, prk, {}, label);
}

std::unique_ptr<KDF> HKDF_Extract::new_object() const {
   return std::make_unique<HKDF_Extract>(m_prf->new_object());
}

std::string HKDF_Extract::name() const {
   return fmt("HKDF-Extract({})", m_prf->name());
}

void HKDF_Extract::perform_kdf(std::span<uint8_t> key,
                               std::span<const uint8_t> secret,
                               std::span<const uint8_t> salt,
                               std::span<const uint8_t> label) const {
   const size_t prf_output_len = m_prf->output_length();
   BOTAN_ARG_CHECK(key.size() <= prf_output_len, "HKDF-Extract maximum output length exceeded");
   BOTAN_ARG_CHECK(label.empty(), "HKDF-Extract does not support a label input");

   if(key.empty()) {
      return;
   }

   if(salt.empty()) {
      m_prf->set_key(std::vector<uint8_t>(prf_output_len));
   } else {
      m_prf->set_key(salt);
   }

   m_prf->update(secret);

   if(key.size() == prf_output_len) {
      m_prf->final(key);
   } else {
      const auto prk = m_prf->final();
      copy_mem(key, std::span{prk}.first(key.size()));
   }
}

std::unique_ptr<KDF> HKDF_Expand::new_object() const {
   return std::make_unique<HKDF_Expand>(m_prf->new_object());
}

std::string HKDF_Expand::name() const {
   return fmt("HKDF-Expand({})", m_prf->name());
}

void HKDF_Expand::perform_kdf(std::span<uint8_t> key,
                              std::span<const uint8_t> secret,
                              std::span<const uint8_t> salt,
                              std::span<const uint8_t> label) const {
   const auto prf_output_length = m_prf->output_length();
   BOTAN_ARG_CHECK(key.size() <= prf_output_length * 255, "HKDF-Expand maximum output length exceeded");

   if(key.empty()) {
      return;
   }

   // Keep a reference to the previous PRF output (empty by default).
   std::span<uint8_t> h = {};

   BufferStuffer k(key);
   m_prf->set_key(secret);
   for(uint8_t counter = 1; !k.full(); ++counter) {
      m_prf->update(h);
      m_prf->update(label);
      m_prf->update(salt);
      m_prf->update(counter);

      // Write straight into the output buffer, except if the PRF output needs
      // a truncation in the final iteration.
      if(k.remaining_capacity() >= prf_output_length) {
         h = k.next(prf_output_length);
         m_prf->final(h);
      } else {
         const auto full_prf_output = m_prf->final();
         h = {};  // this is the final iteration!
         k.append(std::span{full_prf_output}.first(k.remaining_capacity()));
      }
   }
}

secure_vector<uint8_t> hkdf_expand_label(std::string_view hash_fn,
                                         std::span<const uint8_t> secret,
                                         std::string_view label,
                                         std::span<const uint8_t> hash_val,
                                         size_t length) {
   BOTAN_ARG_CHECK(length <= 0xFFFF, "HKDF-Expand-Label requested output too large");
   BOTAN_ARG_CHECK(label.size() <= 0xFF, "HKDF-Expand-Label label too long");
   BOTAN_ARG_CHECK(hash_val.size() <= 0xFF, "HKDF-Expand-Label hash too long");

   const HKDF_Expand hkdf(MessageAuthenticationCode::create_or_throw(fmt("HMAC({})", hash_fn)));

   const auto prefix = concat<std::vector<uint8_t>>(store_be(static_cast<uint16_t>(length)),
                                                    store_be(static_cast<uint8_t>(label.size())),
                                                    as_span_of_bytes(label),
                                                    store_be(static_cast<uint8_t>(hash_val.size())));

   /*
   * We do something a little dirty here to avoid copying the hash_val,
   * making use of the fact that Botan's KDF interface supports label+salt,
   * and knowing that our HKDF hashes first param label then param salt.
   */
   return hkdf.derive_key(length, secret, hash_val, prefix);
}

}  // namespace Botan
