// Copyright 2025 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/xwing.h>

#include <openssl/bytestring.h>
#include <openssl/curve25519.h>
#include <openssl/mlkem.h>
#include <openssl/rand.h>

#include "../fipsmodule/bcm_interface.h"
#include "../fipsmodule/keccak/internal.h"


using namespace bssl;

struct private_key {
  MLKEM768_private_key mlkem_private_key;
  uint8_t x25519_private_key[32];
  uint8_t seed[XWING_PRIVATE_KEY_BYTES];
};

static_assert(sizeof(XWING_private_key) == sizeof(private_key));
static_assert(alignof(XWING_private_key) == alignof(private_key));

static const private_key *private_key_from_external(
    const XWING_private_key *external) {
  return reinterpret_cast<const private_key *>(external);
}
static private_key *private_key_from_external(XWING_private_key *external) {
  return reinterpret_cast<private_key *>(external);
}

static void xwing_expand_private_key(private_key *inout_private_key) {
  BORINGSSL_keccak_st context;
  BORINGSSL_keccak_init(&context, boringssl_shake256);
  BORINGSSL_keccak_absorb(&context, inout_private_key->seed,
                          sizeof(inout_private_key->seed));

  // ML-KEM-768
  uint8_t mlkem_seed[64];
  BORINGSSL_keccak_squeeze(&context, mlkem_seed, sizeof(mlkem_seed));
  MLKEM768_private_key_from_seed(&inout_private_key->mlkem_private_key,
                                 mlkem_seed, sizeof(mlkem_seed));

  // X25519
  BORINGSSL_keccak_squeeze(&context, inout_private_key->x25519_private_key,
                           sizeof(inout_private_key->x25519_private_key));
}

static int xwing_parse_private_key(private_key *out_private_key, CBS *in) {
  if (!CBS_copy_bytes(in, out_private_key->seed,
                      sizeof(out_private_key->seed))) {
    return 0;
  }

  xwing_expand_private_key(out_private_key);
  return 1;
}

static int xwing_marshal_private_key(CBB *out, const private_key *private_key) {
  return CBB_add_bytes(out, private_key->seed, sizeof(private_key->seed));
}

static int xwing_public_from_private(
    uint8_t out_encoded_public_key[XWING_PUBLIC_KEY_BYTES],
    const private_key *private_key) {
  CBB cbb;
  if (!CBB_init_fixed(&cbb, out_encoded_public_key, XWING_PUBLIC_KEY_BYTES)) {
    return 0;
  }

  // ML-KEM-768
  MLKEM768_public_key mlkem_public_key;
  MLKEM768_public_from_private(&mlkem_public_key,
                               &private_key->mlkem_private_key);

  if (!MLKEM768_marshal_public_key(&cbb, &mlkem_public_key)) {
    return 0;
  }

  // X25519
  uint8_t *buf;
  if (!CBB_add_space(&cbb, &buf, 32)) {
    return 0;
  }
  X25519_public_from_private(buf, private_key->x25519_private_key);

  if (CBB_len(&cbb) != XWING_PUBLIC_KEY_BYTES) {
    return 0;
  }
  return 1;
}

static void xwing_combiner(
    uint8_t out_shared_secret[XWING_SHARED_SECRET_BYTES],
    const uint8_t mlkem_shared_secret[MLKEM_SHARED_SECRET_BYTES],
    const uint8_t x25519_shared_secret[32], const uint8_t x25519_ciphertext[32],
    const uint8_t x25519_public_key[32]) {
  BORINGSSL_keccak_st context;
  BORINGSSL_keccak_init(&context, boringssl_sha3_256);

  BORINGSSL_keccak_absorb(&context, mlkem_shared_secret,
                          MLKEM_SHARED_SECRET_BYTES);
  BORINGSSL_keccak_absorb(&context, x25519_shared_secret, 32);
  BORINGSSL_keccak_absorb(&context, x25519_ciphertext, 32);
  BORINGSSL_keccak_absorb(&context, x25519_public_key, 32);

  uint8_t xwing_label[6] = {0x5c, 0x2e, 0x2f, 0x2f, 0x5e, 0x5c};
  BORINGSSL_keccak_absorb(&context, xwing_label, sizeof(xwing_label));

  BORINGSSL_keccak_squeeze(&context, out_shared_secret,
                           XWING_SHARED_SECRET_BYTES);
}

// Public API.

int XWING_parse_private_key(struct XWING_private_key *out_private_key,
                            CBS *in) {
  if (!xwing_parse_private_key(private_key_from_external(out_private_key),
                               in) ||
      CBS_len(in) != 0) {
    return 0;
  }
  return 1;
}

int XWING_marshal_private_key(CBB *out,
                              const struct XWING_private_key *private_key) {
  return xwing_marshal_private_key(out, private_key_from_external(private_key));
}

int XWING_generate_key(uint8_t out_encoded_public_key[XWING_PUBLIC_KEY_BYTES],
                       struct XWING_private_key *out_private_key) {
  private_key *private_key = private_key_from_external(out_private_key);
  RAND_bytes(private_key->seed, sizeof(private_key->seed));

  xwing_expand_private_key(private_key);

  return XWING_public_from_private(out_encoded_public_key, out_private_key);
}

int XWING_public_from_private(
    uint8_t out_encoded_public_key[XWING_PUBLIC_KEY_BYTES],
    const struct XWING_private_key *private_key) {
  return xwing_public_from_private(out_encoded_public_key,
                                   private_key_from_external(private_key));
}

int XWING_encap(uint8_t out_ciphertext[XWING_CIPHERTEXT_BYTES],
                uint8_t out_shared_secret[XWING_SHARED_SECRET_BYTES],
                const uint8_t encoded_public_key[XWING_PUBLIC_KEY_BYTES]) {
  uint8_t eseed[64];
  RAND_bytes(eseed, sizeof(eseed));

  return XWING_encap_external_entropy(out_ciphertext, out_shared_secret,
                                      encoded_public_key, eseed);
}

int XWING_encap_external_entropy(
    uint8_t out_ciphertext[XWING_CIPHERTEXT_BYTES],
    uint8_t out_shared_secret[XWING_SHARED_SECRET_BYTES],
    const uint8_t encoded_public_key[XWING_PUBLIC_KEY_BYTES],
    const uint8_t eseed[64]) {
  // X25519
  static_assert(XWING_PUBLIC_KEY_BYTES >= MLKEM768_PUBLIC_KEY_BYTES + 32);
  const uint8_t *x25519_public_key =
      encoded_public_key + MLKEM768_PUBLIC_KEY_BYTES;
  const uint8_t *x25519_ephemeral_private_key = eseed + 32;
  uint8_t *x25519_ciphertext = out_ciphertext + MLKEM768_CIPHERTEXT_BYTES;
  X25519_public_from_private(x25519_ciphertext, x25519_ephemeral_private_key);

  uint8_t x25519_shared_secret[32];
  if (!X25519(x25519_shared_secret, x25519_ephemeral_private_key,
              x25519_public_key)) {
    return 0;
  }

  // ML-KEM-768
  CBS mlkem_cbs;
  static_assert(MLKEM768_PUBLIC_KEY_BYTES <= XWING_PUBLIC_KEY_BYTES);
  CBS_init(&mlkem_cbs, encoded_public_key, MLKEM768_PUBLIC_KEY_BYTES);

  MLKEM768_public_key mlkem_public_key;
  if (!MLKEM768_parse_public_key(&mlkem_public_key, &mlkem_cbs)) {
    return 0;
  }

  uint8_t *mlkem_ciphertext = out_ciphertext;
  uint8_t mlkem_shared_secret[MLKEM_SHARED_SECRET_BYTES];
  BCM_mlkem768_encap_external_entropy(mlkem_ciphertext, mlkem_shared_secret,
                                      &mlkem_public_key, eseed);

  // Combine the shared secrets
  xwing_combiner(out_shared_secret, mlkem_shared_secret, x25519_shared_secret,
                 x25519_ciphertext, x25519_public_key);
  return 1;
}

static int xwing_decap(uint8_t out_shared_secret[XWING_SHARED_SECRET_BYTES],
                       const uint8_t ciphertext[XWING_CIPHERTEXT_BYTES],
                       const private_key *private_key) {
  static_assert(XWING_CIPHERTEXT_BYTES >= MLKEM768_CIPHERTEXT_BYTES + 32);
  const uint8_t *mlkem_ciphertext = ciphertext;
  const uint8_t *x25519_ciphertext = ciphertext + MLKEM768_CIPHERTEXT_BYTES;

  // ML-KEM-768
  uint8_t mlkem_shared_secret[MLKEM_SHARED_SECRET_BYTES];
  if (!MLKEM768_decap(mlkem_shared_secret, mlkem_ciphertext,
                      MLKEM768_CIPHERTEXT_BYTES,
                      &private_key->mlkem_private_key)) {
    goto err;
  }

  // X25519
  uint8_t x25519_public_key[32];
  X25519_public_from_private(x25519_public_key,
                             private_key->x25519_private_key);

  uint8_t x25519_shared_secret[32];
  if (!X25519(x25519_shared_secret, private_key->x25519_private_key,
              x25519_ciphertext)) {
    goto err;
  }

  // Combine the shared secrets
  xwing_combiner(out_shared_secret, mlkem_shared_secret, x25519_shared_secret,
                 x25519_ciphertext, x25519_public_key);
  return 1;

err:
  // In case of error, fill the shared secret with random bytes so that if the
  // caller forgets to check the return code:
  // - no intermediate information leaks,
  // - the shared secret is unpredictable, so for example any data encrypted
  //   with it wouldn't be trivially decryptable by an attacker.
  RAND_bytes(out_shared_secret, XWING_SHARED_SECRET_BYTES);
  return 0;
}

int XWING_decap(uint8_t out_shared_secret[XWING_SHARED_SECRET_BYTES],
                const uint8_t ciphertext[XWING_CIPHERTEXT_BYTES],
                const struct XWING_private_key *private_key) {
  return xwing_decap(out_shared_secret, ciphertext,
                     private_key_from_external(private_key));
}