ChaCha20.h

/**
 * @file chacha20.h
 * @brief ChaCha20 stream cipher implementation for Windows Kernel Drivers.
 *
 * This implementation provides a ChaCha20 encryption/decryption interface suitable for use
 * within a Windows Kernel driver environment. It does not rely on user-mode functions or
 * dynamic allocation. All functions are designed to operate at IRQL <= DISPATCH_LEVEL.
 *
 * The code follows the ChaCha20 specification described in RFC 8439.
 *
 * Reference: https://www.rfc-editor.org/rfc/rfc8439
 */

#include <ntifs.h>
#include <ntddk.h>

#pragma warning(push)
#pragma warning(disable: 4214)

#ifdef _MSC_VER
#pragma warning(disable: 4201)
#endif

 /**
  * @def CHACHA20_BLOCK_SIZE
  * @brief The size of a single ChaCha20 block in bytes.
  */
#define CHACHA20_BLOCK_SIZE 64

  /**
   * @def CHACHA20_KEY_SIZE
   * @brief The size of a ChaCha20 key in bytes (256-bit key).
   */
#define CHACHA20_KEY_SIZE 32

   /**
    * @def CHACHA20_NONCE_SIZE
    * @brief The size of a ChaCha20 nonce (IV) in bytes.
    */
#define CHACHA20_NONCE_SIZE 12

    /**
     * @def CHACHA20_ROUNDS
     * @brief Number of rounds in the ChaCha20 block function.
     *
     * RFC 8439 specifies 20 rounds.
     */
#define CHACHA20_ROUNDS 20

     /**
      * @struct CHACHA20_CONTEXT
      * @brief Context structure holding ChaCha20 state.
      *
      * The context includes the key setup, initial block counter, and nonce.
      * The state is initialized once and can be used to encrypt/decrypt multiple
      * blocks by updating the counter as data is processed.
      */
typedef struct _CHACHA20_CONTEXT
{
    ULONG State[16]; /**< Internal state words of the ChaCha20 cipher */
} CHACHA20_CONTEXT, * PCHACHA20_CONTEXT;

//
// Internal Functions (static) Prototypes
//

/**
 * @brief Perform a ChaCha20 quarter round on four 32-bit words.
 *
 * @param a Reference to the first word.
 * @param b Reference to the second word.
 * @param c Reference to the third word.
 * @param d Reference to the fourth word.
 */
static __forceinline VOID
ChaCha20QuarterRound(
    _Inout_ ULONG* a,
    _Inout_ ULONG* b,
    _Inout_ ULONG* c,
    _Inout_ ULONG* d
)
{
    *a += *b; *d ^= *a; *d = _rotl(*d, 16);
    *c += *d; *b ^= *c; *b = _rotl(*b, 12);
    *a += *b; *d ^= *a; *d = _rotl(*d, 8);
    *c += *d; *b ^= *c; *b = _rotl(*b, 7);
}

/**
 * @brief Generate one 64-byte block of keystream from the internal state.
 *
 * This function takes the current state, performs the ChaCha20 rounds,
 * and produces a 64-byte keystream block. The internal state should not
 * be modified directly by this function; the caller increments the counter
 * after processing a block as needed.
 *
 * @param ctx Pointer to a CHACHA20_CONTEXT initialized with a key, nonce, and counter.
 * @param keystream Buffer of at least 64 bytes to receive the generated keystream.
 */
static __forceinline VOID
ChaCha20BlockFunction(
    _In_  PCHACHA20_CONTEXT ctx,
    _Out_writes_all_(CHACHA20_BLOCK_SIZE) PUCHAR keystream
)
{
    ULONG workingState[16];
    ULONG i;

    for (i = 0; i < 16; i++)
    {
        workingState[i] = ctx->State[i];
    }

    for (i = 0; i < CHACHA20_ROUNDS; i += 2)
    {
        ChaCha20QuarterRound(&workingState[0], &workingState[4], &workingState[8], &workingState[12]);
        ChaCha20QuarterRound(&workingState[1], &workingState[5], &workingState[9], &workingState[13]);
        ChaCha20QuarterRound(&workingState[2], &workingState[6], &workingState[10], &workingState[14]);
        ChaCha20QuarterRound(&workingState[3], &workingState[7], &workingState[11], &workingState[15]);

        ChaCha20QuarterRound(&workingState[0], &workingState[5], &workingState[10], &workingState[15]);
        ChaCha20QuarterRound(&workingState[1], &workingState[6], &workingState[11], &workingState[12]);
        ChaCha20QuarterRound(&workingState[2], &workingState[7], &workingState[8], &workingState[13]);
        ChaCha20QuarterRound(&workingState[3], &workingState[4], &workingState[9], &workingState[14]);
    }

    for (i = 0; i < 16; i++)
    {
        ULONG outWord = workingState[i] + ctx->State[i];
        keystream[(i * 4) + 0] = (UCHAR)(outWord & 0xFF);
        keystream[(i * 4) + 1] = (UCHAR)((outWord >> 8) & 0xFF);
        keystream[(i * 4) + 2] = (UCHAR)((outWord >> 16) & 0xFF);
        keystream[(i * 4) + 3] = (UCHAR)((outWord >> 24) & 0xFF);
    }
}

/**
 * @brief Increment the 64-bit block counter in the context.
 *
 * The counter is located at State[12] and State[13] in the ChaCha20 state.
 * This function increments the block counter after each processed block.
 *
 * @param ctx Pointer to the CHACHA20_CONTEXT.
 */
static __forceinline VOID
ChaCha20IncrementCounter(
    _Inout_ PCHACHA20_CONTEXT ctx
)
{
    ctx->State[12]++;
    if (ctx->State[12] == 0)
    {
        ctx->State[13]++;
    }
}

//
// Public Funcs
//

/**
 * @brief Initialize a ChaCha20 context with a given key, nonce, and initial block counter.
 *
 * This function sets up the state array from the given 256-bit key, 96-bit nonce,
 * and 32-bit initial counter. The initial counter can be used to start encryption
 * at a non-zero block index if desired.
 *
 * @param ctx Pointer to a CHACHA20_CONTEXT to initialize.
 * @param key A 32-byte (256-bit) key.
 * @param nonce A 12-byte (96-bit) nonce.
 * @param initialCounter A 32-bit initial block counter, often starting at 0.
 *
 * @return NTSTATUS indicating success or failure.
 *
 * @retval STATUS_SUCCESS if the context is successfully initialized.
 * @retval STATUS_INVALID_PARAMETER if any parameter is invalid.
 */
__forceinline
NTSTATUS
Chacha20Initialize(
    PCHACHA20_CONTEXT ctx,
    PCUCHAR key,
    PCUCHAR nonce,
    ULONG initialCounter
)
{
    static UCHAR sigma[16] = { 0, };
    __movsb(sigma, "expand 32-byte k", 16);

    ULONG i;

    if (ctx == NULL || key == NULL || nonce == NULL)
    {
        return STATUS_INVALID_PARAMETER;
    }

    ctx->State[0] = ((ULONG)sigma[0]) | ((ULONG)sigma[1] << 8) | ((ULONG)sigma[2] << 16) | ((ULONG)sigma[3] << 24);
    ctx->State[1] = ((ULONG)sigma[4]) | ((ULONG)sigma[5] << 8) | ((ULONG)sigma[6] << 16) | ((ULONG)sigma[7] << 24);
    ctx->State[2] = ((ULONG)sigma[8]) | ((ULONG)sigma[9] << 8) | ((ULONG)sigma[10] << 16) | ((ULONG)sigma[11] << 24);
    ctx->State[3] = ((ULONG)sigma[12]) | ((ULONG)sigma[13] << 8) | ((ULONG)sigma[14] << 16) | ((ULONG)sigma[15] << 24);

    for (i = 0; i < 8; i++)
    {
        ctx->State[4 + i] = ((ULONG)key[i * 4 + 0]) |
            ((ULONG)key[i * 4 + 1] << 8) |
            ((ULONG)key[i * 4 + 2] << 16) |
            ((ULONG)key[i * 4 + 3] << 24);
    }

    ctx->State[12] = initialCounter;

    ctx->State[14] = ((ULONG)nonce[0]) | ((ULONG)nonce[1] << 8) | ((ULONG)nonce[2] << 16) | ((ULONG)nonce[3] << 24);
    ctx->State[15] = ((ULONG)nonce[4]) | ((ULONG)nonce[5] << 8) | ((ULONG)nonce[6] << 16) | ((ULONG)nonce[7] << 24);

    ctx->State[13] = ((ULONG)nonce[0]) | ((ULONG)nonce[1] << 8) | ((ULONG)nonce[2] << 16) | ((ULONG)nonce[3] << 24);
    ctx->State[14] = ((ULONG)nonce[4]) | ((ULONG)nonce[5] << 8) | ((ULONG)nonce[6] << 16) | ((ULONG)nonce[7] << 24);
    ctx->State[15] = ((ULONG)nonce[8]) | ((ULONG)nonce[9] << 8) | ((ULONG)nonce[10] << 16) | ((ULONG)nonce[11] << 24);

    return STATUS_SUCCESS;
}

/**
 * @brief Encrypt or decrypt data using the ChaCha20 cipher.
 *
 * ChaCha20 is a stream cipher; encryption and decryption are identical operations:
 * ciphertext = plaintext XOR keystream
 * plaintext  = ciphertext XOR keystream
 *
 * This function will process a given buffer in place. The caller is responsible for
 * ensuring the buffer is non-paged or locked in memory as necessary for kernel-level access.
 *
 * The context is updated as data is processed (the internal counter increments for each block).
 * As a result, subsequent calls continue from where the previous call left off, allowing for
 * large data streams to be processed incrementally.
 *
 * @param ctx Pointer to an initialized CHACHA20_CONTEXT.
 * @param data Pointer to the data buffer to encrypt/decrypt in-place.
 * @param dataLength Size of the data buffer in bytes.
 *
 * @return NTSTATUS indicating success or failure.
 *
 * @retval STATUS_SUCCESS if operation completed successfully.
 * @retval STATUS_INVALID_PARAMETER if input parameters are invalid.
 */
__forceinline
NTSTATUS
Chacha20Crypt(
    PCHACHA20_CONTEXT ctx,
    PUCHAR data,
    SIZE_T dataLength
)
{
    UCHAR block[CHACHA20_BLOCK_SIZE];
    SIZE_T offset = 0;
    SIZE_T bytesToProcess;

    if (ctx == NULL || data == NULL)
    {
        return STATUS_INVALID_PARAMETER;
    }

    while (offset < dataLength)
    {
        SIZE_T bytesRemaining = dataLength - offset;
        bytesToProcess = (bytesRemaining < CHACHA20_BLOCK_SIZE) ? bytesRemaining : CHACHA20_BLOCK_SIZE;

        ChaCha20BlockFunction(ctx, block);

        for (SIZE_T i = 0; i < bytesToProcess; i++)
        {
            data[offset + i] ^= block[i];
        }

        ChaCha20IncrementCounter(ctx);

        offset += bytesToProcess;
    }

    return STATUS_SUCCESS;
}

#pragma warning(pop)