/****************************************************************************
*																			*
*					cryptlib ML-KEM Key Exchange Routines					*
*						Copyright Peter Gutmann 2025-2026					*
*																			*
****************************************************************************/

#define PKC_CONTEXT		/* Indicate that we're working with PKC contexts */
#include "crypt.h"
#if defined( INC_ALL )
  #include "context.h"
  #include "kem.h"
#else
  #include "context/context.h"
  #include "crypt/kem.h"
#endif /* Compiler-specific includes */

/* Defines from mlkem_native.h */

#define MLKEM768_CIPHERTEXTBYTES	1088
#define MLKEM768_SYMBYTES			32
#define MLKEM768_BYTES				32

/* The ML-KEM key exchange process is pretty peculiar because it acts like a
   kind of reverse-RSA, with the client providing the public key and the
   server encrypting a session key with it.  However since a new public key
   has to be generated each time it's also partly like Diffie-Hellman.  Or,
   to put it another way, like anything PQC it doesn't look like anything 
   normal.
   
   The sequence of operations is:

	A.load:		generate privK, pubK
				write pubK

	B.wrap		read pubK
				generate session-key
				wrap session-key using pubK

	A.unwrap	read wrapped key
				unwrap session-key using privK

   See FIPS 2.30 section 3 for details */

#ifdef USE_MLKEM

/****************************************************************************
*																			*
*								Utility Routines							*
*																			*
****************************************************************************/

/* The lower-level ML-KEM code uses a callback for its entropy generation, 
   this is just a wrapper for its use in keygen, where it requests 2 * 
   MLKEM_SYMBYTES, and encapsulation, where it requests MLKEM_SYMBYTES.
   The return value is checked for == 0 on success, != 0 on error, so it
   matches cryptlib's handling of error values */

int mlk_randombytes( uint8_t *out, size_t outlen )
	{
	MESSAGE_DATA msgData;

	assert( isWritePtr( out, outlen ) );

	ENSURES( outlen == MLKEM768_SYMBYTES || \
			 outlen == MLKEM768_SYMBYTES * 2 );
			 
	setMessageData( &msgData, out, outlen );
	return( krnlSendMessage( SYSTEM_OBJECT_HANDLE, IMESSAGE_GETATTRIBUTE_S,
							 &msgData, CRYPT_IATTRIBUTE_RANDOM ) );
	}

/****************************************************************************
*																			*
*								Algorithm Self-test							*
*																			*
****************************************************************************/

/* Perform a pairwise consistency test on a public/private key pair */

CHECK_RETVAL_BOOL STDC_NONNULL_ARG( ( 1 ) ) \
static BOOLEAN pairwiseConsistencyTest( CONTEXT_INFO *contextInfoPtr )
	{
	const CAPABILITY_INFO *capabilityInfoPtr;
	KEYAGREE_PARAMS keyAgreeParams;
	BYTE buffer[ MLKEM768_BYTES + 8 ];
	int status;

	assert( isWritePtr( contextInfoPtr, sizeof( CONTEXT_INFO ) ) );

	REQUIRES_B( sanityCheckContext( contextInfoPtr ) );

	capabilityInfoPtr = DATAPTR_GET( contextInfoPtr->capabilityInfo );
	REQUIRES_B( capabilityInfoPtr != NULL );

	/* Encrypt with the public key */
	memset( &keyAgreeParams, 0, sizeof( KEYAGREE_PARAMS ) );
	status = capabilityInfoPtr->encryptFunction( contextInfoPtr, 
					( BYTE * ) &keyAgreeParams, sizeof( KEYAGREE_PARAMS ) );
	if( cryptStatusError( status ) )
		return( FALSE );

	/* Take a copy of the shared key, which is generated by the encrypt 
	   function, and clear it from the key-agreement parameters */
	memcpy( buffer, keyAgreeParams.wrappedKey, MLKEM768_BYTES );
	memset( keyAgreeParams.wrappedKey, 0, MLKEM768_BYTES );

	/* Decrypt with the private key and make sure that we've recovered the 
	   shared key */
	status = capabilityInfoPtr->decryptFunction( contextInfoPtr, 
					( BYTE * ) &keyAgreeParams, sizeof( KEYAGREE_PARAMS ) );
	if( cryptStatusOK( status ) && \
		memcmp( buffer, keyAgreeParams.wrappedKey, MLKEM768_BYTES ) )
		status = CRYPT_ERROR_FAILED;

	/* Clean up */
	zeroise( &keyAgreeParams, sizeof( KEYAGREE_PARAMS ) );

	return( cryptStatusOK( status ) ? TRUE : FALSE );
	}

#ifndef CONFIG_NO_SELFTEST

/* Test the ML-KEM implementation.  Because a lot of the high-level 
   encryption routines don't exist yet we cheat a bit and set up a dummy 
   encryption context with just enough information for the following code to 
   work */

CHECK_RETVAL \
static int selfTest( void )
	{
	const CAPABILITY_INFO *capabilityInfoPtr;
	CONTEXT_INFO contextInfo;
	PKC_INFO contextData, *pkcInfo = &contextData;
	MLKEM_KEY_INFO *mlkemKey;
	int mlkStatus, status;

	static_assert( sizeof( MLKEM_KEY_INFO ) <= sizeof( BN_CTX ),
				   "ML-KEM storage exceeds allocated space" );
	static_assert( MLKEM768_CIPHERTEXTBYTES <= MLKEM768_PUBLICKEYBYTES,
				   "ML-KEM ciphertext size exceeds public key size" );
	static_assert( MLKEM768_PUBLICKEYBYTES <= MAX_PKCSIZE_PQC,
				   "ML-KEM public-key size exceeds MAX_PKCSIZE_PQC" );
	
	/* Initialise the context */
	status = staticInitContext( &contextInfo, CONTEXT_PKC, 
								getMLKEMCapability(), &contextData, 
								sizeof( PKC_INFO ), NULL );
	if( cryptStatusError( status ) )
		return( CRYPT_ERROR_FAILED );
	mlkemKey = pkcInfo->mlkemKey;
	
	/* Generate a public/private key pair and check that it's OK.  This 
	   checking is also performed by initKeyFunction() but we make it 
	   explicit here to catch any problems at this stage */
	mlkStatus = mlk_kem_keypair( mlkemKey->pubKey, mlkemKey->privKey, 0 );
	if( mlkStatus == 0 )
		mlkStatus = mlk_kem_check_sk( mlkemKey->privKey, 0 );
	if( mlkStatus == 0 )
		mlkStatus = mlk_kem_check_pk( mlkemKey->pubKey, 0 );
	if( mlkStatus != 0 )
		{
		staticDestroyContext( &contextInfo );
		retIntError();
		}
	mlkemKey->pubKeySize = MLKEM768_PUBLICKEYBYTES;
	mlkemKey->privKeySize = MLKEM768_SECRETKEYBYTES;

	ENSURES( sanityCheckPKCInfo( pkcInfo ) );

	/* Perform the test key exchange on a block of data */
	capabilityInfoPtr = DATAPTR_GET( contextInfo.capabilityInfo );
	REQUIRES( capabilityInfoPtr != NULL );
	status = capabilityInfoPtr->initKeyFunction( &contextInfo, NULL, 0 );
	if( cryptStatusOK( status ) && \
		!pairwiseConsistencyTest( &contextInfo ) )
		status = CRYPT_ERROR_FAILED;

	/* Clean up */
	staticDestroyContext( &contextInfo );

	return( status );
	}
#else
	#define selfTest	NULL
#endif /* !CONFIG_NO_SELFTEST */

/****************************************************************************
*																			*
*						Diffie-Hellman Key Exchange Routines				*
*																			*
****************************************************************************/

/* Perform ML-KEM encapsulation.  We have to append the distinguisher 'Fn' 
   to the name since some systems already have 'encrypt' and 'decrypt' in 
   their standard headers */

CHECK_RETVAL STDC_NONNULL_ARG( ( 1, 2 ) ) \
static int encryptFn( INOUT_PTR CONTEXT_INFO *contextInfoPtr, 
					  INOUT_BUFFER_FIXED( noBytes ) BYTE *buffer, 
					  IN_LENGTH_FIXED( sizeof( KEYAGREE_PARAMS ) ) int noBytes )
	{
	PKC_INFO *pkcInfo = contextInfoPtr->ctxPKC;
	MLKEM_KEY_INFO *mlkemKey = pkcInfo->mlkemKey;
	KEYAGREE_PARAMS *keyAgreeParams = ( KEYAGREE_PARAMS * ) buffer;

	assert( isWritePtr( contextInfoPtr, sizeof( CONTEXT_INFO ) ) );
	assert( isWritePtr( keyAgreeParams, sizeof( KEYAGREE_PARAMS ) ) );

	REQUIRES( sanityCheckContext( contextInfoPtr ) );
	REQUIRES( noBytes == sizeof( KEYAGREE_PARAMS ) );

	/* Clear return value */
	memset( keyAgreeParams, 0, sizeof( KEYAGREE_PARAMS ) );

	/* Generate the secret key and wrap it using ML-KEM.  Like everything 
	   PQC this works in a strange way, generating both the secret key being
	   wrapped and the wrapped key so that the operation has two outputs 
	   instead of one.  Since there's no normal way to communicate this to 
	   the caller we put the secret key in the wrappedKey storage */
	if( mlk_kem_enc( keyAgreeParams->publicValue, 
					 keyAgreeParams->wrappedKey, 
					 mlkemKey->pubKey, 0 ) != 0 )
		return( CRYPT_ERROR_FAILED );
	keyAgreeParams->publicValueLen = MLKEM768_CIPHERTEXTBYTES;
	keyAgreeParams->wrappedKeyLen = MLKEM768_BYTES;
	
	ENSURES( sanityCheckPKCInfo( pkcInfo ) );

	return( CRYPT_OK );
	}

/* Perform ML-KEM decapsulation */

CHECK_RETVAL STDC_NONNULL_ARG( ( 1, 2 ) ) \
static int decryptFn( INOUT_PTR CONTEXT_INFO *contextInfoPtr, 
					  INOUT_BUFFER_FIXED( noBytes ) BYTE *buffer, 
					  IN_LENGTH_FIXED( sizeof( KEYAGREE_PARAMS ) ) int noBytes )
	{
	PKC_INFO *pkcInfo = contextInfoPtr->ctxPKC;
	MLKEM_KEY_INFO *mlkemKey = pkcInfo->mlkemKey;
	KEYAGREE_PARAMS *keyAgreeParams = ( KEYAGREE_PARAMS * ) buffer;

	assert( isWritePtr( contextInfoPtr, sizeof( CONTEXT_INFO ) ) );
	assert( isWritePtr( keyAgreeParams, sizeof( KEYAGREE_PARAMS ) ) );
	assert( isReadPtrDynamic( keyAgreeParams->publicValue, 
							  keyAgreeParams->publicValueLen ) );

	REQUIRES( sanityCheckContext( contextInfoPtr ) );
	REQUIRES( noBytes == sizeof( KEYAGREE_PARAMS ) );

	/* Check that the input is valid.  There's no explicit check that the
	   data we're being passed is valid, see FIPS 203 section 7.3, however 
	   mlk_kem_dec() performs an mlk_kem_enc() after it's finished and 
	   checks that the result matches the input.  Unfortunately this is an 
	   ambulance-at-the-bottom-of-the-cliff approach that allows an attacker 
	   unfettered access to the internals of the ML-KEM code, but it's all 
	   that we've got.
	   
	   In particular ML-KEM has so many side-channels that one analysis,
	   admittedly of ML-DSA rather than ML-KEM, ("Migration to Post-Quantum
	   Cryptography: From ECDSA to ML-DSA", Daniel Dinu) says that "It is 
	   simpler to enumerate the operations and intermediate variables that 
	   are not sensitive to attacks" and "almost all operations of the 
	   ML-DSA algorithm are susceptible to these attacks".
	   
	   See the cryptlib threat model documentation for how we deal with 
	   this */
	if( keyAgreeParams->publicValueLen != MLKEM768_CIPHERTEXTBYTES )
		return( CRYPT_ERROR_BADDATA );

	/* Unwrap the secret key using ML-KEM */
	if( mlk_kem_dec( keyAgreeParams->wrappedKey, 
					 keyAgreeParams->publicValue, 
					 mlkemKey->privKey, 0 ) != 0 )
		return( CRYPT_ERROR_FAILED );
	keyAgreeParams->wrappedKeyLen = MLKEM768_BYTES;

	/* Make sure that the result looks random */
	if( !checkEntropy( keyAgreeParams->wrappedKey, 
					   keyAgreeParams->wrappedKeyLen ) )
		{
		zeroise( keyAgreeParams->wrappedKey, keyAgreeParams->wrappedKeyLen );
		return( CRYPT_ERROR_NOSECURE );
		}

	ENSURES( sanityCheckPKCInfo( pkcInfo ) );

	return( CRYPT_OK );
	}

/****************************************************************************
*																			*
*								Key Management								*
*																			*
****************************************************************************/

/* Load key components into an encryption context */

CHECK_RETVAL STDC_NONNULL_ARG( ( 1 ) ) \
static int initKey( INOUT_PTR CONTEXT_INFO *contextInfoPtr, 
					IN_BUFFER_OPT( keyLength ) const void *key,
					IN_LENGTH_SHORT_OPT const int keyLength )
	{
	PKC_INFO *pkcInfo = contextInfoPtr->ctxPKC;
	MLKEM_KEY_INFO *mlkemKey = pkcInfo->mlkemKey;

	assert( isWritePtr( contextInfoPtr, sizeof( CONTEXT_INFO ) ) );
	assert( ( key == NULL && keyLength == 0 ) || \
			( isReadPtrDynamic( key, keyLength ) && \
			  keyLength == sizeof( CRYPT_PKCINFO_PQC ) ) );

	REQUIRES( sanityCheckContext( contextInfoPtr ) );
	REQUIRES( ( key == NULL && keyLength == 0 ) || \
			  ( key != NULL && keyLength == sizeof( CRYPT_PKCINFO_PQC ) ) );

#ifndef USE_FIPS140
	/* Load the key component from the external representation into the
	   internal storage unless we're doing an internal load */
	if( key != NULL )
		{
		const CRYPT_PKCINFO_PQC *pqcKey = ( CRYPT_PKCINFO_PQC * ) key;

		if( pqcKey->isPublicKey )
			SET_FLAG( contextInfoPtr->flags, CONTEXT_FLAG_ISPUBLICKEY );
		else
			SET_FLAG( contextInfoPtr->flags, CONTEXT_FLAG_PBO );

		if( pqcKey->pubLen != bytesToBits( MLKEM768_PUBLICKEYBYTES ) )
			return( CRYPT_ERROR_BADDATA );
		REQUIRES( rangeCheck( bitsToBytes( pqcKey->pubLen ),
							  MLKEM768_PUBLICKEYBYTES, 
							  MLKEM768_PUBLICKEYBYTES ) );
		memcpy( mlkemKey->pubKey, pqcKey->pub, MLKEM768_PUBLICKEYBYTES );
		mlkemKey->pubKeySize = MLKEM768_PUBLICKEYBYTES;
		if( !pqcKey->isPublicKey )
			{
			if( pqcKey->privLen != bytesToBits( MLKEM768_SECRETKEYBYTES ) )
				return( CRYPT_ERROR_BADDATA );
			REQUIRES( rangeCheck( bitsToBytes( pqcKey->privLen ),
								  MLKEM768_SECRETKEYBYTES, 
								  MLKEM768_SECRETKEYBYTES ) );
			memcpy( mlkemKey->privKey, pqcKey->priv, MLKEM768_SECRETKEYBYTES );
			mlkemKey->privKeySize = MLKEM768_SECRETKEYBYTES;
			}

		ENSURES( sanityCheckPKCInfo( pkcInfo ) );
		}
#endif /* USE_FIPS140 */
	pkcInfo->keySizeBits = bytesToBits( MLKEM768_PUBLICKEYBYTES );

	/* Checksum the key data try and detect fault attacks.  Since we're
	   setting the checksum at this point there's no need to check the 
	   return value */
	( void ) checksumContextData( pkcInfo, TRUE );

	/* Complete the key checking and setup.  Since everything is handled by 
	   the low-level ML-KEM code we can do this directly rather than using 
	   a separate keygen module */
	if( mlk_kem_check_pk( mlkemKey->pubKey, 0 ) != 0 )
		return( CRYPT_ARGERROR_STR1 );
	if( !TEST_FLAG( contextInfoPtr->flags, CONTEXT_FLAG_ISPUBLICKEY ) && \
		mlk_kem_check_sk( mlkemKey->privKey, 0 ) != 0 )
		return( CRYPT_ARGERROR_STR1 );

	return( CRYPT_OK );
	}

/* Generate a key into an encryption context */

CHECK_RETVAL STDC_NONNULL_ARG( ( 1 ) ) \
static int generateKey( INOUT_PTR CONTEXT_INFO *contextInfoPtr, 
						IN_LENGTH_SHORT_MIN( MIN_PKCSIZE * 8 ) \
							const int keySizeBits )
	{
	PKC_INFO *pkcInfo = contextInfoPtr->ctxPKC;
	MLKEM_KEY_INFO *mlkemKey = pkcInfo->mlkemKey;
	int mlkStatus;

	assert( isWritePtr( contextInfoPtr, sizeof( CONTEXT_INFO ) ) );

	REQUIRES( sanityCheckContext( contextInfoPtr ) );
	REQUIRES( keySizeBits == bytesToBits( MLKEM768_PUBLICKEYBYTES ) );

	/* Generate and consistency-check a key.  Since everything is handled by 
	   the low-level ML-KEM code we can do this directly rather than using 
	   a separate keygen module */
	mlkStatus = mlk_kem_keypair( mlkemKey->pubKey, mlkemKey->privKey, 0 );
	if( mlkStatus == 0 )
		mlkStatus = mlk_kem_check_sk( mlkemKey->privKey, 0 );
	if( mlkStatus == 0 )
		mlkStatus = mlk_kem_check_pk( mlkemKey->pubKey, 0 );
	if( mlkStatus != 0 )
		return( CRYPT_ERROR_FAILED );
	mlkemKey->pubKeySize = MLKEM768_PUBLICKEYBYTES;
	mlkemKey->privKeySize = MLKEM768_SECRETKEYBYTES;
	if( !pairwiseConsistencyTest( contextInfoPtr ) )
		{
		DEBUG_DIAG(( "Consistency check of freshly-generated ML-KEM key "
					 "failed" ));
		assert( DEBUG_WARN );
		return( CRYPT_ERROR_FAILED );
		}
	pkcInfo->keySizeBits = bytesToBits( MLKEM768_PUBLICKEYBYTES );

	/* Checksum the key data try and detect fault attacks.  Since we're
	   setting the checksum at this point there's no need to check the 
	   return value */
	( void ) checksumContextData( pkcInfo, TRUE );

	return( CRYPT_OK );
	}

/****************************************************************************
*																			*
*						Capability Access Routines							*
*																			*
****************************************************************************/

static const CAPABILITY_INFO capabilityInfo = {
	CRYPT_ALGO_MLKEM, bitsToBytes( 0 ), "ML-KEM", 6,
	MLKEM768_PUBLICKEYBYTES, MLKEM768_PUBLICKEYBYTES, MLKEM768_PUBLICKEYBYTES,
	selfTest, getDefaultInfo, NULL, NULL, initKey, generateKey, 
	encryptFn, decryptFn, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, 
	readPublicKeyPQCFunction, writePublicKeyPQCFunction, 
	NULL, NULL
	};

CHECK_RETVAL_PTR_NONNULL \
const CAPABILITY_INFO *getMLKEMCapability( void )
	{
	return( &capabilityInfo );
	}
#endif /* USE_MLKEM */
