/****************************************************************************
*																			*
*					cryptlib ChaCha20 Encryption Routines					*
*					  Copyright Peter Gutmann 2017-2025						*
*																			*
****************************************************************************/

#include "crypt.h"
#if defined( INC_ALL )
  #include "context.h"
  #include "djb.h"
#else
  #include "context/context.h"
  #include "crypt/djb.h"
#endif /* Compiler-specific includes */

#ifdef USE_CHACHA20

/* The size of the ChaCha20 key and cipher block */

#define CHACHA20_KEY_SIZE		32
#define CHACHA20_IV_SIZE		16
#define CHACHA20_BLOCK_SIZE		64

/* The size of the ChaCha20 state.  Due to the fact that it's actually a 
   block cipher pretending to be a stream cipher, we have to store up to a
   block's worth of keystream alongside the cipher state, and the fact that 
   it has a 512-bit block size means there's a lot of state to store */

typedef struct {
	chacha_ctx keyInfo;
	BYTE keystream[ CHACHA20_BLOCK_SIZE + 8 ];
	int keystreamPos;
	} CHACHA20_STATE;

#define CHACHA20_STATE_SIZE		sizeof( CHACHA20_STATE )

/****************************************************************************
*																			*
*							ChaCha20 Self-test Routines						*
*																			*
****************************************************************************/

#ifndef CONFIG_NO_SELFTEST

typedef struct {
	const BOOLEAN isValid;
	const BYTE key[ CHACHA20_KEY_SIZE ];
	const BYTE nonce[ CHACHA20_IV_SIZE ];
	const BYTE keyStream[ CHACHA20_BLOCK_SIZE ];
	} CHACHA20_TEST;

/* ChaCha20 test vectors from draft-agl-tls-chacha20poly1305-04.  We can't
   use these because they use a 64-bit IV, not the weirdo 96:32 one 
   standardised by the IETF */

#if 0

static const CHACHA20_TEST testCHACHA20[] = {
	{ TRUE,
	  { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 },
	  { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 },
	  { 0x76, 0xB8, 0xE0, 0xAD, 0xA0, 0xF1, 0x3D, 0x90,
		0x40, 0x5D, 0x6A, 0xE5, 0x53, 0x86, 0xBD, 0x28,
		0xBD, 0xD2, 0x19, 0xB8, 0xA0, 0x8D, 0xED, 0x1A,
		0xA8, 0x36, 0xEF, 0xCC, 0x8B, 0x77, 0x0D, 0xC7,
		0xDA, 0x41, 0x59, 0x7C, 0x51, 0x57, 0x48, 0x8D,
		0x77, 0x24, 0xE0, 0x3F, 0xB8, 0xD8, 0x4A, 0x37,
		0x6A, 0x43, 0xB8, 0xF4, 0x15, 0x18, 0xA1, 0x1C,
		0xC3, 0x87, 0xB6, 0x69, 0xB2, 0xEE, 0x65, 0x86 } },
	{ TRUE,
	  { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01 },
	  { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 },
	  { 0x45, 0x40, 0xF0, 0x5A, 0x9F, 0x1F, 0xB2, 0x96,
		0xD7, 0x73, 0x6E, 0x7B, 0x20, 0x8E, 0x3C, 0x96,
		0xEB, 0x4F, 0xE1, 0x83, 0x46, 0x88, 0xD2, 0x60,
		0x4F, 0x45, 0x09, 0x52, 0xED, 0x43, 0x2D, 0x41,
		0xBB, 0xE2, 0xA0, 0xB6, 0xEA, 0x75, 0x66, 0xD2,
		0xA5, 0xD1, 0xE7, 0xE2, 0x0D, 0x42, 0xAF, 0x2C,
		0x53, 0xD7, 0x92, 0xB1, 0xC4, 0x3F, 0xEA, 0x81,
		0x7E, 0x9A, 0xD2, 0x75, 0xAE, 0x54, 0x69, 0x63 } },
	{ TRUE,
	  { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 },
	  { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01 },
	  { 0xDE, 0x9C, 0xBA, 0x7B, 0xF3, 0xD6, 0x9E, 0xF5,
		0xE7, 0x86, 0xDC, 0x63, 0x97, 0x3F, 0x65, 0x3A,
		0x0B, 0x49, 0xE0, 0x15, 0xAD, 0xBF, 0xF7, 0x13,
		0x4F, 0xCB, 0x7D, 0xF1, 0x37, 0x82, 0x10, 0x31,
		0xE8, 0x5A, 0x05, 0x02, 0x78, 0xA7, 0x08, 0x45,
		0x27, 0x21, 0x4F, 0x73, 0xEF, 0xC7, 0xFA, 0x5B,
		0x52, 0x77, 0x06, 0x2E, 0xB7, 0xA0, 0x43, 0x3E,
		0x44, 0x5F, 0x41, 0xE3 /*, 0x	, 0x  , 0x	, 0x*/ } },
	{ TRUE,
	  { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 },
	  { 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 },
	  { 0xEF, 0x3F, 0xDF, 0xD6, 0xC6, 0x15, 0x78, 0xFB,
		0xF5, 0xCF, 0x35, 0xBD, 0x3D, 0xD3, 0x3B, 0x80,
		0x09, 0x63, 0x16, 0x34, 0xD2, 0x1E, 0x42, 0xAC,
		0x33, 0x96, 0x0B, 0xD1, 0x38, 0xE5, 0x0D, 0x32,
		0x11, 0x1E, 0x4C, 0xAF, 0x23, 0x7E, 0xE5, 0x3C,
		0xA8, 0xAD, 0x64, 0x26, 0x19, 0x4A, 0x88, 0x54,
		0x5D, 0xDC, 0x49, 0x7A, 0x0B, 0x46, 0x6E, 0x7D,
		0x6B, 0xBD, 0xB0, 0x04, 0x1B, 0x2F, 0x58, 0x6B } },
	{ TRUE,
	  { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
		0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 
		0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 
		0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F },
	  { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07 },
	  { 0xF7, 0x98, 0xA1, 0x89, 0xF1, 0x95, 0xE6, 0x69,
		0x82, 0x10, 0x5F, 0xFB, 0x64, 0x0B, 0xB7, 0x75,
		0x7F, 0x57, 0x9D, 0xA3, 0x16, 0x02, 0xFC, 0x93,
		0xEC, 0x01, 0xAC, 0x56, 0xF8, 0x5A, 0xC3, 0xC1,
		0x34, 0xA4, 0x54, 0x7B, 0x73, 0x3B, 0x46, 0x41,
		0x30, 0x42, 0xC9, 0x44, 0x00, 0x49, 0x17, 0x69,
		0x05, 0xD3, 0xBE, 0x59, 0xEA, 0x1C, 0x53, 0xF1,
		0x59, 0x16, 0x15, 0x5C, 0x2B, 0xE8, 0x24, 0x1A } }
	};
#endif /* 0 */

/* ChaCha20 test vectors from RFC 8439, the first from section 2.4.2, the
   remainder from Appendix A.2 */

static const CHACHA20_TEST testCHACHA20[] = {
	{ TRUE,
	  /* 2.4.2 Key */
	  { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
		0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 
		0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 
		0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F },
	  /* Counter + Nonce */
	  { 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
	    0x00, 0x00, 0x00, 0x4A, 0x00, 0x00, 0x00, 0x00 },
	  /* Keystream */
	  { 0x22, 0x4F, 0x51, 0xF3, 0x40, 0x1B, 0xD9, 0xE1, 
		0x2F, 0xDE, 0x27, 0x6F, 0xB8, 0x63, 0x1D, 0xED, 
		0x8C, 0x13, 0x1F, 0x82, 0x3D, 0x2C, 0x06, 0xE2, 
		0x7E, 0x4F, 0xCA, 0xEC, 0x9E, 0xF3, 0xCF, 0x78, 
		0x8A, 0x3B, 0x0A, 0xA3, 0x72, 0x60, 0x0A, 0x92, 
		0xB5, 0x79, 0x74, 0xCD, 0xED, 0x2B, 0x93, 0x34, 
		0x79, 0x4C, 0xBA, 0x40, 0xC6, 0x3E, 0x34, 0xCD, 
		0xEA, 0x21, 0x2C, 0x4C, 0xF0, 0x7D, 0x41, 0xB7 } },
	{ TRUE,
	  /* A.2 Key */
	  { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 },
	  /* Counter + Nonce */
	  { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 },
	  /* Ciphertext */
	  { 0x76, 0xB8, 0xE0, 0xAD, 0xA0, 0xF1, 0x3D, 0x90, 
		0x40, 0x5D, 0x6A, 0xE5, 0x53, 0x86, 0xBD, 0x28,
		0xBD, 0xD2, 0x19, 0xB8, 0xA0, 0x8D, 0xED, 0x1A, 
		0xA8, 0x36, 0xEF, 0xCC, 0x8B, 0x77, 0x0D, 0xC7,
		0xDA, 0x41, 0x59, 0x7C, 0x51, 0x57, 0x48, 0x8D, 
		0x77, 0x24, 0xE0, 0x3F, 0xB8, 0xD8, 0x4A, 0x37,
		0x6A, 0x43, 0xB8, 0xF4, 0x15, 0x18, 0xA1, 0x1C, 
		0xC3, 0x87, 0xB6, 0x69, 0xB2, 0xEE, 0x65, 0x86 } },
	{ FALSE, { 0 }, { 0 }, { 0 } },
		{ FALSE, { 0 }, { 0 }, { 0 } }
	};

static const BYTE chacha20Key[] = {
	0x80, 0x81, 0x82, 0x83, 0x84, 0x85, 0x86, 0x87,
	0x88, 0x89, 0x8A, 0x8B, 0x8C, 0x8D, 0x8E, 0x8F,
	0x90, 0x91, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97,
	0x98, 0x99, 0x9A, 0x9B, 0x9C, 0x9D, 0x9E, 0x9F
	};
static const BYTE chacha20IV[] = {
	0x07, 0x00, 0x00, 0x00, 0x40, 0x41, 0x42, 0x43,
	0x44, 0x45, 0x46, 0x47
	};
static const BYTE chacha20PT[] = {
	0x4C, 0x61, 0x64, 0x69, 0x65, 0x73, 0x20, 0x61,
	0x6E, 0x64, 0x20, 0x47, 0x65, 0x6E, 0x74, 0x6C,
	0x65, 0x6D, 0x65, 0x6E, 0x20, 0x6F, 0x66, 0x20,
	0x74, 0x68, 0x65, 0x20, 0x63, 0x6C, 0x61, 0x73,
	0x73, 0x20, 0x6F, 0x66, 0x20, 0x27, 0x39, 0x39,
	0x3A, 0x20, 0x49, 0x66, 0x20, 0x49, 0x20, 0x63,
	0x6F, 0x75, 0x6C, 0x64, 0x20, 0x6F, 0x66, 0x66,
	0x65, 0x72, 0x20, 0x79, 0x6F, 0x75, 0x20, 0x6F,
	0x6E, 0x6C, 0x79, 0x20, 0x6F, 0x6E, 0x65, 0x20,
	0x74, 0x69, 0x70, 0x20, 0x66, 0x6F, 0x72, 0x20,
	0x74, 0x68, 0x65, 0x20, 0x66, 0x75, 0x74, 0x75,
	0x72, 0x65, 0x2C, 0x20, 0x73, 0x75, 0x6E, 0x73,
	0x63, 0x72, 0x65, 0x65, 0x6E, 0x20, 0x77, 0x6F,
	0x75, 0x6C, 0x64, 0x20, 0x62, 0x65, 0x20, 0x69,
	0x74, 0x2E
	};
static const BYTE chacha20CT[] = {
	0xD3, 0x1A, 0x8D, 0x34, 0x64, 0x8E, 0x60, 0xDB,
	0x7B, 0x86, 0xAF, 0xBC, 0x53, 0xEF, 0x7E, 0xC2,
	0xA4, 0xAD, 0xED, 0x51, 0x29, 0x6E, 0x08, 0xFE,
	0xA9, 0xE2, 0xB5, 0xA7, 0x36, 0xEE, 0x62, 0xD6,
	0x3D, 0xBE, 0xA4, 0x5E, 0x8C, 0xA9, 0x67, 0x12,
	0x82, 0xFA, 0xFB, 0x69, 0xDA, 0x92, 0x72, 0x8B,
	0x1A, 0x71, 0xDE, 0x0A, 0x9E, 0x06, 0x0B, 0x29,
	0x05, 0xD6, 0xA5, 0xB6, 0x7E, 0xCD, 0x3B, 0x36,
	0x92, 0xDD, 0xBD, 0x7F, 0x2D, 0x77, 0x8B, 0x8C,
	0x98, 0x03, 0xAE, 0xE3, 0x28, 0x09, 0x1B, 0x58,
	0xFA, 0xB3, 0x24, 0xE4, 0xFA, 0xD6, 0x75, 0x94,
	0x55, 0x85, 0x80, 0x8B, 0x48, 0x31, 0xD7, 0xBC,
	0x3F, 0xF4, 0xDE, 0xF0, 0x8E, 0x4B, 0x7A, 0x9D,
	0xE5, 0x76, 0xD2, 0x65, 0x86, 0xCE, 0xC6, 0x4B,
	0x61, 0x16
	};

static const BYTE poly1305Key[] = {
	0x80, 0x81, 0x82, 0x83, 0x84, 0x85, 0x86, 0x87,
	0x88, 0x89, 0x8A, 0x8B, 0x8C, 0x8D, 0x8E, 0x8F,
	0x90, 0x91, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97,
	0x98, 0x99, 0x9A, 0x9B, 0x9C, 0x9D, 0x9E, 0x9F
	};
static const BYTE poly1305Nonce[] = {
	0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x02, 0x03,
	0x04, 0x05, 0x06, 0x07
	};
static const BYTE poly1305Output[] = {
	0x8A, 0xD5, 0xA0, 0x8B, 0x90, 0x5F, 0x81, 0xCC,
	0x81, 0x50, 0x40, 0x27, 0x4A, 0xB2, 0x94, 0x71,
	0xA8, 0x33, 0xB6, 0x37, 0xE3, 0xFD, 0x0D, 0xA5,
	0x08, 0xDB, 0xB8, 0xE2, 0xFD, 0xD1, 0xA6, 0x46
	};

/* Test the ChaCha20 code */

CHECK_RETVAL \
static int chacha20Test( IN_BUFFER_C( CHACHA20_KEY_SIZE ) const BYTE *key, 
						 IN_BUFFER_C( CHACHA20_IV_SIZE ) const BYTE *iv,
						 IN_BUFFER_OPT( length ) const BYTE *plaintext,
						 IN_BUFFER( length ) const BYTE *ciphertext, 
						 IN_LENGTH_SHORT const int length )
	{
	const CAPABILITY_INFO *capabilityInfo = getChaCha20Capability();
	CONTEXT_INFO contextInfo;
	CONV_INFO contextData;
	ALIGN_DATA( keyData, CHACHA20_STATE_SIZE, 16 );
	void *keyDataPtr = ALIGN_GET_PTR( keyData, 16 );
	BYTE temp[ ( CHACHA20_BLOCK_SIZE * 2 ) + 8 ];
	int status;

	assert( isReadPtr( key, CHACHA20_KEY_SIZE ) );
	assert( isReadPtr( iv, 8 ) );
	assert( plaintext == NULL || isReadPtrDynamic( ciphertext, length ) );
	assert( isReadPtrDynamic( ciphertext, length ) );

	REQUIRES( rangeCheck( length, 1, ( CHACHA20_BLOCK_SIZE * 2 ) ) );

	memset( keyDataPtr, 0, CHACHA20_STATE_SIZE );	/* Keep static analysers happy */
	status = staticInitContext( &contextInfo, CONTEXT_CONV, capabilityInfo,
								&contextData, sizeof( CONV_INFO ), 
								keyDataPtr );
	if( cryptStatusError( status ) )
		return( status );
	if( plaintext != NULL )
		{
		REQUIRES( rangeCheck( length, 1, ( CHACHA20_BLOCK_SIZE * 2 ) ) );
		memcpy( temp, plaintext, length );
		}
	else
		memset( temp, 0, ( CHACHA20_BLOCK_SIZE * 2 ) );
	status = capabilityInfo->initKeyFunction( &contextInfo, key, 
											  CHACHA20_KEY_SIZE );
	if( cryptStatusOK( status ) )
		{
		/* Since ChaCha20 changes its keying state on every en/decrypt, the 
		   encryption function re-checksums the key data after the state has
		   been updated.  Normally the checksumming is handled by higher-
		   level code but since we're calling directly into internal 
		   functions we have to do it explicitly here */
		contextData.keyDataSize = CHACHA20_STATE_SIZE;
		contextData.keyDataChecksum = checksumData( contextData.key, 
													contextData.keyDataSize );
		status = capabilityInfo->initParamsFunction( &contextInfo, 
													 KEYPARAM_IV, 
													 iv, CHACHA20_IV_SIZE );
		}
	if( cryptStatusOK( status ) )
		{
		status = capabilityInfo->encryptCFBFunction( &contextInfo, temp,
													 length );
		}
	staticDestroyContext( &contextInfo );
	if( cryptStatusError( status ) || \
		memcmp( temp, ciphertext, length ) )
		return( CRYPT_ERROR_FAILED );

	return( CRYPT_OK );
	}

CHECK_RETVAL \
static int selfTest( void )
	{
	BYTE iv[ CHACHA20_IV_SIZE + 8 ];
	int i, LOOP_ITERATOR, status;

	/* Test the ChaCha20 algorithm */
	LOOP_SMALL( i = 0, 
				i < FAILSAFE_ARRAYSIZE( testCHACHA20, CHACHA20_TEST ) && \
					testCHACHA20[ i ].isValid,
				i++ )
		{
		ENSURES( LOOP_INVARIANT_SMALL( i, 0, 
									   FAILSAFE_ARRAYSIZE( testCHACHA20, \
														   CHACHA20_TEST ) - 1 ) );

		status = chacha20Test( testCHACHA20[i].key, testCHACHA20[ i ].nonce,
							   NULL, testCHACHA20[ i ].keyStream, 
							   CHACHA20_BLOCK_SIZE );
		if( cryptStatusError( status ) )
			return( status );
		}
	ENSURES( LOOP_BOUND_OK );

	/* Test the ChaCha20 algorithm as used in ChaCha20-Poly1305.  Note that
	   this requires that we construct a 128-bit ChaCha20 IV from the given
	   96-bit nonce and the 32-bit counter that we initialise to 1 in little-
	   endian form */
	REQUIRES( boundsCheck( bitsToBytes( 32 ), bitsToBytes( 96 ), 
						   CHACHA20_IV_SIZE ) );
	memcpy( iv, "\x01\x00\x00\x00", 4 );
	memcpy( iv + bitsToBytes( 32 ), chacha20IV, bitsToBytes( 96 ) );
	status = chacha20Test( chacha20Key, iv, chacha20PT, chacha20CT, 114 );
	if( cryptStatusError( status ) )
		return( status );

	/* Finally, test ChaCha20 when used for Poly1305 key generation */
	memset( iv, 0, CHACHA20_IV_SIZE );
	REQUIRES( boundsCheck( bitsToBytes( 32 ), bitsToBytes( 96 ), 
			  CHACHA20_IV_SIZE ) );
	memcpy( iv + bitsToBytes( 32 ), poly1305Nonce, bitsToBytes( 96 ) );
	return( chacha20Test( poly1305Key, iv, NULL, poly1305Output, 32 ) );
	}
#else
	#define selfTest	NULL
#endif /* !CONFIG_NO_SELFTEST */

/****************************************************************************
*																			*
*								Control Routines							*
*																			*
****************************************************************************/

/* Return context subtype-specific information */

CHECK_RETVAL STDC_NONNULL_ARG( ( 3 ) ) \
static int getInfo( IN_ENUM( CAPABILITY_INFO ) const CAPABILITY_INFO_TYPE type, 
					INOUT_PTR_OPT CONTEXT_INFO *contextInfoPtr,
					OUT_PTR void *data, 
					IN_INT_Z const int length )
	{
	int *valuePtr = ( int * ) data;

	assert( contextInfoPtr == NULL || \
			isWritePtr( contextInfoPtr, sizeof( CONTEXT_INFO ) ) );
	assert( ( length == 0 && isWritePtr( data, sizeof( int ) ) ) || \
			( length > 0 && isWritePtrDynamic( data, length ) ) );

	REQUIRES( isEnumRange( type, CAPABILITY_INFO ) );
	REQUIRES( ( contextInfoPtr == NULL ) || \
			  sanityCheckContext( contextInfoPtr ) );

	switch( type )
		{
		case CAPABILITY_INFO_STATESIZE:
			{
			*valuePtr = CHACHA20_STATE_SIZE;

			return( CRYPT_OK );
			}

		case CAPABILITY_INFO_STATEALIGNTYPE:
			{
			/* The ChaCha20 code requires alignment to 128-bit boundaries */
			*valuePtr = bitsToBytes( 128 );

			return( CRYPT_OK );
			}

		default:
			return( getDefaultInfo( type, contextInfoPtr, data, length ) );
		}

	retIntError();
	}

/****************************************************************************
*																			*
*						ChaCha20 En/Decryption Routines						*
*																			*
****************************************************************************/

/* Encrypt/decrypt data.  Since ChaCha20 pretends to be a stream cipher, 
   encryption and decryption are the same operation.  However, it's actually
   a block cipher in counter mode, which means that using it as if it was a
   stream cipher doesn't work unless everything falls on block boundaries.  
   To deal with this we store the generated keystream in an internal buffer 
   and use that as required.
   
   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 int noBytes )
	{
	CONV_INFO *convInfo = contextInfoPtr->ctxConv;
	CHACHA20_STATE *stateInfo = convInfo->key;
	static const BYTE zeroes[ CHACHA20_BLOCK_SIZE ] = { 0 };

	assert( isWritePtr( contextInfoPtr, sizeof( CONTEXT_INFO ) ) );
	assert( isWritePtrDynamic( buffer, noBytes ) );

	REQUIRES( sanityCheckContext( contextInfoPtr ) );
	REQUIRES( isIntegerRangeNZ( noBytes ) );

	/* We're about to modify the keying data, make sure that it's still 
	   valid before we start */
	if( checksumData( convInfo->key, \
					  convInfo->keyDataSize ) != convInfo->keyDataChecksum )
		retIntError();

	while( noBytes > 0 )
		{
		int bytesToUse, i, LOOP_ITERATOR;

		ENSURES( rangeCheck( stateInfo->keystreamPos, 
							 0, CHACHA20_BLOCK_SIZE - 1 ) );

		/* If there's any keystream material left, use it now */
		if( stateInfo->keystreamPos > 0 )
			{
			/* Find out how much keystream material we can use */
			REQUIRES( !checkOverflowSub( CHACHA20_BLOCK_SIZE,
										 stateInfo->keystreamPos ) );
			bytesToUse = CHACHA20_BLOCK_SIZE - stateInfo->keystreamPos;
			if( noBytes < bytesToUse )
				bytesToUse = noBytes;
			REQUIRES( rangeCheck( bytesToUse, 1, CHACHA20_BLOCK_SIZE ) );

			/* Encrypt the data */
			LOOP_LARGE( i = 0, i < bytesToUse, i++ )
				{
				ENSURES( LOOP_INVARIANT_LARGE( i, 0, bytesToUse - 1 ) );

				buffer[ i ] ^= \
					stateInfo->keystream[ i + stateInfo->keystreamPos ];
				}
			ENSURES( LOOP_BOUND_OK );

			/* Adjust the byte count and buffer position */
			REQUIRES( !checkOverflowSub( noBytes, bytesToUse ) );
			noBytes -= bytesToUse;
			buffer += bytesToUse;
			REQUIRES( !checkOverflowAdd( stateInfo->keystreamPos, 
										 bytesToUse ) );
			stateInfo->keystreamPos += bytesToUse;

			/* If we've satisfied the request from the existing keystream, 
			   we're done */
			if( stateInfo->keystreamPos < CHACHA20_BLOCK_SIZE )
				break;

			/* We've consumed all of the keystream, reset the position 
			   indicator */
			stateInfo->keystreamPos = 0;

			/* If the remaining amount of keystream exactly filled the 
			   request, we're done */
			if( noBytes <= 0 )
				break;
			}

		/* Generate the next block of keystream */
		chacha20_encrypt_bytes( &stateInfo->keyInfo, zeroes, 
								stateInfo->keystream, CHACHA20_BLOCK_SIZE );

		/* Encrypt the data */
		bytesToUse = min( noBytes, CHACHA20_BLOCK_SIZE );
		LOOP_LARGE( i = 0, i < bytesToUse, i++ )
			{
			ENSURES( LOOP_INVARIANT_LARGE( i, 0, bytesToUse - 1 ) );

			buffer[ i ] ^= stateInfo->keystream[ i ];
			}
		ENSURES( LOOP_BOUND_OK );
		REQUIRES( !checkOverflowSub( noBytes, bytesToUse ) );
		noBytes -= bytesToUse;
		buffer += bytesToUse;

		/* If we've done the last block and there's any keystream left over,
		   remember what's left */
		if( noBytes <= 0 && bytesToUse < CHACHA20_BLOCK_SIZE )
			stateInfo->keystreamPos = bytesToUse;
		}

	/* This en/decryption process updates changes the key data, so we have 
	   to update the checksum before we return to the caller */
	convInfo->keyDataChecksum = checksumData( convInfo->key, 
											  convInfo->keyDataSize );

	return( CRYPT_OK );
	}

/****************************************************************************
*																			*
*						ChaCha20 Key Management Routines					*
*																			*
****************************************************************************/

/* Initialise crypto parameters such as the IV/counter */

CHECK_RETVAL STDC_NONNULL_ARG( ( 1 ) ) \
static int initParams( INOUT_PTR CONTEXT_INFO *contextInfoPtr, 
					   IN_ENUM( KEYPARAM ) const KEYPARAM_TYPE paramType,
					   IN_PTR_OPT const void *data, 
					   IN_INT const int dataLength )
	{
	CONV_INFO *convInfo = contextInfoPtr->ctxConv;
	const BYTE *dataPtr = data;

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

	REQUIRES( contextInfoPtr->type == CONTEXT_CONV );
	REQUIRES( isEnumRange( paramType, KEYPARAM ) );

	if( paramType == KEYPARAM_IV )
		{
		CHACHA20_STATE *stateInfo = convInfo->key;

		ENSURES( dataLength == bitsToBytes( 128 ) );

		/* We're about to modify the keying data, make sure that it's still 
		   valid before we start.  Since ChaCha20 is an awkward algorithm 
		   that makes the IV part of the security-relevant cryptovariables, 
		   normally if we load a key and then set the IV these are already 
		   set up but if we load the IV first they aren't so we only do the
		   checksumming if they've been set up */
		if( convInfo->keyDataSize > 0 && \
			checksumData( convInfo->key, \
						  convInfo->keyDataSize ) != convInfo->keyDataChecksum )
			retIntError();

		/* Update the ChaCha20 state with the IV.  The original ChaCha20 
		   splits the 128-bit value 64:64 counter and IV, the IETF form 
		   changes the split to 32:96 counter and IV */
		chacha_ietf_ivsetup( &stateInfo->keyInfo, 
							 dataPtr + bitsToBytes( 32 ), dataPtr );

		/* Since we've now changed the IV, we also need to reset any stored
		   keystream */
		memset( stateInfo->keystream, 0, CHACHA20_BLOCK_SIZE );
		stateInfo->keystreamPos = 0;

		/* This process updates internal state which changes the key data 
		   checksum, so we have to update the checksum before we return to 
		   the caller */
		if( convInfo->keyDataSize > 0 )
			{
			convInfo->keyDataChecksum = checksumData( convInfo->key, 
													  convInfo->keyDataSize );
			}

		/* Fall through to load the IV into the context data */
		}

	/* Pass the call on down to the global parameter-handling function */	
	return( initGenericParams( contextInfoPtr, paramType, data, 
							   dataLength ) );
	}

/* Key schedule a ChaCha20 key */

CHECK_RETVAL STDC_NONNULL_ARG( ( 1, 2 ) ) \
static int initKey( INOUT_PTR CONTEXT_INFO *contextInfoPtr, 
					IN_BUFFER( keyLength ) const void *key, 
					IN_LENGTH_SHORT const int keyLength )
	{
	CONV_INFO *convInfo = contextInfoPtr->ctxConv;
	CHACHA20_STATE *stateInfo = convInfo->key;

	assert( isWritePtr( contextInfoPtr, sizeof( CONTEXT_INFO ) ) );
	assert( isReadPtrDynamic( key, keyLength ) );

	REQUIRES( sanityCheckContext( contextInfoPtr ) );
	REQUIRES( keyLength == CHACHA20_KEY_SIZE );

	/* Copy the key to internal storage */
	if( convInfo->userKey != key )
		{
		REQUIRES( rangeCheck( keyLength, 1, CRYPT_MAX_KEYSIZE ) );
		memcpy( convInfo->userKey, key, keyLength );
		convInfo->userKeyLength = keyLength;
		}

	chacha_keysetup( &stateInfo->keyInfo, key );
	memset( stateInfo->keystream, 0, CHACHA20_BLOCK_SIZE );
	stateInfo->keystreamPos = 0;

	return( CRYPT_OK );
	}

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

/* The block size, which also acts as the IV size, is a problem due to the
   special-snowflake form of ChaCha20.  If you have a block cipher in stream
   mode, e.g. AES-CFB, then everything works fine, you get an IV of 128 bits 
   as expected.  The problem is that ChaCha20 is a block cipher pretending 
   to be a stream cipher, being a replacement for the pure stream cipher RC4 
   via Salsa20.  It also doesn't really have an IV but some portion of the 
   value being a nonce and then the rest being an internal counter, with the 
   split varying depending on which standard you follow.

   Internally where cryptlib uses ChaCha20 it has to encode whatever nonce 
   size the protocol requires.  For example for TLS it knows that it needs 
   to load a 96-bit value (actually not even that, it's an XOR of the TLS 
   sequence number with an IV for a total of 96 bits with the internal 
   counter adding another 32 bits to the mix).  For SSH the split is 64 bits 
   of big-endian packet sequence number and 64 bits of internal little-
   endian counter, so two counters of different endianness and no actual 
   nonce, because why not?

   The user-supplied value for the IV can therefore be either 64, 96, or 
   possibly 128 bits depending on which protocol you're using it with.  One 
   option would be to return the IV size as 96 bits on the assumption that 
   the IETF and not the SSH interpretation is being used, but that's still 
   going to cause problems if someone does want to use the SSH 
   interpretation.  Setting the IV is really a write-only quantity where you 
   set the nonce part of the value, the ChaCha20 code uses the bytes it 
   wants and runs the internal-counter part itself, and you can't read the 
   value back because its interpretation is situation-dependant.  Poly1305 
   (which goes with ChaCha20) is even more problematic, you can't just key 
   it and use it like a normal MAC but have to go through all sorts of 
   calisthenics to use it appropriately */

static const CAPABILITY_INFO capabilityInfo = {
	CRYPT_ALGO_CHACHA20, bitsToBytes( 8 ), "ChaCha20", 8,
	CHACHA20_KEY_SIZE, CHACHA20_KEY_SIZE, CHACHA20_KEY_SIZE,
	selfTest, getInfo, NULL, initParams, initKey, NULL,
	NULL, NULL, NULL, NULL, encryptFn, encryptFn 
	};						/* Pseudo-CFB mode */

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