1
2
3
4
5 package mlkem
6
7 import (
8 "crypto/internal/fips140/sha3"
9 "crypto/internal/fips140deps/byteorder"
10 "errors"
11 )
12
13
14 type fieldElement uint16
15
16
17 func fieldCheckReduced(a uint16) (fieldElement, error) {
18 if a >= q {
19 return 0, errors.New("unreduced field element")
20 }
21 return fieldElement(a), nil
22 }
23
24
25 func fieldReduceOnce(a uint16) fieldElement {
26 x := a - q
27
28 x += (x >> 15) * q
29 return fieldElement(x)
30 }
31
32 func fieldAdd(a, b fieldElement) fieldElement {
33 x := uint16(a + b)
34 return fieldReduceOnce(x)
35 }
36
37 func fieldSub(a, b fieldElement) fieldElement {
38 x := uint16(a - b + q)
39 return fieldReduceOnce(x)
40 }
41
42 const (
43 barrettMultiplier = 5039
44 barrettShift = 24
45 )
46
47
48
49 func fieldReduce(a uint32) fieldElement {
50 quotient := uint32((uint64(a) * barrettMultiplier) >> barrettShift)
51 return fieldReduceOnce(uint16(a - quotient*q))
52 }
53
54 func fieldMul(a, b fieldElement) fieldElement {
55 x := uint32(a) * uint32(b)
56 return fieldReduce(x)
57 }
58
59
60
61 func fieldMulSub(a, b, c fieldElement) fieldElement {
62 x := uint32(a) * uint32(b-c+q)
63 return fieldReduce(x)
64 }
65
66
67
68 func fieldAddMul(a, b, c, d fieldElement) fieldElement {
69 x := uint32(a) * uint32(b)
70 x += uint32(c) * uint32(d)
71 return fieldReduce(x)
72 }
73
74
75
76 func compress(x fieldElement, d uint8) uint16 {
77
78
79
80
81
82 dividend := uint32(x) << d
83 quotient := uint32(uint64(dividend) * barrettMultiplier >> barrettShift)
84 remainder := dividend - quotient*q
85
86
87
88
89
90
91
92
93
94
95
96
97
98 quotient += (q/2 - remainder) >> 31 & 1
99 quotient += (q + q/2 - remainder) >> 31 & 1
100
101
102 var mask uint32 = (1 << d) - 1
103 return uint16(quotient & mask)
104 }
105
106
107
108 func decompress(y uint16, d uint8) fieldElement {
109
110
111
112 dividend := uint32(y) * q
113 quotient := dividend >> d
114
115
116
117
118 quotient += dividend >> (d - 1) & 1
119
120
121 return fieldElement(quotient)
122 }
123
124
125
126 type ringElement [n]fieldElement
127
128
129 func polyAdd[T ~[n]fieldElement](a, b T) (s T) {
130 for i := range s {
131 s[i] = fieldAdd(a[i], b[i])
132 }
133 return s
134 }
135
136
137 func polySub[T ~[n]fieldElement](a, b T) (s T) {
138 for i := range s {
139 s[i] = fieldSub(a[i], b[i])
140 }
141 return s
142 }
143
144
145
146
147 func polyByteEncode[T ~[n]fieldElement](b []byte, f T) []byte {
148 out, B := sliceForAppend(b, encodingSize12)
149 for i := 0; i < n; i += 2 {
150 x := uint32(f[i]) | uint32(f[i+1])<<12
151 B[0] = uint8(x)
152 B[1] = uint8(x >> 8)
153 B[2] = uint8(x >> 16)
154 B = B[3:]
155 }
156 return out
157 }
158
159
160
161
162
163
164 func polyByteDecode[T ~[n]fieldElement](b []byte) (T, error) {
165 if len(b) != encodingSize12 {
166 return T{}, errors.New("mlkem: invalid encoding length")
167 }
168 var f T
169 for i := 0; i < n; i += 2 {
170 d := uint32(b[0]) | uint32(b[1])<<8 | uint32(b[2])<<16
171 const mask12 = 0b1111_1111_1111
172 var err error
173 if f[i], err = fieldCheckReduced(uint16(d & mask12)); err != nil {
174 return T{}, errors.New("mlkem: invalid polynomial encoding")
175 }
176 if f[i+1], err = fieldCheckReduced(uint16(d >> 12)); err != nil {
177 return T{}, errors.New("mlkem: invalid polynomial encoding")
178 }
179 b = b[3:]
180 }
181 return f, nil
182 }
183
184
185
186
187
188 func sliceForAppend(in []byte, n int) (head, tail []byte) {
189 if total := len(in) + n; cap(in) >= total {
190 head = in[:total]
191 } else {
192 head = make([]byte, total)
193 copy(head, in)
194 }
195 tail = head[len(in):]
196 return
197 }
198
199
200
201
202
203
204 func ringCompressAndEncode1(s []byte, f ringElement) []byte {
205 s, b := sliceForAppend(s, encodingSize1)
206 clear(b)
207 for i := range f {
208 b[i/8] |= uint8(compress(f[i], 1) << (i % 8))
209 }
210 return s
211 }
212
213
214
215
216
217
218 func ringDecodeAndDecompress1(b *[encodingSize1]byte) ringElement {
219 var f ringElement
220 for i := range f {
221 b_i := b[i/8] >> (i % 8) & 1
222 const halfQ = (q + 1) / 2
223 f[i] = fieldElement(b_i) * halfQ
224 }
225 return f
226 }
227
228
229
230
231
232
233 func ringCompressAndEncode4(s []byte, f ringElement) []byte {
234 s, b := sliceForAppend(s, encodingSize4)
235 for i := 0; i < n; i += 2 {
236 b[i/2] = uint8(compress(f[i], 4) | compress(f[i+1], 4)<<4)
237 }
238 return s
239 }
240
241
242
243
244
245
246 func ringDecodeAndDecompress4(b *[encodingSize4]byte) ringElement {
247 var f ringElement
248 for i := 0; i < n; i += 2 {
249 f[i] = fieldElement(decompress(uint16(b[i/2]&0b1111), 4))
250 f[i+1] = fieldElement(decompress(uint16(b[i/2]>>4), 4))
251 }
252 return f
253 }
254
255
256
257
258
259
260 func ringCompressAndEncode10(s []byte, f ringElement) []byte {
261 s, b := sliceForAppend(s, encodingSize10)
262 for i := 0; i < n; i += 4 {
263 var x uint64
264 x |= uint64(compress(f[i], 10))
265 x |= uint64(compress(f[i+1], 10)) << 10
266 x |= uint64(compress(f[i+2], 10)) << 20
267 x |= uint64(compress(f[i+3], 10)) << 30
268 b[0] = uint8(x)
269 b[1] = uint8(x >> 8)
270 b[2] = uint8(x >> 16)
271 b[3] = uint8(x >> 24)
272 b[4] = uint8(x >> 32)
273 b = b[5:]
274 }
275 return s
276 }
277
278
279
280
281
282
283 func ringDecodeAndDecompress10(bb *[encodingSize10]byte) ringElement {
284 b := bb[:]
285 var f ringElement
286 for i := 0; i < n; i += 4 {
287 x := uint64(b[0]) | uint64(b[1])<<8 | uint64(b[2])<<16 | uint64(b[3])<<24 | uint64(b[4])<<32
288 b = b[5:]
289 f[i] = fieldElement(decompress(uint16(x>>0&0b11_1111_1111), 10))
290 f[i+1] = fieldElement(decompress(uint16(x>>10&0b11_1111_1111), 10))
291 f[i+2] = fieldElement(decompress(uint16(x>>20&0b11_1111_1111), 10))
292 f[i+3] = fieldElement(decompress(uint16(x>>30&0b11_1111_1111), 10))
293 }
294 return f
295 }
296
297
298
299
300
301
302 func ringCompressAndEncode(s []byte, f ringElement, d uint8) []byte {
303 var b byte
304 var bIdx uint8
305 for i := 0; i < n; i++ {
306 c := compress(f[i], d)
307 var cIdx uint8
308 for cIdx < d {
309 b |= byte(c>>cIdx) << bIdx
310 bits := min(8-bIdx, d-cIdx)
311 bIdx += bits
312 cIdx += bits
313 if bIdx == 8 {
314 s = append(s, b)
315 b = 0
316 bIdx = 0
317 }
318 }
319 }
320 if bIdx != 0 {
321 panic("mlkem: internal error: bitsFilled != 0")
322 }
323 return s
324 }
325
326
327
328
329
330
331 func ringDecodeAndDecompress(b []byte, d uint8) ringElement {
332 var f ringElement
333 var bIdx uint8
334 for i := 0; i < n; i++ {
335 var c uint16
336 var cIdx uint8
337 for cIdx < d {
338 c |= uint16(b[0]>>bIdx) << cIdx
339 c &= (1 << d) - 1
340 bits := min(8-bIdx, d-cIdx)
341 bIdx += bits
342 cIdx += bits
343 if bIdx == 8 {
344 b = b[1:]
345 bIdx = 0
346 }
347 }
348 f[i] = fieldElement(decompress(c, d))
349 }
350 if len(b) != 0 {
351 panic("mlkem: internal error: leftover bytes")
352 }
353 return f
354 }
355
356
357
358
359
360
361 func ringCompressAndEncode5(s []byte, f ringElement) []byte {
362 return ringCompressAndEncode(s, f, 5)
363 }
364
365
366
367
368
369
370 func ringDecodeAndDecompress5(bb *[encodingSize5]byte) ringElement {
371 return ringDecodeAndDecompress(bb[:], 5)
372 }
373
374
375
376
377
378
379 func ringCompressAndEncode11(s []byte, f ringElement) []byte {
380 return ringCompressAndEncode(s, f, 11)
381 }
382
383
384
385
386
387
388 func ringDecodeAndDecompress11(bb *[encodingSize11]byte) ringElement {
389 return ringDecodeAndDecompress(bb[:], 11)
390 }
391
392
393
394
395 func samplePolyCBD(s []byte, b byte) ringElement {
396 prf := sha3.NewShake256()
397 prf.Write(s)
398 prf.Write([]byte{b})
399 B := make([]byte, 64*2)
400 prf.Read(B)
401
402
403
404
405 var f ringElement
406 for i := 0; i < n; i += 2 {
407 b := B[i/2]
408 b_7, b_6, b_5, b_4 := b>>7, b>>6&1, b>>5&1, b>>4&1
409 b_3, b_2, b_1, b_0 := b>>3&1, b>>2&1, b>>1&1, b&1
410 f[i] = fieldSub(fieldElement(b_0+b_1), fieldElement(b_2+b_3))
411 f[i+1] = fieldSub(fieldElement(b_4+b_5), fieldElement(b_6+b_7))
412 }
413 return f
414 }
415
416
417
418 type nttElement [n]fieldElement
419
420
421
422 var gammas = [128]fieldElement{17, 3312, 2761, 568, 583, 2746, 2649, 680, 1637, 1692, 723, 2606, 2288, 1041, 1100, 2229, 1409, 1920, 2662, 667, 3281, 48, 233, 3096, 756, 2573, 2156, 1173, 3015, 314, 3050, 279, 1703, 1626, 1651, 1678, 2789, 540, 1789, 1540, 1847, 1482, 952, 2377, 1461, 1868, 2687, 642, 939, 2390, 2308, 1021, 2437, 892, 2388, 941, 733, 2596, 2337, 992, 268, 3061, 641, 2688, 1584, 1745, 2298, 1031, 2037, 1292, 3220, 109, 375, 2954, 2549, 780, 2090, 1239, 1645, 1684, 1063, 2266, 319, 3010, 2773, 556, 757, 2572, 2099, 1230, 561, 2768, 2466, 863, 2594, 735, 2804, 525, 1092, 2237, 403, 2926, 1026, 2303, 1143, 2186, 2150, 1179, 2775, 554, 886, 2443, 1722, 1607, 1212, 2117, 1874, 1455, 1029, 2300, 2110, 1219, 2935, 394, 885, 2444, 2154, 1175}
423
424
425
426
427 func nttMul(f, g nttElement) nttElement {
428 var h nttElement
429
430 for i := 0; i < 256; i += 2 {
431 a0, a1 := f[i], f[i+1]
432 b0, b1 := g[i], g[i+1]
433 h[i] = fieldAddMul(a0, b0, fieldMul(a1, b1), gammas[i/2])
434 h[i+1] = fieldAddMul(a0, b1, a1, b0)
435 }
436 return h
437 }
438
439
440
441 var zetas = [128]fieldElement{1, 1729, 2580, 3289, 2642, 630, 1897, 848, 1062, 1919, 193, 797, 2786, 3260, 569, 1746, 296, 2447, 1339, 1476, 3046, 56, 2240, 1333, 1426, 2094, 535, 2882, 2393, 2879, 1974, 821, 289, 331, 3253, 1756, 1197, 2304, 2277, 2055, 650, 1977, 2513, 632, 2865, 33, 1320, 1915, 2319, 1435, 807, 452, 1438, 2868, 1534, 2402, 2647, 2617, 1481, 648, 2474, 3110, 1227, 910, 17, 2761, 583, 2649, 1637, 723, 2288, 1100, 1409, 2662, 3281, 233, 756, 2156, 3015, 3050, 1703, 1651, 2789, 1789, 1847, 952, 1461, 2687, 939, 2308, 2437, 2388, 733, 2337, 268, 641, 1584, 2298, 2037, 3220, 375, 2549, 2090, 1645, 1063, 319, 2773, 757, 2099, 561, 2466, 2594, 2804, 1092, 403, 1026, 1143, 2150, 2775, 886, 1722, 1212, 1874, 1029, 2110, 2935, 885, 2154}
442
443
444
445
446 func ntt(f ringElement) nttElement {
447 k := 1
448 for len := 128; len >= 2; len /= 2 {
449 for start := 0; start < 256; start += 2 * len {
450 zeta := zetas[k]
451 k++
452
453 f, flen := f[start:start+len], f[start+len:start+len+len]
454 for j := 0; j < len; j++ {
455 t := fieldMul(zeta, flen[j])
456 flen[j] = fieldSub(f[j], t)
457 f[j] = fieldAdd(f[j], t)
458 }
459 }
460 }
461 return nttElement(f)
462 }
463
464
465
466
467 func inverseNTT(f nttElement) ringElement {
468 k := 127
469 for len := 2; len <= 128; len *= 2 {
470 for start := 0; start < 256; start += 2 * len {
471 zeta := zetas[k]
472 k--
473
474 f, flen := f[start:start+len], f[start+len:start+len+len]
475 for j := 0; j < len; j++ {
476 t := f[j]
477 f[j] = fieldAdd(t, flen[j])
478 flen[j] = fieldMulSub(zeta, flen[j], t)
479 }
480 }
481 }
482 for i := range f {
483 f[i] = fieldMul(f[i], 3303)
484 }
485 return ringElement(f)
486 }
487
488
489
490
491 func sampleNTT(rho []byte, ii, jj byte) nttElement {
492 B := sha3.NewShake128()
493 B.Write(rho)
494 B.Write([]byte{ii, jj})
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520 var a nttElement
521 var j int
522 var buf [24]byte
523 off := len(buf)
524 for {
525 if off >= len(buf) {
526 B.Read(buf[:])
527 off = 0
528 }
529 d1 := byteorder.LEUint16(buf[off:]) & 0b1111_1111_1111
530 d2 := byteorder.LEUint16(buf[off+1:]) >> 4
531 off += 3
532 if d1 < q {
533 a[j] = fieldElement(d1)
534 j++
535 }
536 if j >= len(a) {
537 break
538 }
539 if d2 < q {
540 a[j] = fieldElement(d2)
541 j++
542 }
543 if j >= len(a) {
544 break
545 }
546 }
547 return a
548 }
549
View as plain text