> [!CAUTION]
> This ML-KEM implementation is conformant with ML-KEM standard https://doi.org/10.6028/NIST.FIPS.203 and I also *try* to make it timing leakage free, but be informed that this implementation is not *yet* audited. **If you consider using it in production, please be careful !**

# ML-KEM (formerly known as Kyber)
Module-Lattice -based Key Encapsulation Mechanism Standard by NIST i.e. FIPS 203.

## Motivation

ML-KEM has been standardized by NIST as post-quantum secure key encapsulation mechanism (KEM), which can be used for key establishment, between two parties, communicating over insecure channel.

ML-KEM offers an *IND-CCA-secure* Key Encapsulation Mechanism - its security is based on the hardness of solving the learning-with-errors (LWE) problem in module (i.e. structured) lattices.

ML-KEM is built on top of *IND-CPA-secure K-PKE*, where two communicating parties, both generating their key pairs, while publishing only their public keys to each other, can encrypt fixed length ( = 32 -bytes ) message using peer's public key. Cipher text can be decrypted by corresponding secret key ( which is private to the keypair owner ) and 32 -bytes message can be recovered back. Then a slightly tweaked Fujisaki–Okamoto (FO) transform is applied on *IND-CPA-secure K-PKE* - giving us the *IND-CCA-secure ML-KEM* construction. In KEM scheme, two parties interested in establishing a secure communication channel, over public & insecure channel, can generate a 32 -bytes shared secret key. Now they can be use this 32 -bytes shared secret key in any symmetric key primitive, either for encrypting their communication (in much faster way) or deriving new/ longer keys.

Algorithm | Input | Output
--- | :-: | --:
KeyGen | - | Public Key and Secret Key
Encapsulation | Public Key | Cipher Text and 32B Shared Secret
Decapsulation | Secret Key and Cipher Text | 32B Shared Secret

Here I'm maintaining `ml-kem` - a C++20 header-only fully `constexpr` library, implementing ML-KEM, supporting ML-KEM-{512, 768, 1024} parameter sets, as defined in table 2 of ML-KEM standard. It's pretty easy to use, see [usage](#usage). It shows following performance characteristics on desktop and server grade CPUs.

ML-KEM-768 Algorithm | Time taken on "12th Gen Intel(R) Core(TM) i7-1260P" | Time taken on "AWS EC2 Instance c8g.large"
--- | --: | --:
keygen | 22.3us | 31.5us
encaps | 25.6us | 35.9us
decaps | 30.1us | 43.7us

> [!NOTE]
> Find ML-KEM standard @ https://doi.org/10.6028/NIST.FIPS.203 - this is the document that I followed when implementing ML-KEM. I suggest you go through the specification to get an in-depth understanding of the scheme.

## Prerequisites

- A C++ compiler such as `clang++`/ `g++`, with support for compiling C++20 programs.

```bash
$ clang++ --version
Ubuntu clang version 17.0.6 (9ubuntu1)
Target: x86_64-pc-linux-gnu
Thread model: posix
InstalledDir: /usr/bin
```

- Build tools such as `make`, `cmake`.
- For testing ML-KEM implementation, you need to globally install `google-test` library and headers. Follow guide @ https://github.com/google/googletest/tree/main/googletest#standalone-cmake-project, if you don't have it installed.
- For benchmarking ML-KEM implementation, you'll need to have `google-benchmark` header and library globally installed. I found guide @ https://github.com/google/benchmark#installation helpful.

> [!NOTE]
> If you are on a machine running GNU/Linux kernel and you want to obtain *CPU cycle* count for ML-KEM routines, you should consider building `google-benchmark` library with `libPFM` support, following https://gist.github.com/itzmeanjan/05dc3e946f635d00c5e0b21aae6203a7, a step-by-step guide. Find more about libPFM @ https://perfmon2.sourceforge.net.

> [!TIP]
> Git submodule based dependencies will normally be imported automatically, but in case that doesn't work, you can manually initialize and update them by issuing `$ git submodule update --init --recursive` from inside the root of this repository.

## Testing

For testing functional correctness of this implementation and conformance with ML-KEM standard, you have to issue

> [!NOTE]
> Known Answer Test (KAT) files living in [this](./kats/) directory are generated by following (reproducible) steps, described in https://gist.github.com/itzmeanjan/c8f5bc9640d0f0bdd2437dfe364d7710.

```bash
make test -j               # Run tests without any sort of sanitizers, with default C++ compiler.
CXX=clang++ make test -j   # Switch to non-default compiler, by setting variable `CXX`.

make debug_asan_test -j    # Run tests with AddressSanitizer enabled, with `-O1`.
make release_asan_test -j  # Run tests with AddressSanitizer enabled, with `-O3 -march=native`.
make debug_ubsan_test -j   # Run tests with UndefinedBehaviourSanitizer enabled, with `-O1`.
make release_ubsan_test -j # Run tests with UndefinedBehaviourSanitizer enabled, with `-O3 -march=native`.
```

```bash
PASSED TESTS (15/15):
       1 ms: build/test/test.out ML_KEM.ML_KEM_1024_DecapsFailureDueToBitFlippedCipherText
       1 ms: build/test/test.out ML_KEM.ML_KEM_512_DecapsFailureDueToBitFlippedCipherText
       1 ms: build/test/test.out ML_KEM.PolynomialSerialization
       2 ms: build/test/test.out ML_KEM.ML_KEM_512_EncapsFailureDueToNonReducedPubKey
       2 ms: build/test/test.out ML_KEM.ML_KEM_512_KeygenEncapsDecaps
       2 ms: build/test/test.out ML_KEM.ML_KEM_768_KeygenEncapsDecaps
       2 ms: build/test/test.out ML_KEM.ML_KEM_768_DecapsFailureDueToBitFlippedCipherText
       2 ms: build/test/test.out ML_KEM.ML_KEM_768_EncapsFailureDueToNonReducedPubKey
       2 ms: build/test/test.out ML_KEM.ML_KEM_1024_EncapsFailureDueToNonReducedPubKey
       3 ms: build/test/test.out ML_KEM.ML_KEM_1024_KeygenEncapsDecaps
      15 ms: build/test/test.out ML_KEM.ML_KEM_512_KnownAnswerTests
      24 ms: build/test/test.out ML_KEM.ML_KEM_768_KnownAnswerTests
      30 ms: build/test/test.out ML_KEM.ML_KEM_1024_KnownAnswerTests
     111 ms: build/test/test.out ML_KEM.CompressDecompressZq
     136 ms: build/test/test.out ML_KEM.ArithmeticOverZq
```

> [!NOTE]
> There is a help menu, which introduces you to all available commands; just run `make` from the root directory of this project.

## Benchmarking

For benchmarking ML-KEM public functions such as keygen, encaps and decaps, for various suggested parameter sets, you have to issue.

```bash
make benchmark -j  # If you haven't built google-benchmark library with libPFM support.
make perf -j       # If you have built google-benchmark library with libPFM support.
```

> [!CAUTION]
> When benchmarking, ensure that you've disabled CPU frequency scaling, by following guide @ https://github.com/google/benchmark/blob/main/docs/reducing_variance.md.

### On 12th Gen Intel(R) Core(TM) i7-1260P
Benchmark results are in JSON format @ [bench_result_on_Linux_6.11.0-19-generic_x86_64_with_g++_14](./bench_result_on_Linux_6.11.0-19-generic_x86_64_with_g++_14.json).

### On AWS EC2 Instance `c8g.large` i.e. AWS Graviton4
Benchmark results are in JSON format @ [bench_result_on_Linux_6.8.0-1021-aws_aarch64_with_g++_13](./bench_result_on_Linux_6.8.0-1021-aws_aarch64_with_g++_13.json).

## Usage

`ml-kem` is written as a header-only C++20 fully `constexpr` library, majorly targeting 64 -bit desktop/ server grade platforms and it's pretty easy to get started with. All you need to do is following.

- Clone `ml-kem` repository.

```bash
cd

# Single step cloning and importing of submodules
git clone https://github.com/itzmeanjan/ml-kem.git --recurse-submodules
# Or clone and then run tests, which will automatically bring in dependencies
git clone https://github.com/itzmeanjan/ml-kem.git && pushd ml-kem && make test -j && popd
```

- Write your program while including proper header files ( based on which variant of ML-KEM you want to use, see [include](./include/ml_kem/) directory ), which includes declarations ( and definitions ) of all required ML-KEM routines and constants ( such as byte length of public/ private key, cipher text etc. ).

```cpp
// main.cpp

#include "ml_kem/ml_kem_512.hpp"
#include "randomshake/randomshake.hpp"
#include <algorithm>
#include <array>
#include <cassert>

int
main()
{
  std::array<uint8_t, ml_kem_512::SEED_D_BYTE_LEN> d{};
  std::array<uint8_t, ml_kem_512::SEED_Z_BYTE_LEN> z{};

  std::array<uint8_t, ml_kem_512::PKEY_BYTE_LEN> pkey{};
  std::array<uint8_t, ml_kem_512::SKEY_BYTE_LEN> skey{};

  std::array<uint8_t, ml_kem_512::SEED_M_BYTE_LEN> m{};
  std::array<uint8_t, ml_kem_512::CIPHER_TEXT_BYTE_LEN> cipher{};

  std::array<uint8_t, ml_kem_512::SHARED_SECRET_BYTE_LEN> sender_key{};
  std::array<uint8_t, ml_kem_512::SHARED_SECRET_BYTE_LEN> receiver_key{};

  randomshake::randomshake_t<128> csprng;

  csprng.generate(d);
  csprng.generate(z);
  csprng.generate(m);

  ml_kem_512::keygen(d, z, pkey, skey);
  assert(ml_kem_512::encapsulate(m, pkey, cipher, sender_key)); // Key Encapsulation might fail, if input public key is malformed
  ml_kem_512::decapsulate(skey, cipher, receiver_key);

  assert(sender_key == receiver_key);
  return 0;
}
```

- When compiling your program, let your compiler know where it can find `ml-kem`, `sha3`, `RandomShake` and `subtle` headers, which includes their definitions ( all of them are header-only libraries ) too.

```bash
# Assuming `ml-kem` was cloned just under $HOME

ML_KEM_HEADERS=~/ml-kem/include
SHA3_HEADERS=~/ml-kem/sha3/include
RANDOMSHAKE_HEADERS=~/ml-kem/RandomShake/include
SUBTLE_HEADERS=~/ml-kem/subtle/include

g++ -std=c++20 -Wall -Wextra -Wpedantic -O3 -march=native -I $ML_KEM_HEADERS -I $SHA3_HEADERS -I $RANDOMSHAKE_HEADERS -I $SUBTLE_HEADERS main.cpp
```

ML-KEM Variant | Namespace | Header
:-- | :-: | --:
ML-KEM-512 Routines | `ml_kem_512::` | `include/ml_kem/ml_kem_512.hpp`
ML-KEM-768 Routines | `ml_kem_768::` | `include/ml_kem/ml_kem_768.hpp`
ML-KEM-1024 Routines | `ml_kem_1024::` | `include/ml_kem/ml_kem_1024.hpp`

> [!NOTE]
> ML-KEM parameter sets are taken from table 2 of ML-KEM standard @ https://doi.org/10.6028/NIST.FIPS.203.

All the functions, in this ML-KEM header-only library, are implemented as `constexpr` functions. Hence you should be able to evaluate ML-KEM key generation, encapsulation or decapsulation at compile-time itself, given that all inputs are known at compile-time. I present you with following demonstration program, which generates a ML-KEM-512 keypair and encapsulates a message, producing a ML-KEM-512 cipher text and a fixed size shared secret, given `seed_{d, z, m}` as input - all at program compile-time. Notice, the *static assertion*.

```cpp
// compile-time-ml-kem-512.cpp
//
// Compile and run this program with
// $ g++ -std=c++20 -Wall -Wextra -Wpedantic -I include -I sha3/include -I subtle/include compile-time-ml-kem-512.cpp && ./a.out
// or
// $ clang++ -std=c++20 -Wall -Wextra -Wpedantic -fconstexpr-steps=4000000 -I include -I sha3/include -I subtle/include compile-time-ml-kem-512.cpp && ./a.out

#include "ml_kem/ml_kem_512.hpp"

// Compile-time evaluation of ML-KEM-512 key generation and encapsulation, using NIST official KAT no. (1).
constexpr auto
eval_ml_kem_768_encaps() -> auto
{
  using seed_t = std::array<uint8_t, ml_kem_512::SEED_D_BYTE_LEN>;

  // 7c9935a0b07694aa0c6d10e4db6b1add2fd81a25ccb148032dcd739936737f2d
  constexpr seed_t seed_d = { 124, 153, 53, 160, 176, 118, 148, 170, 12, 109, 16,  228, 219, 107, 26,  221, 47,  216, 26, 37,  204, 177, 72,  3,   45, 205, 115, 153, 54,  115, 127, 45 };
  // b505d7cfad1b497499323c8686325e4792f267aafa3f87ca60d01cb54f29202a
  constexpr seed_t seed_z = {181, 5, 215, 207, 173, 27, 73, 116, 153, 50, 60, 134, 134, 50, 94, 71, 146, 242, 103, 170, 250, 63, 135, 202, 96, 208, 28, 181, 79, 41, 32, 42};
  // eb4a7c66ef4eba2ddb38c88d8bc706b1d639002198172a7b1942eca8f6c001ba
  constexpr seed_t seed_m = {235, 74, 124, 102, 239, 78, 186, 45, 219, 56, 200, 141, 139, 199, 6, 177, 214, 57, 0, 33, 152, 23, 42, 123, 25, 66, 236, 168, 246, 192, 1, 186};

  std::array<uint8_t, ml_kem_512::PKEY_BYTE_LEN> pubkey{};
  std::array<uint8_t, ml_kem_512::SKEY_BYTE_LEN> seckey{};
  std::array<uint8_t, ml_kem_512::CIPHER_TEXT_BYTE_LEN> cipher{};

  std::array<uint8_t, ml_kem_512::SHARED_SECRET_BYTE_LEN> shared_secret{};

  ml_kem_512::keygen(seed_d, seed_z, pubkey, seckey);
  (void)ml_kem_512::encapsulate(seed_m, pubkey, cipher, shared_secret);

  return shared_secret;
}

int
main()
{
  // This step is being evaluated at compile-time, thanks to the fact that my ML-KEM implementation is `constexpr`.
  static constexpr auto computed_shared_secret = eval_ml_kem_768_encaps();
  // b4c8e3c4115f9511f2fddb288c4b78c5cd7c89d2d4d321f46b4edc54ddf0eb36
  constexpr std::array<uint8_t, ml_kem_512::SHARED_SECRET_BYTE_LEN> expected_shared_secret = { 180, 200, 227, 196, 17, 95, 149, 17, 242, 253, 219, 40, 140, 75, 120, 197, 205, 124, 137, 210, 212, 211, 33, 244, 107, 78, 220, 84, 221, 240, 235, 54 };

  // Notice static_assert, yay !
  static_assert(computed_shared_secret == expected_shared_secret, "Must be able to compute shared secret at compile-time !");
  return 0;
}
```

See example [program](./examples/ml_kem_768.cpp), where I show how to use ML-KEM-768 API. Issue following command to build and execute example.

```bash
make example -j
```

```bash
ML-KEM-768
Pubkey         :  
Seckey         :  
Encapsulated ? : true
Cipher         :  
Shared secret  : e6a9fc79df8a91733c7f385bc66602a526b54bbf78ed2ac11029a42a2a56f515
```

> [!NOTE]
> Looking at API documentation, in header files, can give you good idea of how to use ML-KEM API. Note, this library doesn't expose any raw pointer based interface, rather everything is wrapped under statically defined `std::span` - which one can easily create from `std::{array, vector}`. I opt for using statically defined `std::span` based function interfaces because we always know, at compile-time, how many bytes the seeds/ keys/ cipher-texts/ shared-secrets are, for various different ML-KEM parameters. This gives much better type safety and compile-time error reporting.
