/*
* Implementation of ML-KEM, previously known as 'Crystals: Kyber'.
*/
#include <stdio.h>
#include <stdarg.h>
#include <stdlib.h>
#include <assert.h>
#include "putty.h"
#include "ssh.h"
#include "mlkem.h"
#include "smallmoduli.h"
/* ----------------------------------------------------------------------
* General definitions.
*/
/*
* Arithmetic in this system works mod 3329, which is prime, and
* congruent to 1 mod 256 (in fact it's 13*256 + 1), meaning that
* 256th roots of unity exist.
*/
#define Q 3329
/*
* Parameter structure describing a particular instance of ML-KEM.
*/
struct mlkem_params {
int k; /* dimensions of the matrices used */
int eta_1, eta_2; /* parameters for mlkem_matrix_poly_cbd calls */
int d_u, d_v; /* bit counts to use in lossy compressed encoding */
};
/*
* Specific parameter sets.
*/
const mlkem_params mlkem_params_512 = {
.k = 2, .eta_1 = 3, .eta_2 = 2, .d_u = 10, .d_v = 4,
};
const mlkem_params mlkem_params_768 = {
.k = 3, .eta_1 = 2, .eta_2 = 2, .d_u = 10, .d_v = 4,
};
const mlkem_params mlkem_params_1024 = {
.k = 4, .eta_1 = 2, .eta_2 = 2, .d_u = 11, .d_v = 5,
};
#define KMAX 4
/* ----------------------------------------------------------------------
* Number-theoretic transform on ring elements.
*
* The ring R used by ML-KEM is (Z/qZ)[X] / <X^256+1> (where q=3329 as
* above). If the quotient polynomial were X^256-1 then it would split
* into 256 linear factors, so that R could be expressed as the direct
* sum of 256 rings (Z/qZ)[X] / <X-zeta^i> (where zeta is some fixed
* primitive 256th root of unity mod q), each isomorphic to Z/qZ
* itself. But X^256+1 only splits into 128 _quadratic_ factors, and
* hence we can only decompose R as the direct sum of rings of the
* form (Z/qZ)[X] / <X^2-zeta^j> for odd j, each a quadratic extension
* of Z/qZ, and all mutually nonisomorphic. This means the NTT runs
* one pass fewer than you'd "normally" expect, and also, multiplying
* two elements of R in their NTT representation is not quite as
* trivial as it would normally be - within each component ring of the
* direct sum you have to do the multiplication slightly differently
* depending on the power of zeta in its quotient polynomial.
*
* We take zeta=17 to be the canonical primitive 256th root of unity
* for NTT purposes.
*/
/*
* First 128 powers of zeta, reordered by bit-reversing the 7-bit
* index. That is, the nth element of this array contains
* zeta^(bitrev7(n)). Used by the NTT itself.
*/
static const uint16_t powers_reversed_order[128] = {
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,
};
/*
* First 128 _odd_ powers of zeta: the nth element is
* zeta^(2*bitrev7(n)+1). Each of these is used for multiplication in
* one of the 128 quadratic-extension rings in the NTT decomposition.
*/
static const uint16_t powers_odd_reversed_order[128] = {
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,
};
/*
* Convert a ring element into NTT representation.
*
* The input v is an array of 256 uint16_t, giving the coefficients of
* a polynomial in X, with v[i] being the coefficient of X^i.
*
* v is modified in place. On output, adjacent pairs of elements of v
* give the coefficients of a smaller polynomial in X, with the pair
* v[2i],v[2i+1] being the coefficients of X^0 and X^1 respectively in
* the ring (Z/qZ)[X] / <X^2 - k>, where k = powers_odd_reversed_order[i].
*/
static void mlkem_ntt(uint16_t *v)
{
const uint64_t Qrecip = reciprocal_for_reduction(Q);
size_t next_power = 1;
for (size_t len = 128; len >= 2; len /= 2) {
for (size_t start = 0; start < 256; start += 2*len) {
uint16_t mult = powers_reversed_order[next_power++];
for (size_t j = start; j < start + len; j++) {
uint16_t t = reduce(mult * v[j + len], Q, Qrecip);
v[j + len] = reduce(v[j] + Q - t, Q, Qrecip);
v[j] = reduce(v[j] + t, Q, Qrecip);
}
}
}
}
/*
* Convert back from NTT representation. Exactly inverts mlkem_ntt().
*/
static void mlkem_inverse_ntt(uint16_t *v)
{
const uint64_t Qrecip = reciprocal_for_reduction(Q);
size_t next_power = 127;
for (size_t len = 2; len <= 128; len *= 2) {
for (size_t start = 0; start < 256; start += 2*len) {
uint16_t mult = powers_reversed_order[next_power--];
for (size_t j = start; j < start + len; j++) {
uint16_t t = v[j];
v[j] = reduce(t + v[j + len], Q, Qrecip);
v[j + len] = reduce(mult * (v[j + len] + Q - t), Q, Qrecip);
}
}
}
for (size_t i = 0; i < 256; i++)
v[i] = reduce(v[i] * 3303, Q, Qrecip);
}
/*
* Multiply two elements of R in NTT representation.
*
* The output can alias an input completely, but mustn't alias one
* partially.
*/
static void mlkem_multiply_ntts(
uint16_t *out, const uint16_t *a, const uint16_t *b)
{
const uint64_t Qrecip = reciprocal_for_reduction(Q);
for (size_t i = 0; i < 128; i++) {
uint16_t a0 = a[2*i], a1 = a[2*i+1];
uint16_t b0 = b[2*i], b1 = b[2*i+1];
uint16_t mult = powers_odd_reversed_order[i];
uint16_t a1b1 = reduce(a1 * b1, Q, Qrecip);
out[2*i] = reduce(a0 * b0 + a1b1 * mult, Q, Qrecip);
out[2*i+1] = reduce(a0 * b1 + a1 * b0, Q, Qrecip);
}
}
/* ----------------------------------------------------------------------
* Operations on matrices over the ring R.
*
* Most of these don't mind whether the matrix contains ring elements
* represented directly as polynomials, or in NTT form. The exception
* is that mlkem_matrix_mul requires it to be in NTT form (because
* multiplying is a huge pain in the ordinary representation).
*/
typedef struct mlkem_matrix mlkem_matrix;
struct mlkem_matrix {
unsigned nrows, ncols;
/*
* (nrows * ncols * 256) 16-bit integers. Each 256-word block
* contains an element of R; the blocks are in in row-major order,
* so that (data + 256*(ncols*y + x)) points at the start of the
* element in row y column x.
*/
uint16_t *data;
};
/* Storage used for multiple matrices, to free all at once afterwards */
typedef struct mlkem_matrix_storage mlkem_matrix_storage;
struct mlkem_matrix_storage {
uint16_t *data;
size_t n; /* number of ring elements */
};
/*
* Allocate space for multiple matrices. All the arrays of uint16_t
* are allocated as a single big array. This makes it easy to free the
* whole lot in one go afterwards.
*
* It also means that the arrays have a fixed memory relationship to
* each other, which matters not at all during live use, but
* eliminates spurious control-flow divergences in testsc based on
* accidents of memory allocation when vectorised code checks two
* memory regions to see if they alias. (The compiler-generated
* aliasing check must do two comparisons, one for each direction, and
* the order of those two regions in memory affects whether the first
* comparison decides the second one is necessary.)
*
* The variadic arguments for this function consist of a sequence of
* triples (mlkem_matrix *m, int nrows, int ncols), terminated by a
* null matrix pointer.
*/
static void mlkem_matrix_alloc(mlkem_matrix_storage *storage, ...)
{
va_list ap;
mlkem_matrix *m;
storage->n = 0;
va_start(ap, storage);
while ((m = va_arg(ap, mlkem_matrix *)) != NULL) {
int nrows = va_arg(ap, int), ncols = va_arg(ap, int);
storage->n += nrows * ncols;
}
va_end(ap);
storage->data = snewn(256 * storage->n, uint16_t);
size_t pos = 0;
va_start(ap, storage);
while ((m = va_arg(ap, mlkem_matrix *)) != NULL) {
int nrows = va_arg(ap, int), ncols = va_arg(ap, int);
m->nrows = nrows;
m->ncols = ncols;
m->data = storage->data + 256 * pos;
pos += nrows * ncols;
}
va_end(ap);
}
/* Clear and free the storage allocated by mlkem_matrix_alloc. */
static void mlkem_matrix_storage_free(mlkem_matrix_storage *storage)
{
smemclr(storage->data, 256 * storage->n * sizeof(uint16_t));
sfree(storage->data);
}
/* Add two matrices. */
static void mlkem_matrix_add(mlkem_matrix *out, const mlkem_matrix *left,
const mlkem_matrix *right)
{
const uint64_t Qrecip = reciprocal_for_reduction(Q);
assert(out->nrows == left->nrows);
assert(out->ncols == left->ncols);
assert(out->nrows == right->nrows);
assert(out->ncols == right->ncols);
for (size_t i = 0; i < out->nrows; i++) {
for (size_t j = 0; j < out->ncols; j++) {
const uint16_t *lv = left->data + 256*(i * left->ncols + j);
const uint16_t *rv = right->data + 256*(i * right->ncols + j);
uint16_t *ov = out->data + 256*(i * out->ncols + j);
for (size_t p = 0; p < 256; p++)
ov[p] = reduce(lv[p] + rv[p] , Q, Qrecip);
}
}
}
/* Subtract matrices. */
static void mlkem_matrix_sub(mlkem_matrix *out, const mlkem_matrix *left,
const mlkem_matrix *right)
{
const uint64_t Qrecip = reciprocal_for_reduction(Q);
assert(out->nrows == left->nrows);
assert(out->ncols == left->ncols);
assert(out->nrows == right->nrows);
assert(out->ncols == right->ncols);
for (size_t i = 0; i < out->nrows; i++) {
for (size_t j = 0; j < out->ncols; j++) {
const uint16_t *lv = left->data + 256*(i * left->ncols + j);
const uint16_t *rv = right->data + 256*(i * right->ncols + j);
uint16_t *ov = out->data + 256*(i * out->ncols + j);
for (size_t p = 0; p < 256; p++)
ov[p] = reduce(lv[p] + Q - rv[p] , Q, Qrecip);
}
}
}
/* Convert every element of a matrix into NTT representation. */
static void mlkem_matrix_ntt(mlkem_matrix *m)
{
for (size_t i = 0; i < m->nrows * m->ncols; i++)
mlkem_ntt(m->data + i * 256);
}
/* Convert every element of a matrix out of NTT representation. */
static void mlkem_matrix_inverse_ntt(mlkem_matrix *m)
{
for (size_t i = 0; i < m->nrows * m->ncols; i++)
mlkem_inverse_ntt(m->data + i * 256);
}
/*
* Multiply two matrices, assuming their elements to be currently in
* NTT representation.
*
* The left input must have the same number of columns as the right
* has rows, in the usual fashion. The output matrix is overwritten.
*
* If 'left_transposed' is true then the left matrix is used as if
* transposed.
*/
static void mlkem_matrix_mul(mlkem_matrix *out, const mlkem_matrix *left,
const mlkem_matrix *right, bool left_transposed)
{
const uint64_t Qrecip = reciprocal_for_reduction(Q);
size_t left_nrows = (left_transposed ? left->ncols : left->nrows);
size_t left_ncols = (left_transposed ? left->nrows : left->ncols);
assert(out->nrows == left_nrows);
assert(left_ncols == right->nrows);
assert(right->ncols == out->ncols);
uint16_t work[256];
for (size_t i = 0; i < out->nrows; i++) {
for (size_t j = 0; j < out->ncols; j++) {
uint16_t *thisout = out->data + 256 * (i * out->ncols + j);
memset(thisout, 0, 256 * sizeof(uint16_t));
for (size_t k = 0; k < right->nrows; k++) {
size_t left_index = left_transposed ?
k * left->ncols + i : i * left->ncols + k;
const uint16_t *lv = left->data + 256*left_index;
const uint16_t *rv = right->data + 256*(k * right->ncols + j);
mlkem_multiply_ntts(work, lv, rv);
for (size_t p = 0; p < 256; p++)
thisout[p] = reduce(thisout[p] + work[p], Q, Qrecip);
}
}
}
smemclr(work, sizeof(work));
}
/* ----------------------------------------------------------------------
* Random sampling functions to make up various kinds of randomised
* matrix and vector.
*/
static void mlkem_sample_ntt(uint16_t *output, ptrlen seed); /* forward ref */
/*
* Invent a matrix based on a 32-bit random seed rho.
*
* This matrix is logically part of the public (encryption) key: it's
* not transmitted explicitly, but the seed is, so that the receiver
* can reconstruct the same matrix. As a result, this function
* _doesn't_ have to worry about side channel resistance, or even
* leaving data lying around in arrays.
*/
static void mlkem_matrix_from_seed(mlkem_matrix *m, const void *rho)
{
for (unsigned r = 0; r < m->nrows; r++) {
for (unsigned c = 0; c < m->ncols; c++) {
unsigned char seedbuf[34];
memcpy(seedbuf, rho, 32);
seedbuf[32] = c;
seedbuf[33] = r;
mlkem_sample_ntt(m->data + 256 * (r * m->nrows + c),
make_ptrlen(seedbuf, sizeof(seedbuf)));
}
}
}
/*
* Invent a single element of the ring R, uniformly at random, derived
* in a specified way from the input random seed.
*
* Used as a subroutine of mlkem_matrix_from_seed() above. So, for the
* same reasons, this doesn't have to worry about side channels,
* making the 'rejection sampling' generation technique easy.
*
* The name SampleNTT (in the official spec) reflects the fact that
* the output elements are regarded as being in NTT representation.
* But since the NTT is a bijection, and the sampling is from the
* uniform probability distribution over R, nothing in this function
* actually needs to worry about that.
*/
static void mlkem_sample_ntt(uint16_t *output, ptrlen seed)
{
ShakeXOF *sx = shake128_xof_from_input(seed);
unsigned char bytebuf[4];
bytebuf[3] = '\0';
for (size_t pos = 0; pos < 256 ;) {
/* Read 3 bytes into the low-order end of bytebuf. The fourth
* byte is always 0, so this gives us a random 24-bit integer. */
shake_xof_read(sx, &bytebuf, 3);
uint32_t random24 = GET_32BIT_LSB_FIRST(bytebuf);
/*
* Split that integer up into two 12-bit ones, and use each
* one if it's in range (taking care for the second one that
* we didn't just reach the end of the buffer).
*
* This function is only used for generating matrices from an
* element of the public key, so we can use data-dependent
* control flow here without worrying about giving away
* secrets.
*/
uint16_t d1 = random24 & 0xFFF;
uint16_t d2 = random24 >> 12;
if (d1 < Q)
output[pos++] = d1;
if (d2 < Q && pos < 256)
output[pos++] = d2;
}
shake_xof_free(sx);
}
/*
* Invent a random vector, with its elements _not_ in NTT
* representation, and all the coefficients very small integers (a lot
* smaller than q) of one sign or the other.
*
* eta is a parameter of the probability distribution, sigma is an
* input 32-byte random seed. Each element of the vector is made by a
* separate hash operation based on sigma plus a distinguishing
* integer suffix; 'offset' indicates the starting point for those
* suffixes, so that the ith output value has suffix (offset+i).
*/
static void mlkem_matrix_poly_cbd(
mlkem_matrix *v, int eta, const void *sigma, int offset)
{
const uint64_t Qrecip = reciprocal_for_reduction(Q);
unsigned char seedbuf[33];
memcpy(seedbuf, sigma, 32);
unsigned char *randombuf = snewn(eta * 64, unsigned char);
for (unsigned r = 0; r < v->nrows * v->ncols; r++) {
seedbuf[32] = r + offset;
ShakeXOF *sx = shake256_xof_from_input(make_ptrlen(seedbuf, 33));
shake_xof_read(sx, randombuf, eta * 64);
shake_xof_free(sx);
for (size_t i = 0; i < 256; i++) {
unsigned x = 0, y = 0;
for (size_t j = 0; j < eta; j++) {
size_t bitpos = 2 * i * eta + j;
x += 1 & ((randombuf[bitpos >> 3]) >> (bitpos & 7));
}
for (size_t j = 0; j < eta; j++) {
size_t bitpos = 2 * i * eta + eta + j;
y += 1 & ((randombuf[bitpos >> 3]) >> (bitpos & 7));
}
v->data[256 * r + i] = reduce(x + Q - y, Q, Qrecip);
}
}
smemclr(seedbuf, sizeof(seedbuf));
smemclr(randombuf, eta * 64);
sfree(randombuf);
}
/* ----------------------------------------------------------------------
* Byte-encoding and decoding functions.
*/
/*
* Losslessly encode one or more elements of the ring R.
*
* Each polynomial coefficient, in the range [0,q), is represented as
* a 12-bit integer. So encoding an entire ring element requires
* (256*12)/8 = 384 bytes, and if that 384-byte string were
* interpreted as a little-endian 3072-bit integer D, then the
* coefficient of X^i could be recovered as (D >> (12*i)) & 0xFFF.
*
* The input is expected to be an array of 256*n uint16_t (often the
* 'data' pointer in an mlkem_matrix). The output is 384*n bytes.
*/
static void mlkem_byte_encode_lossless(
void *outv, const uint16_t *in, size_t n)
{
unsigned char *out = (unsigned char *)outv;
uint32_t buffer = 0, bufbits = 0;
for (size_t i = 0; i < 256*n; i++) {
buffer |= (uint32_t) in[i] << bufbits;
bufbits += 12;
while (bufbits >= 8) {
*out++ = buffer & 0xFF;
buffer >>= 8;
bufbits -= 8;
}
}
}
/*
* Decode a string written by mlkem_byte_encode_lossless.
*
* Each 12-bit value extracted from the input data is checked to make
* sure it's in the range [0,q); if it's out of range, the whole
* function fails and returns false. (But it need not do so in
* constant time, because that's an "abandon the whole connection"
* error, not a "subtly make things not work for the attacker" error.)
*/
static bool mlkem_byte_decode_lossless(
uint16_t *out, const void *inv, size_t n)
{
const unsigned char *in = (const unsigned char *)inv;
uint32_t buffer = 0, bufbits = 0;
for (size_t i = 0; i < 384*n; i++) {
buffer |= (uint32_t) in[i] << bufbits;
bufbits += 8;
while (bufbits >= 12) {
uint16_t value = buffer & 0xFFF;
if (value >= Q)
return false;
*out++ = value;
buffer >>= 12;
bufbits -= 12;
}
}
return true;
}
/*
* Lossily encode one or more elements of R, using d bits for each
* polynomial coefficient, for some d < 12. Each output d-bit value is
* obtained as if by regarding the input coefficient as an integer in
* the range [0,q), multiplying by 2^d/q, and rounding to the nearest
* integer. (Since q is odd, 'round to nearest' can't have a tie.)
*
* This means that a large enough input coefficient can round up to
* 2^d itself. In that situation the output d-bit value is 0.
*/
static void mlkem_byte_encode_compressed(
void *outv, const uint16_t *in, unsigned d, size_t n)
{
const uint64_t Qrecip = reciprocal_for_reduction(2*Q);
unsigned char *out = (unsigned char *)outv;
uint32_t buffer = 0, bufbits = 0;
for (size_t i = 0; i < 256*n; i++) {
uint32_t dividend = ((uint32_t)in[i] << (d+1)) + Q;
uint32_t quotient;
reduce_with_quot(dividend, "ient, 2*Q, Qrecip);
buffer |= (uint32_t) (quotient & ((1 << d) - 1)) << bufbits;
bufbits += d;
while (bufbits >= 8) {
*out++ = buffer & 0xFF;
buffer >>= 8;
bufbits -= 8;
}
}
}
/*
* Decode the lossily encoded output of mlkem_byte_encode_compressed.
*
* Each d-bit chunk of the encoding is converted back into a
* polynomial coefficient as if by multiplying by q/2^d and then
* rounding to nearest. Unlike the rounding in the encode step, this
* _can_ have a tie when an unrounded value is half way between two
* integers. Ties are broken by rounding up (as if the whole rounding
* were performed by the simple rounding method of adding 1/2 and then
* truncating).
*
* Unlike the lossless decode function, this one can't fail input
* validation, because any d-bit value generates some legal
* coefficient.
*/
static void mlkem_byte_decode_compressed(
uint16_t *out, const void *inv, unsigned d, size_t n)
{
const unsigned char *in = (const unsigned char *)inv;
uint32_t buffer = 0, bufbits = 0;
for (size_t i = 0; i < 32*d*n; i++) {
buffer |= (uint32_t) in[i] << bufbits;
bufbits += 8;
while (bufbits >= d) {
uint32_t value = buffer & ((1 << d) - 1);
*out++ = (value * (2*Q) + (1 << d)) >> (d + 1);;
buffer >>= d;
bufbits -= d;
}
}
}
/* ----------------------------------------------------------------------
* The top-level ML-KEM functions.
*/
/*
* Innermost keygen function, exposed for side-channel testing, with
* separate random values rho (public) and sigma (private), so that
* testsc can vary sigma while leaving rho the same.
*/
void mlkem_keygen_rho_sigma(
BinarySink *ek_out, BinarySink *dk_out, const mlkem_params *params,
const void *rho, const void *sigma, const void *z)
{
mlkem_matrix_storage storage[1];
mlkem_matrix a[1], s[1], e[1], t[1];
mlkem_matrix_alloc(storage,
a, params->k, params->k,
s, params->k, 1,
e, params->k, 1,
t, params->k, 1,
(mlkem_matrix *)NULL);
/*
* Make a random k x k matrix A (regarded as in NTT form).
*/
mlkem_matrix_from_seed(a, rho);
/*
* Make two column vectors s and e, with all components having
* small polynomial coefficients, and then convert them _into_ NTT
* form.
*/
mlkem_matrix_poly_cbd(s, params->eta_1, sigma, 0);
mlkem_matrix_poly_cbd(e, params->eta_1, sigma, params->k);
mlkem_matrix_ntt(s);
mlkem_matrix_ntt(e);
/*
* Compute the vector t = As + e.
*/
mlkem_matrix_mul(t, a, s, false);
mlkem_matrix_add(t, t, e);
/*
* The encryption key is the vector t, plus the random seed rho
* from which anyone can reconstruct the matrix A.
*/
unsigned char ek[1568];
mlkem_byte_encode_lossless(ek, t->data, params->k);
memcpy(ek + 384 * params->k, rho, 32);
size_t eklen = 384 * params->k + 32;
put_data(ek_out, ek, eklen);
/*
* The decryption key (for the internal "K-PKE" public-key system)
* is the vector s.
*/
unsigned char dk[1536];
mlkem_byte_encode_lossless(dk, s->data, params->k);
size_t dklen = 384 * params->k;
/*
* The decapsulation key, for the full ML-KEM, consists of
* - the decryption key as above
* - the encryption key
* - an extra hash of the encryption key
* - the random value z used for "implicit rejection", aka
* constructing a useless output value if tampering is
* detected. (I think so an attacker can't tell the difference
* between "I was rumbled" and "I was undetected but my attempt
* didn't generate the right key">)
*/
put_data(dk_out, dk, dklen);
put_data(dk_out, ek, eklen);
ssh_hash *h = ssh_hash_new(&ssh_sha3_256);
put_data(h, ek, eklen);
unsigned char ekhash[32];
ssh_hash_final(h, ekhash);
put_data(dk_out, ekhash, 32);
put_data(dk_out, z, 32);
mlkem_matrix_storage_free(storage);
smemclr(ek, sizeof(ek));
smemclr(ekhash, sizeof(ekhash));
smemclr(dk, sizeof(dk));
}
/*
* Internal keygen function as described in the official spec, taking
* random values d and z and deterministically constructing a key from
* them. The test vectors are expressed in terms of this.
*/
void mlkem_keygen_internal(
BinarySink *ek, BinarySink *dk, const mlkem_params *params,
const void *d, const void *z)
{
/* Hash the input randomness d to make two 32-byte values rho and sigma */
unsigned char rho_sigma[64];
ssh_hash *h = ssh_hash_new(&ssh_sha3_512);
put_data(h, d, 32);
put_byte(h, params->k);
ssh_hash_final(h, rho_sigma);
mlkem_keygen_rho_sigma(ek, dk, params, rho_sigma, rho_sigma + 32, z);
smemclr(rho_sigma, sizeof(rho_sigma));
}
/*
* Keygen function for live use, making up the values at random.
*/
void mlkem_keygen(
BinarySink *ek, BinarySink *dk, const mlkem_params *params)
{
unsigned char dz[64];
random_read(dz, 64);
mlkem_keygen_internal(ek, dk, params, dz, dz + 32);
smemclr(dz, sizeof(dz));
}
/*
* Internal encapsulation function from the official spec, taking a
* random value m as input and behaving deterministically. Again used
* for test vectors.
*/
bool mlkem_encaps_internal(
BinarySink *c_out, BinarySink *k_out,
const mlkem_params *params, ptrlen ek, const void *m)
{
mlkem_matrix_storage storage[1];
mlkem_matrix t[1], a[1], y[1], e1[1], e2[1], mu[1], u[1], v[1];
mlkem_matrix_alloc(storage,
t, params->k, 1,
a, params->k, params->k,
y, params->k, 1,
e1, params->k, 1,
e2, 1, 1,
mu, 1, 1,
u, params->k, 1,
v, 1, 1,
(mlkem_matrix *)NULL);
/*
* Validate input: ek must be the correct length, and its encoded
* ring elements must not include any 16-bit integer intended to
* represent a value mod q which is not in fact in the range [0,q).
*
* We test the latter property by decoding the matrix t, and
* checking the success status returned by the decode.
*/
if (ek.len != 384 * params->k + 32 ||
!mlkem_byte_decode_lossless(t->data, ek.ptr, params->k)) {
mlkem_matrix_storage_free(storage);
return false;
}
/*
* Regenerate the same matrix A used by key generation, from the
* seed string rho at the end of ek.
*/
mlkem_matrix_from_seed(a, (const unsigned char *)ek.ptr + 384 * params->k);
/*
* Hash the input randomness m, to get the value k we'll use as
* the output shared secret, plus some randomness for making up
* the vectors below.
*/
unsigned char kr[64];
unsigned char ekhash[32];
ssh_hash *h;
/* Hash the encryption key */
h = ssh_hash_new(&ssh_sha3_256);
put_datapl(h, ek);
ssh_hash_final(h, ekhash);
/* Hash the input randomness m with that hash */
h = ssh_hash_new(&ssh_sha3_512);
put_data(h, m, 32);
put_data(h, ekhash, 32);
ssh_hash_final(h, kr);
const unsigned char *k = kr, *r = kr + 32;
/*
* Invent random k-element vectors y and e1, and a random scalar
* e2 (here represented as a 1x1 matrix for the sake of not
* proliferating internal helper functions). All are generated by
* poly_cbd (i.e. their ring elements have polynomial coefficients
* of small magnitude). y needs to be in NTT form.
*
* These generations all use r as their seed, which was the second
* half of the 64-byte hash of the input m. We pass different
* 'offset' values to mlkem_matrix_poly_cbd() to ensure the
* generations are probabilistically independent.
*/
mlkem_matrix_poly_cbd(y, params->eta_1, r, 0);
mlkem_matrix_ntt(y);
mlkem_matrix_poly_cbd(e1, params->eta_2, r, params->k);
mlkem_matrix_poly_cbd(e2, params->eta_2, r, 2 * params->k);
/*
* Invent a random scalar mu (again imagined as a 1x1 matrix),
* this time by doing lossy decompression of the random value m at
* 1 bit per polynomial coefficient. That is, all the polynomial
* coefficients of mu are either 0 or 1665 = (q+1)/2.
*
* This generation reuses the _input_ random value m, not either
* half of the hash we made of it.
*/
mlkem_byte_decode_compressed(mu->data, m, 1, 1);
/*
* Calculate a k-element vector u = A^T y + e1.
*
* A and y are in NTT representation, but e1 is not, and we don't
* want the output to be in NTT form either. So we perform an
* inverse NTT after the multiplication.
*/
mlkem_matrix_mul(u, a, y, true); /* regard a as transposed */
mlkem_matrix_inverse_ntt(u);
mlkem_matrix_add(u, u, e1);
/*
* Calculate a scalar v = t^T y + e2 + mu.
*
* (t and y are column vectors, so t^T y is just a scalar - you
* could think of it as the dot product t.y if you preferred.)
*
* Similarly to above, we multiply t and y which are in NTT
* representation, and then perform an inverse NTT before adding
* e2 and mu, which aren't.
*/
mlkem_matrix_mul(v, t, y, true); /* regard t as transposed */
mlkem_matrix_inverse_ntt(v);
mlkem_matrix_add(v, v, e2);
mlkem_matrix_add(v, v, mu);
/*
* The ciphertext consists of u and v, both encoded lossily, with
* different numbers of bits retained per element.
*/
char c[1568];
mlkem_byte_encode_compressed(c, u->data, params->d_u, params->k);
mlkem_byte_encode_compressed(c + 32 * params->k * params->d_u,
v->data, params->d_v, 1);
put_data(c_out, c, 32 * (params->k * params->d_u + params->d_v));
/*
* The output shared secret is just half of the hash of m (the
* first half, which we didn't use for generating vectors above).
*/
put_data(k_out, k, 32);
smemclr(kr, sizeof(kr));
mlkem_matrix_storage_free(storage);
return true;
}
/*
* Encapsulation function for live use, using the real RNG..
*/
bool mlkem_encaps(BinarySink *ciphertext, BinarySink *kout,
const mlkem_params *params, ptrlen ek)
{
unsigned char m[32];
random_read(m, 32);
bool success = mlkem_encaps_internal(ciphertext, kout, params, ek, m);
smemclr(m, sizeof(m));
return success;
}
/*
* Decapsulation.
*/
bool mlkem_decaps(BinarySink *k_out, const mlkem_params *params,
ptrlen dk, ptrlen c)
{
/*
* Validation: check the input strings are the right lengths.
*/
if (dk.len != 768 * params->k + 96)
return false;
if (c.len != 32 * (params->d_u * params->k + params->d_v))
return false;
/*
* Further validation: extract the encryption key from the middle
* of dk, hash it, and check the hash matches.
*/
const unsigned char *dkp = (const unsigned char *)dk.ptr;
const unsigned char *cp = (const unsigned char *)c.ptr;
ptrlen ek = make_ptrlen(dkp + 384*params->k, 384*params->k + 32);
ssh_hash *h;
unsigned char ekhash[32];
h = ssh_hash_new(&ssh_sha3_256);
put_datapl(h, ek);
ssh_hash_final(h, ekhash);
if (!smemeq(ekhash, dkp + 768*params->k + 32, 32))
return false;
mlkem_matrix_storage storage[1];
mlkem_matrix u[1], v[1], s[1], w[1];
mlkem_matrix_alloc(storage,
u, params->k, 1,
v, 1, 1,
s, params->k, 1,
w, 1, 1,
(mlkem_matrix *)NULL);
/*
* Decode the vector u and the scalar v from the ciphertext. These
* won't come out exactly the same as the originals, because of
* the lossy compression.
*/
mlkem_byte_decode_compressed(u->data, cp, params->d_u, params->k);
mlkem_matrix_ntt(u);
mlkem_byte_decode_compressed(v->data, cp + 32 * params->d_u * params->k,
params->d_v, 1);
/*
* Decode the vector s from the private key.
*/
mlkem_byte_decode_lossless(s->data, dkp, params->k);
/*
* Calculate the scalar w = v - s^T u.
*
* s and u are in NTT representation, but v isn't, so we
* inverse-NTT the product before doing the subtraction. Therefore
* w is not in NTT form either.
*/
mlkem_matrix_mul(w, s, u, true); /* regard s as transposed */
mlkem_matrix_inverse_ntt(w);
mlkem_matrix_sub(w, v, w);
/*
* The aim is that this reconstructs something close enough to the
* random vector mu that was made from the input secret m to
* encapsulation, on the grounds that mu's polynomial coefficients
* were very widely separated (on opposite sides of the cyclic
* additive group of Z/qZ) and the noise added during encryption
* all had _small_ polynomial coefficients.
*
* So we now re-encode this lossily at 1 bit per polynomial
* coefficient, and hope that it reconstructs the actual string m.
*
* However, this _is_ only a hope! The ML-KEM decryption is not a
* true mathematical inverse to encryption. With extreme bad luck,
* the noise can add up enough that it flips a bit of m, and
* everything fails. The parameters are chosen to make this happen
* with negligible probability (the same kind of low probability
* that makes you not worry about spontaneous hash collisions),
* but it's not actually impossible.
*/
unsigned char m[32];
mlkem_byte_encode_compressed(m, w->data, 1, 1);
/*
* Now do the key _encapsulation_ again from scratch, using that
* secret m as input, and check that it generates the identical
* ciphertext. This should catch the above theoretical failure,
* but also, it's a defence against malicious intervention in the
* key exchange.
*
* This is also where we get the output secret k from: the
* encapsulation function creates it as half of the hash of m.
*/
unsigned char c_regen[1568], k[32];
buffer_sink c_sink[1], k_sink[1];
buffer_sink_init(c_sink, c_regen, sizeof(c_regen));
buffer_sink_init(k_sink, k, sizeof(k));
bool success = mlkem_encaps_internal(
BinarySink_UPCAST(c_sink), BinarySink_UPCAST(k_sink), params, ek, m);
/* If any application of ML-KEM uses a dk given to it by someone
* else, then perhaps they have to worry about being given an
* invalid one? But in our application we always expect this to
* succeed, because dk is generated and used at the same end of
* the SSH connection, within the same process, and nobody is
* interfering with it. */
assert(success && "We generated this dk ourselves, how can it be bad?");
/*
* If mlkem_encaps_internal returned success but delivered the
* wrong ciphertext, that's a failure, but we must be careful not
* to let the attacker know exactly what went wrong. So we
* generate a plausible but wrong substitute output secret.
*
* k_reject is that secret; for constant-time reasons we generate
* it unconditionally.
*/
unsigned char k_reject[32];
h = ssh_hash_new(&ssh_shake256_32bytes);
put_data(h, dkp + 768 * params->k + 64, 32);
put_datapl(h, c);
ssh_hash_final(h, k_reject);
/*
* Now replace k with k_reject if the ciphertexts didn't match.
*/
assert((void *)c_sink->out == (void *)(c_regen + c.len));
unsigned match = smemeq(c.ptr, c_regen, c.len);
unsigned mask = match - 1;
for (size_t i = 0; i < 32; i++)
k[i] ^= mask & (k[i] ^ k_reject[i]);
/*
* And we're done! Free everything and return whichever secret we
* chose.
*/
put_data(k_out, k, 32);
mlkem_matrix_storage_free(storage);
smemclr(m, sizeof(m));
smemclr(c_regen, sizeof(c_regen));
smemclr(k, sizeof(k));
smemclr(k_reject, sizeof(k_reject));
return true;
}
/* ----------------------------------------------------------------------
* Implement the pq_kemalg vtable in terms of the above functions.
*/
struct mlkem_dk {
strbuf *encoded;
pq_kem_dk dk;
};
static pq_kem_dk *mlkem_vt_keygen(const pq_kemalg *alg, BinarySink *ek)
{
struct mlkem_dk *mdk = snew(struct mlkem_dk);
mdk->dk.vt = alg;
mdk->encoded = strbuf_new_nm();
mlkem_keygen(ek, BinarySink_UPCAST(mdk->encoded), alg->extra);
return &mdk->dk;
}
static bool mlkem_vt_encaps(const pq_kemalg *alg, BinarySink *c, BinarySink *k,
ptrlen ek)
{
return mlkem_encaps(c, k, alg->extra, ek);
}
static bool mlkem_vt_decaps(pq_kem_dk *dk, BinarySink *k, ptrlen c)
{
struct mlkem_dk *mdk = container_of(dk, struct mlkem_dk, dk);
return mlkem_decaps(k, mdk->dk.vt->extra,
ptrlen_from_strbuf(mdk->encoded), c);
}
static void mlkem_vt_free_dk(pq_kem_dk *dk)
{
struct mlkem_dk *mdk = container_of(dk, struct mlkem_dk, dk);
strbuf_free(mdk->encoded);
sfree(mdk);
}
const pq_kemalg ssh_mlkem512 = {
.keygen = mlkem_vt_keygen,
.encaps = mlkem_vt_encaps,
.decaps = mlkem_vt_decaps,
.free_dk = mlkem_vt_free_dk,
.extra = &mlkem_params_512,
.description = "ML-KEM-512",
.ek_len = 384 * 2 + 32,
.c_len = 32 * (10 * 2 + 4),
};
const pq_kemalg ssh_mlkem768 = {
.keygen = mlkem_vt_keygen,
.encaps = mlkem_vt_encaps,
.decaps = mlkem_vt_decaps,
.free_dk = mlkem_vt_free_dk,
.extra = &mlkem_params_768,
.description = "ML-KEM-768",
.ek_len = 384 * 3 + 32,
.c_len = 32 * (10 * 3 + 4),
};
const pq_kemalg ssh_mlkem1024 = {
.keygen = mlkem_vt_keygen,
.encaps = mlkem_vt_encaps,
.decaps = mlkem_vt_decaps,
.free_dk = mlkem_vt_free_dk,
.extra = &mlkem_params_1024,
.description = "ML-KEM-1024",
.ek_len = 384 * 4 + 32,
.c_len = 32 * (11 * 4 + 5),
};