Chacha20 scalar, SSE2, NEON, wasm SIMD implementation (C/C++ comaptible)

chacha.c

#include "chacha.h"

#include <stdint.h>
#include <stddef.h>
#include <string.h>
#include <assert.h>
#if defined(__SSE2__) || defined(__AVX__)
    #include <immintrin.h>
    #define Vec4u32 __m128i

    #define VLOADU(p)    _mm_loadu_si128((__m128i *)(void *)(p))
    #define VSTOREU(p,x) _mm_storeu_si128((__m128i *)(void *)(p), (x))
    #define VADD(a,b)    _mm_add_epi32((a), (b))
    #define VXOR(a,b)    _mm_xor_si128((a), (b))

    #define VROTL(x, n)  _mm_or_si128(_mm_slli_epi32((x), n), _mm_srli_epi32((x), 32 - n))
    #define VLANEROT1(x) _mm_shuffle_epi32((x), _MM_SHUFFLE(0, 3, 2, 1)) /* [1,2,3,0] */
    #define VLANEROT2(x) _mm_shuffle_epi32((x), _MM_SHUFFLE(1, 0, 3, 2)) /* [2,3,0,1] */
    #define VLANEROT3(x) _mm_shuffle_epi32((x), _MM_SHUFFLE(2, 1, 0, 3)) /* [3,0,1,2] */
#elif defined(__ARM_NEON)
    #include <arm_neon.h>
    #define Vec4u32 uint32x4_t

    #define VLOADU(p)   vld1q_u32((uint32_t *)(void *)(p))
    #define VSTOREU(p,x) vst1q_u32((uint32_t *)(void *)(p), (x))
    #define VADD(a,b)  vaddq_u32((a), (b))
    #define VXOR(a,b)  veorq_u32((a), (b))

    #define VROTL(x, n) vorrq_u32(vshlq_n_u32((x), n), vshrq_n_u32((x), 32 - n))
    #define VLANEROT1(x) vextq_u32((x), (x), 1)
    #define VLANEROT2(x) vextq_u32((x), (x), 2)
    #define VLANEROT3(x) vextq_u32((x), (x), 3)
#elif defined(__wasm_simd128__)
    #include <wasm_simd128.h>
    #define Vec4u32 v128_t

    #define VLOADU(p)   wasm_v128_load((void *)(p))
    #define VSTOREU(p,x) wasm_v128_store((void *)(p), (x))
    #define VADD(a,b)  wasm_i32x4_add((a), (b))
    #define VXOR(a,b)  wasm_v128_xor((a), (b))

    #define VROTL(x, n) wasm_v128_or(wasm_i32x4_shl((x), n), wasm_u32x4_shr((x), 32-n))
    #define VLANEROT1(x) wasm_i32x4_shuffle((x), (x), 1, 2, 3, 0)
    #define VLANEROT2(x) wasm_i32x4_shuffle((x), (x), 2, 3, 0, 1)
    #define VLANEROT3(x) wasm_i32x4_shuffle((x), (x), 3, 0, 1, 2)
#else
    typedef struct { uint32_t x, y, z, w; } _Chacha_Vec4u32;
    #define Vec4u32 _Chacha_Vec4u32
    #if defined __cplusplus
        #define SINIT(T) T
    #else
        #define SINIT(T) (T)
    #endif

    static Vec4u32 _chacha_vloadu(void* p) {
        Vec4u32 v = {0};
        memcpy(&v, p, sizeof v);
        return v;
    }
    #define VLOADU(p)   _chacha_vloadu(p)
    #define VSTOREU(p, v) do {Vec4u32 _v = (v); memcpy((p), &_v, sizeof _v); } while(0)
    #define VADD(a, b) SINIT(Vec4u32){(a).x + (b).x, (a).y + (b).y, (a).z + (b).z, (a).w + (b).w}
    #define VXOR(a, b) SINIT(Vec4u32){(a).x ^ (b).x, (a).y ^ (b).y, (a).z ^ (b).z, (a).w ^ (b).w}

    #define SROTL32(x, n) (x << n) | (x >> (32u - n))
    #define VROTL(v, n) SINIT(Vec4u32){ \
        SROTL32((v).x, (n)), \
        SROTL32((v).y, (n)), \
        SROTL32((v).z, (n)), \
        SROTL32((v).w, (n))  \
    } 
    #define VLANEROT1(v) SINIT(Vec4u32){(v).y, (v).z, (v).w, (v).x}
    #define VLANEROT2(v) SINIT(Vec4u32){(v).z, (v).w, (v).x, (v).y}
    #define VLANEROT3(v) SINIT(Vec4u32){(v).w, (v).x, (v).y, (v).z}
#endif

void chacha_init_state(Chacha64* state, const Chacha32* key, uint32_t counter, const Chacha12* nonce) {
    uint8_t* i8 = (uint8_t*) (void*) state;
    uint32_t test = 0x11223344;
    uint8_t test8; memcpy(&test8, &test, 1);
    assert(test8 == 0x44 && "Big endian found while expecting little endian! Change the code!");

    memcpy(i8+0, "expand 32-byte k", 16);
    memcpy(i8+16, key, 32);
    memcpy(i8+48, &counter, 4);
    memcpy(i8+52, nonce, 12);
}

void chacha_init_state64(Chacha64* state, const Chacha32* key, uint64_t counter, const Chacha8* nonce) {
    uint8_t* i8 = (uint8_t*) (void*) state;
    uint32_t test = 0x11223344;
    uint8_t test8; memcpy(&test8, &test, 1);
    assert(test8 == 0x44 && "Big endian found while expecting little endian! Change the code!");

    memcpy(i8+0, "expand 32-byte k", 16);
    memcpy(i8+16, key, 32);
    memcpy(i8+48, &counter, 8);
    memcpy(i8+56, nonce, 8);
}

void chacha_block_xor(Chacha64* out_state, const Chacha64* plaintext, const Chacha64* in_state, uint32_t rounds) {
    uint8_t* i8 = (uint8_t*) (void*) in_state;
    uint8_t* p8 = (uint8_t*) (void*) plaintext;
    uint8_t* o8 = (uint8_t*) (void*) out_state;

    Vec4u32 a = VLOADU(i8 + 4*0);
    Vec4u32 b = VLOADU(i8 + 4*4);
    Vec4u32 c = VLOADU(i8 + 4*8);
    Vec4u32 d = VLOADU(i8 + 4*12);

    Vec4u32 oa = a;
    Vec4u32 ob = b;
    Vec4u32 oc = c;
    Vec4u32 od = d;

    #define QUARTER_ROUND(a, b, c, d) do { \
        a = VADD(a, b); d = VXOR(d, a); d = VROTL(d, 16); \
        c = VADD(c, d); b = VXOR(b, c); b = VROTL(b, 12); \
        a = VADD(a, b); d = VXOR(d, a); d = VROTL(d,  8); \
        c = VADD(c, d); b = VXOR(b, c); b = VROTL(b,  7); \
    } while (0)

    uint32_t doublerounds = rounds >> 1;
    for (uint32_t i = 0; i < doublerounds; i ++) {
        QUARTER_ROUND(a, b, c, d);

        b = VLANEROT1(b);
        c = VLANEROT2(c);
        d = VLANEROT3(d);

        QUARTER_ROUND(a, b, c, d);

        b = VLANEROT3(b);
        c = VLANEROT2(c);
        d = VLANEROT1(d);
    }

    a = VADD(a, oa);
    b = VADD(b, ob);
    c = VADD(c, oc);
    d = VADD(d, od);
    if (plaintext) {
        a = VXOR(a, VLOADU(p8 + 4*0));
        b = VXOR(b, VLOADU(p8 + 4*4));
        c = VXOR(c, VLOADU(p8 + 4*8));
        d = VXOR(d, VLOADU(p8 + 4*12));
    }

    VSTOREU(o8 + 4*0, a);
    VSTOREU(o8 + 4*4, b);
    VSTOREU(o8 + 4*8, c);
    VSTOREU(o8 + 4*12, d);
}

uint32_t chacha_xor_mutate(void* out, const void* in, size_t len, Chacha64* state, uint32_t rounds) {
    uint8_t *o8 = (uint8_t*) out;
    uint8_t *i8 = (uint8_t*) (void*) in;

    uint32_t counter = 0;
    memcpy(&counter, (uint8_t*) state + 48, 4);

    size_t full_to = len & ~(size_t) 63;
    size_t i = 0;
    for (; i < full_to; i += 64) {
        chacha_block_xor(o8 + i, i8 + i, state, rounds);
        counter += 1; memcpy((uint8_t*) state + 48, &counter, 4);
    }

    if (i < len) {
        size_t n = len - full_to;
        uint8_t keystream[64]; 
        chacha_block_xor(keystream, i8 + i, state, rounds);
        memcpy(o8 + i, keystream, n);
        counter += 1; memcpy((uint8_t*) state + 48, &counter, 4);
    }
    return counter;
}

uint32_t chacha_xor(void* out, const void* in, size_t len, const Chacha64* state, uint32_t rounds) {
    uint8_t state_copy[64];
    memcpy(state_copy, state, sizeof state_copy);
    return chacha_xor_mutate(out, in, len, state_copy, rounds);
}

void chacha_block(Chacha64* out_state, const Chacha64* in_state, uint32_t rounds) {
    chacha_block_xor(out_state, NULL, in_state, rounds);
}


#undef QUARTER_ROUND
#undef VLOADU
#undef VSTOREU
#undef VADD
#undef VXOR
#undef VROTL
#undef VLANEROT1
#undef VLANEROT2
#undef VLANEROT3
#undef Vec4u32
#ifdef SROTL32
#undef SROTL32
#endif

chacha.h

#pragma once

#include <stdint.h>

//Note: all buffers may be unaligned / point to any type, little endian assumed.
// typedefs to indicate fixed sized buffers without u8 
// which is annoying reinterpret to cast to
typedef void Chacha64;
typedef void Chacha32;
typedef void Chacha12;
typedef void Chacha8;

//Initializes chacha state with the given key, counter and nonce. 
// Each combination of key + counter + nonce should be used only once!
void chacha_init_state(Chacha64* state, const Chacha32* key, uint32_t counter, const Chacha12* nonce);
void chacha_init_state64(Chacha64* state, const Chacha32* key, uint64_t counter, const Chacha8* nonce);

//perform chacha on a block. The _xor variant is equivalent to calling chacha_block and then xoring the out_state with plaintext.
// rounds should be 20 according to spec but you can use lower value like 12 when using chacha as RNG.
void chacha_block(Chacha64* out_state, const Chacha64* in_state, uint32_t rounds);
void chacha_block_xor(Chacha64* out_state, const Chacha64* plaintext, const Chacha64* in_state, uint32_t rounds);

//xor encrypts input into output using the provided chacha state. Returns the modified counter value of the state
// The _mutate variant is same as above except mutates state, incrementing the counter.
uint32_t chacha_xor(void* output, const void* input, size_t len, const Chacha64* state, uint32_t rounds);
uint32_t chacha_xor_mutate(void* output, const void* input, size_t len, Chacha64* state, uint32_t rounds);

//define one of the following when compiling the .c file to get simd implementation:
//#define __SSE2__
//#define __AVX__
//#define __ARM_NEON
//#define __wasm_simd128__

main.c

#include "chacha.h"

#include <assert.h>
#include <string.h>
#include <stdio.h>

int main()
{
    //Comapre with test cases from https://www.rfc-editor.org/rfc/rfc8439#page-10
    {
        uint32_t key[8] = {
            0x03020100, 0x07060504, 0x0b0a0908, 0x0f0e0d0c,
            0x13121110, 0x17161514, 0x1b1a1918, 0x1f1e1d1c
        };

        uint32_t nonce[3] = {0x09000000, 0x4a000000, 0x00000000};
        uint32_t counter = 1;

        uint32_t state[16] = {0};
        chacha_init_state(state, key, counter, nonce);

        uint32_t output_block[16] = {0};
        chacha_block(output_block, state, 20);

        uint32_t expected_block[16] = {
            0xe4e7f110, 0x15593bd1, 0x1fdd0f50, 0xc47120a3,
            0xc7f4d1c7, 0x0368c033, 0x9aaa2204, 0x4e6cd4c3,
            0x466482d2, 0x09aa9f07, 0x05d7c214, 0xa2028bd9,
            0xd19c12b5, 0xb94e16de, 0xe883d0cb, 0x4e3c50a2,
        };
        assert(memcmp(output_block, expected_block, sizeof(output_block)) == 0);
    }

    {
        uint32_t key[8] = {
            0x03020100, 0x07060504, 0x0b0a0908, 0x0f0e0d0c,
            0x13121110, 0x17161514, 0x1b1a1918, 0x1f1e1d1c
        };

        uint32_t nonce[3] = {0x00000000, 0x4a000000, 0x00000000};
        uint32_t counter = 1;
        uint32_t state[16] = {0};
        chacha_init_state(state, key, counter, nonce);

        const char input_text[] = 
            "Ladies and Gentlemen of the class of '99: If I could offer you"
            " only one tip for the future, sunscreen would be it.";

        uint8_t  output_text[128] = {0};
        uint32_t len_text = sizeof(input_text) - 1;
        chacha_xor(output_text, input_text, len_text, state, 20);

        uint8_t expected_text[] = {
            0x6e, 0x2e, 0x35, 0x9a, 0x25, 0x68, 0xf9, 0x80,
            0x41, 0xba, 0x07, 0x28, 0xdd, 0x0d, 0x69, 0x81,
            0xe9, 0x7e, 0x7a, 0xec, 0x1d, 0x43, 0x60, 0xc2,
            0x0a, 0x27, 0xaf, 0xcc, 0xfd, 0x9f, 0xae, 0x0b,
            0xf9, 0x1b, 0x65, 0xc5, 0x52, 0x47, 0x33, 0xab,
            0x8f, 0x59, 0x3d, 0xab, 0xcd, 0x62, 0xb3, 0x57,
            0x16, 0x39, 0xd6, 0x24, 0xe6, 0x51, 0x52, 0xab,
            0x8f, 0x53, 0x0c, 0x35, 0x9f, 0x08, 0x61, 0xd8,
            0x07, 0xca, 0x0d, 0xbf, 0x50, 0x0d, 0x6a, 0x61,
            0x56, 0xa3, 0x8e, 0x08, 0x8a, 0x22, 0xb6, 0x5e,
            0x52, 0xbc, 0x51, 0x4d, 0x16, 0xcc, 0xf8, 0x06,
            0x81, 0x8c, 0xe9, 0x1a, 0xb7, 0x79, 0x37, 0x36,
            0x5a, 0xf9, 0x0b, 0xbf, 0x74, 0xa3, 0x5b, 0xe6,
            0xb4, 0x0b, 0x8e, 0xed, 0xf2, 0x78, 0x5e, 0x42,
            0x87, 0x4d
        };
        assert(memcmp(output_text, expected_text, len_text) == 0);
    }

    printf("tests passed");
}