From f23a540f3da63be234fb0a76feafb6725bc51060 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Thomas=20Pr=C3=A9vost?= Date: Wed, 14 Feb 2024 16:43:21 +0100 Subject: [PATCH] Can retrieve Shannon's entropy of all stored keys --- src/entropy/mod.rs | 63 ++++++++++++++++++++++++++++ src/lib.rs | 5 +-- src/qkd_manager/http_response_obj.rs | 16 +++++++ src/qkd_manager/mod.rs | 58 +++++++++++++++++++++---- src/routes/keys/mod.rs | 3 ++ src/routes/keys/route_entropy.rs | 30 +++++++++++++ tests/data/total_entropy.json | 3 ++ tests/entropy.rs | 22 ++++++++++ 8 files changed, 189 insertions(+), 11 deletions(-) create mode 100644 src/entropy/mod.rs create mode 100644 src/routes/keys/route_entropy.rs create mode 100644 tests/data/total_entropy.json create mode 100644 tests/entropy.rs diff --git a/src/entropy/mod.rs b/src/entropy/mod.rs new file mode 100644 index 0000000..9380482 --- /dev/null +++ b/src/entropy/mod.rs @@ -0,0 +1,63 @@ +pub(crate) trait EntropyAccumulator { + fn add_bytes(&mut self, bytes: &[u8]); + fn get_entropy(&self) -> f64; +} + +pub(crate) struct ShannonEntropyAccumulator { + /// Counter for each byte value + bytes_counter: [u64; 256], + /// Total received bytes + total_bytes: u64, +} + +impl ShannonEntropyAccumulator { + pub(crate) fn new() -> Self { + Self { + bytes_counter: [0; 256], + total_bytes: 0, + } + } +} + +impl EntropyAccumulator for ShannonEntropyAccumulator { + fn add_bytes(&mut self, bytes: &[u8]) { + for byte in bytes { + self.bytes_counter[*byte as usize] += 1; + } + self.total_bytes += bytes.len() as u64; + } + + fn get_entropy(&self) -> f64 { + let mut entropy = 0.0; + for count in self.bytes_counter.iter() { + if *count == 0 { + continue; + } + let symbol_probability = *count as f64 / self.total_bytes as f64; + entropy -= symbol_probability * symbol_probability.log2(); + } + entropy + } +} + +#[cfg(test)] +mod tests { + use crate::entropy::EntropyAccumulator; + + #[test] + fn test_shannon_entropy_accumulator() { + let mut entropy_accumulator_1 = super::ShannonEntropyAccumulator::new(); + entropy_accumulator_1.add_bytes(&[0, 0, 0, 0, 0, 0]); + assert_eq!(entropy_accumulator_1.get_entropy(), 0.0); + + let mut entropy_accumulator_2 = super::ShannonEntropyAccumulator::new(); + entropy_accumulator_2.add_bytes(&[0x00, 0x00, 0x01, 0x01, 0x02]); + assert_eq!(entropy_accumulator_2.get_entropy(), 1.5219280948873621); + + let mut entropy_accumulator_3 = super::ShannonEntropyAccumulator::new(); + entropy_accumulator_3.add_bytes("Souvent sur la montagne, à l’ombre du vieux chêne,\n".as_bytes()); + assert_eq!(entropy_accumulator_3.get_entropy(), 4.465641023041018); + entropy_accumulator_3.add_bytes("Au coucher du soleil, tristement je m’assieds ;\n".as_bytes()); + assert_eq!(entropy_accumulator_3.get_entropy(), 4.507894683096287); + } +} \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index 1fa9e21..17a6c8a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -15,6 +15,7 @@ pub mod server; pub mod routes; pub mod qkd_manager; pub mod config; +pub(crate) mod entropy; /// Cast a string to an io::Error @@ -82,6 +83,4 @@ mod test { assert_eq!(err.kind(), std::io::ErrorKind::Other); assert_eq!(err.to_string(), "test"); } -} - -// TODO shannon entropy \ No newline at end of file +} \ No newline at end of file diff --git a/src/qkd_manager/http_response_obj.rs b/src/qkd_manager/http_response_obj.rs index a282d27..b9fe048 100644 --- a/src/qkd_manager/http_response_obj.rs +++ b/src/qkd_manager/http_response_obj.rs @@ -98,6 +98,13 @@ pub(crate) struct ResponseQkdKeysList { } impl HttpResponseBody for ResponseQkdKeysList {} +#[derive(serde::Serialize, Debug, PartialEq)] +pub(crate) struct ResponseTotalKeysEntropy { + pub(crate) total_entropy: f64, +} + +impl HttpResponseBody for ResponseTotalKeysEntropy {} + #[cfg(test)] mod test { use crate::qkd_manager::http_response_obj::HttpResponseBody; @@ -151,4 +158,13 @@ mod test { let response_qkd_sae_info_json = response_qkd_sae_info.to_json().unwrap(); assert_eq!(response_qkd_sae_info_json.replace("\r", ""), "{\n \"SAE_ID\": 1,\n \"KME_ID\": 1\n}"); } + + #[test] + fn test_serialize_response_total_keys_entropy() { + let response_total_keys_entropy = super::ResponseTotalKeysEntropy { + total_entropy: 1.0, + }; + let response_total_keys_entropy_json = response_total_keys_entropy.to_json().unwrap(); + assert_eq!(response_total_keys_entropy_json.replace("\r", ""), "{\n \"total_entropy\": 1.0\n}"); + } } \ No newline at end of file diff --git a/src/qkd_manager/mod.rs b/src/qkd_manager/mod.rs index f56b4a0..99b73da 100644 --- a/src/qkd_manager/mod.rs +++ b/src/qkd_manager/mod.rs @@ -12,7 +12,8 @@ use log::error; use sha1::Digest; use crate::qkd_manager::http_response_obj::ResponseQkdKeysList; use crate::qkd_manager::QkdManagerResponse::TransmissionError; -use crate::{KmeId, QkdEncKey, SaeClientCertSerial, SaeId}; +use crate::{io_err, KmeId, QkdEncKey, SaeClientCertSerial, SaeId}; +use crate::entropy::{EntropyAccumulator, ShannonEntropyAccumulator}; /// QKD manager interface, can be cloned for instance in each request handler task #[derive(Clone)] @@ -25,6 +26,8 @@ pub struct QkdManager { pub(crate) dir_watcher: Arc>>, /// The ID of the KME this QKD manager belongs to pub kme_id: KmeId, + /// Shannon's entropy calculator for keys stored in the database + shannon_entropy_calculator: Arc>, } impl QkdManager { @@ -60,6 +63,7 @@ impl QkdManager { response_rx, dir_watcher, kme_id: this_kme_id, + shannon_entropy_calculator: Arc::new(Mutex::new(ShannonEntropyAccumulator::new())), } } @@ -74,20 +78,27 @@ impl QkdManager { } /// Add a new QKD key to the database + /// Increases the total entropy of all keys in the database /// # Arguments /// * `key` - The QKD key to add (key + origin SAE ID + target SAE ID) /// # Returns /// Ok if the key was added successfully, an error otherwise pub fn add_pre_init_qkd_key(&self, key: PreInitQkdKeyWrapper) -> Result { - self.command_tx.send(QkdManagerCommand::AddPreInitKey(key)).map_err(|_| { + const EXPECTED_QKD_MANAGER_RESPONSE: QkdManagerResponse = QkdManagerResponse::Ok; + + self.command_tx.send(QkdManagerCommand::AddPreInitKey(key.to_owned())).map_err(|_| { TransmissionError })?; - match self.response_rx.recv().map_err(|_| { + let add_key_status = self.response_rx.recv().map_err(|_| { TransmissionError - })? { - QkdManagerResponse::Ok => Ok(QkdManagerResponse::Ok), // Ok is the QkdManagerResponse expected here - qkd_response_error => Err(qkd_response_error), + })?; + if add_key_status != EXPECTED_QKD_MANAGER_RESPONSE { + return Err(add_key_status); } + self.shannon_entropy_calculator.lock().map_err(|_| { + TransmissionError + })?.add_bytes(&key.key); + Ok(EXPECTED_QKD_MANAGER_RESPONSE) } /// Get a QKD key from the database (shall be called by the master SAE) @@ -254,6 +265,20 @@ impl QkdManager { qkd_response_error => Err(qkd_response_error), } } + + /// Get the Shannon entropy of all stored keys + /// # Returns + /// The total Shannon entropy of all stored keys, an error in case of concurrency issues + pub async fn get_total_keys_shannon_entropy(&self) -> Result { + let entropy_calculator = Arc::clone(&self.shannon_entropy_calculator); + Ok(tokio::task::spawn_blocking(move || { + Ok::(entropy_calculator + .lock() + .map_err(|_| io_err("Mutex locking error"))? + .get_entropy()) + }).await + .map_err(|_| io_err("Async task error"))??) + } } /// A Pre-init QKD key, with its origin and target KME IDs @@ -383,17 +408,34 @@ pub enum QkdManagerResponse { #[cfg(test)] mod test { use serial_test::serial; + use crate::QkdEncKey; const CLIENT_CERT_SERIAL_SIZE_BYTES: usize = 20; - #[test] - fn test_add_qkd_key() { + #[tokio::test] + async fn test_add_qkd_key() { const SQLITE_DB_PATH: &'static str = ":memory:"; let qkd_manager = super::QkdManager::new(SQLITE_DB_PATH, 1); let key = super::PreInitQkdKeyWrapper::new(1, &[0; crate::QKD_KEY_SIZE_BYTES]).unwrap(); let response = qkd_manager.add_pre_init_qkd_key(key); assert!(response.is_ok()); assert_eq!(response.unwrap(), super::QkdManagerResponse::Ok); + assert_eq!(qkd_manager.get_total_keys_shannon_entropy().await.unwrap(), 0.0); + } + + #[tokio::test] + async fn test_stored_keys_entropy() { + const SQLITE_DB_PATH: &'static str = ":memory:"; + let first_key: QkdEncKey = <[u8; crate::QKD_KEY_SIZE_BYTES]>::try_from("ABCDEFGHIJKLMNOPQRSTUVWXYZ012345".as_bytes()).unwrap(); + let second_key: QkdEncKey = <[u8; crate::QKD_KEY_SIZE_BYTES]>::try_from("6789+-abcdefghijklmnopqrstuvwxyz".as_bytes()).unwrap(); + + let qkd_manager = super::QkdManager::new(SQLITE_DB_PATH, 1); + let key = super::PreInitQkdKeyWrapper::new(1, &first_key).unwrap(); + qkd_manager.add_pre_init_qkd_key(key).unwrap(); + assert_eq!(qkd_manager.get_total_keys_shannon_entropy().await.unwrap(), 5.0); + let key = super::PreInitQkdKeyWrapper::new(1, &second_key).unwrap(); + qkd_manager.add_pre_init_qkd_key(key).unwrap(); + assert_eq!(qkd_manager.get_total_keys_shannon_entropy().await.unwrap(), 6.0); } #[test] diff --git a/src/routes/keys/mod.rs b/src/routes/keys/mod.rs index c042424..fc7c753 100644 --- a/src/routes/keys/mod.rs +++ b/src/routes/keys/mod.rs @@ -7,6 +7,7 @@ use hyper::body::Bytes; use crate::routes::{EtsiSaeQkdRoutesV1, RequestContext}; mod get_key; +mod route_entropy; /// Dispatches the request to the correct function pub(super) async fn key_handler(rcx: &RequestContext<'_>, req: Request, uri_segments: &[&str]) -> Result>, Infallible> { @@ -17,6 +18,8 @@ pub(super) async fn key_handler(rcx: &RequestContext<'_>, req: Request get_key::route_get_key(rcx, req, slave_sae_id).await, // Get key(s) from a slave SAE, with ID provided by the master SAE ([slave_sae_id, "dec_keys"], &hyper::Method::POST) => get_key::route_get_key_with_id(rcx, req, slave_sae_id).await, + // Retrieve Shannon's entropy for all stored keys in KME database + (["entropy", "total"], &hyper::Method::GET) => route_entropy::route_get_entropy_total(rcx, req).await, // Route not found _ => EtsiSaeQkdRoutesV1::not_found(), } diff --git a/src/routes/keys/route_entropy.rs b/src/routes/keys/route_entropy.rs new file mode 100644 index 0000000..2f94827 --- /dev/null +++ b/src/routes/keys/route_entropy.rs @@ -0,0 +1,30 @@ +use std::convert::Infallible; +use http_body_util::Full; +use hyper::{body, Request, Response}; +use hyper::body::Bytes; +use log::error; +use crate::qkd_manager::http_response_obj::HttpResponseBody; +use crate::routes::request_context::RequestContext; + +pub(in crate::routes) async fn route_get_entropy_total(rcx: &RequestContext<'_>, _req: Request) -> Result>, Infallible> { + // Get the total entropy from stored keys + let entropy = match rcx.qkd_manager.get_total_keys_shannon_entropy().await { + Ok(entropy) => entropy, + Err(e) => { + error!("Error getting total entropy: {}", e.to_string()); + return super::EtsiSaeQkdRoutesV1::internal_server_error(); + } + }; + let total_entropy_response_obj = crate::qkd_manager::http_response_obj::ResponseTotalKeysEntropy { + total_entropy: entropy, + }; + // Return the total entropy as a JSON object response + let total_entropy_response_obj_json = match total_entropy_response_obj.to_json() { + Ok(json) => json, + Err(_) => { + error!("Error serializing total entropy object"); + return super::EtsiSaeQkdRoutesV1::internal_server_error(); + } + }; + Ok(crate::routes::EtsiSaeQkdRoutesV1::json_response_from_str(&total_entropy_response_obj_json)) +} \ No newline at end of file diff --git a/tests/data/total_entropy.json b/tests/data/total_entropy.json new file mode 100644 index 0000000..c271a73 --- /dev/null +++ b/tests/data/total_entropy.json @@ -0,0 +1,3 @@ +{ + "total_entropy": 3.6123008763572906 +} \ No newline at end of file diff --git a/tests/entropy.rs b/tests/entropy.rs new file mode 100644 index 0000000..3023521 --- /dev/null +++ b/tests/entropy.rs @@ -0,0 +1,22 @@ +use const_format::concatcp; +use serial_test::serial; +use crate::common::util::assert_string_equal; + +mod common; + +#[tokio::test] +#[serial] +async fn get_total_entropy() { + const EXPECTED_BODY: &'static str = include_str!("data/total_entropy.json"); + const REQUEST_URL: &'static str = concatcp!("https://", common::HOST_PORT ,"/api/v1/keys/entropy/total"); + + common::setup(); + let reqwest_client = common::setup_cert_auth_reqwest_client(); + + let response = reqwest_client.get(REQUEST_URL).send().await; + assert!(response.is_ok()); + let response = response.unwrap(); + assert_eq!(response.status(), 200); + let response_body = response.text().await.unwrap(); + assert_string_equal(&response_body, EXPECTED_BODY); +} \ No newline at end of file