Add helper functions to compute CPA

Ledger-Donjon · Apr 5, 2024 · 8399fa2 · 8399fa2
1 parent 60a11e6
commit 8399fa2
Show file tree

Hide file tree

Showing 3 changed files with 114 additions and 64 deletions.
diff --git a/benches/cpa.rs b/benches/cpa.rs
@@ -1,12 +1,11 @@
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
-use muscat::cpa::Cpa;
+use muscat::cpa::{self, Cpa};
 use muscat::cpa_normal;
 use muscat::leakage::{hw, sbox};
 use ndarray::{Array2, ArrayView1, Axis};
 use ndarray_rand::rand::{rngs::StdRng, SeedableRng};
 use ndarray_rand::rand_distr::Uniform;
 use ndarray_rand::RandomExt;
-use rayon::iter::{ParallelBridge, ParallelIterator};
 use std::iter::zip;
 
 pub fn leakage_model(value: usize, guess: usize) -> usize {
@@ -28,36 +27,6 @@ fn cpa_sequential(leakages: &Array2<f64>, plaintexts: &Array2<u8>) -> Cpa {
     cpa
 }
 
-fn cpa_parallel(leakages: &Array2<f64>, plaintexts: &Array2<u8>) -> Cpa {
-    let chunk_size = 500;
-
-    let mut cpa = zip(
-        leakages.axis_chunks_iter(Axis(0), chunk_size),
-        plaintexts.axis_chunks_iter(Axis(0), chunk_size),
-    )
-    .par_bridge()
-    .map(|(leakages_chunk, plaintexts_chunk)| {
-        let mut cpa = Cpa::new(leakages.shape()[1], 256, 0, leakage_model);
-
-        for i in 0..leakages_chunk.shape()[0] {
-            cpa.update(
-                leakages_chunk.row(i).map(|&x| x as usize),
-                plaintexts_chunk.row(i).map(|&y| y as usize),
-            );
-        }
-
-        cpa
-    })
-    .reduce(
-        || Cpa::new(leakages.shape()[1], 256, 0, leakage_model),
-        |a: Cpa, b| a + b,
-    );
-
-    cpa.finalize();
-
-    cpa
-}
-
 pub fn leakage_model_normal(value: ArrayView1<usize>, guess: usize) -> usize {
     hw(sbox((value[1] ^ guess) as u8) as usize)
 }
@@ -82,33 +51,6 @@ fn cpa_normal_sequential(leakages: &Array2<f64>, plaintexts: &Array2<u8>) -> cpa
     cpa
 }
 
-fn cpa_normal_parallel(leakages: &Array2<f64>, plaintexts: &Array2<u8>) -> cpa_normal::Cpa {
-    let chunk_size = 500;
-
-    let mut cpa = zip(
-        leakages.axis_chunks_iter(Axis(0), chunk_size),
-        plaintexts.axis_chunks_iter(Axis(0), chunk_size),
-    )
-    .par_bridge()
-    .map(|(leakages_chunk, plaintexts_chunk)| {
-        let mut cpa =
-            cpa_normal::Cpa::new(leakages.shape()[1], chunk_size, 256, leakage_model_normal);
-        cpa.update(
-            leakages_chunk.map(|&x| x as f32),
-            plaintexts_chunk.to_owned(),
-        );
-        cpa
-    })
-    .reduce(
-        || cpa_normal::Cpa::new(leakages.shape()[1], chunk_size, 256, leakage_model_normal),
-        |x, y| x + y,
-    );
-
-    cpa.finalize();
-
-    cpa
-}
-
 fn bench_cpa(c: &mut Criterion) {
     // Seed rng to get the same output each run
     let mut rng = StdRng::seed_from_u64(0);
@@ -133,7 +75,18 @@ fn bench_cpa(c: &mut Criterion) {
         group.bench_with_input(
             BenchmarkId::new("cpa_parallel", nb_traces),
             &(&leakages, &plaintexts),
-            |b, (leakages, plaintexts)| b.iter(|| cpa_parallel(leakages, plaintexts)),
+            |b, (leakages, plaintexts)| {
+                b.iter(|| {
+                    cpa::cpa(
+                        &leakages.map(|&x| x as usize),
+                        &plaintexts.map(|&x| x as usize),
+                        256,
+                        0,
+                        leakage_model,
+                        500,
+                    )
+                })
+            },
         );
         // For 25000 traces, 60s of measurement_time is too low
         if nb_traces <= 10000 {
@@ -146,7 +99,17 @@ fn bench_cpa(c: &mut Criterion) {
         group.bench_with_input(
             BenchmarkId::new("cpa_normal_parallel", nb_traces),
             &(&leakages, &plaintexts),
-            |b, (leakages, plaintexts)| b.iter(|| cpa_normal_parallel(leakages, plaintexts)),
+            |b, (leakages, plaintexts)| {
+                b.iter(|| {
+                    cpa_normal::cpa(
+                        &leakages.map(|&x| x as f32),
+                        plaintexts,
+                        256,
+                        leakage_model_normal,
+                        500,
+                    )
+                })
+            },
         );
     }
 

diff --git a/src/cpa.rs b/src/cpa.rs
@@ -1,6 +1,52 @@
 use ndarray::{concatenate, s, Array1, Array2, ArrayView1, ArrayView2, Axis};
-use rayon::prelude::{IntoParallelIterator, ParallelIterator};
-use std::ops::Add;
+use rayon::{
+    iter::ParallelBridge,
+    prelude::{IntoParallelIterator, ParallelIterator},
+};
+use std::{iter::zip, ops::Add};
+
+/// Computes the [`Cpa`] of the given traces.
+///
+/// # Panics
+/// - Panic if `leakages.shape()[0] != plaintexts.shape()[0]`
+/// - Panic if `chunk_size` is 0.
+pub fn cpa(
+    leakages: &Array2<usize>,
+    plaintexts: &Array2<usize>,
+    guess_range: usize,
+    target_byte: usize,
+    leakage_func: fn(usize, usize) -> usize,
+    chunk_size: usize,
+) -> Cpa {
+    assert_eq!(leakages.shape()[0], plaintexts.shape()[0]);
+    assert!(chunk_size > 0);
+
+    let mut cpa = zip(
+        leakages.axis_chunks_iter(Axis(0), chunk_size),
+        plaintexts.axis_chunks_iter(Axis(0), chunk_size),
+    )
+    .par_bridge()
+    .map(|(leakages_chunk, plaintexts_chunk)| {
+        let mut cpa = Cpa::new(leakages.shape()[1], guess_range, target_byte, leakage_func);
+
+        for i in 0..leakages_chunk.shape()[0] {
+            cpa.update(
+                leakages_chunk.row(i).to_owned(),
+                plaintexts_chunk.row(i).to_owned(),
+            );
+        }
+
+        cpa
+    })
+    .reduce(
+        || Cpa::new(leakages.shape()[1], guess_range, target_byte, leakage_func),
+        |a: Cpa, b| a + b,
+    );
+
+    cpa.finalize();
+
+    cpa
+}
 
 pub struct Cpa {
     sum_leakages: Array1<usize>,

diff --git a/src/cpa_normal.rs b/src/cpa_normal.rs
@@ -1,5 +1,46 @@
 use ndarray::{concatenate, Array1, Array2, ArrayView1, ArrayView2, Axis};
-use std::ops::Add;
+use rayon::iter::{ParallelBridge, ParallelIterator};
+use std::{iter::zip, ops::Add};
+
+/// Computes the [`Cpa`] of the given traces.
+///
+/// # Panics
+/// - Panic if `leakages.shape()[0] != plaintexts.shape()[0]`
+/// - Panic if `chunk_size` is 0.
+pub fn cpa<T, U>(
+    leakages: &Array2<T>,
+    plaintexts: &Array2<U>,
+    guess_range: usize,
+    leakage_func: fn(ArrayView1<usize>, usize) -> usize,
+    chunk_size: usize,
+) -> Cpa
+where
+    T: Into<f32> + Copy + Sync,
+    U: Into<usize> + Copy + Sync,
+{
+    assert_eq!(leakages.shape()[0], plaintexts.shape()[0]);
+    assert!(chunk_size > 0);
+
+    let mut cpa = zip(
+        leakages.axis_chunks_iter(Axis(0), chunk_size),
+        plaintexts.axis_chunks_iter(Axis(0), chunk_size),
+    )
+    .par_bridge()
+    .map(|(leakages_chunk, plaintexts_chunk)| {
+        let mut cpa = Cpa::new(leakages.shape()[1], chunk_size, guess_range, leakage_func);
+        cpa.update(leakages_chunk.to_owned(), plaintexts_chunk.to_owned());
+        cpa
+    })
+    .reduce(
+        || Cpa::new(leakages.shape()[1], chunk_size, guess_range, leakage_func),
+        |x, y| x + y,
+    );
+
+    cpa.finalize();
+
+    cpa
+}
+
 pub struct Cpa {
     /* List of internal class variables */
     sum_leakages: Array1<f32>,