/src/SymCrypt/lib/mlkem.c

Source (jump to first uncovered line)
//
// mlkem.c   ML-KEM related functionality
//
// Copyright (c) Microsoft Corporation. Licensed under the MIT license.
//

#include "precomp.h"

const SYMCRYPT_MLKEM_INTERNAL_PARAMS SymCryptMlKemInternalParamsMlKem512 =
{
    .params         = SYMCRYPT_MLKEM_PARAMS_MLKEM512,
    .cbPolyElement  = SYMCRYPT_INTERNAL_MLKEM_SIZEOF_POLYRINGELEMENT,
    .nRows          = 2,
    .cbVector       = sizeof(SYMCRYPT_MLKEM_VECTOR) + (2*SYMCRYPT_INTERNAL_MLKEM_SIZEOF_POLYRINGELEMENT),
    .cbMatrix       = sizeof(SYMCRYPT_MLKEM_MATRIX) + (2*2*SYMCRYPT_INTERNAL_MLKEM_SIZEOF_POLYRINGELEMENT),
    .nEta1          = 3,
    .nEta2          = 2,
    .nBitsOfU       = 10,
    .nBitsOfV       = 4,
};

const SYMCRYPT_MLKEM_INTERNAL_PARAMS SymCryptMlKemInternalParamsMlKem768 =
{
    .params         = SYMCRYPT_MLKEM_PARAMS_MLKEM768,
    .cbPolyElement  = SYMCRYPT_INTERNAL_MLKEM_SIZEOF_POLYRINGELEMENT,
    .nRows          = 3,
    .cbVector       = sizeof(SYMCRYPT_MLKEM_VECTOR) + (3*SYMCRYPT_INTERNAL_MLKEM_SIZEOF_POLYRINGELEMENT),
    .cbMatrix       = sizeof(SYMCRYPT_MLKEM_MATRIX) + (3*3*SYMCRYPT_INTERNAL_MLKEM_SIZEOF_POLYRINGELEMENT),
    .nEta1          = 2,
    .nEta2          = 2,
    .nBitsOfU       = 10,
    .nBitsOfV       = 4,
};

const SYMCRYPT_MLKEM_INTERNAL_PARAMS SymCryptMlKemInternalParamsMlKem1024 =
{
    .params         = SYMCRYPT_MLKEM_PARAMS_MLKEM1024,
    .cbPolyElement  = SYMCRYPT_INTERNAL_MLKEM_SIZEOF_POLYRINGELEMENT,
    .nRows          = 4,
    .cbVector       = sizeof(SYMCRYPT_MLKEM_VECTOR) + (4*SYMCRYPT_INTERNAL_MLKEM_SIZEOF_POLYRINGELEMENT),
    .cbMatrix       = sizeof(SYMCRYPT_MLKEM_MATRIX) + (4*4*SYMCRYPT_INTERNAL_MLKEM_SIZEOF_POLYRINGELEMENT),
    .nEta1          = 2,
    .nEta2          = 2,
    .nBitsOfU       = 11,
    .nBitsOfV       = 5,
};

static
SYMCRYPT_ERROR
SYMCRYPT_CALL
SymCryptMlKemkeyGetInternalParamsFromParams(
            SYMCRYPT_MLKEM_PARAMS           params,
    _Out_   PSYMCRYPT_MLKEM_INTERNAL_PARAMS pInternalParams )
{
    SYMCRYPT_ERROR scError = SYMCRYPT_NO_ERROR;

    switch( params )
    {
        case SYMCRYPT_MLKEM_PARAMS_MLKEM512:
            *pInternalParams = SymCryptMlKemInternalParamsMlKem512;
            break;
        case SYMCRYPT_MLKEM_PARAMS_MLKEM768:
            *pInternalParams = SymCryptMlKemInternalParamsMlKem768;
            break;
        case SYMCRYPT_MLKEM_PARAMS_MLKEM1024:
            *pInternalParams = SymCryptMlKemInternalParamsMlKem1024;
            break;
        default:
            scError = SYMCRYPT_INVALID_ARGUMENT;
            goto cleanup;
    }

cleanup:
    return scError;
}

static
PSYMCRYPT_MLKEMKEY
SYMCRYPT_CALL
SymCryptMlKemkeyInitialize(
    _In_                        PCSYMCRYPT_MLKEM_INTERNAL_PARAMS    pInternalParams,
    _Out_writes_bytes_(cbKey)   PBYTE                               pbKey,
                                UINT32                              cbKey )
{
    PSYMCRYPT_MLKEMKEY pRes = NULL;
    PSYMCRYPT_MLKEMKEY pKey = (PSYMCRYPT_MLKEMKEY)pbKey;
    PBYTE pbCurr = pbKey + sizeof(SYMCRYPT_MLKEMKEY);

    SymCryptWipeKnownSize( pbKey, cbKey );

    pKey->fAlgorithmInfo = 0;
    pKey->params = *pInternalParams;
    pKey->cbTotalSize = cbKey;
    pKey->hasPrivateSeed = FALSE;
    pKey->hasPrivateKey = FALSE;

    pKey->pmAtranspose = SymCryptMlKemMatrixCreate( pbCurr, pInternalParams->cbMatrix, pInternalParams->nRows );
    if( pKey->pmAtranspose == NULL )
    {
        goto cleanup;
    }
    pbCurr += pInternalParams->cbMatrix;

    pKey->pvt = SymCryptMlKemVectorCreate( pbCurr, pInternalParams->cbVector, pInternalParams->nRows );
    if( pKey->pvt == NULL )
    {
        goto cleanup;
    }
    pbCurr += pInternalParams->cbVector;

    pKey->pvs = SymCryptMlKemVectorCreate( pbCurr, pInternalParams->cbVector, pInternalParams->nRows );
    if( pKey->pvs == NULL )
    {
        goto cleanup;
    }
    pbCurr += pInternalParams->cbVector;

    SYMCRYPT_ASSERT( pbCurr == (pbKey + cbKey) );

    SYMCRYPT_SET_MAGIC( pKey );

    pRes = pKey;

cleanup:
    return pRes;
}

PSYMCRYPT_MLKEMKEY
SYMCRYPT_CALL
SymCryptMlKemkeyAllocate(
    SYMCRYPT_MLKEM_PARAMS   params )
{
    SYMCRYPT_ERROR scError = SYMCRYPT_NO_ERROR;
    PBYTE  pbKey = NULL;
    UINT32 cbKey;
    SYMCRYPT_MLKEM_INTERNAL_PARAMS internalParams;

    PSYMCRYPT_MLKEMKEY pKey = NULL;

    scError = SymCryptMlKemkeyGetInternalParamsFromParams(params, &internalParams);
    if( scError != SYMCRYPT_NO_ERROR )
    {
        goto cleanup;
    }

    cbKey = sizeof(SYMCRYPT_MLKEMKEY) + internalParams.cbMatrix + (2*internalParams.cbVector);

    pbKey = SymCryptCallbackAlloc( cbKey );
    if ( pbKey == NULL )
    {
        goto cleanup;
    }

    pKey = SymCryptMlKemkeyInitialize( &internalParams, pbKey, cbKey );
    if ( pKey == NULL )
    {
        goto cleanup;
    }

    pbKey = NULL;

cleanup:
    if ( pbKey != NULL )
    {
        SymCryptCallbackFree( pbKey );
    }

    return pKey;
}

VOID
SYMCRYPT_CALL
SymCryptMlKemkeyFree(
    _Inout_ PSYMCRYPT_MLKEMKEY  pkMlKemkey )
{
    SYMCRYPT_CHECK_MAGIC( pkMlKemkey );

    SymCryptWipe( (PBYTE) pkMlKemkey, pkMlKemkey->cbTotalSize );

    SymCryptCallbackFree( pkMlKemkey );
}

#define SYMCRYPT_MLKEM_SIZEOF_ENCODED_UNCOMPRESSED_VECTOR(_nRows)   (384UL * _nRows)

// d and z are each 32 bytes
#define SYMCRYPT_MLKEM_SIZEOF_FORMAT_PRIVATE_SEED               (2*32)
// s and t are encoded uncompressed vectors
// public seed, H(encapsulation key) and z are each 32 bytes
#define SYMCRYPT_MLKEM_SIZEOF_FORMAT_DECAPSULATION_KEY(_nRows)  ((2*SYMCRYPT_MLKEM_SIZEOF_ENCODED_UNCOMPRESSED_VECTOR(_nRows)) + (3*32))
// t is encoded uncompressed vector
// public seed is 32 bytes
#define SYMCRYPT_MLKEM_SIZEOF_FORMAT_ENCAPSULATION_KEY(_nRows)  (SYMCRYPT_MLKEM_SIZEOF_ENCODED_UNCOMPRESSED_VECTOR(_nRows) + 32)

SYMCRYPT_ERROR
SYMCRYPT_CALL
SymCryptMlKemSizeofKeyFormatFromParams(
            SYMCRYPT_MLKEM_PARAMS       params,
            SYMCRYPT_MLKEMKEY_FORMAT    mlKemkeyFormat,
    _Out_   SIZE_T*                     pcbKeyFormat )
{
    SYMCRYPT_ERROR scError = SYMCRYPT_NO_ERROR;
    SYMCRYPT_MLKEM_INTERNAL_PARAMS internalParams;

    if( mlKemkeyFormat == SYMCRYPT_MLKEMKEY_FORMAT_NULL )
    {
        scError = SYMCRYPT_INCOMPATIBLE_FORMAT;
        goto cleanup;
    }

    scError = SymCryptMlKemkeyGetInternalParamsFromParams(params, &internalParams);
    if( scError != SYMCRYPT_NO_ERROR )
    {
        goto cleanup;
    }

    switch( mlKemkeyFormat )
    {
        case SYMCRYPT_MLKEMKEY_FORMAT_PRIVATE_SEED:
            *pcbKeyFormat = SYMCRYPT_MLKEM_SIZEOF_FORMAT_PRIVATE_SEED;
            break;

        case SYMCRYPT_MLKEMKEY_FORMAT_DECAPSULATION_KEY:
            *pcbKeyFormat = SYMCRYPT_MLKEM_SIZEOF_FORMAT_DECAPSULATION_KEY(internalParams.nRows);
            break;

        case SYMCRYPT_MLKEMKEY_FORMAT_ENCAPSULATION_KEY:
            *pcbKeyFormat = SYMCRYPT_MLKEM_SIZEOF_FORMAT_ENCAPSULATION_KEY(internalParams.nRows);
            break;

        default:
            scError = SYMCRYPT_INVALID_ARGUMENT;
            goto cleanup;
    }

cleanup:
    return scError;
}

SYMCRYPT_ERROR
SYMCRYPT_CALL
SymCryptMlKemSizeofCiphertextFromParams(
            SYMCRYPT_MLKEM_PARAMS       params,
    _Out_   SIZE_T*                     pcbCiphertext )
{
    SYMCRYPT_ERROR scError = SYMCRYPT_NO_ERROR;
    SYMCRYPT_MLKEM_INTERNAL_PARAMS internalParams;
    SIZE_T cbU, cbV;

    scError = SymCryptMlKemkeyGetInternalParamsFromParams(params, &internalParams);
    if( scError != SYMCRYPT_NO_ERROR )
    {
        goto cleanup;
    }

    // u vector encoded with nBitsOfU * SYMCRYPT_MLWE_POLYNOMIAL_COEFFICIENTS bits per polynomial
    cbU = ((SIZE_T)internalParams.nRows) * internalParams.nBitsOfU * (SYMCRYPT_MLWE_POLYNOMIAL_COEFFICIENTS / 8);
    // v polynomial encoded with nBitsOfV * SYMCRYPT_MLWE_POLYNOMIAL_COEFFICIENTS bits
    cbV = ((SIZE_T)internalParams.nBitsOfV) * (SYMCRYPT_MLWE_POLYNOMIAL_COEFFICIENTS / 8);
    *pcbCiphertext = cbU + cbV;

    SYMCRYPT_ASSERT( (internalParams.params != SYMCRYPT_MLKEM_PARAMS_MLKEM512)  || ((cbU + cbV) == SYMCRYPT_MLKEM_CIPHERTEXT_SIZE_MLKEM512)  );
    SYMCRYPT_ASSERT( (internalParams.params != SYMCRYPT_MLKEM_PARAMS_MLKEM768)  || ((cbU + cbV) == SYMCRYPT_MLKEM_CIPHERTEXT_SIZE_MLKEM768)  );
    SYMCRYPT_ASSERT( (internalParams.params != SYMCRYPT_MLKEM_PARAMS_MLKEM1024) || ((cbU + cbV) == SYMCRYPT_MLKEM_CIPHERTEXT_SIZE_MLKEM1024) );

cleanup:
    return scError;
}

static
VOID
SYMCRYPT_CALL
SymCryptMlKemkeyExpandPublicMatrixFromPublicSeed(
    _Inout_ PSYMCRYPT_MLKEMKEY                                  pkMlKemkey,
    _Inout_ PSYMCRYPT_MLKEM_INTERNAL_COMPUTATION_TEMPORARIES    pCompTemps )
{
    UINT32 i, j;
    BYTE coordinates[2];

    PSYMCRYPT_SHAKE128_STATE pShakeStateBase = &pCompTemps->hashState0.shake128State;
    PSYMCRYPT_SHAKE128_STATE pShakeStateWork = &pCompTemps->hashState1.shake128State;
    const UINT32 nRows = pkMlKemkey->params.nRows;

    SymCryptShake128Init( pShakeStateBase );
    SymCryptShake128Append( pShakeStateBase, pkMlKemkey->publicSeed, sizeof(pkMlKemkey->publicSeed) );

    for( i=0; i<nRows; i++ )
    {
        coordinates[1] = (BYTE)i;
        for( j=0; j<nRows; j++ )
        {
            coordinates[0] = (BYTE)j;
            SymCryptShake128StateCopy( pShakeStateBase, pShakeStateWork );
            SymCryptShake128Append( pShakeStateWork, coordinates, sizeof(coordinates) );

            SymCryptMlKemPolyElementSampleNTTFromShake128( pShakeStateWork, pkMlKemkey->pmAtranspose->apPolyElements[(i*nRows)+j] );
        }
    }

    // no need to wipe; everything computed here is always public
}

static
VOID
SYMCRYPT_CALL
SymCryptMlKemkeyComputeEncapsulationKeyHash(
    _Inout_ PSYMCRYPT_MLKEMKEY                                  pkMlKemkey,
    _Inout_ PSYMCRYPT_MLKEM_INTERNAL_COMPUTATION_TEMPORARIES    pCompTemps,
            SIZE_T                                              cbEncodedVector )
{
    PSYMCRYPT_SHA3_256_STATE pState = &pCompTemps->hashState0.sha3_256State;

    SymCryptSha3_256Init( pState );
    SymCryptSha3_256Append( pState, pkMlKemkey->encodedT, cbEncodedVector );
    SymCryptSha3_256Append( pState, pkMlKemkey->publicSeed, sizeof(pkMlKemkey->publicSeed) );
    SymCryptSha3_256Result( pState, pkMlKemkey->encapsKeyHash );
}

static
VOID
SYMCRYPT_CALL
SymCryptMlKemkeyExpandFromPrivateSeed(
    _Inout_ PSYMCRYPT_MLKEMKEY                                  pkMlKemkey,
    _Inout_ PSYMCRYPT_MLKEM_INTERNAL_COMPUTATION_TEMPORARIES    pCompTemps )
{
    BYTE privateSeedHash[SYMCRYPT_SHA3_512_RESULT_SIZE];
    BYTE CBDSampleBuffer[3*64 + 1];
    PSYMCRYPT_MLKEM_VECTOR pvTmp;
    PSYMCRYPT_MLKEM_POLYELEMENT_ACCUMULATOR paTmp;
    PSYMCRYPT_SHAKE256_STATE pShakeStateBase = &pCompTemps->hashState0.shake256State;
    PSYMCRYPT_SHAKE256_STATE pShakeStateWork = &pCompTemps->hashState1.shake256State;
    UINT32 i;
    const UINT32 nRows = pkMlKemkey->params.nRows;
    const UINT32 nEta1 = pkMlKemkey->params.nEta1;
    const SIZE_T cbEncodedVector = SYMCRYPT_MLKEM_SIZEOF_ENCODED_UNCOMPRESSED_VECTOR(nRows);
    const UINT32 cbPolyElement = pkMlKemkey->params.cbPolyElement;
    const UINT32 cbVector = pkMlKemkey->params.cbVector;

    SYMCRYPT_ASSERT( pkMlKemkey->hasPrivateSeed );
    SYMCRYPT_ASSERT( (nEta1 == 2) || (nEta1 == 3) );
    SYMCRYPT_ASSERT( cbEncodedVector <= sizeof(pkMlKemkey->encodedT) );

    pvTmp = SymCryptMlKemVectorCreate( pCompTemps->abVectorBuffer0, cbVector, nRows );
    SYMCRYPT_ASSERT( pvTmp != NULL );
    paTmp = SymCryptMlKemPolyElementAccumulatorCreate( pCompTemps->abPolyElementAccumulatorBuffer, 2*cbPolyElement );
    SYMCRYPT_ASSERT( paTmp != NULL );

    // (rho || sigma) = G(d || k)
    // use CBDSampleBuffer to concatenate the private seed and encoding of nRows
    memcpy( CBDSampleBuffer, pkMlKemkey->privateSeed, sizeof(pkMlKemkey->privateSeed) );
    CBDSampleBuffer[sizeof(pkMlKemkey->privateSeed)] = (BYTE) nRows;
    SymCryptSha3_512( CBDSampleBuffer, sizeof(pkMlKemkey->privateSeed)+1, privateSeedHash );

    // copy public seed
    memcpy( pkMlKemkey->publicSeed, privateSeedHash, sizeof(pkMlKemkey->publicSeed) );

    // generate A from public seed
    SymCryptMlKemkeyExpandPublicMatrixFromPublicSeed( pkMlKemkey, pCompTemps );

    // Initialize pShakeStateBase with sigma
    SymCryptShake256Init( pShakeStateBase );
    SymCryptShake256Append( pShakeStateBase, privateSeedHash+sizeof(pkMlKemkey->publicSeed), 32 );

    // Expand s in place
    for( i=0; i<nRows; i++ )
    {
        CBDSampleBuffer[0] = (BYTE) i;
        SymCryptShake256StateCopy( pShakeStateBase, pShakeStateWork );
        SymCryptShake256Append( pShakeStateWork, CBDSampleBuffer, 1 );

        SymCryptShake256Extract( pShakeStateWork, CBDSampleBuffer, 64ul*nEta1, FALSE );

        SymCryptMlKemPolyElementSampleCBDFromBytes( CBDSampleBuffer, nEta1, SYMCRYPT_INTERNAL_MLKEM_VECTOR_ELEMENT(i, pkMlKemkey->pvs) );
    }
    // Expand e in t, ready for multiply-add
    for( i=0; i<nRows; i++ )
    {
        CBDSampleBuffer[0] = (BYTE) (nRows+i);
        SymCryptShake256StateCopy( pShakeStateBase, pShakeStateWork );
        SymCryptShake256Append( pShakeStateWork, CBDSampleBuffer, 1 );

        SymCryptShake256Extract( pShakeStateWork, CBDSampleBuffer, 64ul*nEta1, FALSE );

        SymCryptMlKemPolyElementSampleCBDFromBytes( CBDSampleBuffer, nEta1, SYMCRYPT_INTERNAL_MLKEM_VECTOR_ELEMENT(i, pkMlKemkey->pvt) );
    }

    // Perform NTT on s and e
    SymCryptMlKemVectorNTT( pkMlKemkey->pvs );
    SymCryptMlKemVectorNTT( pkMlKemkey->pvt );

    // pvTmp = s .* R
    SymCryptMlKemVectorMulR( pkMlKemkey->pvs, pvTmp );

    // t = ((A o (s .* R)) ./ R) + e = A o s + e
    SymCryptMlKemMatrixVectorMontMulAndAdd( pkMlKemkey->pmAtranspose, pvTmp, pkMlKemkey->pvt, paTmp );

    // transpose A
    SymCryptMlKemMatrixTranspose( pkMlKemkey->pmAtranspose );

    // precompute byte-encoding of public vector t
    SymCryptMlKemVectorCompressAndEncode( pkMlKemkey->pvt, 12, pkMlKemkey->encodedT, cbEncodedVector );

    // precompute hash of encapsulation key blob
    SymCryptMlKemkeyComputeEncapsulationKeyHash( pkMlKemkey, pCompTemps, cbEncodedVector );

    // Cleanup!
    SymCryptWipeKnownSize( privateSeedHash, sizeof(privateSeedHash) );
    SymCryptWipeKnownSize( CBDSampleBuffer, sizeof(CBDSampleBuffer) );
}

SYMCRYPT_ERROR
SYMCRYPT_CALL
SymCryptMlKemkeySetValue(
    _In_reads_bytes_( cbSrc )   PCBYTE                      pbSrc,
                                SIZE_T                      cbSrc,
                                SYMCRYPT_MLKEMKEY_FORMAT    mlKemkeyFormat,
                                UINT32                      flags,
    _Inout_                     PSYMCRYPT_MLKEMKEY          pkMlKemkey )
{
    SYMCRYPT_ERROR scError = SYMCRYPT_NO_ERROR;
    PCBYTE pbCurr = pbSrc;
    PSYMCRYPT_MLKEM_INTERNAL_COMPUTATION_TEMPORARIES pCompTemps = NULL;
    const UINT32 nRows = pkMlKemkey->params.nRows;
    const SIZE_T cbEncodedVector = SYMCRYPT_MLKEM_SIZEOF_ENCODED_UNCOMPRESSED_VECTOR( nRows );

    // Ensure only allowed flags are specified
    UINT32 allowedFlags = SYMCRYPT_FLAG_KEY_NO_FIPS | SYMCRYPT_FLAG_KEY_MINIMAL_VALIDATION;

    if ( ( flags & ~allowedFlags ) != 0 )
    {
        scError = SYMCRYPT_INVALID_ARGUMENT;
        goto cleanup;
    }

    // Check that minimal validation flag only specified with no fips
    if ( ( ( flags & SYMCRYPT_FLAG_KEY_NO_FIPS ) == 0 ) &&
         ( ( flags & SYMCRYPT_FLAG_KEY_MINIMAL_VALIDATION ) != 0 ) )
    {
        scError = SYMCRYPT_INVALID_ARGUMENT;
        goto cleanup;
    }

    if( mlKemkeyFormat == SYMCRYPT_MLKEMKEY_FORMAT_NULL )
    {
        scError = SYMCRYPT_INVALID_ARGUMENT;
        goto cleanup;
    }

    if( ( flags & SYMCRYPT_FLAG_KEY_NO_FIPS ) == 0 )
    {
        // Ensure ML-KEM algorithm selftest is run before first use of ML-KEM algorithms;
        // notably _before_ first full KeyGen
        SYMCRYPT_RUN_SELFTEST_ONCE(
            SymCryptMlKemSelftest,
            SYMCRYPT_SELFTEST_ALGORITHM_MLKEM);
    }

    pCompTemps = SymCryptCallbackAlloc( sizeof(SYMCRYPT_MLKEM_INTERNAL_COMPUTATION_TEMPORARIES) );
    if( pCompTemps == NULL )
    {
        scError = SYMCRYPT_MEMORY_ALLOCATION_FAILURE;
        goto cleanup;
    }

    if( mlKemkeyFormat == SYMCRYPT_MLKEMKEY_FORMAT_PRIVATE_SEED )
    {
        if( cbSrc != SYMCRYPT_MLKEM_SIZEOF_FORMAT_PRIVATE_SEED )
        {
            scError = SYMCRYPT_WRONG_KEY_SIZE;
            goto cleanup;
        }

        pkMlKemkey->hasPrivateSeed = TRUE;
        memcpy( pkMlKemkey->privateSeed, pbCurr, sizeof(pkMlKemkey->privateSeed) );
        pbCurr += sizeof(pkMlKemkey->privateSeed);

        pkMlKemkey->hasPrivateKey = TRUE;
        memcpy( pkMlKemkey->privateRandom, pbCurr, sizeof(pkMlKemkey->privateRandom) );
        pbCurr += sizeof(pkMlKemkey->privateRandom);

        SymCryptMlKemkeyExpandFromPrivateSeed( pkMlKemkey, pCompTemps );
    }
    else if( mlKemkeyFormat == SYMCRYPT_MLKEMKEY_FORMAT_DECAPSULATION_KEY )
    {
        if( cbSrc != SYMCRYPT_MLKEM_SIZEOF_FORMAT_DECAPSULATION_KEY( nRows ) )
        {
            scError = SYMCRYPT_WRONG_KEY_SIZE;
            goto cleanup;
        }

        // decode s
        scError = SymCryptMlKemVectorDecodeAndDecompress( pbCurr, cbEncodedVector, 12, pkMlKemkey->pvs );
        if( scError != SYMCRYPT_NO_ERROR )
        {
            goto cleanup;
        }
        pbCurr += cbEncodedVector;

        // copy t and decode t
        memcpy( pkMlKemkey->encodedT, pbCurr, cbEncodedVector );
        pbCurr += cbEncodedVector;
        scError = SymCryptMlKemVectorDecodeAndDecompress( pkMlKemkey->encodedT, cbEncodedVector, 12, pkMlKemkey->pvt );
        if( scError != SYMCRYPT_NO_ERROR )
        {
            goto cleanup;
        }

        // copy public seed and expand public matrix
        memcpy( pkMlKemkey->publicSeed, pbCurr, sizeof(pkMlKemkey->publicSeed) );
        pbCurr += sizeof(pkMlKemkey->publicSeed);
        SymCryptMlKemkeyExpandPublicMatrixFromPublicSeed( pkMlKemkey, pCompTemps );

        // transpose A
        SymCryptMlKemMatrixTranspose( pkMlKemkey->pmAtranspose );

        // copy hash of encapsulation key
        memcpy( pkMlKemkey->encapsKeyHash, pbCurr, sizeof(pkMlKemkey->encapsKeyHash) );
        pbCurr += sizeof(pkMlKemkey->encapsKeyHash);

        // copy private random
        memcpy( pkMlKemkey->privateRandom, pbCurr, sizeof(pkMlKemkey->privateRandom) );
        pbCurr += sizeof(pkMlKemkey->privateRandom);

        pkMlKemkey->hasPrivateSeed = FALSE;
        pkMlKemkey->hasPrivateKey  = TRUE;
    }
    else if( mlKemkeyFormat == SYMCRYPT_MLKEMKEY_FORMAT_ENCAPSULATION_KEY )
    {
        if( cbSrc != SYMCRYPT_MLKEM_SIZEOF_FORMAT_ENCAPSULATION_KEY( nRows ) )
        {
            scError = SYMCRYPT_WRONG_KEY_SIZE;
            goto cleanup;
        }

        // copy t and decode t
        memcpy( pkMlKemkey->encodedT, pbCurr, cbEncodedVector );
        pbCurr += cbEncodedVector;
        scError = SymCryptMlKemVectorDecodeAndDecompress( pkMlKemkey->encodedT, cbEncodedVector, 12, pkMlKemkey->pvt );
        if( scError != SYMCRYPT_NO_ERROR )
        {
            goto cleanup;
        }

        // copy public seed and expand public matrix
        memcpy( pkMlKemkey->publicSeed, pbCurr, sizeof(pkMlKemkey->publicSeed) );
        pbCurr += sizeof(pkMlKemkey->publicSeed);
        SymCryptMlKemkeyExpandPublicMatrixFromPublicSeed( pkMlKemkey, pCompTemps );

        // transpose A
        SymCryptMlKemMatrixTranspose( pkMlKemkey->pmAtranspose );

        // precompute hash of encapsulation key blob
        SymCryptMlKemkeyComputeEncapsulationKeyHash( pkMlKemkey, pCompTemps, cbEncodedVector );

        pkMlKemkey->hasPrivateSeed = FALSE;
        pkMlKemkey->hasPrivateKey  = FALSE;
    }
    else
    {
        scError = SYMCRYPT_NOT_IMPLEMENTED;
        goto cleanup;
    }

    SYMCRYPT_ASSERT( pbCurr == pbSrc + cbSrc );

cleanup:
    if( pCompTemps != NULL )
    {
        SymCryptWipe( pCompTemps, sizeof(*pCompTemps) );
        SymCryptCallbackFree( pCompTemps );
    }

    return scError;
}


SYMCRYPT_ERROR
SYMCRYPT_CALL
SymCryptMlKemkeyGetValue(
    _In_                        PCSYMCRYPT_MLKEMKEY         pkMlKemkey,
    _Out_writes_bytes_( cbDst ) PBYTE                       pbDst,
                                SIZE_T                      cbDst,
                                SYMCRYPT_MLKEMKEY_FORMAT    mlKemkeyFormat,
                                UINT32                      flags )
{
    SYMCRYPT_ERROR scError = SYMCRYPT_NO_ERROR;
    PBYTE pbCurr = pbDst;
    const UINT32 nRows = pkMlKemkey->params.nRows;
    const SIZE_T cbEncodedVector = SYMCRYPT_MLKEM_SIZEOF_ENCODED_UNCOMPRESSED_VECTOR( nRows );

    UNREFERENCED_PARAMETER( flags );

    if( mlKemkeyFormat == SYMCRYPT_MLKEMKEY_FORMAT_NULL )
    {
        scError = SYMCRYPT_INVALID_ARGUMENT;
        goto cleanup;
    }

    if( mlKemkeyFormat == SYMCRYPT_MLKEMKEY_FORMAT_PRIVATE_SEED )
    {
        if( cbDst != SYMCRYPT_MLKEM_SIZEOF_FORMAT_PRIVATE_SEED )
        {
            scError = SYMCRYPT_WRONG_KEY_SIZE;
            goto cleanup;
        }

        if( !pkMlKemkey->hasPrivateSeed )
        {
            scError = SYMCRYPT_INCOMPATIBLE_FORMAT;
            goto cleanup;
        }

        memcpy( pbCurr, pkMlKemkey->privateSeed, sizeof(pkMlKemkey->privateSeed) );
        pbCurr += sizeof(pkMlKemkey->privateSeed);

        memcpy( pbCurr, pkMlKemkey->privateRandom, sizeof(pkMlKemkey->privateRandom) );
        pbCurr += sizeof(pkMlKemkey->privateRandom);
    }
    else if( mlKemkeyFormat == SYMCRYPT_MLKEMKEY_FORMAT_DECAPSULATION_KEY )
    {
        if( cbDst != SYMCRYPT_MLKEM_SIZEOF_FORMAT_DECAPSULATION_KEY( nRows ) )
        {
            scError = SYMCRYPT_INVALID_ARGUMENT;
            goto cleanup;
        }

        if( !pkMlKemkey->hasPrivateKey )
        {
            scError = SYMCRYPT_INVALID_ARGUMENT;
            goto cleanup;
        }

        // We don't precompute byte-encoding of private key as exporting decapsulation key is not a critical path operation
        // All other fields are kept in memory
        SymCryptMlKemVectorCompressAndEncode( pkMlKemkey->pvs, 12, pbCurr, cbEncodedVector );
        pbCurr += cbEncodedVector;

        memcpy( pbCurr, pkMlKemkey->encodedT, cbEncodedVector );
        pbCurr += cbEncodedVector;

        memcpy( pbCurr, pkMlKemkey->publicSeed, sizeof(pkMlKemkey->publicSeed) );
        pbCurr += sizeof(pkMlKemkey->publicSeed);

        memcpy( pbCurr, pkMlKemkey->encapsKeyHash, sizeof(pkMlKemkey->encapsKeyHash) );
        pbCurr += sizeof(pkMlKemkey->encapsKeyHash);

        memcpy( pbCurr, pkMlKemkey->privateRandom, sizeof(pkMlKemkey->privateRandom) );
        pbCurr += sizeof(pkMlKemkey->privateRandom);
    }
    else if( mlKemkeyFormat == SYMCRYPT_MLKEMKEY_FORMAT_ENCAPSULATION_KEY )
    {
        if( cbDst != SYMCRYPT_MLKEM_SIZEOF_FORMAT_ENCAPSULATION_KEY( nRows ) )
        {
            scError = SYMCRYPT_INVALID_ARGUMENT;
            goto cleanup;
        }

        memcpy( pbCurr, pkMlKemkey->encodedT, cbEncodedVector );
        pbCurr += cbEncodedVector;

        memcpy( pbCurr, pkMlKemkey->publicSeed, sizeof(pkMlKemkey->publicSeed) );
        pbCurr += sizeof(pkMlKemkey->publicSeed);
    }
    else
    {
        scError = SYMCRYPT_NOT_IMPLEMENTED;
        goto cleanup;
    }

    SYMCRYPT_ASSERT( pbCurr == pbDst + cbDst );

cleanup:
    return scError;
}


SYMCRYPT_ERROR
SYMCRYPT_CALL
SymCryptMlKemkeyGenerate(
    _Inout_                     PSYMCRYPT_MLKEMKEY  pkMlKemkey,
                                UINT32              flags )
{
    SYMCRYPT_ERROR scError = SYMCRYPT_NO_ERROR;
    BYTE privateSeed[SYMCRYPT_MLKEM_SIZEOF_FORMAT_PRIVATE_SEED];

    // Ensure only allowed flags are specified
    UINT32 allowedFlags = SYMCRYPT_FLAG_KEY_NO_FIPS;

    if ( ( flags & ~allowedFlags ) != 0 )
    {
        scError = SYMCRYPT_INVALID_ARGUMENT;
        goto cleanup;
    }

    scError = SymCryptCallbackRandom( privateSeed, sizeof(privateSeed) );
    if( scError != SYMCRYPT_NO_ERROR )
    {
        goto cleanup;
    }

    scError = SymCryptMlKemkeySetValue( privateSeed, sizeof(privateSeed), SYMCRYPT_MLKEMKEY_FORMAT_PRIVATE_SEED, flags, pkMlKemkey );
    if( scError != SYMCRYPT_NO_ERROR )
    {
        goto cleanup;
    }

    // SymCryptMlKemkeySetValue ensures the self-test is run before
    // first operational use of MlKem

    // Awaiting feedback from NIST for discussion from PQC forum and CMUF
    // before implementing costly PCT on ML-KEM key generation which is
    // not expected by FIPS 203

cleanup:
    SymCryptWipeKnownSize( privateSeed, sizeof(privateSeed) );

    return scError;
}

SYMCRYPT_ERROR
SYMCRYPT_CALL
SymCryptMlKemEncapsulateInternal(
    _In_    PCSYMCRYPT_MLKEMKEY                                 pkMlKemkey,
    _Out_writes_bytes_( cbAgreedSecret )
            PBYTE                                               pbAgreedSecret,
            SIZE_T                                              cbAgreedSecret,
    _Out_writes_bytes_( cbCiphertext )
            PBYTE                                               pbCiphertext,
            SIZE_T                                              cbCiphertext,
    _In_reads_bytes_( SYMCRYPT_MLKEM_SIZEOF_ENCAPS_RANDOM )
            PCBYTE                                              pbRandom,
    _Inout_ PSYMCRYPT_MLKEM_INTERNAL_COMPUTATION_TEMPORARIES    pCompTemps )
{
    BYTE CBDSampleBuffer[3*64 + 1];
    SYMCRYPT_ERROR scError = SYMCRYPT_NO_ERROR;
    PSYMCRYPT_MLKEM_VECTOR pvrInner;
    PSYMCRYPT_MLKEM_VECTOR pvTmp;
    PSYMCRYPT_MLKEM_POLYELEMENT peTmp0, peTmp1;
    PSYMCRYPT_MLKEM_POLYELEMENT_ACCUMULATOR paTmp;
    PSYMCRYPT_SHA3_512_STATE pHashState = &pCompTemps->hashState0.sha3_512State;
    PSYMCRYPT_SHAKE256_STATE pShakeBaseState = &pCompTemps->hashState0.shake256State;
    PSYMCRYPT_SHAKE256_STATE pShakeWorkState = &pCompTemps->hashState1.shake256State;
    SIZE_T cbU, cbV;
    UINT32 i;
    const UINT32 nRows = pkMlKemkey->params.nRows;
    const UINT32 nBitsOfU = pkMlKemkey->params.nBitsOfU;
    const UINT32 nBitsOfV = pkMlKemkey->params.nBitsOfV;
    const UINT32 nEta1 = pkMlKemkey->params.nEta1;
    const UINT32 nEta2 = pkMlKemkey->params.nEta2;
    const UINT32 cbPolyElement = pkMlKemkey->params.cbPolyElement;
    const UINT32 cbVector = pkMlKemkey->params.cbVector;

    // u vector encoded with nBitsOfU * SYMCRYPT_MLWE_POLYNOMIAL_COEFFICIENTS bits per polynomial
    cbU = nRows * nBitsOfU * (SYMCRYPT_MLWE_POLYNOMIAL_COEFFICIENTS / 8);
    // v polynomial encoded with nBitsOfV * SYMCRYPT_MLWE_POLYNOMIAL_COEFFICIENTS bits
    cbV = nBitsOfV * (SYMCRYPT_MLWE_POLYNOMIAL_COEFFICIENTS / 8);

    if( (cbAgreedSecret != SYMCRYPT_MLKEM_SIZEOF_AGREED_SECRET) ||
        (cbCiphertext != cbU + cbV) )
    {
        scError = SYMCRYPT_INVALID_ARGUMENT;
        goto cleanup;
    }

    pvrInner = SymCryptMlKemVectorCreate( pCompTemps->abVectorBuffer0, cbVector, nRows );
    SYMCRYPT_ASSERT( pvrInner != NULL );
    pvTmp = SymCryptMlKemVectorCreate( pCompTemps->abVectorBuffer1, cbVector, nRows );
    SYMCRYPT_ASSERT( pvTmp != NULL );
    peTmp0 = SymCryptMlKemPolyElementCreate( pCompTemps->abPolyElementBuffer0, cbPolyElement );
    SYMCRYPT_ASSERT( peTmp0 != NULL );
    peTmp1 = SymCryptMlKemPolyElementCreate( pCompTemps->abPolyElementBuffer1, cbPolyElement );
    SYMCRYPT_ASSERT( peTmp1 != NULL );
    paTmp = SymCryptMlKemPolyElementAccumulatorCreate( pCompTemps->abPolyElementAccumulatorBuffer, 2*cbPolyElement );
    SYMCRYPT_ASSERT( paTmp != NULL );

    // CBDSampleBuffer = (K || rOuter) = SHA3-512(pbRandom || encapsKeyHash)
    SymCryptSha3_512Init( pHashState );
    SymCryptSha3_512Append( pHashState, pbRandom, SYMCRYPT_MLKEM_SIZEOF_ENCAPS_RANDOM );
    SymCryptSha3_512Append( pHashState, pkMlKemkey->encapsKeyHash, sizeof(pkMlKemkey->encapsKeyHash) );
    SymCryptSha3_512Result( pHashState, CBDSampleBuffer );

    // Write K to pbAgreedSecret
    memcpy( pbAgreedSecret, CBDSampleBuffer, SYMCRYPT_MLKEM_SIZEOF_AGREED_SECRET );

    // Initialize pShakeStateBase with rOuter
    SymCryptShake256Init( pShakeBaseState );
    SymCryptShake256Append( pShakeBaseState, CBDSampleBuffer+cbAgreedSecret, 32 );

    // Expand rInner vector
    for( i=0; i<nRows; i++ )
    {
        CBDSampleBuffer[0] = (BYTE) i;
        SymCryptShake256StateCopy( pShakeBaseState, pShakeWorkState );
        SymCryptShake256Append( pShakeWorkState, CBDSampleBuffer, 1 );

        SymCryptShake256Extract( pShakeWorkState, CBDSampleBuffer, 64ul*nEta1, FALSE );

        SymCryptMlKemPolyElementSampleCBDFromBytes( CBDSampleBuffer, nEta1, SYMCRYPT_INTERNAL_MLKEM_VECTOR_ELEMENT(i, pvrInner) );
    }

    // Perform NTT on rInner
    SymCryptMlKemVectorNTT( pvrInner );

    // Set pvTmp to 0
    SymCryptMlKemVectorSetZero( pvTmp );

    // pvTmp = (Atranspose o rInner) ./ R
    SymCryptMlKemMatrixVectorMontMulAndAdd( pkMlKemkey->pmAtranspose, pvrInner, pvTmp, paTmp );

    // pvTmp = INTT(Atranspose o rInner)
    SymCryptMlKemVectorINTTAndMulR( pvTmp );

    // Expand e1 and add it to pvTmp - do addition PolyElement-wise to reduce memory usage
    for( i=0; i<nRows; i++ )
    {
        CBDSampleBuffer[0] = (BYTE) (nRows+i);
        SymCryptShake256StateCopy( pShakeBaseState, pShakeWorkState );
        SymCryptShake256Append( pShakeWorkState, CBDSampleBuffer, 1 );

        SymCryptShake256Extract( pShakeWorkState, CBDSampleBuffer, 64ul*nEta2, FALSE );

        SymCryptMlKemPolyElementSampleCBDFromBytes( CBDSampleBuffer, nEta2, peTmp0 );

        SymCryptMlKemPolyElementAdd( SYMCRYPT_INTERNAL_MLKEM_VECTOR_ELEMENT(i, pvTmp), peTmp0, SYMCRYPT_INTERNAL_MLKEM_VECTOR_ELEMENT(i, pvTmp) );
    }

    // pvTmp = u = INTT(Atranspose o rInner) + e1
    // Compress and encode u into prefix of ciphertext
    SymCryptMlKemVectorCompressAndEncode( pvTmp, nBitsOfU, pbCiphertext, cbU );

    // peTmp0 = (t o r) ./ R
    SymCryptMlKemVectorMontDotProduct( pkMlKemkey->pvt, pvrInner, peTmp0, paTmp );

    // peTmp0 = INTT(t o r)
    SymCryptMlKemPolyElementINTTAndMulR( peTmp0 );

    // Expand e2 polynomial in peTmp1
    CBDSampleBuffer[0] = (BYTE) (2*nRows);
    SymCryptShake256StateCopy( pShakeBaseState, pShakeWorkState );
    SymCryptShake256Append( pShakeWorkState, CBDSampleBuffer, 1 );

    SymCryptShake256Extract( pShakeWorkState, CBDSampleBuffer, 64ul*nEta2, FALSE );

    SymCryptMlKemPolyElementSampleCBDFromBytes( CBDSampleBuffer, nEta2, peTmp1 );

    // peTmp = INTT(t o r) + e2
    SymCryptMlKemPolyElementAdd( peTmp0, peTmp1, peTmp0 );

    // peTmp1 = mu
    SymCryptMlKemPolyElementDecodeAndDecompress( pbRandom, 1, peTmp1 );

    // peTmp0 = v = INTT(t o r) + e2 + mu
    SymCryptMlKemPolyElementAdd( peTmp0, peTmp1, peTmp0 );

    // Compress and encode v into remainder of ciphertext
    SymCryptMlKemPolyElementCompressAndEncode( peTmp0, nBitsOfV, pbCiphertext+cbU );

cleanup:
    SymCryptWipeKnownSize( CBDSampleBuffer, sizeof(CBDSampleBuffer) );

    return scError;
}


SYMCRYPT_ERROR
SYMCRYPT_CALL
SymCryptMlKemEncapsulateEx(
    _In_                                    PCSYMCRYPT_MLKEMKEY pkMlKemkey,
    _In_reads_bytes_( cbRandom )            PCBYTE              pbRandom,
                                            SIZE_T              cbRandom,
    _Out_writes_bytes_( cbAgreedSecret )    PBYTE               pbAgreedSecret,
                                            SIZE_T              cbAgreedSecret,
    _Out_writes_bytes_( cbCiphertext )      PBYTE               pbCiphertext,
                                            SIZE_T              cbCiphertext )
{
    SYMCRYPT_ERROR scError = SYMCRYPT_NO_ERROR;
    PSYMCRYPT_MLKEM_INTERNAL_COMPUTATION_TEMPORARIES pCompTemps = NULL;

    if( cbRandom != SYMCRYPT_MLKEM_SIZEOF_ENCAPS_RANDOM )
    {
        scError = SYMCRYPT_INVALID_ARGUMENT;
        goto cleanup;
    }

    pCompTemps = SymCryptCallbackAlloc( sizeof(SYMCRYPT_MLKEM_INTERNAL_COMPUTATION_TEMPORARIES) );
    if( pCompTemps == NULL )
    {
        scError = SYMCRYPT_MEMORY_ALLOCATION_FAILURE;
        goto cleanup;
    }

    scError = SymCryptMlKemEncapsulateInternal(
        pkMlKemkey,
        pbAgreedSecret, cbAgreedSecret,
        pbCiphertext, cbCiphertext,
        pbRandom,
        pCompTemps );

cleanup:
    if( pCompTemps != NULL )
    {
        SymCryptWipe( pCompTemps, sizeof(*pCompTemps) );
        SymCryptCallbackFree( pCompTemps );
    }

    return scError;
}

SYMCRYPT_ERROR
SYMCRYPT_CALL
SymCryptMlKemEncapsulate(
    _In_                                    PCSYMCRYPT_MLKEMKEY pkMlKemkey,
    _Out_writes_bytes_( cbAgreedSecret )    PBYTE               pbAgreedSecret,
                                            SIZE_T              cbAgreedSecret,
    _Out_writes_bytes_( cbCiphertext )      PBYTE               pbCiphertext,
                                            SIZE_T              cbCiphertext )
{
    SYMCRYPT_ERROR scError = SYMCRYPT_NO_ERROR;
    BYTE pbm[SYMCRYPT_MLKEM_SIZEOF_ENCAPS_RANDOM];

    scError = SymCryptCallbackRandom( pbm, sizeof(pbm) );
    if( scError != SYMCRYPT_NO_ERROR )
    {
        goto cleanup;
    }

    scError = SymCryptMlKemEncapsulateEx(
        pkMlKemkey,
        pbm, sizeof(pbm),
        pbAgreedSecret, cbAgreedSecret,
        pbCiphertext, cbCiphertext );

cleanup:
    SymCryptWipeKnownSize( pbm, sizeof(pbm) );

    return scError;
}

SYMCRYPT_ERROR
SYMCRYPT_CALL
SymCryptMlKemDecapsulate(
    _In_                                    PCSYMCRYPT_MLKEMKEY pkMlKemkey,
    _In_reads_bytes_( cbCiphertext )        PCBYTE              pbCiphertext,
                                            SIZE_T              cbCiphertext,
    _Out_writes_bytes_( cbAgreedSecret )    PBYTE               pbAgreedSecret,
                                            SIZE_T              cbAgreedSecret )
{
    PSYMCRYPT_MLKEM_INTERNAL_COMPUTATION_TEMPORARIES pCompTemps = NULL;
    BYTE pbDecryptedRandom[SYMCRYPT_MLKEM_SIZEOF_ENCAPS_RANDOM];
    BYTE pbDecapsulatedSecret[SYMCRYPT_MLKEM_SIZEOF_AGREED_SECRET];
    BYTE pbImplicitRejectionSecret[SYMCRYPT_MLKEM_SIZEOF_AGREED_SECRET];
    PBYTE pbReadCiphertext, pbReencapsulatedCiphertext;
    BOOLEAN successfulReencrypt;

    PBYTE pbCurr;
    PBYTE pbAlloc = NULL;
    const SIZE_T cbAlloc = sizeof(SYMCRYPT_MLKEM_INTERNAL_COMPUTATION_TEMPORARIES) + (2*cbCiphertext);

    SYMCRYPT_ERROR scError = SYMCRYPT_NO_ERROR;
    SIZE_T cbU, cbV, cbCopy;
    PSYMCRYPT_MLKEM_VECTOR pvu;
    PSYMCRYPT_MLKEM_POLYELEMENT peTmp0, peTmp1;
    PSYMCRYPT_MLKEM_POLYELEMENT_ACCUMULATOR paTmp;
    PSYMCRYPT_SHAKE256_STATE pShakeState;
    const UINT32 nRows = pkMlKemkey->params.nRows;
    const UINT32 nBitsOfU = pkMlKemkey->params.nBitsOfU;
    const UINT32 nBitsOfV = pkMlKemkey->params.nBitsOfV;
    const UINT32 cbPolyElement = pkMlKemkey->params.cbPolyElement;
    const UINT32 cbVector = pkMlKemkey->params.cbVector;

    // u vector encoded with nBitsOfU * SYMCRYPT_MLWE_POLYNOMIAL_COEFFICIENTS bits per polynomial
    cbU = nRows * nBitsOfU * (SYMCRYPT_MLWE_POLYNOMIAL_COEFFICIENTS / 8);
    // v polynomial encoded with nBitsOfV * SYMCRYPT_MLWE_POLYNOMIAL_COEFFICIENTS bits
    cbV = nBitsOfV * (SYMCRYPT_MLWE_POLYNOMIAL_COEFFICIENTS / 8);

    if( (cbAgreedSecret != SYMCRYPT_MLKEM_SIZEOF_AGREED_SECRET) ||
        (cbCiphertext != cbU + cbV) ||
        !pkMlKemkey->hasPrivateKey )
    {
        scError = SYMCRYPT_INVALID_ARGUMENT;
        goto cleanup;
    }

    pbAlloc = SymCryptCallbackAlloc( cbAlloc );
    if( pbAlloc == NULL )
    {
        scError = SYMCRYPT_MEMORY_ALLOCATION_FAILURE;
        goto cleanup;
    }
    pbCurr = pbAlloc;

    pCompTemps = (PSYMCRYPT_MLKEM_INTERNAL_COMPUTATION_TEMPORARIES) pbCurr;
    pbCurr += sizeof(SYMCRYPT_MLKEM_INTERNAL_COMPUTATION_TEMPORARIES);

    pbReadCiphertext = pbCurr;
    pbCurr += cbCiphertext;

    pbReencapsulatedCiphertext = pbCurr;
    pbCurr += cbCiphertext;

    SYMCRYPT_ASSERT( pbCurr == (pbAlloc + cbAlloc) );

    // Read the input ciphertext once to local pbReadCiphertext to ensure our view of ciphertext consistent
    memcpy( pbReadCiphertext, pbCiphertext, cbCiphertext );

    pvu = SymCryptMlKemVectorCreate( pCompTemps->abVectorBuffer0, cbVector, nRows );
    SYMCRYPT_ASSERT( pvu != NULL );
    peTmp0 = SymCryptMlKemPolyElementCreate( pCompTemps->abPolyElementBuffer0, cbPolyElement );
    SYMCRYPT_ASSERT( peTmp0 != NULL );
    peTmp1 = SymCryptMlKemPolyElementCreate( pCompTemps->abPolyElementBuffer1, cbPolyElement );
    SYMCRYPT_ASSERT( peTmp1 != NULL );
    paTmp = SymCryptMlKemPolyElementAccumulatorCreate( pCompTemps->abPolyElementAccumulatorBuffer, 2*cbPolyElement );
    SYMCRYPT_ASSERT( paTmp != NULL );

    // Decode and decompress u
    scError = SymCryptMlKemVectorDecodeAndDecompress( pbReadCiphertext, cbU, nBitsOfU, pvu );
    SYMCRYPT_ASSERT( scError == SYMCRYPT_NO_ERROR );

    // Perform NTT on u
    SymCryptMlKemVectorNTT( pvu );

    // peTmp0 = (s o NTT(u)) ./ R
    SymCryptMlKemVectorMontDotProduct( pkMlKemkey->pvs, pvu, peTmp0, paTmp );

    // peTmp0 = INTT(s o NTT(u))
    SymCryptMlKemPolyElementINTTAndMulR( peTmp0 );

    // Decode and decompress v
    scError = SymCryptMlKemPolyElementDecodeAndDecompress( pbReadCiphertext+cbU, nBitsOfV, peTmp1 );
    SYMCRYPT_ASSERT( scError == SYMCRYPT_NO_ERROR );

    // peTmp0 = w = v - INTT(s o NTT(u))
    SymCryptMlKemPolyElementSub( peTmp1, peTmp0, peTmp0 );

    // pbDecryptedRandom = m' = Encoding of w
    SymCryptMlKemPolyElementCompressAndEncode( peTmp0, 1, pbDecryptedRandom );

    // Compute:
    //  pbDecapsulatedSecret = K' = Decapsulated secret (without implicit rejection)
    //  pbReencapsulatedCiphertext = c' = Ciphertext from re-encapsulating decrypted random value
    scError = SymCryptMlKemEncapsulateInternal(
        pkMlKemkey,
        pbDecapsulatedSecret, sizeof(pbDecapsulatedSecret),
        pbReencapsulatedCiphertext, cbCiphertext,
        pbDecryptedRandom,
        pCompTemps );
    SYMCRYPT_ASSERT( scError == SYMCRYPT_NO_ERROR );

    // Compute the secret we will return if using implicit rejection
    // pbImplicitRejectionSecret = K_bar = SHAKE256( z || c )
    pShakeState = &pCompTemps->hashState0.shake256State;
    SymCryptShake256Init( pShakeState );
    SymCryptShake256Append( pShakeState, pkMlKemkey->privateRandom, sizeof(pkMlKemkey->privateRandom) );
    SymCryptShake256Append( pShakeState, pbReadCiphertext, cbCiphertext );
    SymCryptShake256Extract( pShakeState, pbImplicitRejectionSecret, sizeof(pbImplicitRejectionSecret), FALSE );

    // Constant time test if re-encryption successful
    successfulReencrypt = SymCryptEqual( pbReencapsulatedCiphertext, pbReadCiphertext, cbCiphertext );

    // If not successful, perform side-channel-safe copy of Implicit Rejection secret over Decapsulated secret
    cbCopy = (((SIZE_T)successfulReencrypt)-1) & SYMCRYPT_MLKEM_SIZEOF_AGREED_SECRET;
    SymCryptScsCopy( pbImplicitRejectionSecret, cbCopy, pbDecapsulatedSecret, SYMCRYPT_MLKEM_SIZEOF_AGREED_SECRET );

    // Write agreed secret (with implicit rejection) to pbAgreedSecret
    memcpy( pbAgreedSecret, pbDecapsulatedSecret, SYMCRYPT_MLKEM_SIZEOF_AGREED_SECRET );

cleanup:
    if( pbAlloc != NULL )
    {
        SymCryptWipe( pbAlloc, cbAlloc );
        SymCryptCallbackFree( pbAlloc );
    }

    SymCryptWipeKnownSize( pbDecryptedRandom, sizeof(pbDecryptedRandom) );
    SymCryptWipeKnownSize( pbDecapsulatedSecret, sizeof(pbDecapsulatedSecret) );
    SymCryptWipeKnownSize( pbImplicitRejectionSecret, sizeof(pbImplicitRejectionSecret) );

    return scError;
}

Coverage Report

Created: 2024-11-21 07:03