From cce93a5b2580f22f50c839125b9e57665ea24241 Mon Sep 17 00:00:00 2001 From: TrAyZeN <1810leo@gmail.com> Date: Fri, 5 Apr 2024 14:11:12 +0200 Subject: [PATCH] Return errors from utils instead of panicking --- Cargo.toml | 2 ++ examples/cpa.rs | 57 +++++++++++++++++++++++---------------- examples/cpa_partioned.rs | 30 +++++++++++---------- examples/rank.rs | 24 ++++++++++------- examples/snr.rs | 9 ++++--- src/quicklog.rs | 38 +++++++++++--------------- src/util.rs | 34 +++++++++++------------ 7 files changed, 104 insertions(+), 90 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 53b0e27..cc0aa33 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,10 +18,12 @@ rayon = "1.10.0" indicatif = "0.17.8" ndarray-npy ="0.8.1" itertools = "0.12.1" +thiserror = "1.0.58" [dev-dependencies] criterion = "0.5.1" ndarray-rand = "0.14.0" +anyhow = "1.0.81" [[bench]] name = "cpa" diff --git a/examples/cpa.rs b/examples/cpa.rs index a5eb6c3..b0837d4 100644 --- a/examples/cpa.rs +++ b/examples/cpa.rs @@ -1,10 +1,11 @@ +use anyhow::Result; use indicatif::ProgressIterator; use muscat::cpa_normal::*; use muscat::leakage::{hw, sbox}; -use muscat::util::{progress_bar, read_array_2_from_npy_file, save_array2}; +use muscat::util::{progress_bar, read_array2_from_npy_file, save_array2}; use ndarray::*; use rayon::iter::{ParallelBridge, ParallelIterator}; -use std::time::{self}; +use std::time; // leakage model pub fn leakage_model(value: ArrayView1, guess: usize) -> usize { @@ -16,21 +17,21 @@ type FormatTraces = f64; type FormatMetadata = u8; #[allow(dead_code)] -fn cpa() { - let start_sample: usize = 0; - let end_sample: usize = 5000; - let size: usize = end_sample - start_sample; // Number of samples - let patch: usize = 500; +fn cpa() -> Result<()> { + let start_sample = 0; + let end_sample = 5000; + let size = end_sample - start_sample; // Number of samples + let patch = 500; let guess_range = 256; // 2**(key length) let folder = String::from("../../data/cw"); let dir_l = format!("{folder}/leakages.npy"); let dir_p = format!("{folder}/plaintexts.npy"); - let leakages: Array2 = read_array_2_from_npy_file::(&dir_l); - let plaintext: Array2 = read_array_2_from_npy_file::(&dir_p); + let leakages = read_array2_from_npy_file::(&dir_l)?; + let plaintext = read_array2_from_npy_file::(&dir_p)?; let len_traces = leakages.shape()[0]; + let mut cpa_parallel = ((0..len_traces).step_by(patch)) .progress_with(progress_bar(len_traces)) - .map(|row| row) .par_bridge() .map(|row_number| { let mut cpa = Cpa::new(size, patch, guess_range, leakage_model); @@ -48,29 +49,33 @@ fn cpa() { || Cpa::new(size, patch, guess_range, leakage_model), |x, y| x + y, ); + cpa_parallel.finalize(); println!("Guessed key = {}", cpa_parallel.pass_guess()); - save_array2("results/corr.npy", cpa_parallel.pass_corr_array().view()); + + save_array2("results/corr.npy", cpa_parallel.pass_corr_array().view())?; + + Ok(()) } #[allow(dead_code)] -fn success() { - let start_sample: usize = 0; - let end_sample: usize = 5000; - let size: usize = end_sample - start_sample; // Number of samples - let patch: usize = 500; +fn success() -> Result<()> { + let start_sample = 0; + let end_sample = 5000; + let size = end_sample - start_sample; // Number of samples + let patch = 500; let guess_range = 256; // 2**(key length) let folder = String::from("../data/log_584012"); // "../../../intenship/scripts/log_584012" let nfiles = 13; // Number of files in the directory. TBD: Automating this value - let rank_traces: usize = 1000; + let rank_traces = 1000; + let mut cpa = Cpa::new(size, patch, guess_range, leakage_model); cpa.success_traces(rank_traces); for i in (0..nfiles).progress() { let dir_l = format!("{folder}/l/{i}.npy"); let dir_p = format!("{folder}/p/{i}.npy"); - let leakages: Array2 = read_array_2_from_npy_file::(&dir_l); - let plaintext: Array2 = - read_array_2_from_npy_file::(&dir_p); + let leakages = read_array2_from_npy_file::(&dir_l)?; + let plaintext = read_array2_from_npy_file::(&dir_p)?; let len_leakages = leakages.shape()[0]; for row in (0..len_leakages).step_by(patch) { let range_samples = start_sample..end_sample; @@ -84,14 +89,20 @@ fn success() { cpa.update_success(sample_traces, sample_metadata); } } + cpa.finalize(); println!("Guessed key = {}", cpa.pass_guess()); + // save corr key curves in npy - save_array2("results/success.npy", cpa.pass_rank().view()); + save_array2("results/success.npy", cpa.pass_rank().view())?; + + Ok(()) } -fn main() { +fn main() -> Result<()> { let t = time::Instant::now(); - cpa(); + cpa()?; println!("{:?}", t.elapsed()); + + Ok(()) } diff --git a/examples/cpa_partioned.rs b/examples/cpa_partioned.rs index 814363e..a6dae2d 100644 --- a/examples/cpa_partioned.rs +++ b/examples/cpa_partioned.rs @@ -1,11 +1,9 @@ -// use simple_bar::ProgressBar; +use anyhow::Result; use indicatif::ProgressIterator; use muscat::cpa::*; use muscat::leakage::{hw, sbox}; -use muscat::util::{progress_bar, read_array_2_from_npy_file, save_array}; -use ndarray::*; +use muscat::util::{progress_bar, read_array2_from_npy_file, save_array}; use rayon::prelude::{ParallelBridge, ParallelIterator}; -use std::time::Instant; // traces format type FormatTraces = i16; @@ -17,8 +15,8 @@ pub fn leakage_model(value: usize, guess: usize) -> usize { } // multi-threading cpa -fn cpa() { - let size: usize = 5000; // Number of samples +fn cpa() -> Result<()> { + let size = 5000; // Number of samples let guess_range = 256; // 2**(key length) let target_byte = 1; let folder = String::from("../../data"); // Directory of leakages and metadata @@ -26,19 +24,17 @@ fn cpa() { /* Parallel operation using multi-threading on patches */ let mut cpa = (0..nfiles) - .into_iter() .progress_with(progress_bar(nfiles)) .map(|n| { let dir_l = format!("{folder}/l{n}.npy"); let dir_p = format!("{folder}/p{n}.npy"); - let leakages: Array2 = read_array_2_from_npy_file(&dir_l); - let plaintext: Array2 = read_array_2_from_npy_file(&dir_p); + let leakages = read_array2_from_npy_file::(&dir_l).unwrap(); + let plaintext = read_array2_from_npy_file::(&dir_p).unwrap(); (leakages, plaintext) }) - .into_iter() .par_bridge() .map(|patch| { - let mut c: Cpa = Cpa::new(size, guess_range, target_byte, leakage_model); + let mut c = Cpa::new(size, guess_range, target_byte, leakage_model); let len_leakage = patch.0.shape()[0]; for i in 0..len_leakage { c.update( @@ -52,12 +48,18 @@ fn cpa() { || Cpa::new(size, guess_range, target_byte, leakage_model), |a: Cpa, b| a + b, ); + cpa.finalize(); println!("Guessed key = {}", cpa.pass_guess()); + // save corr key curves in npy - save_array("../results/corr.npy", &cpa.pass_corr_array()); + save_array("../results/corr.npy", &cpa.pass_corr_array())?; + + Ok(()) } -fn main() { - cpa(); +fn main() -> Result<()> { + cpa()?; + + Ok(()) } diff --git a/examples/rank.rs b/examples/rank.rs index 9dccb16..aaa66eb 100644 --- a/examples/rank.rs +++ b/examples/rank.rs @@ -1,10 +1,10 @@ +use anyhow::Result; use indicatif::ProgressIterator; use muscat::cpa::*; use muscat::leakage::{hw, sbox}; -use muscat::util::{progress_bar, read_array_2_from_npy_file, save_array}; +use muscat::util::{progress_bar, read_array2_from_npy_file, save_array}; use ndarray::*; use rayon::prelude::{ParallelBridge, ParallelIterator}; -use std::time::Instant; // traces format type FormatTraces = i16; @@ -15,8 +15,8 @@ pub fn leakage_model(value: usize, guess: usize) -> usize { hw(sbox((value ^ guess) as u8) as usize) } -fn rank() { - let size: usize = 5000; // Number of samples +fn rank() -> Result<()> { + let size = 5000; // Number of samples let guess_range = 256; // 2**(key length) let target_byte = 1; let folder = String::from("../../data"); @@ -26,8 +26,8 @@ fn rank() { for file in (0..nfiles).progress_with(progress_bar(nfiles)) { let dir_l = format!("{folder}/l{file}.npy"); let dir_p = format!("{folder}/p{file}.npy"); - let leakages: Array2 = read_array_2_from_npy_file(&dir_l); - let plaintext: Array2 = read_array_2_from_npy_file(&dir_p); + let leakages = read_array2_from_npy_file::(&dir_l)?; + let plaintext = read_array2_from_npy_file::(&dir_p)?; let len_file = leakages.shape()[0]; for sample in (0..len_file).step_by(chunk) { let l_sample: ndarray::ArrayBase< @@ -36,7 +36,6 @@ fn rank() { > = leakages.slice(s![sample..sample + chunk, ..]); let p_sample = plaintext.slice(s![sample..sample + chunk, ..]); let x = (0..chunk) - .into_iter() .par_bridge() .fold( || Cpa::new(size, guess_range, target_byte, leakage_model), @@ -56,10 +55,15 @@ fn rank() { rank.finalize(); } } + // save rank key curves in npy - save_array("../results/rank.npy", &rank.pass_rank()); + save_array("../results/rank.npy", &rank.pass_rank())?; + + Ok(()) } -fn main() { - rank(); +fn main() -> Result<()> { + rank()?; + + Ok(()) } diff --git a/examples/snr.rs b/examples/snr.rs index 8597de2..a36ec75 100644 --- a/examples/snr.rs +++ b/examples/snr.rs @@ -1,13 +1,14 @@ +use anyhow::Result; use indicatif::ProgressIterator; use muscat::processors::Snr; use muscat::quicklog::{BatchIter, Log}; use muscat::util::{progress_bar, save_array}; use rayon::prelude::{ParallelBridge, ParallelIterator}; -fn main() { +fn main() -> Result<()> { // Open log file // This uses logs from the python quicklog library. - let log = Log::::new("log").unwrap(); + let log = Log::::new("log")?; let leakage_size = log.leakage_size(); let trace_count = log.len(); @@ -37,5 +38,7 @@ fn main() { .reduce(|| Snr::new(leakage_size, 256), |a, b| a + b); // Save the resulting SNR trace to a numpy file - save_array("result.npy", &result.snr()); + save_array("result.npy", &result.snr())?; + + Ok(()) } diff --git a/src/quicklog.rs b/src/quicklog.rs index a983f6d..9997d2b 100644 --- a/src/quicklog.rs +++ b/src/quicklog.rs @@ -4,12 +4,13 @@ use ndarray::Array1; use npyz::{Deserialize, NpyFile}; use std::{ fs::File, - io::{BufRead, BufReader, Lines, Seek, SeekFrom}, + io::{self, BufRead, BufReader, Lines, Seek, SeekFrom}, marker::PhantomData, path::Path, }; +use thiserror::Error; -use crate::{trace::Trace, util::read_array_1_from_npy_file}; +use crate::{trace::Trace, util::read_array1_from_npy_file}; /// Returns traces database directory from `TRACESDIR` environment variable, or `None` if it is not /// defined. @@ -45,25 +46,16 @@ pub fn guess_leakages_size(path: &str) -> usize { } /// Parsing records log can produce to types of errors, wrapped into this single error type. -#[derive(Debug)] +#[derive(Error, Debug)] pub enum LogError { - IoError(std::io::Error), - JsonError(serde_json::Error), + #[error("IO error")] + IoError(#[from] io::Error), + #[error("Failed to deserialize json")] + JsonError(#[from] serde_json::Error), + #[error("No records")] NoRecords, } -impl From for LogError { - fn from(error: std::io::Error) -> Self { - Self::IoError(error) - } -} - -impl From for LogError { - fn from(error: serde_json::Error) -> Self { - Self::JsonError(error) - } -} - /// Opens a log file and allows iterating over the records. /// /// `T` specifies the type of the elements in the leakages. @@ -178,7 +170,7 @@ impl Record { f.seek(SeekFrom::Start(toff)).unwrap(); let buf = BufReader::new(f); let npy = NpyFile::new(buf).unwrap(); - Ok(read_array_1_from_npy_file(npy)) + Ok(read_array1_from_npy_file(npy)) } else if let Some(_tid) = self.tid() { // Trace is stored in a single file todo!() @@ -195,7 +187,7 @@ pub struct FileRecordIterator { } impl FileRecordIterator { - pub fn new(path: &str) -> Result { + pub fn new(path: &str) -> Result { let file = File::open(path)?; let reader = BufReader::new(file); Ok(Self { @@ -218,10 +210,10 @@ impl Iterator for FileRecordIterator { data: value, phantom: PhantomData, })), - Err(err) => Some(Err(LogError::JsonError(err))), + Err(err) => Some(Err(err.into())), } } - Err(err) => Some(Err(LogError::IoError(err))), + Err(err) => Some(Err(err.into())), } } else { None @@ -255,7 +247,7 @@ impl CachedLoader { let toff = record.toff(); let chunk = &self.current_data.as_slice()[toff as usize..]; let npy = NpyFile::new(chunk).unwrap(); - Ok(read_array_1_from_npy_file(npy)) + Ok(read_array1_from_npy_file(npy)) } else { record.load_trace() } @@ -415,5 +407,5 @@ impl Iterator for BatchTraceIterator { pub fn array_from_bytes(bytes: &[u8], toff: usize) -> Array1 { let chunk = &bytes[toff..]; let npy = NpyFile::new(chunk).unwrap(); - read_array_1_from_npy_file(npy) + read_array1_from_npy_file(npy) } diff --git a/src/util.rs b/src/util.rs index f61d438..dc8316f 100644 --- a/src/util.rs +++ b/src/util.rs @@ -1,20 +1,22 @@ //! Convenient utility functions. -use std::{fs::File, io::BufWriter, time::Duration}; +use std::{io::Read, path::Path, time::Duration}; use indicatif::{ProgressBar, ProgressStyle}; use ndarray::{Array, Array1, Array2, ArrayView2}; -use ndarray_npy::{write_npy, ReadNpyExt, ReadableElement, WriteNpyExt}; +use ndarray_npy::{read_npy, write_npy, ReadNpyError, ReadableElement, WriteNpyError}; use npyz::{Deserialize, NpyFile}; /// Reads a [`NpyFile`] as a [`Array1`] /// /// This does the same as [`NpyFile.into_vec`] but faster, as this method reserves the resulting -/// vector to the final size directly instead of relying on `collect`. It however panics in -/// case of IO error. -pub fn read_array_1_from_npy_file(npy: NpyFile) -> Array1 { +/// vector to the final size directly instead of relying on `collect`. +/// +/// # Panics +/// This function panics in case of IO error. +pub fn read_array1_from_npy_file(npy: NpyFile) -> Array1 { let mut v: Vec = Vec::new(); - v.reserve_exact(npy.shape()[0] as usize); + v.reserve_exact(npy.shape()[0].try_into().unwrap()); v.extend(npy.data().unwrap().map(|x| x.unwrap())); Array::from_vec(v) } @@ -30,11 +32,10 @@ pub fn save_array< S: ndarray::Data, D: ndarray::Dimension, >( - path: &str, + path: impl AsRef, array: &ndarray::ArrayBase, -) { - // let dir = BufWriter::new(File::create(path).unwrap()); - write_npy(path, array).unwrap(); +) -> Result<(), WriteNpyError> { + write_npy(path, array) } /// Creates a [`ProgressBar`] with a predefined default style. @@ -46,13 +47,12 @@ pub fn progress_bar(len: usize) -> ProgressBar { progress_bar } -pub fn read_array_2_from_npy_file(dir: &str) -> Array2 { - let reader: File = File::open(dir).unwrap(); - let arr: Array2 = Array2::::read_npy(reader).unwrap(); - arr +pub fn read_array2_from_npy_file( + path: impl AsRef, +) -> Result, ReadNpyError> { + read_npy(path) } -pub fn save_array2(path: &str, ar: ArrayView2) { - let writer = BufWriter::new(File::create(path).unwrap()); - ar.write_npy(writer).unwrap(); +pub fn save_array2(path: impl AsRef, array: ArrayView2) -> Result<(), WriteNpyError> { + write_npy(path, &array) }