feat: Module-Lattice KEM

src/algebra/field/prime/mod.rs

@@ -36,7 +36,7 @@ pub type AESField = PrimeField<2>;
 
 /// The [`PrimeField`] struct represents elements of a field with prime order. The field is defined
 /// by a prime number `P`, and the elements are integers modulo `P`.
-#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Default, PartialOrd)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Default, PartialOrd, Ord)]
 pub struct PrimeField<const P: usize> {
   pub(crate) value: usize,
 }

src/algebra/ring.rs

@@ -0,0 +1,30 @@
+use std::{
+  cmp::{Eq, PartialEq},
+  ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign},
+};
+
+use super::Finite;
+
+pub trait Ring:
+  std::fmt::Debug
+  + Default
+  + Sized
+  + Copy
+  + Clone
+  + PartialEq
+  + Eq
+  + Add<Output = Self>
+  + AddAssign
+  + Sub<Output = Self>
+  + SubAssign
+  + Mul<Output = Self>
+  + MulAssign
+  + Neg<Output = Self>
+  + 'static {
+  const ZERO: Self;
+  const ONE: Self;
+}
+
+// pub trait FiniteRing: Finite + Ring {
+//   const PRIMITIVE_ELEMENT: Self;
+// }

src/encryption/asymmetric/lwe/mod.rs

@@ -0,0 +1,186 @@
+//! This module implements the Learning With Errors (LWE) public-key encryption scheme.
+//!
+//! LWE is a lattice-based cryptosystem whose security is based on the hardness of the
+//! Learning With Errors problem, first introduced by Regev in 2005. The scheme encrypts
+//! single bits and provides security against quantum computers.
+//!
+//! The implementation uses three type parameters:
+//! - Q: The modulus (must be prime)
+//! - N: The dimension of the lattice
+//! - K: The parameter for the binomial distribution used for error sampling
+
+use rand::Rng;
+
+use super::AsymmetricEncryption;
+use crate::algebra::field::{prime::PrimeField, Field};
+
+/// An implementation of Learning With Errors (LWE) encryption
+pub struct LWE<const Q: usize, const N: usize, const K: usize> {
+  private_key: PrivateKey<Q, N>,
+  public_key:  PublicKey<Q, N>,
+}
+
+/// The public key consisting of a random matrix A and vector b = As + e
+#[derive(Debug, Clone)]
+pub struct PublicKey<const Q: usize, const N: usize> {
+  /// Random matrix in Z_q^{n×n}
+  a: [[PrimeField<Q>; N]; N],
+  /// Vector b = As + e where s is secret and e is error
+  b: [PrimeField<Q>; N],
+}
+
+/// The private key consisting of the secret vector s
+#[derive(Debug, Clone)]
+pub struct PrivateKey<const Q: usize, const N: usize> {
+  /// Secret vector with small coefficients
+  s: [PrimeField<Q>; N],
+}
+
+/// A ciphertext consisting of vector u and scalar v
+#[derive(Debug, Clone)]
+pub struct Ciphertext<const Q: usize, const N: usize> {
+  /// First component u = A^T r
+  u: [PrimeField<Q>; N],
+  /// Second component v = b^T r + ⌊q/2⌋m
+  v: PrimeField<Q>,
+}
+
+/// Sample from a centered binomial distribution with parameter K
+///
+/// Returns a value in the range [-K, K] following a discrete approximation
+/// of a Gaussian distribution.
+pub fn sample_binomial<const Q: usize, const K: usize>() -> PrimeField<Q> {
+  let mut rng = rand::thread_rng();
+  let mut sum = PrimeField::<Q>::ZERO;
+
+  for _ in 0..K {
+    if rng.gen::<bool>() {
+      sum += PrimeField::<Q>::ONE;
+    }
+    if rng.gen::<bool>() {
+      sum -= PrimeField::<Q>::ONE;
+    }
+  }
+  sum
+}
+
+impl<const Q: usize, const N: usize, const K: usize> LWE<Q, N, K> {
+  pub fn new() -> LWE<Q, N, K> {
+    let mut rng = rand::thread_rng();
+
+    // Generate random matrix A
+    let mut a = [[PrimeField::<Q>::ZERO; N]; N];
+    for i in 0..N {
+      for j in 0..N {
+        a[i][j] = PrimeField::from(rng.gen_range(0..Q));
+      }
+    }
+
+    // Sample secret vector s with small coefficients
+    let mut s = [PrimeField::<Q>::ZERO; N];
+    for i in 0..N {
+      s[i] = sample_binomial::<Q, K>();
+    }
+
+    // Generate error vector e
+    let mut e = [PrimeField::<Q>::ZERO; N];
+    for i in 0..N {
+      e[i] = sample_binomial::<Q, K>();
+    }
+
+    // Compute b = As + e
+    let mut b = [PrimeField::<Q>::ZERO; N];
+    for i in 0..N {
+      let mut sum = PrimeField::<Q>::ZERO;
+      for j in 0..N {
+        sum += a[i][j] * s[j];
+      }
+      sum += e[i];
+      b[i] = sum;
+    }
+
+    Self { public_key: PublicKey { a, b }, private_key: PrivateKey { s } }
+  }
+}
+
+impl<const Q: usize, const N: usize, const K: usize> AsymmetricEncryption for LWE<Q, N, K> {
+  type Ciphertext = Ciphertext<Q, N>;
+  type Plaintext = bool;
+  type PrivateKey = PrivateKey<Q, N>;
+  type PublicKey = PublicKey<Q, N>;
+
+  fn encrypt(&self, plaintext: &Self::Plaintext) -> Self::Ciphertext {
+    let mut rng = rand::thread_rng();
+
+    // Sample random vector r (binary)
+    let mut r = [PrimeField::<Q>::ZERO; N];
+    for i in 0..N {
+      r[i] = PrimeField::from(rng.gen_range(0..2));
+    }
+
+    // Compute u = A^T r
+    let mut u = [PrimeField::<Q>::ZERO; N];
+    for i in 0..N {
+      for j in 0..N {
+        u[i] += self.public_key.a[j][i] * r[j];
+      }
+    }
+
+    // Compute v = b^T r + ⌊q/2⌋m
+    let mut v = PrimeField::<Q>::ZERO;
+    for i in 0..N {
+      v += self.public_key.b[i] * r[i];
+    }
+
+    if *plaintext {
+      v += PrimeField::from(Q / 2);
+    }
+
+    Ciphertext { u, v }
+  }
+
+  fn decrypt(&self, ct: &Self::Ciphertext) -> Self::Plaintext {
+    // Compute v - s^T u
+    let mut result = ct.v;
+    for i in 0..N {
+      result -= self.private_key.s[i] * ct.u[i];
+    }
+
+    // Get q/2 as a field element
+    let q_half = PrimeField::<Q>::from(Q / 2);
+
+    // For distance to zero, we need min(x, q-x)
+    let dist_to_zero = result.min(-result);
+
+    // For distance to q/2, we need min(|x - q/2|, |q/2 - x|)
+    let dist_to_q_half = if result >= q_half {
+      // If result ≥ q/2, distance is min(result - q/2, q - result + q/2)
+      (result - q_half).min(-result + q_half)
+    } else {
+      // If result < q/2, distance is min(q/2 - result, result + q/2)
+      (q_half - result).min(result + q_half)
+    };
+
+    dist_to_q_half < dist_to_zero
+  }
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+
+  #[test]
+  fn test_encryption_decryption() {
+    for _ in 0..100 {
+      let lwe = LWE::<97, 4, 2>::new();
+
+      // Test encryption and decryption of 0
+      let ct_zero = lwe.encrypt(&false);
+      assert_eq!(lwe.decrypt(&ct_zero), false, "Failed decrypting 0");
+
+      // Test encryption and decryption of 1
+      let ct_one = lwe.encrypt(&true);
+      assert_eq!(lwe.decrypt(&ct_one), true, "Failed decrypting 1");
+    }
+  }
+}

src/encryption/asymmetric/mod.rs

@@ -1,2 +1,16 @@
 //! Contains implementation of asymmetric cryptographic primitives like RSA encryption.
+pub mod lwe;
 pub mod rsa;
+
+pub trait AsymmetricEncryption {
+  type PublicKey;
+  type PrivateKey;
+  type Plaintext;
+  type Ciphertext;
+
+  /// Encrypts plaintext using key and returns ciphertext
+  fn encrypt(&self, plaintext: &Self::Plaintext) -> Self::Ciphertext;
+
+  /// Decrypts ciphertext using key and returns plaintext
+  fn decrypt(&self, ciphertext: &Self::Ciphertext) -> Self::Plaintext;
+}

src/hashes/sha3.rs

@@ -41,6 +41,19 @@ const RHO: [[u32; 5]; 5] =
     27, 20, 39, 8, 14,
   ]];
 
+/// Type alias for SHA3-224.
+pub type Sha3_224 = Sha3<28>;
+/// Type alias for SHA3-256.
+pub type Sha3_256 = Sha3<32>;
+/// Type alias for SHA3-384.
+pub type Sha3_384 = Sha3<48>;
+/// Type alias for SHA3-512.
+pub type Sha3_512 = Sha3<64>;
+/// Type alias for SHAKE128.
+pub type Shake128 = Shake<128>;
+/// Type alias for SHAKE256.
+pub type Shake256 = Shake<256>;
+
 #[derive(Clone, Debug)]
 struct KeccakState {
   lanes: [[u64; 5]; 5],
@@ -277,19 +290,6 @@ impl<const SECURITY_BITS: usize> Shake<SECURITY_BITS> {
   }
 }
 
-/// Type alias for SHA3-224.
-pub type Sha3_224 = Sha3<28>;
-/// Type alias for SHA3-256.
-pub type Sha3_256 = Sha3<32>;
-/// Type alias for SHA3-384.
-pub type Sha3_384 = Sha3<48>;
-/// Type alias for SHA3-512.
-pub type Sha3_512 = Sha3<64>;
-/// Type alias for SHAKE128.
-pub type Shake128 = Shake<128>;
-/// Type alias for SHAKE256.
-pub type Shake256 = Shake<256>;
-
 #[cfg(test)]
 mod tests {
   use hex_literal::hex;

src/kem/README.md

@@ -0,0 +1,23 @@
+# Key Encpasulation Mechanism
+
+## ToC
+- What is KEM?
+- How is it useful?
+- Security
+
+KEM is a cryptographic algorithm, that is used to establish a shared key between two parties on a untrusted channel.
+
+KEM consists of three algorithms:
+- $\text{KeyGen}$: Generates encapsulation and decapsulation keys
+- $\text{Encaps}$: Encapsulates key and ciphertext that is sent to other part to generate shared key
+- $\text{Decaps}$: Generates shared key from decapsulation key
+
+```mermaid
+flowchart TB
+a[Alice]-->k[KeyGen]
+k--"Decapsulation Key"-->d[Decaps]
+k--"Encapsulation Key"-->e[Encaps]
+e--"ciphertext"-->d
+d-->al[Alice's shared key]
+e-->bo[Bob's shared key]
+```

src/kem/kyber/README.md

@@ -0,0 +1,88 @@
+# Module Lattice Based KEM (FIPS 203)
+
+## ToC
+- [ ] Intrduction
+- [ ] Preliminaries
+- [ ] Parameters
+- [ ] Encryption
+- [ ] KEM
+- [ ] Optimisation
+- [ ] Performance
+- [ ] Security
+
+
+## Introduction
+
+> [!NOTE]
+> Implementation and terminilogy follows FIPS 203 that formalises PQ-KEM standard for the internet.
+
+ML-KEM is based on Module Learning with Errors problems introduced by [Reg05][Reg05].
+
+## Preliminaries
+- [ ] Why Polynomial Rings?
+- [ ] Cyclotomic Polynomial Rings
+- [ ] NTT
+- [ ] Lattices
+
+## Implementation Details
+- Polynomial Rings
+  - representation
+  - Arithmetic with Polynomials
+- Matrices and Vectors
+  - representation
+  - Arithmetic
+- cryptographic functions
+  - SHA3-256, SHA3-512
+  - SHAKE128, SHAKE256
+  - PRF: $\textsf{PRF}_{\eta}(s,b):=\text{SHAKE256}(s\|b,8\cdot 64 \cdot \eta)$
+  - Hash function:
+    - $H: H(s) := \text{SHA3-256}(s)$, where $s\in \mathbb{B}^*$
+    - $J: J(s) := \text{SHAKE256}(s,8\cdot 32)$, where $s\in \mathbb{B}^*$
+    - $G:\ G(c \in \mathbb{B}^*) := \text{SHA3-512}(c) \in \mathbb{B}^{32}\times\mathbb{B}^{32}$
+  - XOF: SHAKE128
+- General Functions:
+  - BitsToBytes, BytesToBits
+  - Compression, Decompression
+  - ByteEncode, ByteDecode
+- Sampling algorithms:
+  - SampleNTT
+- NTT, RevNTT
+- MultiplyNTT
+- BaseCaseMultiply
+
+## NTT: Number-Theroretic Transform
+
+- Special case of FFT performed over $\mathbb{Z}_q^*$.
+- Most computationally intensive operation in encryption scheme is the multiplication of polynomials in $\mathcal{R}_{q,f}$ that is $\mathcal{O}(n^2)$, where $n$ is the degree of the polynomial. NTT allows to perform this in $\mathcal{O}(n\log n)$.
+- Polynomial: $\mathbb{Z}_q[X]/(X^d+\alpha)$, where $d=2^k$ and $-\alpha$ is perfect square of $d\in\mathbb{Z}_q$
+  - $X^d+\alpha\equiv(X^{d/2}-r)(X^{d/2}+r)\mod q$
+- Using, Chinese Remainder Theorem, $ab\in\mathbb{Z}_q/(X^d+\alpha)$ can be written as $ab \mod (X^{d/2}+r),\ ab \mod (X^{d/2}-r)$, and can be converted back to $\mod (X^d+\alpha)$.
+- Doing this recursively, asymptotic complexity of the algorithm turns out to be $d\log d$
+
+## Parameters
+- $q=3329=2^8+1$
+- $f=X^{256}+1$
+- $k,\eta_1,\eta_2,d_u,d_v$ vary according to parameter sets.
+
+|            | k   | $\eta_1$ | $\eta_2$ | $d_u$ | $d_v$ | decryption error | pk size | ciphertext size |
+| ---------- | --- | -------- | -------- | ----- | ----- | ---------------- | ------- | --------------- |
+| Kyber-512  | 3   | 3        | 2        | 10    | 4     | $2^{−139}$       | 800B    | 768B            |
+| Kyber-768  | 4   | 2        | 2        | 10    | 4     | $2^{−164}$       | 1184B   | 1088B           |
+| Kyber-1024 | 5   | 2        | 2        | 11    | 5     | $2^{−174}$       | 1568B   | 1568B           |
+
+## CPA-secure Encryption Scheme
+1. $\text{KeyGen}$
+2.
+
+## TODO
+
+- [ ] FIPS 202: SHA3-{256,512}, SHAKE{128,256}
+- [ ] Rings, Polynomial Rings
+  - [ ] arithmetic for rings
+  - [ ]
+
+## References
+- [Module-Lattice-Based Key-Encapsulation Mechanism Standard](https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.203.pdf)
+- [RustCrypto/KEMs](https://github.com/RustCrypto/KEMs): const functions inspiration
+
+[Reg05]: <https://cims.nyu.edu/~regev/papers/lwesurvey.pdf>