/****************************************************************************
*																			*
*					cryptlib Internal Memory Management API					*
*						Copyright Peter Gutmann 1992-2025					*
*																			*
****************************************************************************/

#include <stdarg.h>
#if defined( INC_ALL )
  #include "crypt.h"
#else
  #include "crypt.h"
#endif /* Compiler-specific includes */

/****************************************************************************
*																			*
*						Checked Buffer Management Routines					*
*																			*
****************************************************************************/

/* Add cookies around a buffer, used to check for overwrites.  This converts
   a memory buffer:

	| buffer
	v
	+---------------+
	|	Buffer		|
	+---------------+
	|<-- bufSize -->|

   into:

	|bufPtr	| buffer
	v		v
	+-------+---------------+-------+
	|Cookie |	Buffer		|Cookie	|
	+-------+---------------+-------+
			|<-- bufSize -->|

   To do this we need to convert a request for an allocation of bufSize 
   bytes into one of COOKIE_SIZE + bufSize + COOKIE_SIZE bytes and then 
   return a pointer to the actual buffer data inside the cookie'd buffer */

/* A buffer cookie, used as a canary to check for overwrites, being the XOR 
   of the low 64 bits of the address and a random value */

#ifdef FIXED_SEED
static const BYTE canarySeed[ SAFEBUFFER_COOKIE_SIZE ] = { FIXED_SEED };
#else
static const BYTE canarySeed[ SAFEBUFFER_COOKIE_SIZE ] = \
											{ SAFEBUFFER_COOKIE_DATA };
#endif /* FIXED_SEED */

STDC_NONNULL_ARG( ( 1 ) ) \
static void makeCanary( OUT_BUFFER_FIXED( SAFEBUFFER_COOKIE_SIZE ) \
							BYTE *canary,
						const void *address )
	{
	uintptr_t addressValue = ( uintptr_t ) address;
	const int shiftCountMod = ( sizeof( uintptr_t ) == 4 ) ? 32 : 64;
	LOOP_INDEX i;

	/* The following works for both 32 and 64-bit values because for 32 bits 
	   it stores two copies of the address in the 64-bit cookie if uintptr_t 
	   is 32 bits since the shift count is taken mod the register size, for 
	   a full breakdown see 
	   https://devblogs.microsoft.com/oldnewthing/20230904-00/?p=108704.
	   We don't actually rely on this for anything, it just happens to be
	   a convenient side-effect of how shifts work.
	   
	   However at this point we run into yet more gcc braindamage: Even 
	   though this works just fine, gcc knows that shifting by more than
	   the word size is UB and so goes out of its way to silently emit code 
	   that breaks, stopping the loop once 32 bits have been output and 
	   leaving the remaining 32 bits uninitialised (!!).  Because of this we 
	   have to perform an explicit reduction of the shift amount in order to
	   achieve the same effect that the hardware would produce anyway if gcc
	   didn't go out of its way to make things break */
	LOOP_SMALL( i = 0, i < SAFEBUFFER_COOKIE_SIZE, i++ )
		{
		const int shiftCount = ( i * 8 ) % shiftCountMod;
		
		ENSURES_V( LOOP_INVARIANT_SMALL( i, 0, SAFEBUFFER_COOKIE_SIZE - 1 ) );
		
		canary[ i ] = intToByte( addressValue >> shiftCount ) ^ canarySeed[ i ];
		}
	ENSURES_V( LOOP_BOUND_OK );
	}

/* Initialise a safe buffer with canaries and check that the canaries 
   are still valid.  The annotations here aren't quite correct since 
   we're accessing memory just outside the buffer, but there's no way to 
   correctly annotate what's going on */

STDC_NONNULL_ARG( ( 1 ) ) \
void safeBufferInit( INOUT_BUFFER_FIXED( bufSize ) void *buffer, 
					 IN_DATALENGTH const int bufSize )
	{
	BYTE *startCookiePtr = ( ( BYTE * ) buffer ) - SAFEBUFFER_COOKIE_SIZE;
	BYTE *endCookiePtr = ( ( BYTE * ) buffer ) + bufSize;
	BYTE cookie[ SAFEBUFFER_COOKIE_SIZE + 16 ];

	REQUIRES_V( isBufsizeRangeMin( bufSize, 256 ) );

	/* Insert the cookies, which correspond to the address at which they're
	   stored XOR'd with a fixed magic value */
	makeCanary( cookie, startCookiePtr );
	memcpy( startCookiePtr, cookie, SAFEBUFFER_COOKIE_SIZE );
	makeCanary( cookie, endCookiePtr );
	memcpy( endCookiePtr, cookie, SAFEBUFFER_COOKIE_SIZE );
	}

CHECK_RETVAL_PTR \
void *safeBufferAlloc( IN_DATALENGTH const int bufSize )
	{
	void *bufPtr;
	
	REQUIRES_N( isBufsizeRangeMin( bufSize, MIN_BUFFER_SIZE ) );
	
	/* Allocate the buffer, with extra space for the cookies and overflow 
	   data.  The overflow amount is added implicitly here because the 
	   caller only records the buffer size as being the base size without 
	   the overflow amount, which means that when we perform the cookie 
	   check on { buffer, bufSize } there's additional overflow space 
	   between the purported end of the buffer and the cookie */
	REQUIRES_N( !checkOverflowAdd( SAFEBUFFER_SIZE( bufSize ), 8 ) );
	bufPtr = clAlloc( "safeBufferAlloc", SAFEBUFFER_SIZE( bufSize ) + 8 );
	if( bufPtr == NULL )
		return( NULL );
	safeBufferInit( SAFEBUFFER_PTR( bufPtr ), bufSize );

	return( SAFEBUFFER_PTR( bufPtr ) );
	}

void safeBufferFree( const void *buffer )
	{
	void *startCookiePtr = ( ( BYTE * ) buffer ) - SAFEBUFFER_COOKIE_SIZE;
	BYTE cookie[ SAFEBUFFER_COOKIE_SIZE + 16 ];

	assert( isReadPtr( buffer, SAFEBUFFER_COOKIE_SIZE ) );

	/* We can't check the overall buffer state since we don't know its size
	   in order to locate the end cookie, but we can check at least the 
	   start cookie for validity */
	makeCanary( cookie, startCookiePtr );
	if( memcmp( cookie, startCookiePtr, SAFEBUFFER_COOKIE_SIZE ) )
		{
		/* BUffer corrupted, don't try and free it */
		assert( DEBUG_WARN );
		return;
		}

	clFree( "safeBufferFree", startCookiePtr );
	}

CHECK_RETVAL_BOOL STDC_NONNULL_ARG( ( 1 ) ) \
BOOLEAN safeBufferCheck( IN_BUFFER( bufSize ) const void *buffer, 
						 IN_DATALENGTH const int bufSize )
	{
	const BYTE *startCookiePtr = ( ( const BYTE * ) buffer ) - SAFEBUFFER_COOKIE_SIZE;
	const BYTE *endCookiePtr = ( ( const BYTE * ) buffer ) + bufSize;
	BYTE cookie[ SAFEBUFFER_COOKIE_SIZE + 16 ];

	REQUIRES_B( isBufsizeRangeMin( bufSize, 256 ) );

	/* Check that the cookies haven't been disturbed.  We don't DEBUG_WARN on
	   these because they're used in the self-test, relying on actual checks 
	   being wrapped in REQUIRES_B()/ENSURES_B() */
	makeCanary( cookie, startCookiePtr );
	if( memcmp( cookie, startCookiePtr, SAFEBUFFER_COOKIE_SIZE ) )
		return( FALSE );
	makeCanary( cookie, endCookiePtr );
	if( memcmp( cookie, endCookiePtr, SAFEBUFFER_COOKIE_SIZE ) )
		return( FALSE );

	return( TRUE );
	}

/****************************************************************************
*																			*
*						Dynamic Buffer Management Routines					*
*																			*
****************************************************************************/

/* Dynamic buffer management functions.  When reading variable-length
   object data we can usually fit the data into a small fixed-length buffer 
   but occasionally we have to cope with larger data amounts that require a 
   dynamically-allocated buffer.  The following routines manage this 
   process, dynamically allocating and freeing a larger buffer if required */

CHECK_RETVAL STDC_NONNULL_ARG( ( 1 ) ) \
static int getDynData( OUT_PTR DYNBUF *dynBuf, 
					   IN_HANDLE const CRYPT_HANDLE cryptHandle,
					   IN_MESSAGE const MESSAGE_TYPE message, 
					   IN_INT const int messageParam )
	{
	MESSAGE_DATA msgData;
	void *dataPtr = NULL;
	int status;

	assert( isWritePtr( dynBuf, sizeof( DYNBUF ) ) );

	REQUIRES( isHandleRangeValid( cryptHandle ) );
	REQUIRES( ( message == IMESSAGE_GETATTRIBUTE_S && \
				( isAttribute( messageParam ) || \
				  isInternalAttribute( messageParam ) ) ) || \
			  ( message == IMESSAGE_CRT_EXPORT && \
		 		( messageParam == CRYPT_CERTFORMAT_CERTIFICATE || \
				  messageParam == CRYPT_CERTFORMAT_CERTCHAIN ) ) );

	/* Clear return values.  Note that we don't use the usual memset() to 
	   clear the value since the structure contains the storage for the 
	   fixed-size portion of the buffer appended to it, and using memset() 
	   to clear that is just unnecessary overhead */
	dynBuf->data = dynBuf->dataBuffer;
	dynBuf->length = 0;

	/* Get the data from the object */
	setMessageData( &msgData, NULL, 0 );
	status = krnlSendMessage( cryptHandle, message, &msgData, messageParam );
	if( cryptStatusError( status ) )
		return( status );
	if( msgData.length > DYNBUF_SIZE )
		{
		/* The data is larger than the built-in buffer size, dynamically
		   allocate a larger buffer */
		REQUIRES( isIntegerRangeNZ( msgData.length ) );
		if( ( dataPtr = clDynAlloc( "getDynData", msgData.length ) ) == NULL )
			return( CRYPT_ERROR_MEMORY );
		msgData.data = dataPtr;
		status = krnlSendMessage( cryptHandle, message, &msgData,
								  messageParam );
		if( cryptStatusError( status ) )
			{
			clFree( "getDynData", dataPtr );
			return( status );
			}
		dynBuf->data = dataPtr;
		}
	else
		{
		/* The data will fit into the built-in buffer, read it directly into
		   the buffer */
		msgData.data = dynBuf->data;
		status = krnlSendMessage( cryptHandle, message, &msgData,
								  messageParam );
		if( cryptStatusError( status ) )
			return( status );
		}
	dynBuf->length = msgData.length;

	return( CRYPT_OK );
	}

CHECK_RETVAL STDC_NONNULL_ARG( ( 1 ) ) \
int dynCreate( OUT_PTR DYNBUF *dynBuf, 
			   IN_HANDLE const CRYPT_HANDLE cryptHandle,
			   IN_ATTRIBUTE const CRYPT_ATTRIBUTE_TYPE attributeType )
	{
	assert( isWritePtr( dynBuf, sizeof( DYNBUF ) ) );

	REQUIRES( isHandleRangeValid( cryptHandle ) );
	REQUIRES( isAttribute( attributeType ) || \
			  isInternalAttribute( attributeType ) );

	return( getDynData( dynBuf, cryptHandle, IMESSAGE_GETATTRIBUTE_S,
						attributeType ) );
	}

CHECK_RETVAL STDC_NONNULL_ARG( ( 1 ) ) \
int dynCreateCert( OUT_PTR DYNBUF *dynBuf, 
				   IN_HANDLE const CRYPT_HANDLE cryptHandle,
				   IN_ENUM( CRYPT_CERTFORMAT ) \
				   const CRYPT_CERTFORMAT_TYPE formatType )
	{
	assert( isWritePtr( dynBuf, sizeof( DYNBUF ) ) );

	REQUIRES( isHandleRangeValid( cryptHandle ) );
	REQUIRES( formatType == CRYPT_CERTFORMAT_CERTIFICATE || \
			  formatType == CRYPT_CERTFORMAT_CERTCHAIN );

	return( getDynData( dynBuf, cryptHandle, IMESSAGE_CRT_EXPORT, 
						formatType ) );
	}

STDC_NONNULL_ARG( ( 1 ) ) \
void dynDestroy( INOUT_PTR DYNBUF *dynBuf )
	{
	assert( isWritePtr( dynBuf, sizeof( DYNBUF ) ) );
	assert( isWritePtrDynamic( dynBuf->data, dynBuf->length ) );

	REQUIRES_V( dynBuf->data != NULL );
	REQUIRES_V( isBufsizeRangeNZ( dynBuf->length ) );

	REQUIRES_V( isIntegerRangeNZ( dynBuf->length ) ); 
	zeroise( dynBuf->data, dynBuf->length );
	if( dynBuf->data != dynBuf->dataBuffer )
		clFree( "dynDestroy", dynBuf->data );
	dynBuf->data = NULL;
	dynBuf->length = 0;
	}

/****************************************************************************
*																			*
*						Memory Pool Management Routines						*
*																			*
****************************************************************************/

/* Memory pool management functions.  When allocating many small blocks of
   memory, especially in resource-constrained systems, it's better if we pre-
   allocate a small memory pool ourselves and grab chunks of it as required,
   falling back to dynamically allocating memory later on if we exhaust the
   pool.  The following functions implement the custom memory pool
   management.  Usage is:

	initMemPool( &memPoolState, storage, storageSize );
	newItem = getMemPool( &memPoolState, newItemSize ) */

typedef struct {
	BUFFER( storageSize, storagePos ) 
	void *storage;					/* Memory pool */
	int storageSize, storagePos;	/* Current usage and total size of pool */
	} MEMPOOL_INFO;

#ifndef CONFIG_CONSERVE_MEMORY_EXTRA

CHECK_RETVAL_BOOL STDC_NONNULL_ARG( ( 1 ) ) \
static BOOLEAN sanityCheckMempool( const MEMPOOL_INFO *state )
	{
	/* Make sure that the overall pool size information is in order */
	if( !isShortIntegerRangeMin( state->storageSize, 64 ) )
		{
		DEBUG_PUTS(( "sanityCheckMempool: Storage size" ));
		return( FALSE );
		}

	/* Make sure that the pool allocation information is in order */
	if( !isShortIntegerRange( state->storagePos ) || \
		state->storagePos > state->storageSize )
		{
		DEBUG_PUTS(( "sanityCheckMempool: Storage position" ));
		return( FALSE );
		}

	return( TRUE );
	}
#endif /* !CONFIG_CONSERVE_MEMORY_EXTRA */

CHECK_RETVAL STDC_NONNULL_ARG( ( 1, 2 ) ) \
int initMemPool( OUT_PTR void *statePtr, 
				 IN_BUFFER( memPoolSize ) void *memPool, 
				 IN_LENGTH_SHORT_MIN( 64 ) const int memPoolSize )
	{
	MEMPOOL_INFO *state = ( MEMPOOL_INFO * ) statePtr;

	assert( isWritePtr( state, sizeof( MEMPOOL_INFO ) ) );
	assert( isWritePtrDynamic( memPool, memPoolSize ) );

#if defined( __WIN32__ ) && defined( _MSC_VER )
	#pragma warning( disable: 4127 )	/* Needed for sizeof() in check */
#endif /* VC++ */
	REQUIRES( sizeof( MEMPOOL_STATE ) >= sizeof( MEMPOOL_INFO ) );
	REQUIRES( isShortIntegerRangeMin( memPoolSize, 64 ) );
#if defined( __WIN32__ ) && defined( _MSC_VER )
	#pragma warning( 4: 4127 )
#endif /* VC++ */

	memset( state, 0, sizeof( MEMPOOL_INFO ) );
	state->storage = memPool;
	state->storageSize = memPoolSize;

	return( CRYPT_OK );
	}

CHECK_RETVAL_PTR STDC_NONNULL_ARG( ( 1 ) ) \
void *getMemPool( INOUT_PTR void *statePtr, IN_LENGTH_SHORT const int size )
	{
	MEMPOOL_INFO *state = ( MEMPOOL_INFO * ) statePtr;
	BYTE *allocPtr;
	const int allocSize = roundUp( size, sizeof( int ) );

	assert( isWritePtr( state, sizeof( MEMPOOL_INFO ) ) );
	assert( isWritePtrDynamic( state->storage, state->storageSize ) );

	REQUIRES_N( isShortIntegerRangeNZ( size ) );
	REQUIRES_N( !checkOverflowRoundup( size, sizeof( int ) ) );
	REQUIRES_N( isShortIntegerRangeMin( allocSize, sizeof( int ) ) );
	REQUIRES_N( sanityCheckMempool( state ) );

	/* If we can't satisfy the request from the memory pool we have to
	   allocate the memory block dynamically */
	REQUIRES_N( !checkOverflowAdd( state->storagePos, allocSize ) );
	if( state->storagePos + allocSize > state->storageSize )
		{
		REQUIRES_N( isShortIntegerRangeNZ( size ) );
		return( clDynAlloc( "getMemPool", size ) );
		}

	/* We can satisfy the request from the pool:

	 memPool
		|
		v		 <- size -->
		+-------+-----------+-------+
		|		|			|		|
		+-------+-----------+-------+
				^			^
				|			|
			storagePos	storagePos' */
	allocPtr = ( BYTE * ) state->storage + state->storagePos;
	REQUIRES_N( !checkOverflowAdd( state->storagePos, allocSize ) );
	state->storagePos += allocSize;
	ENSURES_N( sanityCheckMempool( state ) );

	return( allocPtr );
	}

STDC_NONNULL_ARG( ( 1, 2 ) ) \
void freeMemPool( INOUT_PTR void *statePtr, 
				  IN_PTR void *memblock )
	{
	MEMPOOL_INFO *state = ( MEMPOOL_INFO * ) statePtr;

	assert( isWritePtr( state, sizeof( MEMPOOL_INFO ) ) );
	assert( isWritePtrDynamic( state->storage, state->storageSize ) );

	REQUIRES_V( sanityCheckMempool( state ) );

	/* If the memory block to free lies within the pool, there's nothing to 
	   do */
	if( memblock >= state->storage && \
		memblock < ( void * ) ( ( BYTE * ) state->storage + \
										   state->storageSize ) )
		return;

	/* It's outside the pool and therefore dynamically allocated, free it */
	clFree( "freeMemPool", memblock );
	}

/****************************************************************************
*																			*
*							Debugging Malloc Support						*
*																			*
****************************************************************************/

/* Debugging malloc() that dumps memory usage diagnostics to stdout.  Note
   that these functions are only intended to be used during interactive 
   debugging sessions since they trigger assertions under error conditions 
   rather than returning an error status (the fact that they dump 
   diagnostics to stdout during operation should be a clue as to their
   intended status and usage) */

#ifdef CONFIG_DEBUG_MALLOC

#ifdef __WIN32__
  #include <direct.h>
#endif /* __WIN32__ */

#ifdef __WINCE__

CHECK_RETVAL_RANGE( 0, MAX_INTLENGTH_STRING ) STDC_NONNULL_ARG( ( 1 ) ) \
static int wcPrintf( FORMAT_STRING const char *format, ... )
	{
	wchar_t wcBuffer[ 1024 + 8 ];
	char buffer[ 1024 + 8 ];
	va_list argPtr;
	int length;

	va_start( argPtr, format );
	length = vsprintf_s( buffer, 1024, format, argPtr );
	va_end( argPtr );
	if( !rangeCheck( length, 1, 1023 ) )
		return( length );
	mbstowcs( wcBuffer, buffer, length + 1 );
	NKDbgPrintfW( wcBuffer );

	return( length );
	}

#define printf		wcPrintf

#endif /* __WINCE__ */

static int clAllocIndex = 0;

void *clAllocFn( const char *fileName, const char *fnName,
				 const int lineNo, size_t size )
	{
#ifdef CONFIG_MALLOCTEST
	static int mallocCount = 0, mallocFailCount = 0;
#endif /* CONFIG_MALLOCTEST */
	BYTE *memPtr;
	int length;

	assert( fileName != NULL );
	assert( fnName != NULL );
	assert( lineNo > 0 );
	assert( isIntegerRangeNZ( size ) );

	length = DEBUG_PRINT(( "ALLOC: %s:%s:%d", debugGetBasePath( fileName ), 
						   fnName, lineNo ));
	while( length < 56 )
		{
		DEBUG_PRINT(( " " ));
		length++;
		}
	DEBUG_PRINT(( " %4d - %d bytes.\n", clAllocIndex, size ));
#ifdef CONFIG_MALLOCTEST
	/* If we've exceeded the allocation count, make the next attempt to 
	   allocate memory fail */
	if( mallocCount >= mallocFailCount )
		{
		mallocCount = 0;
		mallocFailCount++;

		return( NULL );
		}
	mallocCount++;
#endif /* CONFIG_MALLOCTEST */
	if( ( memPtr = malloc( size + sizeof( LONG ) ) ) == NULL )
		return( NULL );
	mputLong( memPtr, clAllocIndex );	/* Implicit memPtr += sizeof( LONG ) */
	clAllocIndex++;
	return( memPtr );
	}

void clFreeFn( const char *fileName, const char *fnName,
			   const int lineNo, void *memblock )
	{
	BYTE *memPtr = ( BYTE * ) memblock - sizeof( LONG );
	int index, length;

	assert( fileName != NULL );
	assert( fnName != NULL );
	assert( lineNo > 0 );

	index = mget32( memPtr );
	memPtr -= sizeof( LONG );		/* mget32() changes memPtr */
	length = DEBUG_PRINT(( "FREE : %s:%s:%d", debugGetBasePath( fileName ), 
						   fnName, lineNo ));
	while( length < 56 )
		{
		DEBUG_PRINT(( " " ));
		length++;
		}
	DEBUG_PRINT(( " %4d.\n", index ));
	free( memPtr );
	}
#endif /* CONFIG_DEBUG_MALLOC */

/* Fault-testing malloc() that fails after a given number of allocations */

#ifdef CONFIG_FAULT_MALLOC

static int currentAllocCount = 0, failAllocCount = 0;
static BOOLEAN allocFailed = FALSE;

void clFaultAllocSetCount( const int number )
	{
	currentAllocCount = 0;
	failAllocCount = number;
	allocFailed = FALSE;
	}

void *clFaultAllocFn( const char *fileName, const char *fnName, 
					  const int lineNo, size_t size )
	{
	/* If we've failed an allocation we probably shouldn't get here again,
	   however if we're running a multithreaded init then the second thread 
	   could try and allocate memory after the first one has failed */
	if( allocFailed )
		{
#ifdef __WIN32__
		DEBUG_PRINT(( "\n<<< Further allocation call from thread %X after "
					  "previous call failed, called from %s line %d in "
					  "%s.>>>\n", GetCurrentThreadId(), fnName, lineNo, 
					  fileName ));
#else
		DEBUG_PRINT(( "\n<<< Further allocation call after previous call "
					  "failed, called from %s line %d in %s.>>>\n", fnName, 
					  lineNo, fileName ));
#endif /* __WIN32__  */
		if( failAllocCount < 15 )
			{
			DEBUG_PRINT(( "<<<  (This could be because of a multithreaded "
						  "init).>>>\n" ));
			DEBUG_PRINT(( "\n" ));
			}
		return( NULL );
		}

	/* If we haven't reached the failure allocation count, return normally */
	if( currentAllocCount < failAllocCount )
		{
		currentAllocCount++;
		return( malloc( size ) );
		}

	/* We've reached the failure count, fail the allocation */
#ifdef __WIN32__
	DEBUG_PRINT(( "\n<<< Failing allocation call #%d for thread %X, called "
				  "from %s line %d in %s.>>>\n\n", failAllocCount + 1, 
				  GetCurrentThreadId(), fnName, lineNo, fileName ));
#else
	DEBUG_PRINT(( "\n<<< Failing at allocation call #%d, called from %s line "
				  "%d in %s.>>>\n\n", failAllocCount + 1, fnName, lineNo, 
				  fileName ));
#endif /* __WIN32__  */
	allocFailed = TRUE;
	return( NULL );
	}
#endif /* CONFIG_FAULT_MALLOC */
