// Copyright 2024 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/base.h>

#include <assert.h>
#include <stdint.h>
#include <string.h>

#include "../../internal.h"
#include "./address.h"
#include "./params.h"
#include "./thash.h"
#include "./wots.h"


// Implements Algorithm 5: chain function, page 18
static void chain(const slh_dsa_config *config, uint8_t *output,
                  const uint8_t *input, uint32_t start, uint32_t steps,
                  const uint8_t *pub_seed, uint8_t addr[32]) {
  assert(start < config->wots_w);
  assert(steps < config->wots_w);

  OPENSSL_memcpy(output, input, config->n);

  for (uint32_t i = start; i < (start + steps) && i < config->wots_w; ++i) {
    slhdsa_set_hash_addr(config, addr, i);
    slhdsa_thash_f(config, output, output, pub_seed, addr);
  }
}

static void slhdsa_wots_do_chain(const slh_dsa_config *config, uint8_t *out,
                                 uint8_t sk_addr[32], uint8_t addr[32],
                                 uint8_t value, const uint8_t *sk_seed,
                                 const uint8_t *pub_seed,
                                 uint32_t chain_index) {
  uint8_t tmp_sk[SLHDSA_MAX_N];
  slhdsa_set_chain_addr(config, sk_addr, chain_index);
  slhdsa_thash_prf(config, tmp_sk, pub_seed, sk_seed, sk_addr);
  slhdsa_set_chain_addr(config, addr, chain_index);
  chain(config, out, tmp_sk, 0, value, pub_seed, addr);
}

// Implements Algorithm 6: wots_pkGen function, page 18
void slhdsa_wots_pk_gen(const slh_dsa_config *config, uint8_t *pk,
                        const uint8_t *sk_seed, const uint8_t *pub_seed,
                        uint8_t addr[32]) {
  uint8_t wots_pk_addr[32], sk_addr[32];
  OPENSSL_memcpy(wots_pk_addr, addr, sizeof(wots_pk_addr));
  OPENSSL_memcpy(sk_addr, addr, sizeof(sk_addr));
  slhdsa_set_type(config, sk_addr, SLHDSA_ADDR_TYPE_WOTSPRF);
  slhdsa_copy_keypair_addr(config, sk_addr, addr);

  uint8_t tmp[SLHDSA_MAX_WOTS_BYTES];
  const uint32_t n = config->n;
  const uint32_t wots_len = slhdsa_wots_len(config);
  const uint32_t max_chain = config->wots_w - 1;
  for (uint32_t i = 0; i < wots_len; ++i) {
    slhdsa_wots_do_chain(config, tmp + i * n, sk_addr, addr, max_chain, sk_seed,
                         pub_seed, i);
  }

  // Compress pk
  slhdsa_set_type(config, wots_pk_addr, SLHDSA_ADDR_TYPE_WOTSPK);
  slhdsa_copy_keypair_addr(config, wots_pk_addr, addr);
  slhdsa_thash_tl(config, pk, tmp, pub_seed, wots_pk_addr);
}

// Implements Algorithm 7: wots_sign function, page 20
void slhdsa_wots_sign(const slh_dsa_config *config, uint8_t *sig,
                      const uint8_t *msg, const uint8_t *sk_seed,
                      const uint8_t *pub_seed, uint8_t addr[32]) {
  BSSL_CHECK(config->wots_w == 16);
  const size_t n = config->n;
  const uint32_t max_chain = config->wots_w - 1;

  uint16_t csum = 0;
  for (size_t i = 0; i < n; ++i) {
    csum += max_chain - (msg[i] >> 4);
    csum += max_chain - (msg[i] & 15);
  }

  // Compute chains
  uint8_t sk_addr[32];
  OPENSSL_memcpy(sk_addr, addr, sizeof(sk_addr));
  slhdsa_set_type(config, sk_addr, SLHDSA_ADDR_TYPE_WOTSPRF);
  slhdsa_copy_keypair_addr(config, sk_addr, addr);

  uint32_t chain_index = 0;
  for (size_t i = 0; i < n; ++i) {
    slhdsa_wots_do_chain(config, sig, sk_addr, addr, msg[i] >> 4, sk_seed,
                         pub_seed, chain_index++);
    sig += n;

    slhdsa_wots_do_chain(config, sig, sk_addr, addr, msg[i] & 15, sk_seed,
                         pub_seed, chain_index++);
    sig += n;
  }

  // Include the WOTS checksum values (len2 = 3 for the parameter sets we
  // support).
  slhdsa_wots_do_chain(config, sig, sk_addr, addr, (csum >> 8) & 15, sk_seed,
                       pub_seed, chain_index++);
  sig += n;
  slhdsa_wots_do_chain(config, sig, sk_addr, addr, (csum >> 4) & 15, sk_seed,
                       pub_seed, chain_index++);
  sig += n;
  slhdsa_wots_do_chain(config, sig, sk_addr, addr, csum & 15, sk_seed, pub_seed,
                       chain_index++);
}

static void slhdsa_wots_pk_from_sig_do_chain(const slh_dsa_config *config,
                                             uint8_t *out, uint8_t addr[32],
                                             const uint8_t *in, uint8_t value,
                                             const uint8_t *pub_seed,
                                             uint32_t chain_index) {
  slhdsa_set_chain_addr(config, addr, chain_index);
  chain(config, out + chain_index * config->n, in + chain_index * config->n,
        value, (config->wots_w - 1) - value, pub_seed, addr);
}

// Implements Algorithm 8: wots_pkFromSig function, page 21
void slhdsa_wots_pk_from_sig(const slh_dsa_config *config, uint8_t *pk,
                             const uint8_t *sig, const uint8_t *msg,
                             const uint8_t *pub_seed, uint8_t addr[32]) {
  BSSL_CHECK(config->wots_w == 16);
  const size_t n = config->n;
  const uint32_t max_chain = config->wots_w - 1;

  uint16_t csum = 0;
  for (size_t i = 0; i < n; ++i) {
    csum += max_chain - (msg[i] >> 4);
    csum += max_chain - (msg[i] & 15);
  }

  uint8_t tmp[SLHDSA_MAX_WOTS_BYTES];
  uint8_t wots_pk_addr[32];
  OPENSSL_memcpy(wots_pk_addr, addr, sizeof(wots_pk_addr));

  uint32_t chain_index = 0;
  for (size_t i = 0; i < n; ++i) {
    slhdsa_wots_pk_from_sig_do_chain(config, tmp, addr, sig, msg[i] >> 4,
                                     pub_seed, chain_index++);
    slhdsa_wots_pk_from_sig_do_chain(config, tmp, addr, sig, msg[i] & 15,
                                     pub_seed, chain_index++);
  }

  slhdsa_wots_pk_from_sig_do_chain(config, tmp, addr, sig, csum >> 8, pub_seed,
                                   chain_index++);
  slhdsa_wots_pk_from_sig_do_chain(config, tmp, addr, sig, (csum >> 4) & 15,
                                   pub_seed, chain_index++);
  slhdsa_wots_pk_from_sig_do_chain(config, tmp, addr, sig, csum & 15, pub_seed,
                                   chain_index++);

  // Compress pk
  slhdsa_set_type(config, wots_pk_addr, SLHDSA_ADDR_TYPE_WOTSPK);
  slhdsa_copy_keypair_addr(config, wots_pk_addr, addr);
  slhdsa_thash_tl(config, pk, tmp, pub_seed, wots_pk_addr);
}