From d102fa1d289bf3ac4489c76ee6310bde3d0bca34 Mon Sep 17 00:00:00 2001 From: Leo Vasanko Date: Wed, 22 Nov 2023 04:03:48 +0000 Subject: [PATCH] Implement Numpy BitGenerator. --- ,gitignore | 2 + benchmarks/chaopenssl.c | 55 ++++++++++++++++++++ pyproject.toml | 2 +- randquik/cha.py | 2 +- setup.py | 6 +++ src/cha1block.h | 3 -- src/cha4block.h | 3 -- src/cha8block.h | 3 -- src/chacha20.c | 91 --------------------------------- src/chacha20.h | 109 +++++++++++++++++++++++++++++++++------- src/npbitgen.h | 19 +++++++ src/nprand.pyx | 39 ++++++++++++++ 12 files changed, 215 insertions(+), 119 deletions(-) create mode 100644 ,gitignore create mode 100644 benchmarks/chaopenssl.c create mode 100644 setup.py delete mode 100644 src/chacha20.c create mode 100644 src/npbitgen.h create mode 100644 src/nprand.pyx diff --git a/,gitignore b/,gitignore new file mode 100644 index 0000000..6b29aac --- /dev/null +++ b/,gitignore @@ -0,0 +1,2 @@ +__PYCACHE__ +.* diff --git a/benchmarks/chaopenssl.c b/benchmarks/chaopenssl.c new file mode 100644 index 0000000..ef4acf9 --- /dev/null +++ b/benchmarks/chaopenssl.c @@ -0,0 +1,55 @@ +#include +#include +#include + +int main() { + // Key and IV should be appropriately sized for ChaCha20 + unsigned char 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 + }; + unsigned char iv[] = { + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00 + }; + + // Initialize context + EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new(); + if (!ctx) { + perror("EVP_CIPHER_CTX_new failed"); + return 1; + } + + // Initialize the ChaCha20 cipher + if (!EVP_EncryptInit_ex(ctx, EVP_chacha20(), NULL, key, iv)) { + perror("EVP_EncryptInit_ex failed"); + EVP_CIPHER_CTX_free(ctx); + return 1; + } + + // Buffer for the keystream + unsigned char keystream[1000000]; + memset(keystream, 0, sizeof(keystream)); + + for (int i = 0; i < 1000; i++) { + // Generate keystream + int len; + if (!EVP_EncryptUpdate(ctx, keystream, &len, keystream, sizeof keystream)) { + perror("EVP_EncryptUpdate failed"); + EVP_CIPHER_CTX_free(ctx); + return 1; + } + } + // Clean up + EVP_CIPHER_CTX_free(ctx); + + // Print the generated keystream + for (int i = 0; i < 64; i++) { + printf("%02x", keystream[i]); + } + printf("\n"); + + return 0; +} diff --git a/pyproject.toml b/pyproject.toml index 06dfa1f..be3b39a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "randquik" version = "0.1.0" description = "Extremely fast and cryptographically secure random number generator." readme = "README.md" -license = "" +license.text = "Public Domain" authors = [{ name = "Vasanko" }] classifiers = [ "Operating System :: POSIX", diff --git a/randquik/cha.py b/randquik/cha.py index 7f75448..64d45d6 100644 --- a/randquik/cha.py +++ b/randquik/cha.py @@ -46,7 +46,7 @@ def _processKeys(key, iv): iv = bytes(4) + iv else: raise ValueError( - "iv lenth must be 8 (original nonce), 12 (IETF) or 16 (counter in initial 8 bytes)" + "iv lenth must be 8 (ChaCha20 original), 12 (IETF) or 16 (counter in initial 8 bytes)" ) return ffi.from_buffer(key), ffi.from_buffer(iv) diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..6706d91 --- /dev/null +++ b/setup.py @@ -0,0 +1,6 @@ +from distutils.core import setup + +import numpy +from Cython.Build import cythonize + +setup(ext_modules=cythonize("src/nprand.pyx"), include_dirs=[numpy.get_include()]) diff --git a/src/cha1block.h b/src/cha1block.h index fb796af..cd03e92 100644 --- a/src/cha1block.h +++ b/src/cha1block.h @@ -1,6 +1,3 @@ -#pragma once -#include "chacha20.h" - #include #include #include diff --git a/src/cha4block.h b/src/cha4block.h index 84fa0d0..9500f72 100644 --- a/src/cha4block.h +++ b/src/cha4block.h @@ -1,6 +1,3 @@ -#pragma once -#include "chacha20.h" - #define VEC4_ROT(A, IMM) \ _mm_or_si128(_mm_slli_epi32(A, IMM), _mm_srli_epi32(A, (32 - IMM))) diff --git a/src/cha8block.h b/src/cha8block.h index 3d0bc8d..0991163 100644 --- a/src/cha8block.h +++ b/src/cha8block.h @@ -1,6 +1,3 @@ -#pragma once -#include "chacha20.h" - #define VEC8_ROT(A, IMM) \ _mm256_or_si256(_mm256_slli_epi32(A, IMM), _mm256_srli_epi32(A, (32 - IMM))) diff --git a/src/chacha20.c b/src/chacha20.c deleted file mode 100644 index 4b50d39..0000000 --- a/src/chacha20.c +++ /dev/null @@ -1,91 +0,0 @@ -#include "chacha20.h" - -#if defined(__x86_64__) -#ifdef __GNUC__ -#pragma GCC target("sse2") -#pragma GCC target("ssse3") -#pragma GCC target("avx2") -#endif - -#include // SSE2 -#include "cha4block.h" - -#include // AVX2 -#include // SSSE3 -#include "cha8block.h" - -#elif defined(__aarch64__) -#include "sse2neon.h" -#include "cha4block.h" -#endif - -#include "cha1block.h" - -#include -#include -#include -#include -#include - -#include -#include -#include - -void cha_init(cha_ctx* ctx, const uint8_t* key, const uint8_t* iv) { - ctx->state[0] = 0x61707865; - ctx->state[1] = 0x3320646e; - ctx->state[2] = 0x79622d32; - ctx->state[3] = 0x6b206574; - memcpy(ctx->state + 4, key, 32); - memcpy(ctx->state + 12, iv, 16); - memset(ctx->unconsumed, 0, sizeof ctx->unconsumed); - ctx->uncount = 0; -} - -void cha_wipe(cha_ctx* ctx) { memset(ctx, 0, sizeof(cha_ctx)); } - -int cha_update(cha_ctx* ctx, uint8_t* out, uint64_t outlen) { - // The included header will mess with these variables - uint8_t* c = out; - uint8_t* end = out + outlen; - if (ctx->uncount) { - // Deliver stored bytes first - uint64_t N = ctx->uncount >= outlen ? outlen : ctx->uncount; - fprintf(stderr, "%llu, %i, %llu\n", N, ctx->uncount, outlen); - memcpy(c, ctx->unconsumed, N); - ctx->uncount -= N; - c += N; - if (ctx->uncount) { - memmove(ctx->unconsumed, ctx->unconsumed + N, ctx->uncount); - } - if (c == out + outlen) - return 0; - } -#if defined(__x86_64__) - // TODO: Handle resume if we are not at block boundary - if (__builtin_cpu_supports("ssse3")) { - if (__builtin_cpu_supports("avx2")) { - c += _cha_8block(ctx, c, end); - assert(end - c < 512); - } - c += _cha_4block(ctx, c, end); - assert(end - c < 256); - } -#elif defined(__aarch64__) - c += _cha_4block(ctx, c, end); -#endif - c += _cha_block(ctx, c, end); - assert(c == end); - return 0; -} - -// ChaCha20 -int cha_generate( - uint8_t* out, uint64_t outlen, const uint8_t key[32], const uint8_t iv[16] -) { - cha_ctx ctx; - cha_init(&ctx, key, iv); - cha_update(&ctx, out, outlen); - cha_wipe(&ctx); - return 0; -} diff --git a/src/chacha20.h b/src/chacha20.h index 1a69914..91be67f 100644 --- a/src/chacha20.h +++ b/src/chacha20.h @@ -1,4 +1,3 @@ -#pragma once #include #include @@ -6,33 +5,109 @@ static const uint64_t CHA_BLOCK_SIZE = 64; typedef struct cha_ctx { uint32_t state[16]; - uint8_t unconsumed[64]; + uint8_t unconsumed[CHA_BLOCK_SIZE]; uint8_t uncount; } cha_ctx; +#if defined(__x86_64__) +#ifdef __GNUC__ +#pragma GCC target("sse2") +#pragma GCC target("ssse3") +#pragma GCC target("avx2") +#endif + +#include // SSE2 +#include "cha4block.h" + +#include // AVX2 +#include // SSSE3 +#include "cha8block.h" + +#elif defined(__aarch64__) +#include "sse2neon.h" +#include "cha4block.h" +#endif + +#include "cha1block.h" + +#include +#include +#include +#include +#include + +#include +#include +#include + + /// @brief Initialize cha_ctx /// @param ctx holds ChaCha20 state /// @param key 32 byte key /// @param iv 16 bytes, where normally initial 4-8 bytes are zeroes and the rest /// nonce -void cha_init(cha_ctx* ctx, const uint8_t* key, const uint8_t* iv); +void cha_init(cha_ctx* ctx, const uint8_t* key, const uint8_t* iv) { + ctx->state[0] = 0x61707865; + ctx->state[1] = 0x3320646e; + ctx->state[2] = 0x79622d32; + ctx->state[3] = 0x6b206574; + memcpy(ctx->state + 4, key, 32); + memcpy(ctx->state + 12, iv, 16); + memset(ctx->unconsumed, 0, sizeof ctx->unconsumed); + ctx->uncount = 0; +} /// Dispose of sensitive data within the context void cha_wipe(cha_ctx* ctx); -/// @brief Incremental upgrade -/// @param ctx Gets updated -/// @param out -/// @param outlen -/// @return -int cha_update(cha_ctx* ctx, uint8_t* out, uint64_t outlen); +void cha_wipe(cha_ctx* ctx) { memset(ctx, 0, sizeof(cha_ctx)); } -/// @brief Produce a requested number of random bytes of the stream. -/// @param out -/// @param outlen -/// @param key -/// @param iv -/// @return -int cha_generate( +/// @brief Incremental generation, keeps state between calls +/// @param ctx ChaCha20 context +/// @param out output buffer +/// @param outlen output buffer length +void cha_update(cha_ctx* ctx, uint8_t* out, uint64_t outlen) { + // The included header will mess with these variables + uint8_t* c = out; + uint8_t* end = out + outlen; + if (ctx->uncount) { + // Deliver stored bytes first + uint64_t N = ctx->uncount >= outlen ? outlen : ctx->uncount; + memcpy(c, ctx->unconsumed, N); + ctx->uncount -= N; + c += N; + if (ctx->uncount) { + memmove(ctx->unconsumed, ctx->unconsumed + N, ctx->uncount); + } + if (c == out + outlen) return; + } +#if defined(__x86_64__) + // TODO: Handle resume if we are not at block boundary + if (__builtin_cpu_supports("ssse3")) { + if (__builtin_cpu_supports("avx2")) { + c += _cha_8block(ctx, c, end); + assert(end - c < 512); + } + c += _cha_4block(ctx, c, end); + assert(end - c < 256); + } +#elif defined(__aarch64__) + c += _cha_4block(ctx, c, end); +#endif + c += _cha_block(ctx, c, end); + assert(c == end); +} + +/// @brief Produce a requested number of random bytes of the stream, one shot. +/// @param out output buffer +/// @param outlen output buffer length +/// @param key 32 byte key +/// @param iv 16 bytes, where normally initial 4-8 bytes are zeroes (counter) +void cha_generate( uint8_t* out, uint64_t outlen, const uint8_t key[32], const uint8_t iv[16] -); +) { + cha_ctx ctx; + cha_init(&ctx, key, iv); + cha_update(&ctx, out, outlen); + cha_wipe(&ctx); +} diff --git a/src/npbitgen.h b/src/npbitgen.h new file mode 100644 index 0000000..51a7eb5 --- /dev/null +++ b/src/npbitgen.h @@ -0,0 +1,19 @@ +#include "chacha20.h" + + +static uint64_t cha_uint64(void *st) { + cha_ctx *ctx = (cha_ctx *)st; + uint64_t ret; + cha_update(ctx, (uint8_t *)&ret, sizeof ret); + return ret; +} +static uint32_t cha_uint32(void *st) { + cha_ctx *ctx = (cha_ctx *)st; + uint32_t ret; + cha_update(ctx, (uint8_t *)&ret, sizeof ret); + return ret; + +} +static double cha_double(void *st) { + return cha_uint64(st) / (UINT64_MAX + 1.0); +} diff --git a/src/nprand.pyx b/src/nprand.pyx new file mode 100644 index 0000000..5fc1ae0 --- /dev/null +++ b/src/nprand.pyx @@ -0,0 +1,39 @@ +#cython: language_level=3 + +from libc.stdint cimport uint32_t, uint8_t, uint64_t +from cpython.pycapsule cimport PyCapsule_IsValid, PyCapsule_GetPointer + +import numpy as np +cimport numpy as np +cimport cython +import secrets + +from numpy.random cimport BitGenerator + +np.import_array() + +cdef extern from "npbitgen.h": + struct cha_ctx: + uint32_t state[16] + uint8_t unconsumed[64] + uint8_t uncount + + void cha_init(cha_ctx* ctx, const uint8_t* key, const uint8_t* iv) nogil + uint64_t cha_uint64(void *state) nogil + uint32_t cha_uint32(void *state) nogil + double cha_double(void *state) nogil + +cdef class Cha(BitGenerator): + + cdef cha_ctx rng_state + + def __init__(self, seed=None): + BitGenerator.__init__(self, seed) + self._bitgen.state = &self.rng_state + self._bitgen.next_uint64 = &cha_uint64 + self._bitgen.next_uint32 = &cha_uint32 + self._bitgen.next_double = &cha_double + self._bitgen.next_raw = &cha_uint64 + # Generated state is ChaCha20 key + key = self._seed_seq.generate_state(4, np.uint64) + cha_init(&self.rng_state, np.PyArray_DATA(key), b"NumpRand")