/*
 * Copyright (C) 2022 - 2025, Stephan Mueller <smueller@chronox.de>
 *
 * License: see LICENSE file in root directory
 *
 * THIS SOFTWARE IS PROVIDED ``AS IS'' AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE, ALL OF
 * WHICH ARE HEREBY DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT
 * OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
 * USE OF THIS SOFTWARE, EVEN IF NOT ADVISED OF THE POSSIBILITY OF SUCH
 * DAMAGE.
 */

#ifndef KYBER_X25519_KDF_H
#define KYBER_X25519_KDF_H

#include "kyber_type.h"
#include "lc_hash.h"
#include "lc_kmac.h"
#include "lc_sha3.h"

#ifdef __cplusplus
extern "C" {
#endif

/**
 * @brief kyber_double_ss_kdf - KDF to derive arbitrary sized SS from Kyber SS
 *
 *	SS <- KMAC256(K = Kyber-SS || X25519-SS,
 *		      X = Kyber-CT || X25519-ephemeral-PK,
 *		      L = requested SS length, S = "Kyber X25519 KEM SS")
 *
 * This KDF is is consistent with SP800-108 rev 1.
 */
static inline void kyber_x25519_ss_kdf(uint8_t *ss, size_t ss_len,
				       const struct lc_kyber_x25519_ct *ct,
				       const struct lc_kyber_x25519_ss *calc_ss)
{
	static const uint8_t kyber_ss_label[] = "Kyber X25519 KEM SS";

	/*
	 * NOTE: this only works because struct lc_kyber_x25519_ss contains
	 * Kyber SS || X25519 SS in memory. Also, lc_kyber_x25519_ct contains
	 * Kyber CT || X25519 ephemeral PK in memory. If either structure
	 * changes, change this KDF invocation.
	 */
	lc_kmac(lc_cshake256, (uint8_t *)calc_ss, sizeof(struct lc_kyber_ss),
		kyber_ss_label, sizeof(kyber_ss_label) - 1, (uint8_t *)ct,
		sizeof(struct lc_kyber_x25519_ct), ss, ss_len);
}

/**
 * @brief Kyber-X25519 KDF with 3 input values
 *
 * SS <- KMAC256(K = Kyber-SS || X25519-SS from ephemeral key 1,
 *		 X = Kyber-SS || X25519-SS from static key || Nonce,
 *		 L = requested SS length,
 *		 S = "Kyber X25519 KEM 3-way SS")
 *
 * @param [in] ss0 SS0
 * @param [in] ss1 SS1
 * @param [in] in3 input buffer 3
 * @param [in] inlen3 length of input buffer 3
 * @param [out] out output buffer of size
 * @param [in] outlen output buffer length
 */
static inline void kyber_x25519_kdf3(const struct lc_kyber_x25519_ss *ss0,
				     const struct lc_kyber_x25519_ss *ss1,
				     const uint8_t *in3, size_t inlen3,
				     uint8_t *out, size_t outlen)
{
	static const uint8_t kyber_x25519_ss_label[] =
		"Kyber X25519 KEM 3-way SS";
	LC_KMAC_CTX_ON_STACK(kmac_ctx, lc_cshake256);

	/*
	 * NOTE: this only works because struct lc_kyber_x25519_ss contains
	 * Kyber SS || X25519 SS in memory. If this structure changes,
	 * change this KDF invocation.
	 */
	lc_kmac_init(kmac_ctx, (uint8_t *)ss0,
		     sizeof(struct lc_kyber_x25519_ss), kyber_x25519_ss_label,
		     sizeof(kyber_x25519_ss_label) - 1);
	lc_kmac_update(kmac_ctx, (uint8_t *)ss1,
		       sizeof(struct lc_kyber_x25519_ss));
	lc_kmac_update(kmac_ctx, in3, inlen3);
	lc_kmac_final(kmac_ctx, out, outlen);

	lc_kmac_zero(kmac_ctx);
}

/**
 * @brief Kyber-X25519 KDF with 4 input values
 *
 * SS <- KMAC256(K = Kyber-SS || X25519-SS from ephemeral key 1,
 *		 X = Kyber-SS || X25519-SS from ephemeral key 2 ||
 *		     Kyber-SS || X25519-SS from static key || Nonce,
 *		 L = requested SS length,
 *		 S = "Kyber X25519 KEM 4-way SS")
 *
 * @param [in] ss0 SS0
 * @param [in] ss1 SS1
 * @param [in] ss2 SS2
 * @param [in] in4 input buffer 4
 * @param [in] inlen4 length of input buffer 4
 * @param [out] out output buffer of size
 * @param [in] outlen output buffer length
 */
static inline void kyber_x25519_kdf4(const struct lc_kyber_x25519_ss *ss0,
				     const struct lc_kyber_x25519_ss *ss1,
				     const struct lc_kyber_x25519_ss *ss2,
				     const uint8_t *in4, size_t inlen4,
				     uint8_t *out, size_t outlen)
{
	static const uint8_t kyber_x25519_ss_label[] =
		"Kyber X25519 KEM 4-way SS";
	LC_KMAC_CTX_ON_STACK(kmac_ctx, lc_cshake256);

	/*
	 * NOTE: this only works because struct lc_kyber_x25519_ss contqains
	 * Kyber SS || X25519 SS in memory. If this structure changes,
	 * change this KDF invocation.
	 */
	lc_kmac_init(kmac_ctx, (uint8_t *)ss0,
		     sizeof(struct lc_kyber_x25519_ss), kyber_x25519_ss_label,
		     sizeof(kyber_x25519_ss_label) - 1);
	lc_kmac_update(kmac_ctx, (uint8_t *)ss1,
		       sizeof(struct lc_kyber_x25519_ss));
	lc_kmac_update(kmac_ctx, (uint8_t *)ss2,
		       sizeof(struct lc_kyber_x25519_ss));
	lc_kmac_update(kmac_ctx, in4, inlen4);
	lc_kmac_final(kmac_ctx, out, outlen);

	lc_kmac_zero(kmac_ctx);
}

#ifdef __cplusplus
}
#endif

#endif /* KYBER_X25519_KDF_H */
