LeoVasanko/RandQuik
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implement Numpy BitGenerator.

Browse Source
This commit is contained in:
LeoVasanko
2023-11-22 04:03:48 +00:00
parent c573505826
commit d102fa1d28
12 changed files with 215 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cha1block.h
-3
View File
@@ -1,6 +1,3 @@
#pragma once
#include "chacha20.h"
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
src/cha4block.h
-3
View File
@@ -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)))
src/cha8block.h
-3
View File
@@ -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)))
src/chacha20.c
-91
View File
@@ -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 <emmintrin.h> // SSE2
#include "cha4block.h"
#include <immintrin.h> // AVX2
#include <tmmintrin.h> // SSSE3
#include "cha8block.h"
#elif defined(__aarch64__)
#include "sse2neon.h"
#include "cha4block.h"
#endif
#include "cha1block.h"
#include <assert.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <pthread.h>
#include <time.h>
#include <unistd.h>
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;
}
src/chacha20.h
+92 -17
View File
@@ -1,4 +1,3 @@
#pragma once
#include <stdbool.h>
#include <stdint.h>
@@ -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 <emmintrin.h> // SSE2
#include "cha4block.h"
#include <immintrin.h> // AVX2
#include <tmmintrin.h> // SSSE3
#include "cha8block.h"
#elif defined(__aarch64__)
#include "sse2neon.h"
#include "cha4block.h"
#endif
#include "cha1block.h"
#include <assert.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <pthread.h>
#include <time.h>
#include <unistd.h>
/// @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);
}
src/npbitgen.h
+19
View File
@@ -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);
}
src/nprand.pyx
+39
View File
@@ -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 = <void *>&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, <uint8_t *>np.PyArray_DATA(key), b"NumpRand")