Implement parameter validation for ChaCha/Poly modules
diff --git a/library/chacha20.c b/library/chacha20.c
index d14a51e..0757163 100644
--- a/library/chacha20.c
+++ b/library/chacha20.c
@@ -53,6 +53,12 @@
 #define inline __inline
 #endif
 
+/* Parameter validation macros */
+#define CHACHA20_VALIDATE_RET( cond )                                       \
+    MBEDTLS_INTERNAL_VALIDATE_RET( cond, MBEDTLS_ERR_CHACHA20_BAD_INPUT_DATA )
+#define CHACHA20_VALIDATE( cond )                                           \
+    MBEDTLS_INTERNAL_VALIDATE( cond )
+
 #define BYTES_TO_U32_LE( data, offset )                           \
     ( (uint32_t) data[offset]                                     \
           | (uint32_t) ( (uint32_t) data[( offset ) + 1] << 8 )   \
@@ -181,14 +187,13 @@
 
 void mbedtls_chacha20_init( mbedtls_chacha20_context *ctx )
 {
-    if( ctx != NULL )
-    {
-        mbedtls_platform_zeroize( ctx->state, sizeof( ctx->state ) );
-        mbedtls_platform_zeroize( ctx->keystream8, sizeof( ctx->keystream8 ) );
+    CHACHA20_VALIDATE( ctx != NULL );
 
-        /* Initially, there's no keystream bytes available */
-        ctx->keystream_bytes_used = CHACHA20_BLOCK_SIZE_BYTES;
-    }
+    mbedtls_platform_zeroize( ctx->state, sizeof( ctx->state ) );
+    mbedtls_platform_zeroize( ctx->keystream8, sizeof( ctx->keystream8 ) );
+
+    /* Initially, there's no keystream bytes available */
+    ctx->keystream_bytes_used = CHACHA20_BLOCK_SIZE_BYTES;
 }
 
 void mbedtls_chacha20_free( mbedtls_chacha20_context *ctx )
@@ -202,10 +207,8 @@
 int mbedtls_chacha20_setkey( mbedtls_chacha20_context *ctx,
                             const unsigned char key[32] )
 {
-    if( ( ctx == NULL ) || ( key == NULL ) )
-    {
-        return( MBEDTLS_ERR_CHACHA20_BAD_INPUT_DATA );
-    }
+    CHACHA20_VALIDATE_RET( ctx != NULL );
+    CHACHA20_VALIDATE_RET( key != NULL );
 
     /* ChaCha20 constants - the string "expand 32-byte k" */
     ctx->state[0] = 0x61707865;
@@ -230,10 +233,8 @@
                              const unsigned char nonce[12],
                              uint32_t counter )
 {
-    if( ( ctx == NULL ) || ( nonce == NULL ) )
-    {
-        return( MBEDTLS_ERR_CHACHA20_BAD_INPUT_DATA );
-    }
+    CHACHA20_VALIDATE_RET( ctx != NULL );
+    CHACHA20_VALIDATE_RET( nonce != NULL );
 
     /* Counter */
     ctx->state[12] = counter;
@@ -259,15 +260,9 @@
     size_t offset = 0U;
     size_t i;
 
-    if( ctx == NULL )
-    {
-        return( MBEDTLS_ERR_CHACHA20_BAD_INPUT_DATA );
-    }
-    else if( ( size > 0U ) && ( ( input == NULL ) || ( output == NULL ) ) )
-    {
-        /* input and output pointers are allowed to be NULL only if size == 0 */
-        return( MBEDTLS_ERR_CHACHA20_BAD_INPUT_DATA );
-    }
+    CHACHA20_VALIDATE_RET( ctx != NULL );
+    CHACHA20_VALIDATE_RET( size == 0 || input  != NULL );
+    CHACHA20_VALIDATE_RET( size == 0 || output != NULL );
 
     /* Use leftover keystream bytes, if available */
     while( size > 0U && ctx->keystream_bytes_used < CHACHA20_BLOCK_SIZE_BYTES )
@@ -332,6 +327,11 @@
     mbedtls_chacha20_context ctx;
     int ret;
 
+    CHACHA20_VALIDATE_RET( key != NULL );
+    CHACHA20_VALIDATE_RET( nonce != NULL );
+    CHACHA20_VALIDATE_RET( data_len == 0 || input  != NULL );
+    CHACHA20_VALIDATE_RET( data_len == 0 || output != NULL );
+
     mbedtls_chacha20_init( &ctx );
 
     ret = mbedtls_chacha20_setkey( &ctx, key );
diff --git a/library/chachapoly.c b/library/chachapoly.c
index 860f877..ee5cc34 100644
--- a/library/chachapoly.c
+++ b/library/chachapoly.c
@@ -44,6 +44,12 @@
 
 #if !defined(MBEDTLS_CHACHAPOLY_ALT)
 
+/* Parameter validation macros */
+#define CHACHAPOLY_VALIDATE_RET( cond )                                       \
+    MBEDTLS_INTERNAL_VALIDATE_RET( cond, MBEDTLS_ERR_POLY1305_BAD_INPUT_DATA )
+#define CHACHAPOLY_VALIDATE( cond )                                           \
+    MBEDTLS_INTERNAL_VALIDATE( cond )
+
 #define CHACHAPOLY_STATE_INIT       ( 0 )
 #define CHACHAPOLY_STATE_AAD        ( 1 )
 #define CHACHAPOLY_STATE_CIPHERTEXT ( 2 ) /* Encrypting or decrypting */
@@ -90,39 +96,35 @@
 
 void mbedtls_chachapoly_init( mbedtls_chachapoly_context *ctx )
 {
-    if( ctx != NULL )
-    {
-        mbedtls_chacha20_init( &ctx->chacha20_ctx );
-        mbedtls_poly1305_init( &ctx->poly1305_ctx );
-        ctx->aad_len        = 0U;
-        ctx->ciphertext_len = 0U;
-        ctx->state          = CHACHAPOLY_STATE_INIT;
-        ctx->mode           = MBEDTLS_CHACHAPOLY_ENCRYPT;
-    }
+    CHACHAPOLY_VALIDATE( ctx != NULL );
+
+    mbedtls_chacha20_init( &ctx->chacha20_ctx );
+    mbedtls_poly1305_init( &ctx->poly1305_ctx );
+    ctx->aad_len        = 0U;
+    ctx->ciphertext_len = 0U;
+    ctx->state          = CHACHAPOLY_STATE_INIT;
+    ctx->mode           = MBEDTLS_CHACHAPOLY_ENCRYPT;
 }
 
 void mbedtls_chachapoly_free( mbedtls_chachapoly_context *ctx )
 {
     if( ctx != NULL )
-    {
-        mbedtls_chacha20_free( &ctx->chacha20_ctx );
-        mbedtls_poly1305_free( &ctx->poly1305_ctx );
-        ctx->aad_len        = 0U;
-        ctx->ciphertext_len = 0U;
-        ctx->state          = CHACHAPOLY_STATE_INIT;
-        ctx->mode           = MBEDTLS_CHACHAPOLY_ENCRYPT;
-    }
+        return;
+
+    mbedtls_chacha20_free( &ctx->chacha20_ctx );
+    mbedtls_poly1305_free( &ctx->poly1305_ctx );
+    ctx->aad_len        = 0U;
+    ctx->ciphertext_len = 0U;
+    ctx->state          = CHACHAPOLY_STATE_INIT;
+    ctx->mode           = MBEDTLS_CHACHAPOLY_ENCRYPT;
 }
 
 int mbedtls_chachapoly_setkey( mbedtls_chachapoly_context *ctx,
                                const unsigned char key[32] )
 {
     int ret;
-
-    if( ( ctx == NULL ) || ( key == NULL ) )
-    {
-        return( MBEDTLS_ERR_POLY1305_BAD_INPUT_DATA );
-    }
+    CHACHAPOLY_VALIDATE_RET( ctx != NULL );
+    CHACHAPOLY_VALIDATE_RET( key != NULL );
 
     ret = mbedtls_chacha20_setkey( &ctx->chacha20_ctx, key );
 
@@ -135,11 +137,8 @@
 {
     int ret;
     unsigned char poly1305_key[64];
-
-    if( ( ctx == NULL ) || ( nonce == NULL ) )
-    {
-        return( MBEDTLS_ERR_POLY1305_BAD_INPUT_DATA );
-    }
+    CHACHAPOLY_VALIDATE_RET( ctx != NULL );
+    CHACHAPOLY_VALIDATE_RET( nonce != NULL );
 
     /* Set counter = 0, will be update to 1 when generating Poly1305 key */
     ret = mbedtls_chacha20_starts( &ctx->chacha20_ctx, nonce, 0U );
@@ -176,16 +175,10 @@
                                    const unsigned char *aad,
                                    size_t aad_len )
 {
-    if( ctx == NULL )
-    {
-        return( MBEDTLS_ERR_POLY1305_BAD_INPUT_DATA );
-    }
-    else if( ( aad_len > 0U ) && ( aad == NULL ) )
-    {
-        /* aad pointer is allowed to be NULL if aad_len == 0 */
-        return( MBEDTLS_ERR_POLY1305_BAD_INPUT_DATA );
-    }
-    else if( ctx->state != CHACHAPOLY_STATE_AAD )
+    CHACHAPOLY_VALIDATE_RET( ctx != NULL );
+    CHACHAPOLY_VALIDATE_RET( aad_len == 0 || aad != NULL );
+
+    if( ctx->state != CHACHAPOLY_STATE_AAD )
     {
         return( MBEDTLS_ERR_CHACHAPOLY_BAD_STATE );
     }
@@ -201,18 +194,12 @@
                                unsigned char *output )
 {
     int ret;
+    CHACHAPOLY_VALIDATE_RET( ctx != NULL );
+    CHACHAPOLY_VALIDATE_RET( len == 0 || input != NULL );
+    CHACHAPOLY_VALIDATE_RET( len == 0 || output != NULL );
 
-    if( ctx == NULL )
-    {
-        return( MBEDTLS_ERR_POLY1305_BAD_INPUT_DATA );
-    }
-    else if( ( len > 0U ) && ( ( input == NULL ) || ( output == NULL ) ) )
-    {
-        /* input and output pointers are allowed to be NULL if len == 0 */
-        return( MBEDTLS_ERR_POLY1305_BAD_INPUT_DATA );
-    }
-    else if( ( ctx->state != CHACHAPOLY_STATE_AAD ) &&
-              ( ctx->state != CHACHAPOLY_STATE_CIPHERTEXT ) )
+    if( ( ctx->state != CHACHAPOLY_STATE_AAD ) &&
+        ( ctx->state != CHACHAPOLY_STATE_CIPHERTEXT ) )
     {
         return( MBEDTLS_ERR_CHACHAPOLY_BAD_STATE );
     }
@@ -257,12 +244,10 @@
 {
     int ret;
     unsigned char len_block[16];
+    CHACHAPOLY_VALIDATE_RET( ctx != NULL );
+    CHACHAPOLY_VALIDATE_RET( mac != NULL );
 
-    if( ( ctx == NULL ) || ( mac == NULL ) )
-    {
-        return( MBEDTLS_ERR_POLY1305_BAD_INPUT_DATA );
-    }
-    else if( ctx->state == CHACHAPOLY_STATE_INIT )
+    if( ctx->state == CHACHAPOLY_STATE_INIT )
     {
         return( MBEDTLS_ERR_CHACHAPOLY_BAD_STATE );
     }
@@ -350,6 +335,13 @@
                                         unsigned char *output,
                                         unsigned char tag[16] )
 {
+    CHACHAPOLY_VALIDATE_RET( ctx   != NULL );
+    CHACHAPOLY_VALIDATE_RET( nonce != NULL );
+    CHACHAPOLY_VALIDATE_RET( tag   != NULL );
+    CHACHAPOLY_VALIDATE_RET( aad_len == 0 || aad    != NULL );
+    CHACHAPOLY_VALIDATE_RET( length  == 0 || input  != NULL );
+    CHACHAPOLY_VALIDATE_RET( length  == 0 || output != NULL );
+
     return( chachapoly_crypt_and_tag( ctx, MBEDTLS_CHACHAPOLY_ENCRYPT,
                                       length, nonce, aad, aad_len,
                                       input, output, tag ) );
@@ -368,9 +360,12 @@
     unsigned char check_tag[16];
     size_t i;
     int diff;
-
-    if( tag == NULL )
-        return( MBEDTLS_ERR_POLY1305_BAD_INPUT_DATA );
+    CHACHAPOLY_VALIDATE_RET( ctx   != NULL );
+    CHACHAPOLY_VALIDATE_RET( nonce != NULL );
+    CHACHAPOLY_VALIDATE_RET( tag   != NULL );
+    CHACHAPOLY_VALIDATE_RET( aad_len == 0 || aad    != NULL );
+    CHACHAPOLY_VALIDATE_RET( length  == 0 || input  != NULL );
+    CHACHAPOLY_VALIDATE_RET( length  == 0 || output != NULL );
 
     if( ( ret = chachapoly_crypt_and_tag( ctx,
                         MBEDTLS_CHACHAPOLY_DECRYPT, length, nonce,
diff --git a/library/poly1305.c b/library/poly1305.c
index e22d3af..c22a0a1 100644
--- a/library/poly1305.c
+++ b/library/poly1305.c
@@ -49,6 +49,12 @@
 #define inline __inline
 #endif
 
+/* Parameter validation macros */
+#define POLY1305_VALIDATE_RET( cond )                                       \
+    MBEDTLS_INTERNAL_VALIDATE_RET( cond, MBEDTLS_ERR_POLY1305_BAD_INPUT_DATA )
+#define POLY1305_VALIDATE( cond )                                           \
+    MBEDTLS_INTERNAL_VALIDATE( cond )
+
 #define POLY1305_BLOCK_SIZE_BYTES ( 16U )
 
 #define BYTES_TO_U32_LE( data, offset )                           \
@@ -276,27 +282,24 @@
 
 void mbedtls_poly1305_init( mbedtls_poly1305_context *ctx )
 {
-    if( ctx != NULL )
-    {
-        mbedtls_platform_zeroize( ctx, sizeof( mbedtls_poly1305_context ) );
-    }
+    POLY1305_VALIDATE( ctx != NULL );
+
+    mbedtls_platform_zeroize( ctx, sizeof( mbedtls_poly1305_context ) );
 }
 
 void mbedtls_poly1305_free( mbedtls_poly1305_context *ctx )
 {
     if( ctx != NULL )
-    {
-        mbedtls_platform_zeroize( ctx, sizeof( mbedtls_poly1305_context ) );
-    }
+        return;
+
+    mbedtls_platform_zeroize( ctx, sizeof( mbedtls_poly1305_context ) );
 }
 
 int mbedtls_poly1305_starts( mbedtls_poly1305_context *ctx,
                              const unsigned char key[32] )
 {
-    if( ctx == NULL || key == NULL )
-    {
-        return( MBEDTLS_ERR_POLY1305_BAD_INPUT_DATA );
-    }
+    POLY1305_VALIDATE_RET( ctx != NULL );
+    POLY1305_VALIDATE_RET( key != NULL );
 
     /* r &= 0x0ffffffc0ffffffc0ffffffc0fffffff */
     ctx->r[0] = BYTES_TO_U32_LE( key, 0 )  & 0x0FFFFFFFU;
@@ -331,16 +334,8 @@
     size_t remaining = ilen;
     size_t queue_free_len;
     size_t nblocks;
-
-    if( ctx == NULL )
-    {
-        return( MBEDTLS_ERR_POLY1305_BAD_INPUT_DATA );
-    }
-    else if( ( ilen > 0U ) && ( input == NULL ) )
-    {
-        /* input pointer is allowed to be NULL only if ilen == 0 */
-        return( MBEDTLS_ERR_POLY1305_BAD_INPUT_DATA );
-    }
+    POLY1305_VALIDATE_RET( ctx != NULL );
+    POLY1305_VALIDATE_RET( ilen == 0 || input != NULL );
 
     if( ( remaining > 0U ) && ( ctx->queue_len > 0U ) )
     {
@@ -398,10 +393,8 @@
 int mbedtls_poly1305_finish( mbedtls_poly1305_context *ctx,
                              unsigned char mac[16] )
 {
-    if( ( ctx == NULL ) || ( mac == NULL ) )
-    {
-        return( MBEDTLS_ERR_POLY1305_BAD_INPUT_DATA );
-    }
+    POLY1305_VALIDATE_RET( ctx != NULL );
+    POLY1305_VALIDATE_RET( mac != NULL );
 
     /* Process any leftover data */
     if( ctx->queue_len > 0U )
@@ -431,6 +424,9 @@
 {
     mbedtls_poly1305_context ctx;
     int ret;
+    POLY1305_VALIDATE_RET( key != NULL );
+    POLY1305_VALIDATE_RET( mac != NULL );
+    POLY1305_VALIDATE_RET( ilen == 0 || input != NULL );
 
     mbedtls_poly1305_init( &ctx );