Add an ML-KEM implementation #2

bifurcation · 2024-02-07T20:15:39Z

This PR adds an implementation of ML-KEM variant of the Kyber KEM as described
in the initial public draft of FIPS 203. This implementation covers all three
parameter sets described in the specification by making heavy use of Rust
generics and the typenum crate.

In addition to self-compatibility testing, we test correctness by verifying the
test vectors supplied by NIST; see tests/nist.rs.

This PR also includes a benchmark that tests the speed of the ML-KEM-768
variant, the most commonly implemented variant. At least on my MacBook, this
implementation is competitive with / slightly slower than the fastest known implementations (units
are µs/op; only Go and ml_kem measured locally, others from this blog post):

MacOS M1 Pro	`ml_kem`	libcrux	PQClean	BoringSSL	Circl	Go
Key generation	34.487	28.336	30.387	32.051	39.783	42.900
Encapsulation	36.068	29.222	38.688	28.131	44.874	58.831
Decapsulation	42.704	31.635	43.072	33.777	50.025	70.862

tarcieri · 2024-02-07T20:23:09Z

cc @dsprenkels

bifurcation · 2024-02-08T19:03:01Z

Note that the Go version linked above has a very thorough battery of tests. I think it would be a good idea to have that coverage here as well, but it can be handled as a follow-on PR.

ml-kem/.gitignore

ml-kem/Cargo.toml

tarcieri · 2024-02-13T01:40:13Z

ml-kem/src/lib.rs

+/// A key that can be used to decapsulated an encapsulated shared key
+pub trait DecapsulationKey<SharedKey, Ciphertext>: EncodedSizeUser {
+    /// Decapsulate the ciphertext to obtain the encapsulated shared key
+    fn decapsulate(&self, ciphertext: &Ciphertext) -> SharedKey;
+}
+
+/// A key that can be used to encapsulate a shared key such that it can only be decapsulated
+/// by the holder of the corresponding decapsulation key
+pub trait EncapsulationKey<SharedKey, Ciphertext>: EncodedSizeUser {
+    /// Encapsulate a fresh secret to the holder of the corresponding decapsulation key
+    fn encapsulate(&self, rng: &mut impl CryptoRngCore) -> (SharedKey, Ciphertext);
+
+    /// Encapsulate a specific shared key to the holder of the corresponding decapsulation key
+    #[cfg(feature = "deterministic")]
+    fn encapsulate_deterministic(&self, m: &B32) -> (SharedKey, Ciphertext);
+}
+
+/// A generic interface to a Key Encapsulation Method
+pub trait KemCore {
+    /// The size of a shared key generated by this KEM
+    type SharedKeySize: ArrayLength;
+
+    /// The size of a ciphertext encapsulating a shared key
+    type CiphertextSize: ArrayLength;
+
+    /// A decapsulation key for this KEM
+    type DecapsulationKey: DecapsulationKey<SharedKey<Self>, Ciphertext<Self>> + Debug + PartialEq;
+
+    /// An encapsulation key for this KEM
+    type EncapsulationKey: EncapsulationKey<SharedKey<Self>, Ciphertext<Self>> + Debug + PartialEq;
+
+    /// Generate a new (decapsulation, encapsulation) key pair
+    fn generate(rng: &mut impl CryptoRngCore) -> (Self::DecapsulationKey, Self::EncapsulationKey);
+
+    /// Generate a new (decapsulation, encapsulation) key pair deterministically
+    #[cfg(feature = "deterministic")]
+    fn generate_deterministic(d: &B32, z: &B32)
+        -> (Self::DecapsulationKey, Self::EncapsulationKey);
+}


We do have some existing traits for KEMs here:

https://github.com/RustCrypto/traits/tree/master/kem

https://docs.rs/kem/latest/kem/

cc @rozbb

Yup, though its tires haven't been thoroughly kicked so I'm curious if @bifurcation thinks it's an OK fit.

Also it probably needs to be upgraded to hybrid-array

Oh nice re KEM traits, somehow I had missed that. Obviously should use that, was already thinking that that trait should be outside this crate.

@tarcieri if you concur re hybrid_array, happy to switch to that. (Assuming there are no roadblocks, e.g., due to some of the typenum craziness we're going here.) Might do some benchmarking to see if there are perf impacts.

To be clear: I haven't taken a deep look at the KEM traits yet, will do and report back. But it's correct to have the KEM traits outside this crate.

We're migrating to hybrid-array elsewhere. It should be a drop-in replacement for generic-array. There's a migration guide here (it should mostly be find/replace): https://docs.rs/hybrid-array/0.2.0-rc.5/hybrid_array/#migrating-from-generic-array

I can also look at updating the kem crate to use hybrid-array.

On trying to build tests, I find that I need ArraySize implementations for a bunch more sizes: 480, 800, 1088, 1152, 1184, 1536, 1568, 1600, 1632, 2336, 2368, 2400, 3104, 3136, 3168. So increasing the base set size up to 4096 would cover me. Note, however, that typenum doesn't actually define U$n for all of these $n, so you might have different problems.

Given this dependency, we might want to consider migrating to hybrid-array as a follow-on to this PR.

Yeah, that gets a bit tricky and I see typenum::Const<N> doesn't cover that range either. It will require encoding the bit pattern for those sizes, and in such a case I could probably consider e.g. multiples of 32 in the range 1024-4096 which should cover your needs at least.

Given this dependency, we might want to consider migrating to hybrid-array as a follow-on to this PR.

Yes, agreed. Sorry for the unnecessary detour. I just started trying to migrate the crate myself and hit the same large-const issues for Saber.

No worries @rozbb! It does look a bit nicer than generic-array, so I like it as a long-term direction. It's unfortunate that the state of const generics is such that it requires all this hacking around.

This should at least be a stopgap to support the sizes you need with hybrid-array: https://github.com/RustCrypto/hybrid-array/pull/49/files

I'll see if I can come up with a better/consistent/more automated way of implementing these additional sizes.

bifurcation · 2024-02-16T13:56:32Z

Yup, though its tires haven't been thoroughly kicked so I'm curious if @bifurcation thinks it's an OK fit.

@rozbb I filed a bunch of comments on the interface in RustCrypto/traits#1508. Basically, I think it's usable, just could be nicer. It is correct in that it focuses on the encap/decap keys instead of trying to capture the whole system as I do here. I would probably prefer to get that fixed before implementing here.

This adds support for sizes identified as needed for post-quantum KEM/DSA use cases, namely the ones from this comment: RustCrypto/KEMs#2 (comment) These should ideally get expanded into some consistent multiples above 1024, e.g. multiples of 32, and generated in a purely automated manner (e.g. by a script that can break down the bit representation and build the generic syntax), but this should at least be sufficient to unblock these use cases. Note that `UInt<UInt<..<UTerm, B#>, B#...` aliases expressing the explicit bits for a given number are used instead of e.g. `<U1024 as Add<U32>>::Output` because when the latter is used it causes similar errors for conflicting trait impls as we saw with `typenum::U<N>` for whatever reason: error[E0119]: conflicting implementations of trait `traits::ArraySize` for type `UTerm` --> src/sizes.rs:82:13 | 82 | unsafe impl ArraySize for $ty { | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ | | | first implementation here | conflicting implementation for `UTerm`

bifurcation · 2024-03-01T17:15:31Z

Latest commits address the comments so far:

Remove the redundant .gitignore file
Add a workspace-level Cargo.toml file
Convert to hybrid-array

I have not attempted to use the KEM traits, given RustCrypto/traits#1508

The latest commit here builds against the master branch of hybrid-array, since the required changes haven't been released to crates.io yet.

@tarcieri how would you like to proceed here?

tarcieri · 2024-03-01T17:21:02Z

I can cut another hybrid-array release

rozbb · 2024-03-01T17:36:24Z

@bifurcation also would love to get your input on RustCrypto/traits#1509 and see if that addresses your comments. It's very very minimal now. Const-size arrays aren't even needed anymore.

tarcieri · 2024-03-01T17:51:41Z

hybrid-array v0.2.0-rc.6 has been released: RustCrypto/hybrid-array#54

tarcieri · 2024-03-01T18:50:24Z

There's no current CI configuration, although I could potentially follow up with one, since I don't think it would run anyway until this PR is merged.

@bifurcation can you check the code is clean under clippy if you haven't already? I'm noticing some function names with capital letters like PRF and XOF and I'm curious if clippy would flag those

bifurcation · 2024-03-01T22:56:47Z

Thanks for the nudge @rozbb, I had missed that PR. I sent a review. d0945d4 copy/pastes those interfaces into this repo; when that PR gets merged and released, we can switch over to it here.

Re clippy -- Yes, it runs clean. I added an #![allow(non_snake_case)] to the top to allow for variable names to follow the spec. I don't feel strongly, just think it make it easier to see how things line up if someone is reading. 4d83d5d makes a few clippy exceptions more precise. I would be open to doing the same for non_snake_case, since I agree that that's not great to have in general.

Re CI -- I think adding CI in a PR will run it on that PR. @tarcieri do you happen to have a favorite example that could be copy/pasted here?

tarcieri · 2024-03-01T22:58:53Z

I can add it in a separate PR (and since I'm a member, it should run it in that PR)

tarcieri

utACK with the caveat that my review is largely for style and consistency, and not the correctness of the implementation

Add ML-KEM implementation

1809696

bifurcation added 3 commits February 8, 2024 12:44

Pre-compute constants using const fn

60c4890

Use constant-time selection in decapsulation

9fe6e2f

Make determinstic feature non-default

d7ceaf9

tarcieri reviewed Feb 13, 2024

View reviewed changes

ml-kem/.gitignore Outdated Show resolved Hide resolved

tarcieri reviewed Feb 13, 2024

View reviewed changes

ml-kem/Cargo.toml Show resolved Hide resolved

tarcieri reviewed Feb 13, 2024

View reviewed changes

tarcieri mentioned this pull request Feb 13, 2024

Adding Xwing KEM RustCrypto/traits#1506

Closed

Remove redundant .gitignore

20089d7

rozbb mentioned this pull request Feb 16, 2024

No way to use extend supported array lengths RustCrypto/hybrid-array#45

Closed

bifurcation mentioned this pull request Feb 16, 2024

Comments on KEM traits RustCrypto/traits#1508

Open

tarcieri mentioned this pull request Feb 16, 2024

Further expand array size support RustCrypto/hybrid-array#49

Merged

bifurcation added 2 commits March 1, 2024 11:19

Add workspace-level Cargo.toml

3bad718

Convert from generic-array to hybrid-array

4e5404e

Use released version of hybrid-array

48e9e2f

bifurcation added 2 commits March 1, 2024 15:55

Use standard KEM interfaces

d0945d4

Cargo clippy cleanup

4d83d5d

tarcieri approved these changes Mar 1, 2024

View reviewed changes

tarcieri merged commit ea44681 into RustCrypto:master Mar 1, 2024

bifurcation mentioned this pull request May 30, 2024

ml-kem: fix benchmarking code #19

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add an ML-KEM implementation #2

Add an ML-KEM implementation #2

bifurcation commented Feb 7, 2024

tarcieri commented Feb 7, 2024

bifurcation commented Feb 8, 2024

tarcieri Feb 13, 2024 •

edited

Loading

rozbb Feb 13, 2024

bifurcation Feb 13, 2024

bifurcation Feb 13, 2024

tarcieri Feb 13, 2024

bifurcation Feb 16, 2024 •

edited

Loading

tarcieri Feb 16, 2024 •

edited

Loading

rozbb Feb 16, 2024

bifurcation Feb 16, 2024

tarcieri Feb 16, 2024

bifurcation commented Feb 16, 2024 •

edited

Loading

bifurcation commented Mar 1, 2024

tarcieri commented Mar 1, 2024

rozbb commented Mar 1, 2024

tarcieri commented Mar 1, 2024

tarcieri commented Mar 1, 2024

bifurcation commented Mar 1, 2024

tarcieri commented Mar 1, 2024

tarcieri left a comment

Add an ML-KEM implementation #2

Add an ML-KEM implementation #2

Conversation

bifurcation commented Feb 7, 2024

tarcieri commented Feb 7, 2024

bifurcation commented Feb 8, 2024

tarcieri Feb 13, 2024 • edited Loading

Choose a reason for hiding this comment

rozbb Feb 13, 2024

Choose a reason for hiding this comment

bifurcation Feb 13, 2024

Choose a reason for hiding this comment

bifurcation Feb 13, 2024

Choose a reason for hiding this comment

tarcieri Feb 13, 2024

Choose a reason for hiding this comment

bifurcation Feb 16, 2024 • edited Loading

Choose a reason for hiding this comment

tarcieri Feb 16, 2024 • edited Loading

Choose a reason for hiding this comment

rozbb Feb 16, 2024

Choose a reason for hiding this comment

bifurcation Feb 16, 2024

Choose a reason for hiding this comment

tarcieri Feb 16, 2024

Choose a reason for hiding this comment

bifurcation commented Feb 16, 2024 • edited Loading

bifurcation commented Mar 1, 2024

tarcieri commented Mar 1, 2024

rozbb commented Mar 1, 2024

tarcieri commented Mar 1, 2024

tarcieri commented Mar 1, 2024

bifurcation commented Mar 1, 2024

tarcieri commented Mar 1, 2024

tarcieri left a comment

Choose a reason for hiding this comment

tarcieri Feb 13, 2024 •

edited

Loading

bifurcation Feb 16, 2024 •

edited

Loading

tarcieri Feb 16, 2024 •

edited

Loading

bifurcation commented Feb 16, 2024 •

edited

Loading