DeterministicESPAsyncWebServer v6.27.1
Zero-allocation, bounded-execution async HTTP server for ESP32
Loading...
Searching...
No Matches
mlkem.cpp
Go to the documentation of this file.
1// Copyright (C) 2026 Douglas Quigg (dstroy0) <dquigg123@gmail.com>
2// SPDX-License-Identifier: AGPL-3.0-or-later
3
4#include "network_drivers/presentation/pqc/mlkem.h"
5
6#if DETWS_ENABLE_PQC_KEX
7
8#include "network_drivers/presentation/pqc/sha3.h"
9#include <string.h>
10
11// ML-KEM-768 parameters (FIPS 203).
12#define MK_N 256
13#define MK_Q 3329
14#define MK_K 3
15#define MK_ETA 2 // eta1 == eta2 == 2 for ML-KEM-768
16#define MK_DU 10
17#define MK_DV 4
18#define MK_POLYBYTES 384 // 12 bits * 256 / 8
19#define MK_QINV (-3327) // q^-1 mod 2^16, signed
20
21// Twiddle factors zeta^BitRev7(i) in Montgomery form, reduced to (-q/2, q/2] (FIPS 203 / pq-crystals).
22static const int16_t mk_zetas[128] = {
23 -1044, -758, -359, -1517, 1493, 1422, 287, 202, -171, 622, 1577, 182, 962, -1202, -1474, 1468,
24 573, -1325, 264, 383, -829, 1458, -1602, -130, -681, 1017, 732, 608, -1542, 411, -205, -1571,
25 1223, 652, -552, 1015, -1293, 1491, -282, -1544, 516, -8, -320, -666, -1618, -1162, 126, 1469,
26 -853, -90, -271, 830, 107, -1421, -247, -951, -398, 961, -1508, -725, 448, -1065, 677, -1275,
27 -1103, 430, 555, 843, -1251, 871, 1550, 105, 422, 587, 177, -235, -291, -460, 1574, 1653,
28 -246, 778, 1159, -147, -777, 1483, -602, 1119, -1590, 644, -872, 349, 418, 329, -156, -75,
29 817, 1097, 603, 610, 1322, -1285, -1465, 384, -1215, -136, 1218, -1335, -874, 220, -1187, -1659,
30 -1185, -1530, -1278, 794, -1510, -854, -870, 478, -108, -308, 996, 991, 958, -1460, 1522, 1628};
31
32// Montgomery reduction: given a = m*R, return m mod q in (-q, q). R = 2^16.
33static int16_t mont_reduce(int32_t a)
34{
35 int16_t t = (int16_t)((int16_t)a * (int16_t)MK_QINV);
36 t = (int16_t)((a - (int32_t)t * MK_Q) >> 16);
37 return t;
38}
39
40static inline int16_t fqmul(int16_t a, int16_t b)
41{
42 return mont_reduce((int32_t)a * b);
43}
44
45// Barrett reduction: a mod q in (-q/2, q/2].
46static int16_t barrett_reduce(int16_t a)
47{
48 const int16_t v = 20159; // round(2^26 / q)
49 int16_t t = (int16_t)(((int32_t)v * a + (1 << 25)) >> 26);
50 t = (int16_t)(t * MK_Q);
51 return (int16_t)(a - t);
52}
53
54// Forward NTT (in place). Coefficients enter in normal order, leave in bit-reversed NTT order.
55static void ntt(int16_t r[MK_N])
56{
57 unsigned k = 1;
58 for (unsigned len = 128; len >= 2; len >>= 1)
59 {
60 for (unsigned start = 0; start < MK_N; start += (len << 1))
61 {
62 int16_t zeta = mk_zetas[k++];
63 for (unsigned j = start; j < start + len; j++)
64 {
65 int16_t t = fqmul(zeta, r[j + len]);
66 r[j + len] = (int16_t)(r[j] - t);
67 r[j] = (int16_t)(r[j] + t);
68 }
69 }
70 }
71}
72
73// Inverse NTT (in place), with the 1/128 * Montgomery scaling folded into the final multiply.
74static void invntt(int16_t r[MK_N])
75{
76 const int16_t f = 1441; // mont^2 / 128
77 unsigned k = 127;
78 for (unsigned len = 2; len <= 128; len <<= 1)
79 {
80 for (unsigned start = 0; start < MK_N; start += (len << 1))
81 {
82 int16_t zeta = mk_zetas[k--];
83 for (unsigned j = start; j < start + len; j++)
84 {
85 int16_t t = r[j];
86 r[j] = barrett_reduce((int16_t)(t + r[j + len]));
87 r[j + len] = (int16_t)(r[j + len] - t);
88 r[j + len] = fqmul(zeta, r[j + len]);
89 }
90 }
91 }
92 for (unsigned j = 0; j < MK_N; j++)
93 r[j] = fqmul(r[j], f);
94}
95
96// Multiply two degree-1 residues mod (X^2 - zeta) in the NTT domain.
97static void basemul(int16_t r[2], const int16_t a[2], const int16_t b[2], int16_t zeta)
98{
99 r[0] = fqmul(a[1], b[1]);
100 r[0] = fqmul(r[0], zeta);
101 r[0] = (int16_t)(r[0] + fqmul(a[0], b[0]));
102 r[1] = fqmul(a[0], b[1]);
103 r[1] = (int16_t)(r[1] + fqmul(a[1], b[0]));
104}
105
106static void poly_basemul(int16_t r[MK_N], const int16_t a[MK_N], const int16_t b[MK_N])
107{
108 for (unsigned i = 0; i < MK_N / 4; i++)
109 {
110 basemul(&r[4 * i], &a[4 * i], &b[4 * i], mk_zetas[64 + i]);
111 basemul(&r[4 * i + 2], &a[4 * i + 2], &b[4 * i + 2], (int16_t)(-mk_zetas[64 + i]));
112 }
113}
114
115// r = sum_i a[i] o b[i] (pointwise in the NTT domain), then Barrett-reduced.
116static void polyvec_basemul_acc(int16_t r[MK_N], const int16_t a[MK_K][MK_N], const int16_t b[MK_K][MK_N])
117{
118 int16_t t[MK_N];
119 poly_basemul(r, a[0], b[0]);
120 for (unsigned i = 1; i < MK_K; i++)
121 {
122 poly_basemul(t, a[i], b[i]);
123 for (unsigned j = 0; j < MK_N; j++)
124 r[j] = (int16_t)(r[j] + t[j]);
125 }
126 for (unsigned j = 0; j < MK_N; j++)
127 r[j] = barrett_reduce(r[j]);
128}
129
130// ByteDecode_12: 384 octets -> 256 coefficients in [0, 2^12).
131static void poly_frombytes(int16_t r[MK_N], const uint8_t a[MK_POLYBYTES])
132{
133 for (unsigned i = 0; i < MK_N / 2; i++)
134 {
135 r[2 * i] = (int16_t)(((a[3 * i + 0] >> 0) | ((uint16_t)a[3 * i + 1] << 8)) & 0xFFF);
136 r[2 * i + 1] = (int16_t)(((a[3 * i + 1] >> 4) | ((uint16_t)a[3 * i + 2] << 4)) & 0xFFF);
137 }
138}
139
140// Decompress_1: each message bit b -> b ? (q+1)/2 : 0.
141static void poly_frommsg(int16_t r[MK_N], const uint8_t msg[32])
142{
143 for (unsigned i = 0; i < 32; i++)
144 for (unsigned j = 0; j < 8; j++)
145 {
146 int16_t mask = (int16_t)(-(int16_t)((msg[i] >> j) & 1));
147 r[8 * i + j] = (int16_t)(mask & ((MK_Q + 1) / 2));
148 }
149}
150
151static inline uint32_t load32_le(const uint8_t *x)
152{
153 return (uint32_t)x[0] | ((uint32_t)x[1] << 8) | ((uint32_t)x[2] << 16) | ((uint32_t)x[3] << 24);
154}
155
156// Centered binomial distribution, eta = 2: 128 octets of PRF output -> 256 coefficients in [-2, 2].
157static void cbd2(int16_t r[MK_N], const uint8_t buf[128])
158{
159 for (unsigned i = 0; i < MK_N / 8; i++)
160 {
161 uint32_t t = load32_le(buf + 4 * i);
162 uint32_t d = t & 0x55555555u;
163 d += (t >> 1) & 0x55555555u;
164 for (unsigned j = 0; j < 8; j++)
165 {
166 int16_t a = (int16_t)((d >> (4 * j)) & 0x3);
167 int16_t b = (int16_t)((d >> (4 * j + 2)) & 0x3);
168 r[8 * i + j] = (int16_t)(a - b);
169 }
170 }
171}
172
173// PRF_eta(seed, nonce) = SHAKE256(seed || nonce), then sample CBD_eta (eta = 2).
174static void poly_getnoise(int16_t r[MK_N], const uint8_t seed[32], uint8_t nonce)
175{
176 uint8_t extseed[33];
177 memcpy(extseed, seed, 32);
178 extseed[32] = nonce;
179 uint8_t buf[MK_ETA * MK_N / 4]; // 128
180 shake256(buf, sizeof(buf), extseed, sizeof(extseed));
181 cbd2(r, buf);
182}
183
184// One transposed matrix entry: A^T[i][j] = SampleNTT(XOF(rho || i || j)) (FIPS 203 gen with (i,j)).
185static void gen_matrix_entry(int16_t out[MK_N], const uint8_t rho[32], uint8_t i, uint8_t j)
186{
187 uint8_t seed[34];
188 memcpy(seed, rho, 32);
189 seed[32] = i;
190 seed[33] = j;
191 KeccakCtx ctx;
192 shake128_absorb(&ctx, seed, sizeof(seed));
193
194 unsigned count = 0;
195 while (count < MK_N)
196 {
197 uint8_t buf[KECCAK_RATE_SHAKE128]; // 168 = 56*3, no 3-octet group straddles a block
198 keccak_squeeze(&ctx, buf, sizeof(buf));
199 for (unsigned p = 0; p + 3 <= sizeof(buf) && count < MK_N; p += 3)
200 {
201 uint16_t d1 = (uint16_t)(buf[p] | ((uint16_t)(buf[p + 1] & 0xF) << 8));
202 uint16_t d2 = (uint16_t)((buf[p + 1] >> 4) | ((uint16_t)buf[p + 2] << 4));
203 if (d1 < MK_Q)
204 out[count++] = (int16_t)d1;
205 if (count < MK_N && d2 < MK_Q)
206 out[count++] = (int16_t)d2;
207 }
208 }
209}
210
211// Compress_10 + ByteEncode_10 for one polynomial: 256 coefficients -> 320 octets.
212static void poly_compress10(uint8_t r[320], const int16_t a[MK_N])
213{
214 unsigned k = 0;
215 for (unsigned i = 0; i < MK_N / 4; i++)
216 {
217 uint16_t t[4];
218 for (unsigned j = 0; j < 4; j++)
219 {
220 int16_t u = a[4 * i + j];
221 u = (int16_t)(u + ((u >> 15) & MK_Q)); // to [0, q)
222 t[j] = (uint16_t)(((((uint32_t)u << 10) + MK_Q / 2) / MK_Q) & 0x3FF);
223 }
224 r[k + 0] = (uint8_t)(t[0]);
225 r[k + 1] = (uint8_t)((t[0] >> 8) | (t[1] << 2));
226 r[k + 2] = (uint8_t)((t[1] >> 6) | (t[2] << 4));
227 r[k + 3] = (uint8_t)((t[2] >> 4) | (t[3] << 6));
228 r[k + 4] = (uint8_t)(t[3] >> 2);
229 k += 5;
230 }
231}
232
233// Compress_4 + ByteEncode_4 for one polynomial: 256 coefficients -> 128 octets.
234static void poly_compress4(uint8_t r[128], const int16_t a[MK_N])
235{
236 for (unsigned i = 0; i < MK_N / 8; i++)
237 {
238 uint8_t t[8];
239 for (unsigned j = 0; j < 8; j++)
240 {
241 int16_t u = a[8 * i + j];
242 u = (int16_t)(u + ((u >> 15) & MK_Q));
243 t[j] = (uint8_t)(((((uint16_t)u << 4) + MK_Q / 2) / MK_Q) & 15);
244 }
245 r[4 * i + 0] = (uint8_t)(t[0] | (t[1] << 4));
246 r[4 * i + 1] = (uint8_t)(t[2] | (t[3] << 4));
247 r[4 * i + 2] = (uint8_t)(t[4] | (t[5] << 4));
248 r[4 * i + 3] = (uint8_t)(t[6] | (t[7] << 4));
249 }
250}
251
252// FIPS 203 modulus check: ek's decoded coefficients must all be < q.
253static bool check_ek(const uint8_t ek[MLKEM768_EK_BYTES])
254{
255 for (unsigned i = 0; i < MK_K; i++)
256 {
257 int16_t p[MK_N];
258 poly_frombytes(p, ek + i * MK_POLYBYTES);
259 for (unsigned j = 0; j < MK_N; j++)
260 if ((uint16_t)p[j] >= MK_Q)
261 return false;
262 }
263 return true;
264}
265
266// K-PKE.Encrypt(ek, m, r) -> ct. u is streamed and compressed one row at a time to bound stack.
267static void k_pke_encrypt(uint8_t ct[MLKEM768_CT_BYTES], const uint8_t ek[MLKEM768_EK_BYTES], const uint8_t m[32],
268 const uint8_t coins[32])
269{
270 int16_t that[MK_K][MK_N];
271 for (unsigned i = 0; i < MK_K; i++)
272 poly_frombytes(that[i], ek + i * MK_POLYBYTES);
273 const uint8_t *rho = ek + MK_K * MK_POLYBYTES;
274
275 int16_t sp[MK_K][MK_N];
276 for (unsigned i = 0; i < MK_K; i++)
277 {
278 poly_getnoise(sp[i], coins, (uint8_t)i); // y, nonce 0..k-1
279 ntt(sp[i]);
280 }
281
282 // u = NTT^-1(A^T o y) + e1, compressed with du = 10 (320 octets/row).
283 for (unsigned i = 0; i < MK_K; i++)
284 {
285 int16_t at_row[MK_K][MK_N];
286 for (unsigned j = 0; j < MK_K; j++)
287 gen_matrix_entry(at_row[j], rho, (uint8_t)i, (uint8_t)j);
288 int16_t u_row[MK_N];
289 polyvec_basemul_acc(u_row, at_row, sp);
290 invntt(u_row);
291 int16_t e1[MK_N];
292 poly_getnoise(e1, coins, (uint8_t)(MK_K + i)); // e1, nonce k..2k-1
293 for (unsigned x = 0; x < MK_N; x++)
294 u_row[x] = barrett_reduce((int16_t)(u_row[x] + e1[x]));
295 poly_compress10(ct + i * 320, u_row);
296 }
297
298 // v = NTT^-1(t^T o y) + e2 + Decompress_1(m), compressed with dv = 4 (128 octets).
299 int16_t v[MK_N];
300 polyvec_basemul_acc(v, that, sp);
301 invntt(v);
302 int16_t e2[MK_N];
303 poly_getnoise(e2, coins, (uint8_t)(2 * MK_K)); // e2, nonce 2k
304 int16_t mu[MK_N];
305 poly_frommsg(mu, m);
306 for (unsigned x = 0; x < MK_N; x++)
307 v[x] = barrett_reduce((int16_t)(v[x] + e2[x] + mu[x]));
308 poly_compress4(ct + MK_K * 320, v);
309}
310
311bool mlkem768_encaps(const uint8_t ek[MLKEM768_EK_BYTES], const uint8_t m[MLKEM768_MSG_BYTES],
312 uint8_t ct[MLKEM768_CT_BYTES], uint8_t ss[MLKEM768_SS_BYTES])
313{
314 if (!check_ek(ek))
315 return false;
316
317 // (K, r) = G(m || H(ek)); ss = K.
318 uint8_t g_in[64];
319 memcpy(g_in, m, 32);
320 sha3_256(g_in + 32, ek, MLKEM768_EK_BYTES); // H(ek)
321 uint8_t g_out[64];
322 sha3_512(g_out, g_in, sizeof(g_in));
323 memcpy(ss, g_out, 32);
324
325 k_pke_encrypt(ct, ek, m, g_out + 32);
326 return true;
327}
328
329#endif // DETWS_ENABLE_PQC_KEX
mlkem.h
ML-KEM-768 encapsulation (FIPS 203), responder side only.
sha3.h
Keccak-f[1600] sponge: SHA3-256, SHA3-512, SHAKE128, SHAKE256 (FIPS 202).