Niels Möller nisse@lysator.liu.se writes:
Hi,
I've now merge Daiki's ML-KEM implementation, see https://git.lysator.liu.se/nettle/nettle/-/merge_requests/67.
I think there's some further changes I think I'd like to do before release:
I've started to look into this, since I want to have an API i'm happy with prior to release.
- Add more randomized tests, and add assert_maybe for some of the invariants, in particular for arithmetics.
Partly done, at least some more asserts.
- Probably remove the ml_kem_params from the api, and instead use separate functions for ML-KEM 768 and ML-KEM 1024.
I've added specific functions in the most simple way on the branch refactor-ml-kem, but not yet un-exported the ml_kem params things.
- Add functions exposing structs for expanded keys. To avoid having to expand them over and over again if using the same key repeatedly. And for all-in-one functions, these structs could also be used as the types for needed scratch space,
When looking at this, I'm also looking at use of scratch space generally. Some notes...
For the lower-level primitives, the public key is a matrix A and a vector t, and the private key is a vector s. The "scalars" (i.e., an individual element of a matrix or vector) are polynomials mod (Q, X^256+1), represented as 256 uint16_t values.
So for ml-kem-768, the dimension is 3, and we have
A: 3x3x256 uint16_t values (generated from small public seed) t: 3x256 values s: 3x256 values
If we start with key generation, these values are currently all allocated from the provided scratch space (and only encoded versions of s and t are returned to caller). It also uses scratch space for a temporary noise vector e, for a total of 4608 values or 9 KiB. And there's probably another KiB or two allocated on the stack.
It would make sense to me to return A, t and s (the expanded public and private keys), and let functions above deal with encoding. And not encode and decode repeatedly as the same key is used multiple times. The spec hints that at least A can be cached in this way.
Another observation is that need for scratch space can be generally reduced by not computing and storing values long before they are needed. The A matrix is used for one matrix x vector multiplication, A * s. If s is generated upfront, it's sufficient to compute one row of A at a time or even a single scalar element at a time). Similarly, noice vector e could be generated and applied one element at a time.
If we do A row-wise, scratch space is then reduced to 2560 values or 5 KiB. If we do A element-wise, we would save one more KiB. (And relative saving would be larger for ml-kem-1024, where dimension is 4 rather than 3). Maybe still a bit large to unconditionally allocate on the stack, but if we store A, t and s in output areas, I think remaining temporary storage could go in the stack and itch/scratch could be eliminated.
Encapsulation uses a slightly different matrix x vector multiplication, A^T * y, so here it would be natural to generate A one column at a time.
On the other hand, top-level decapsulation uses both low-level decapsulation and encapsulation, so it needs A twice. Here it seems desirable to have all of A allocated in memory (passed in by caller, or generated once from the public seed). So then there's a consistency argument against generating A incrementally in the other functions.
For the encapsulation operation, one could also consider using the output ciphertext area for temporary storage, but it's relatively small so maybe not worth it.
Regards, /Niels