1#include <smp/core.h>
2#include <smp/percpu.h>
3#include <stdint.h>
4#include <string.h>
5
6struct prng_core {
7 uint32_t state[16];
8 uint8_t buffer[64]; // keystream buffer
9 size_t pos;
10};
11
12void prng_build(struct prng_core *this_one, size_t cpu) {
13 (void) this_one, (void) cpu;
14}
15
16PERCPU_DECLARE(pcs, struct prng_core, prng_build);
17#define prng_core_state PERCPU_READ(pcs)
18
19static void prng_seed_core(uint64_t seed) {
20 const char *sigma = "expand 32-byte k";
21 uint8_t key[32];
22
23 uint64_t tsc = smp_core()->last_tsc;
24 uint32_t core_id = smp_core()->id;
25
26 for (int i = 0; i < 4; i++) {
27 key[i * 8 + 0] = (seed >> (i * 8 + 0)) & 0xff;
28 key[i * 8 + 1] = (seed >> (i * 8 + 8)) & 0xff;
29
30 key[i * 8 + 2] = (tsc >> (i * 8 + 32)) & 0xff;
31 key[i * 8 + 3] = (tsc >> (i * 8 + 32)) & 0xff;
32
33 key[i * 8 + 4] = (core_id >> (i * 8 + 0)) & 0xff;
34 key[i * 8 + 5] = (core_id >> (i * 8 + 4)) & 0xff;
35
36 key[i * 8 + 6] = (tsc >> (i * 8 + 16)) & 0xff;
37 key[i * 8 + 7] = (tsc >> (i * 8 + 24)) & 0xff;
38 }
39
40 for (int i = 0; i < 4; i++)
41 prng_core_state.state[i] = ((uint32_t *) sigma)[i];
42 for (int i = 0; i < 8; i++)
43 prng_core_state.state[4 + i] = ((uint32_t *) key)[i];
44
45 prng_core_state.state[12] = 0;
46 prng_core_state.state[13] = 0;
47 prng_core_state.state[14] = 0;
48 prng_core_state.state[15] = 0;
49
50 prng_core_state.pos = 64;
51}
52
53#define QR(a, b, c, d) \
54 a += b; \
55 d ^= a; \
56 d = (d << 16) | (d >> 16); \
57 c += d; \
58 b ^= c; \
59 b = (b << 12) | (b >> 20); \
60 a += b; \
61 d ^= a; \
62 d = (d << 8) | (d >> 24); \
63 c += d; \
64 b ^= c; \
65 b = (b << 7) | (b >> 25)
66
67static void chacha20_generate_block(uint8_t out[64], uint32_t state[16]) {
68 uint32_t x[16];
69 memcpy(x, state, sizeof(x));
70
71 for (int i = 0; i < 10; i++) {
72 QR(x[0], x[4], x[8], x[12]);
73 QR(x[1], x[5], x[9], x[13]);
74 QR(x[2], x[6], x[10], x[14]);
75 QR(x[3], x[7], x[11], x[15]);
76 QR(x[0], x[5], x[10], x[15]);
77 QR(x[1], x[6], x[11], x[12]);
78 QR(x[2], x[7], x[8], x[13]);
79 QR(x[3], x[4], x[9], x[14]);
80 }
81
82 for (int i = 0; i < 16; i++) {
83 x[i] += state[i];
84 out[i * 4 + 0] = x[i] & 0xff;
85 out[i * 4 + 1] = (x[i] >> 8) & 0xff;
86 out[i * 4 + 2] = (x[i] >> 16) & 0xff;
87 out[i * 4 + 3] = (x[i] >> 24) & 0xff;
88 }
89}
90
91uint64_t prng_next(void) {
92 if (prng_core_state.pos >= 64) {
93 chacha20_generate_block(prng_core_state.buffer, prng_core_state.state);
94 prng_core_state.state[12]++;
95 prng_core_state.pos = 0;
96 }
97
98 uint64_t val;
99 memcpy(&val, prng_core_state.buffer + prng_core_state.pos,
100 sizeof(uint64_t));
101 prng_core_state.pos += 8;
102 return val;
103}
104
105void prng_seed(uint64_t seed) {
106 if (seed == 0)
107 seed = smp_core()->last_tsc;
108 prng_seed_core(seed);
109}
110