/*
* Salsa20 / XSalsa20
* (C) 1999-2010,2014 Jack Lloyd
*
* Botan is released under the Simplified BSD License (see license.txt)
*/

#include <botan/internal/salsa20.h>

#include <botan/exceptn.h>
#include <botan/internal/loadstor.h>
#include <botan/internal/rotate.h>

#if defined(BOTAN_HAS_CPUID)
   #include <botan/internal/cpuid.h>
#endif

namespace Botan {

namespace {

inline void salsa20_quarter_round(uint32_t& x1, uint32_t& x2, uint32_t& x3, uint32_t& x4) {
   x2 ^= rotl<7>(x1 + x4);
   x3 ^= rotl<9>(x2 + x1);
   x4 ^= rotl<13>(x3 + x2);
   x1 ^= rotl<18>(x4 + x3);
}

}  // namespace

/*
* Generate HSalsa20 cipher stream (for XSalsa20 IV setup)
*/
//static
void Salsa20::hsalsa20(uint32_t output[8], const uint32_t input[16]) {
   uint32_t x00 = input[0];
   uint32_t x01 = input[1];
   uint32_t x02 = input[2];
   uint32_t x03 = input[3];
   uint32_t x04 = input[4];
   uint32_t x05 = input[5];
   uint32_t x06 = input[6];
   uint32_t x07 = input[7];
   uint32_t x08 = input[8];
   uint32_t x09 = input[9];
   uint32_t x10 = input[10];
   uint32_t x11 = input[11];
   uint32_t x12 = input[12];
   uint32_t x13 = input[13];
   uint32_t x14 = input[14];
   uint32_t x15 = input[15];

   for(size_t i = 0; i != 10; ++i) {
      salsa20_quarter_round(x00, x04, x08, x12);
      salsa20_quarter_round(x05, x09, x13, x01);
      salsa20_quarter_round(x10, x14, x02, x06);
      salsa20_quarter_round(x15, x03, x07, x11);

      salsa20_quarter_round(x00, x01, x02, x03);
      salsa20_quarter_round(x05, x06, x07, x04);
      salsa20_quarter_round(x10, x11, x08, x09);
      salsa20_quarter_round(x15, x12, x13, x14);
   }

   output[0] = x00;
   output[1] = x05;
   output[2] = x10;
   output[3] = x15;
   output[4] = x06;
   output[5] = x07;
   output[6] = x08;
   output[7] = x09;
}

/*
* Generate Salsa20 cipher stream
*/
//static
void Salsa20::salsa_core(uint8_t output[64], const uint32_t input[16], size_t rounds) {
   BOTAN_ASSERT_NOMSG(rounds % 2 == 0);

   uint32_t x00 = input[0];
   uint32_t x01 = input[1];
   uint32_t x02 = input[2];
   uint32_t x03 = input[3];
   uint32_t x04 = input[4];
   uint32_t x05 = input[5];
   uint32_t x06 = input[6];
   uint32_t x07 = input[7];
   uint32_t x08 = input[8];
   uint32_t x09 = input[9];
   uint32_t x10 = input[10];
   uint32_t x11 = input[11];
   uint32_t x12 = input[12];
   uint32_t x13 = input[13];
   uint32_t x14 = input[14];
   uint32_t x15 = input[15];

   for(size_t i = 0; i != rounds / 2; ++i) {
      salsa20_quarter_round(x00, x04, x08, x12);
      salsa20_quarter_round(x05, x09, x13, x01);
      salsa20_quarter_round(x10, x14, x02, x06);
      salsa20_quarter_round(x15, x03, x07, x11);

      salsa20_quarter_round(x00, x01, x02, x03);
      salsa20_quarter_round(x05, x06, x07, x04);
      salsa20_quarter_round(x10, x11, x08, x09);
      salsa20_quarter_round(x15, x12, x13, x14);
   }

   store_le(x00 + input[0], output + 4 * 0);
   store_le(x01 + input[1], output + 4 * 1);
   store_le(x02 + input[2], output + 4 * 2);
   store_le(x03 + input[3], output + 4 * 3);
   store_le(x04 + input[4], output + 4 * 4);
   store_le(x05 + input[5], output + 4 * 5);
   store_le(x06 + input[6], output + 4 * 6);
   store_le(x07 + input[7], output + 4 * 7);
   store_le(x08 + input[8], output + 4 * 8);
   store_le(x09 + input[9], output + 4 * 9);
   store_le(x10 + input[10], output + 4 * 10);
   store_le(x11 + input[11], output + 4 * 11);
   store_le(x12 + input[12], output + 4 * 12);
   store_le(x13 + input[13], output + 4 * 13);
   store_le(x14 + input[14], output + 4 * 14);
   store_le(x15 + input[15], output + 4 * 15);
}

size_t Salsa20::parallelism() {
#if defined(BOTAN_HAS_SALSA20_AVX512)
   if(CPUID::has(CPUID::Feature::AVX512)) {
      return 16;
   }
#endif

#if defined(BOTAN_HAS_SALSA20_AVX2)
   if(CPUID::has(CPUID::Feature::AVX2)) {
      return 8;
   }
#endif

   return 4;
}

std::string Salsa20::provider() const {
#if defined(BOTAN_HAS_SALSA20_AVX512)
   if(auto feat = CPUID::check(CPUID::Feature::AVX512)) {
      return *feat;
   }
#endif

#if defined(BOTAN_HAS_SALSA20_AVX2)
   if(auto feat = CPUID::check(CPUID::Feature::AVX2)) {
      return *feat;
   }
#endif

#if defined(BOTAN_HAS_SALSA20_SIMD32)
   if(auto feat = CPUID::check(CPUID::Feature::SIMD_4X32)) {
      return *feat;
   }
#endif

   return "base";
}

//static
void Salsa20::salsa20(uint8_t output[], size_t output_blocks, uint32_t state[16], size_t rounds) {
   BOTAN_ASSERT(rounds % 2 == 0, "Valid rounds");

#if defined(BOTAN_HAS_SALSA20_AVX512)
   if(CPUID::has(CPUID::Feature::AVX512)) {
      while(output_blocks >= 16) {
         Salsa20::salsa20_avx512_x16(output, state, rounds);
         output += 16 * 64;
         output_blocks -= 16;
      }
   }
#endif

#if defined(BOTAN_HAS_SALSA20_AVX2)
   if(CPUID::has(CPUID::Feature::AVX2)) {
      while(output_blocks >= 8) {
         Salsa20::salsa20_avx2_x8(output, state, rounds);
         output += 8 * 64;
         output_blocks -= 8;
      }
   }
#endif

#if defined(BOTAN_HAS_SALSA20_SIMD32)
   if(CPUID::has(CPUID::Feature::SIMD_4X32)) {
      while(output_blocks >= 4) {
         Salsa20::salsa20_simd32_x4(output, state, rounds);
         output += 4 * 64;
         output_blocks -= 4;
      }
   }
#endif

   for(size_t i = 0; i != output_blocks; ++i) {
      salsa_core(output + 64 * i, state, rounds);

      ++state[8];
      if(state[8] == 0) {
         state[9] += 1;
      }
   }
}

/*
* Combine cipher stream with message
*/
void Salsa20::cipher_bytes(const uint8_t in[], uint8_t out[], size_t length) {
   assert_key_material_set();

   while(length >= m_buffer.size() - m_position) {
      const size_t available = m_buffer.size() - m_position;

      xor_buf(out, in, &m_buffer[m_position], available);
      salsa20(m_buffer.data(), m_buffer.size() / 64, m_state.data(), 20);

      length -= available;
      in += available;
      out += available;

      m_position = 0;
   }

   xor_buf(out, in, &m_buffer[m_position], length);

   m_position += length;
}

void Salsa20::generate_keystream(uint8_t out[], size_t length) {
   assert_key_material_set();

   while(length >= m_buffer.size() - m_position) {
      const size_t available = m_buffer.size() - m_position;

      // TODO: this could write directly to the output buffer
      // instead of bouncing it through m_buffer first
      copy_mem(out, &m_buffer[m_position], available);
      salsa20(m_buffer.data(), m_buffer.size() / 64, m_state.data(), 20);

      length -= available;
      out += available;
      m_position = 0;
   }

   copy_mem(out, &m_buffer[m_position], length);

   m_position += length;
}

void Salsa20::initialize_state() {
   static const uint32_t TAU[] = {0x61707865, 0x3120646e, 0x79622d36, 0x6b206574};

   static const uint32_t SIGMA[] = {0x61707865, 0x3320646e, 0x79622d32, 0x6b206574};

   m_state[1] = m_key[0];
   m_state[2] = m_key[1];
   m_state[3] = m_key[2];
   m_state[4] = m_key[3];

   if(m_key.size() == 4) {
      m_state[0] = TAU[0];
      m_state[5] = TAU[1];
      m_state[10] = TAU[2];
      m_state[15] = TAU[3];
      m_state[11] = m_key[0];
      m_state[12] = m_key[1];
      m_state[13] = m_key[2];
      m_state[14] = m_key[3];
   } else {
      m_state[0] = SIGMA[0];
      m_state[5] = SIGMA[1];
      m_state[10] = SIGMA[2];
      m_state[15] = SIGMA[3];
      m_state[11] = m_key[4];
      m_state[12] = m_key[5];
      m_state[13] = m_key[6];
      m_state[14] = m_key[7];
   }

   m_state[6] = 0;
   m_state[7] = 0;
   m_state[8] = 0;
   m_state[9] = 0;

   m_position = 0;
}

bool Salsa20::has_keying_material() const {
   return !m_state.empty();
}

size_t Salsa20::buffer_size() const {
   return 64;
}

/*
* Salsa20 Key Schedule
*/
void Salsa20::key_schedule(std::span<const uint8_t> key) {
   m_key.resize(key.size() / 4);
   load_le<uint32_t>(m_key.data(), key.data(), m_key.size());

   m_state.resize(16);

   const size_t salsa_block = 64;
   m_buffer.resize(parallelism() * salsa_block);

   set_iv(nullptr, 0);
}

/*
* Set the Salsa IV
*/
void Salsa20::set_iv_bytes(const uint8_t iv[], size_t length) {
   assert_key_material_set();

   if(!valid_iv_length(length)) {
      throw Invalid_IV_Length(name(), length);
   }

   initialize_state();

   if(length == 0) {
      // Salsa20 null IV
      m_state[6] = 0;
      m_state[7] = 0;
   } else if(length == 8) {
      // Salsa20
      m_state[6] = load_le<uint32_t>(iv, 0);
      m_state[7] = load_le<uint32_t>(iv, 1);
   } else {
      // XSalsa20
      m_state[6] = load_le<uint32_t>(iv, 0);
      m_state[7] = load_le<uint32_t>(iv, 1);
      m_state[8] = load_le<uint32_t>(iv, 2);
      m_state[9] = load_le<uint32_t>(iv, 3);

      secure_vector<uint32_t> hsalsa(8);
      hsalsa20(hsalsa.data(), m_state.data());

      m_state[1] = hsalsa[0];
      m_state[2] = hsalsa[1];
      m_state[3] = hsalsa[2];
      m_state[4] = hsalsa[3];
      m_state[6] = load_le<uint32_t>(iv, 4);
      m_state[7] = load_le<uint32_t>(iv, 5);
      m_state[11] = hsalsa[4];
      m_state[12] = hsalsa[5];
      m_state[13] = hsalsa[6];
      m_state[14] = hsalsa[7];
   }

   m_state[8] = 0;
   m_state[9] = 0;

   salsa20(m_buffer.data(), m_buffer.size() / 64, m_state.data(), 20);
   m_position = 0;
}

bool Salsa20::valid_iv_length(size_t iv_len) const {
   return (iv_len == 0 || iv_len == 8 || iv_len == 24);
}

size_t Salsa20::default_iv_length() const {
   return 24;
}

Key_Length_Specification Salsa20::key_spec() const {
   return Key_Length_Specification(16, 32, 16);
}

std::unique_ptr<StreamCipher> Salsa20::new_object() const {
   return std::make_unique<Salsa20>();
}

std::string Salsa20::name() const {
   return "Salsa20";
}

/*
* Clear memory of sensitive data
*/
void Salsa20::clear() {
   zap(m_key);
   zap(m_state);
   zap(m_buffer);
   m_position = 0;
}

void Salsa20::seek(uint64_t offset) {
   assert_key_material_set();

   const uint64_t counter = offset / 64;

   m_state[8] = static_cast<uint32_t>(counter);
   m_state[9] = static_cast<uint32_t>(counter >> 32);

   salsa20(m_buffer.data(), m_buffer.size() / 64, m_state.data(), 20);

   m_position = offset % 64;
}
}  // namespace Botan
