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