Merge pull request #15 from lattice-based-cryptography/structure_KEM_class

refactor KEM struct

Author: jacksonwalters

.github/workflows/basic.yml

@@ -0,0 +1,38 @@
+name: basic
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+env:
+  CARGO_TERM_COLOR: always
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+
+    # Install the Rust toolchain
+    - name: Install Rust toolchain
+      uses: dtolnay/rust-toolchain@stable
+
+    # Cache Cargo dependencies to speed up builds
+    - name: Cache Cargo dependencies
+      uses: actions/cache@v3
+      with:
+        path: ~/.cargo/registry
+        key: cargo-${{ runner.os }}-${{ hashFiles('**/Cargo.lock') }}
+        restore-keys: |
+          cargo-${{ runner.os }}-
+
+    # Build the project
+    - name: Build
+      run: cargo build --verbose
+
+    # Run the tests
+    - name: Run Tests
+      run: cargo test --verbose

.gitignore

@@ -1,3 +1,5 @@
+.DS_Store
+
 # Generated by Cargo
 # will have compiled files and executables
 debug/

src/lib.rs

@@ -1,2 +1,2 @@
-pub mod kem;
+pub mod ml_kem;
 pub mod utils;

src/main.rs

@@ -1,12 +1,13 @@
-use ml_kem::kem::KEM;
+use ml_kem::ml_kem::MLKEM;
 use ml_kem::utils::Parameters;
 mod tests;
 
 fn main() {
     let params = Parameters::default();
 
     // Generate key pair
-    let (public_key, secret_key) = KEM::keygen(&params);
+    let mlkem = MLKEM::new(params); 
+    let (public_key, secret_key) = mlkem.keygen();
 
     // Print keys for verification
     println!("Public Key: {:?}", public_key);

src/ml_kem.rs

@@ -0,0 +1,65 @@
+use module_lwe::utils::{gen_uniform_matrix,mul_mat_vec_simple,gen_small_vector,add_vec};
+use module_lwe::encrypt::encrypt;
+use module_lwe::decrypt::decrypt;
+use ring_lwe::utils::gen_binary_poly;
+use crate::utils::{Parameters, hash};
+use polynomial_ring::Polynomial;
+
+pub struct MLKEM {
+    params: Parameters,
+}
+
+impl MLKEM {
+    // Constructor to initialize MLKEM with parameters
+    pub fn new(params: Parameters) -> Self {
+        MLKEM { params } // Corrected: properly initializes and returns the struct
+    }
+
+    pub fn keygen(&self) -> ((Vec<Vec<Polynomial<i64>>>, Vec<Polynomial<i64>>), Vec<Polynomial<i64>>) {
+        let a = gen_uniform_matrix(self.params.n, self.params.k, self.params.q, None); 
+        
+        let s = gen_small_vector(self.params.n, self.params.k, None);
+        let e = gen_small_vector(self.params.n, self.params.k, None);
+        
+        let b = add_vec(
+            &mul_mat_vec_simple(&a, &s, self.params.q, &self.params.f, self.params.omega), 
+            &e, 
+            self.params.q, 
+            &self.params.f
+        );
+        
+        ((a, b), s)
+    }
+
+    pub fn encapsulate(&self, pk: (Vec<Vec<Polynomial<i64>>>, Vec<Polynomial<i64>>)) -> (String, (Vec<Polynomial<i64>>, Polynomial<i64>)) {
+        let params_mlwe = module_lwe::utils::Parameters { 
+            n: self.params.n, 
+            q: self.params.q, 
+            k: self.params.k, 
+            omega: self.params.omega, 
+            f: self.params.f.clone() 
+        };
+
+        let mut m = gen_binary_poly(self.params.n, None).coeffs().to_vec();
+        m.resize(self.params.n, 0);
+
+        let ct = encrypt(&pk.0, &pk.1, &m, &params_mlwe, None);
+        let k = hash(m);
+        (k, ct)
+    }
+
+    pub fn decapsulate(&self, sk: Vec<Polynomial<i64>>, ct: (Vec<Polynomial<i64>>, Polynomial<i64>)) -> String {
+        let params_mlwe = module_lwe::utils::Parameters { 
+            n: self.params.n, 
+            q: self.params.q, 
+            k: self.params.k, 
+            omega: self.params.omega, 
+            f: self.params.f.clone() 
+        };
+
+        let mut m = decrypt(&sk, &ct.0, &ct.1, &params_mlwe);
+        m.resize(self.params.n, 0);
+
+        hash(m)
+    }
+}

src/tests.rs

@@ -1,15 +1,16 @@
 #[cfg(test)]  // This makes the following module compile only during tests
 mod tests {
     use ml_kem::utils::Parameters;
-	use ml_kem::kem::KEM;
+	use ml_kem::ml_kem::MLKEM;
 
     // Test for basic keygen/encapsulate/decapsulate
     #[test]
     pub fn test_basic() {
         let params = Parameters::default();  // Adjust this if needed
-        let (pk, sk) = KEM::keygen(&params);
-        let (k, ct) = KEM::encapsulate(pk, &params);
-		let k_recovered = KEM::decapsulate(sk, ct, &params);
+        let mlkem = MLKEM::new(params); 
+        let (pk, sk) = mlkem.keygen();
+        let (k, ct) = mlkem.encapsulate(pk);
+		let k_recovered = mlkem.decapsulate(sk, ct);
         assert_eq!(k, k_recovered, "test failed: {} != {}", k, k_recovered);
     }
 }